기술적 요약: PFN-TS

문제 정의

맥락형 밴딧 (Contextual bandits) 은 불확실성 하에서 순차적 의사결정을 다루며, 여기서 학습자는 맥락을 관찰하고 행동을 선택한 후 확률적 보상을 받습니다. 목표는 시간에 따른 누적 후회 (regret) 를 최소화하는 것입니다. **Thompson Sampling(TS)**은 보상 함수에 대한 베이지안 사후분포를 유지하고 이 사후분포에서 행동을 탐욕적으로 샘플링하는 표준 전략입니다. 그러나 TS 의 효과는 근본적인 보상 모델의 품질에 달려 있습니다.

표 형식 (tabular) 응용 분야 (예: 모바일 헬스, 추천 시스템) 에서는 상호작용, 임계값 효과, 이질성과 같은 복잡한 인간 행동을 포착하지 못하기 때문에 선형 보상 모델이 종종 오명세 (misspecified) 됩니다. 더 풍부한 대안들 (커널 방법, 신경망, 트리 앙상블) 이 존재하지만, 확장성 문제, 높은 계산 오버헤드, 또는 경직된 구조적 가정으로 인해 고통받습니다. 또한 이러한 모델들은 일반적으로 온라인 환경에서 조정하기 어려운 하이퍼파라미터를 필요로 합니다.

**Prior-Data Fitted Networks(PFNs, TabPFN 및 TabICL 등)**는 유망한 대안을 제공합니다. 이들은 광범위한 데이터 생성 과정의 사전분포 (prior) 에 걸쳐 베이지안 사후 예측 분포 (PPDs) 를 근사하도록 훈련된 트랜스포머 기반 모델들입니다. 이들은 문맥 내 학습 (In-Context Learning, ICL) 을 통해 작동하며, 파라미터 업데이트나 하이퍼파라미터 조정이 필요 없습니다. 그러나 중요한 간극이 존재합니다: PFNs 는 PPD(노이즈가 있는 미래 보상에 대한 불확실성) 를 출력하는 반면, TS 는 잠재적 평균 보상 함수에 대한 불확실성을 필요로 합니다. PPD 에서 직접 샘플링하는 것은 평균 함수의 사후분포에서 샘플링하는 것과 동일하지 않아, 비최적의 탐색을 초래합니다.

방법론: PFN-TS

저자들은 PFN-TS 를 제안합니다. 이는 PFN 사후 예측을 잠재적 평균 보상 함수의 샘플로 변환함으로써 이 간극을 메우는 Thompson Sampling 알고리즘입니다. 방법론은 세 가지 핵심 구성 요소로 이루어집니다:

1. 부분표본 예측 중심극한정리 (SubCLT)

PFN 의 예측 시퀀스에서 잠재적 평균 $f_0(x)$의 사후분포를 복원하기 위해, 저자들은 예측 중심극한정리 (CLT) 를 활용합니다.

• 이론적 기반: 특정 가정 (점근적 분산, 준-마팅게일 조건, 유계 모멘트) 하에서, 예측 평균 시퀀스의 꼬리는 잠재적 평균을 중심으로 하는 가우시안 분포로 수렴합니다.
• 효율성 혁신: 전통적인 사후분산 추정 접근법은 데이터셋의 모든 $n$개의 접두사 (prefix) 에 대해 예측 평균을 계산해야 하므로, 온라인 밴딧에서는 계산적으로 불가능합니다 ($O(n)$개의 순전파). PFN-TS 는 $O(\log n)$개의 기하급수적으로 간격을 둔 접두사 ($t_0, t_1, \dots, t_J$) 만을 사용하여 분산을 계산하는 부분표본 추정기를 도입합니다.
• 분산 추정: 분산은 $D_j = m_{t_j}(x) - m_{t_{j-1}}(x)$인 블록 증가분을 조화적으로 가중치 ($w_j$) 하여 추정합니다. 저자들은 이 부분표본 추정기가 일관성 ($\hat{V}_{s_n} \xrightarrow{p} V(x)$) 을 가지며 CLT 를 만족함을 증명하여, Thompson Sampling 을 위한 처리 가능한 가우시안 근사 $N(m_{s_n}(x), \hat{V}_{s_n}(x)/s_n)$을 가능하게 합니다.

2. 캐싱 및 계산 최적화

라운드당 비용을 추가로 줄이기 위해 PFN-TS 는 사용된 특정 PFN 인 TabICL 의 아키텍처를 활용합니다:

