CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

기술 요약: CRUMB – 분포 일치형 컨텍스트 배칭을 통한 효율적인 사전 적합 네트워크 추론

문제 정의
TabPFN 및 TabICL과 같은 사전 적합 네트워크(Prior-fitted networks, PFNs)는 인컨텍스트 학습(in-context learning)이 가능한 테이블형 파운데이션 모델의 한 부류를 나타냅니다. 이들은 레이블이 지정된 전체 훈련 세트를 컨텍스트로 받아 단 한 번의 순전파(forward pass)를 통해 테스트 쿼리에 대한 예측을 생성하며, 종종 소규모 내지 중규모 데이터셋에서 전통적인 그래디언트 부스팅 방법보다 뛰어난 성능을 보입니다. 그러나 이러한 트랜스포머 기반 PFN의 셀프 어텐션 메커니즘은 훈련 샘플 수($N$)에 따라 제곱($O(N^2)$)으로 스케일링된다는 결정적인 병목 현상이 존재합니다. 이는 대규모 데이터셋($N > 50,000$)에서 추론을 계산적으로 불가능하게 만듭니다.

기존 솔루션들은 정확도와 효율성 사이의 절충안에 직면해 있습니다. 균등 서브샘플링(Uniform subsampling)은 훈련 데이터를 무작위로 버림으로써 속도를 개선하지만 정보를 희생합니다. $k$-최근접 이웃($k$NN)과 같은 쿼리 의존적 방법은 테스트 포인트에 맞춤화된 컨텍스트를 제공하여 정확도를 높이지만, 배칭(batching)을 방해하여 $T$개의 테스트 포인트에 대해 $T$번의 별도 순전파를 요구합니다. MICP(Mixture of In-Context Prompters)와 같은 다른 접근 방식은 훈련 데이터를 클러스터링하여 쿼리를 배칭하려고 시도하지만, 특히 공변량 드리프트(covariate drift) 상황에서 훈련 컨텍스트를 테스트 데이터의 특정 분포와 정렬하는 데 실패할 수 있습니다.

방법론
저자들은 모델 재학습이 필요 없는 3단계, 아키텍처 불가지론적(architecture-agnostic) 추론 래퍼인 CRUMB(Clustered Retrieval Using Minimised-MMD Batching)를 제안합니다. 이 방법은 다음 단계들을 통해 컨텍스트 품질과 배칭 효율성 사이의 긴장을 해결합니다.

1. 테스트 측 클러스터링(Test-Side Clustering): 테스트 쿼리($D_{test}$)는 표준화된 입력 피처에 대한 $k$-평균(k-means)을 사용하여 $K$개의 클러스터($C_1, \dots, C_K$)로 분할됩니다. 이를 통해 시스템은 공통된 훈련 컨텍스트를 공유할 수 있는 유사한 테스트 포인트들을 그룹화할 수 있습니다.
2. 분포 일치형 컨텍스트 선택(Distributionally Matched Context Selection): 각 테스트 클러스터 $C_k$에 대해, 전체 훈련 풀($D_{train}$)에서 작은 훈련 서브셋 $S_k$(크기 $n \ll N$)가 선택됩니다. 기하학적 근접성 방법(예: 중심점-최근접 이웃)과 달리, CRUMB는 테스트 클러스터의 경험적 분포와 후보 훈련 서브셋 사이의 **최대 평균 불일치(Maximum Mean Discrepancy, MMD)**를 탐욕적으로 최소화함으로써 $S_k$를 선택합니다.
   • 이는 테스트 클러스터와의 근접성(관련성)에 보상을 주고, 선택된 훈련 포인트 간의 중복성(다양성)에 페널티를 주는 **커널 허딩(kernel herding)**이라는 탐욕적 알고리즘을 통해 수행됩니다.
   • 이 과정을 가속화하기 위해, 저자들은 가우시안 RBF 커널을 근사하기 위해 랜덤 푸리에 특징(Random Fourier Features, RFF)을 사용하고 배치형 탐욕 선택(batched greedy selection)을 채택했습니다.
3. 배치 추론(Batched Inference): PFN 모델은 $K$번의 독립적인 순전파를 실행합니다. 각 순전파에서 모델은 선택된 훈련 서브셋 $S_k$를 컨텍스트로 받고, 전체 클러스터 $C_k$를 쿼리로 받습니다. 이는 총 순전파 횟수를 (per-query $k$NN의 경우처럼) $T$에서 $K$로 줄여, 분포 정렬을 유지하면서도 효율적인 GPU 활용을 가능하게 합니다.

주요 기여