• KV-Cache 재사용: TabICL 은 고정된 훈련 세트에 대해 키 - 값 (key-value) 어텐션 행렬을 캐시합니다. PFN-TS 는 이러한 캐시된 표현을 라운드 간에 재사용합니다.
• 새로고침 지점 (Refresh Points): 알고리즘은 SubCLT 그리드에 의해 정의된 기하급수적인 "새로고침 지점"($t_j$) 에서만 새로운 순전파를 수행합니다. 새로고침 지점 사이에서는 최신 캐시된 상태에서 파생된 스냅샷 가우시안 근사를 사용합니다. 이로 인해 라운드당 암 (arm) 당 순전파 복잡도가 $O(n)$에서 $O(\log n)$으로 감소합니다.

3. 적응형 인코딩 전략

이 논문은 입력에서 암 (arm) 의 정체성을 표현하는 방식에 대한 설계 선택을 다룹니다:

• 분리 인코딩 (Disjoint Encoding): 암마다 독립적인 모델을 유지합니다.
• 결합 인코딩 (Joint Encoding): 맥락에 원-핫 벡터로 암의 정체성을 추가하여 암 간 정보 공유를 가능하게 합니다.
• 적응형 선택: 어떤 인코딩도 보편적으로 우세하지 않기 때문에, PFN-TS 는 "이중 인코딩 (dual-encoding)"단계 동안 두 가지를 병렬로 실행합니다. **연속 순위 확률 점수 (CRPS, Continuous Ranked Probability Score)**를 기반으로 활성 인코딩을 선택합니다. CRPS 는 보정 (calibration) 을 측정하는 엄격하게 적절한 점수 규칙입니다. 누적 CRPS 가 더 낮은 인코딩이 유지되고, challenger 는 사전 지정된 전환 시간 이후에 폐기됩니다.

주요 기여

1. 알고리즘적 연결고리: PFN-TS 는 일관된 부분표본 CLT 를 통해 잠재적 평균 사후분산에 대한 일관된 추정기를 도출함으로써, PPD 를 출력하는 PFNs 를 Thompson Sampling 에 성공적으로 적응시킨 최초의 방법입니다.
2. 이론적 보장:
   • 일관성: mild 한 가정 하에서 부분표본 분산 추정기의 일관성을 증명합니다.
   • 후회 상한 (Regret Bound): PFN-TS 에 대한 베이지안 후회 상한을 유도합니다. 이 상한은 PFN 의 사전분포 하에서의 정확한 Thompson Sampling 의 후회와 PFN 및 부분표본 과정의 근사 품질에 의존하는 오차 항으로 분해됩니다.
3. 효율성: $O(\log n)$개의 순전파만으로도 일관된 분산 추정이 충분함을 보여, PFN 기반 밴딧이 온라인 배포에 실현 가능함을 입증합니다.
4. 실증적 성능:
   • 합성 벤치마크: PFN-TS 는 비선형 합성 작업 (Friedman 변형) 에서 평균 순위가 가장 우수하며, 선형 작업에서도 경쟁력을 유지합니다.
   • OpenML 벤치마크: 여덟 개의 분류 - 밴딧 데이터셋 중 여섯 개에서 LinTS, NeuralTS, BFTS, XGBoostTS 와 같은 베이스라인을 능가합니다.
   • 실제 응용: Drink Less 모바일 헬스 시험의 오프라인 평가에서 PFN-TS 는 최종 시점에서 가장 높은 추정 정책 가치 (SNIPS) 를 달성하여 BFTS 및 기타 베이스라인을 능가했습니다.

중요성 및 주장

이 논문은 PFN-TS 가 효율적인 불확실성 추출과 적응형 인코딩이 구현된다면, 표 형식 문맥 내 학습이 순차적 의사결정을 위한 효과적인 보상 모델로 기능할 수 있음을 보여준다고 주장합니다.

• 모델 유연성: 전통적인 모델과 달리 PFN-TS 는 하이퍼파라미터 조정이 필요 없으며, PFN 의 광범위한 사전분포에 내재된 데이터 구조 (희소성, 부드러움, 비선형성) 에 자동으로 적응합니다.
• 실용적 타당성: 데이터셋 크기에 대한 불확실성 추정의 계산 병목 현상을 선형에서 로그arithmic 으로 줄임으로써, 온라인 밴딧 환경에서 파운데이션 모델의 배포를 실용적으로 만듭니다.
• 강건성: 이 방법은 다양한 데이터 생성 과정에 걸쳐 강건하며, 비선형 합성 작업과 실제 모바일 헬스 데이터 모두에서 잘 수행됨이 입증되었습니다.

저자들은 한계에 대해 겸손하게 언급하며, 반복적인 모델 평가로 인해 런타임이 여전히 제약 사항이며, 암의 역사가 짧거나 정보가 약할 때 유한 표본 보정이 poor 할 수 있음을 지적합니다. 그들은 향후 연구로 정규화된 분산 추정과 특정 TabPFN/TabICL 사전분포에 대한 더 날카로운 후회 상한을 제안합니다.