• MMD 기반 컨텍스트 검색: 본 논문은 훈련 서브셋을 선택하는 메커니즘으로서 탐욕적 MMD 최소화(커널 허딩)를 도입했습니다. 저자들은 이러한 분포 매칭 접근 방식이 동일한 클러스터링 프레임워크 내에서 중심점-최근접 이웃 및 Voronoi-균등 검색 전략보다 우수함을 입증했습니다.
• 테스트 측 클러스터링을 통한 배치 컨텍스트 선택: 테스트 쿼리를 클러스터링하고 각 클러스터에 분포적으로 일치하는 단일 훈련 컨텍스트를 할당함으로써, CRUMB는 쿼리 관련성을 희생하지 않으면서 효율적인 배칭을 가능하게 합니다.
• 공변량 드리프트에 대한 강건성: CRUMB는 공변량 드리프트를 자연스럽게 처리합니다. 훈련 컨텍스트가 (고정된 훈련 분포가 아니라) 현재의 테스트 배치 분포에 맞춰 선택되기 때문에, CRUMB는 변화하는 테스트 데이터에 적응합니다. 반면 고정된 훈련 클러스터에 의존하는 MICP와 같은 방법은 테스트 분포가 훈련 분포에서 벗어날 때 크게 성능이 저하됩니다.

실험 결과
저자들은 세 가지 PFN 아키텍처(TabPFNv2, TabICLv1, TabICLv2)에 대해 TabArena 벤치마크(51개의 다양한 테이블 데이터셋)에서 CRUMB를 평가했습니다.

• 베이스라인 대비 성능: TabArena 벤치마크에서 CRUMB는 동일한 컨텍스트 예산 하에서 균등 서봅플링 및 MICP 기술을 유의미하게 능가했습니다. CRUMB는 성능 면에서 per-query $k$NN과 통계적으로 유사하지만, $T$번의 순전파가 아닌 고정된 $K$번의 순전파(예: $K=20$)로 이를 달성하여 추론 시간을 대폭 단축했습니다(보고된 실험에서 361초 대비 35초).
• 대규모 데이터셋: 최대 120,000개의 샘플을 가진 데이터셋(예: GiveMeSomeCredit, Higgs)에 대한 실험에서, 조기 종료(MMD 개선이 정체될 때 컨텍스트 선택을 중단)를 적용한 CRUMB는 특히 TabPFNv2 및 TabICLv1에 대해 MICP를 더 자주 이겼습니다.
• 공변량 드리프트: 테스트 분포가 첫 번째 주성분(first principal component)을 따라 이동하는 통제된 드리프트 실험에서, CRUMB의 MICP 대비 우위는 드리프트 강도가 높아짐에 따라 단조 증가했습니다. 드리프트가 0일 때 CRUMB는 +4.9% 앞섰으며, 최대 드리프트(완전한 out-of-sample) 시 격차는 +17.1%로 벌어졌습니다. 이는 MICP의 고정된 훈련 클러스터가 정렬을 잃는 반면, CR음의 훈련 컨텍스트를 이동된 테스트 분포에 재정렬하는 능력 덕분입니다.

의의 및 주장
본 논문은 CRUMB가 아키텍처 변경이나 모델 미세 조정 없이 대규모 테이블 데이터셋에 대한 PFN 추론을 실용적으로 만든다고 주장합니다. 테스트 측 클러스터링과 분포 일치형 컨텍스트 검색을 결합함으로써, CRUMB는 per-query 검색의 정확도와 균등 배칭의 효율성 사이의 "스윗 스팟(sweet spot)"을 달えます.

저자들은 핵심 혁신이 **테스트 적응형 설계(test-adaptive design)**에 있다고 강조합니다. 즉, 기하학적 근접성만이 아니라 분포적 정렬(MMD)을 기반으로 테스트 포인트를 클러스터링하고 컨텍스트를 선택하는 것입니다. 이 설계는 기존의 컨텍스트 선택 방법들이 흔히 실패하는 영역인 공변량 드리프트에 대해 자연스러운 강건성 메커즘을 제공합니다. 결론적으로, CRUMB는 기본 PFN의 사전 학습된 프라이어(prior)가 가진 성능 상한을 그대로 물려받지만, PFN을 대규모 데이터셋에 적용하는 데 있어 계산적 장벽을 효과적으로 제거합니다.

한계점
저자들은 CRUMB가 클러스터링을 수행하기 위해 전체 테스트 쿼리를 사전에 필요로 한다는 점을 인정하며, 이는 수정 없이는 진정한 온라인 스트리밍 환경에 적합하지 않음을 시사합니다(예: 센트로이드를 캐싱하는 방식은 사실상 MICP 스타일의 라우팅으로 축소됨). 또한, 유클리드 $k$-평균 및 등방성 가우시안 커널에 대한 의존도는 고차원 또는 혼합 유형 피처 공간에서 성능을 제한할 수 있습니다. 마지막으로, 순수 추론 래퍼로서 CRUMB는 PFN의 사전 학습된 프라이어와 특정 데이터셋 간의 불일치를 보완할 수는 없습니다.