기술 요약: 포켓 파운데이션 모델

문제 제기

TabICLv2, TabPFNv2.6, LimiX 와 같은 표 형식 파운데이션 모델 (TFM) 은 대규모 트랜스포머를 통한 컨텍스트 학습 (ICL) 을 활용하여 작은 구조화된 데이터셋에서 최첨단 정확도를 입증해 왔습니다. 그러나 이러한 모델은 배포에 있어 치명적인 장벽에 직면해 있습니다. 즉, 쿼리 시 전체 훈련 세트를 주의 깊게 살펴봐야 하므로 GPU 에서 추론 지연 시간이 151ms 에서 1,275ms 사이에 발생합니다. 이 지연 시간은 사기 탐지, 신용 점수 산정, 환자 분류와 같은 실시간 애플리케이션과 호환되지 않으며, 이러한 애플리케이션은 종종 2ms 미만의 응답을 요구합니다. 더 나아가 GPU 비용과 CPU 전용 생산 제약 조건으로 인해 TFM 의 직접적인 배포는 비현실적입니다.

지식 증류는 이러한 모델을 더 빠른 경량 부스팅 결정 트리 (GBDT) 학생 모델로 압축하는 길을 제공하지만, ICL 기반 교사 모델에는 특정 기술적 장애물이 존재합니다. ICL 교사 모델이 이미 자신의 컨텍스트 (즉, 자신의 훈련 세트) 에 포함된 예제에 점수를 매길 때, 그 예측은 엔트로피가 거의 0 인 거의 원-핫 (one-hot) 벡터로 붕괴됩니다. 이러한 "누출"은 효과적인 증류에 필요한 클래스 간 구조 (소프트 레이블) 를 파괴하여, 학생 모델이 하드 레이블을 통해 학습했을 때보다 더 이상 학습하지 못하게 하거나, 종종 표준 훈련보다 더 나쁜 성능을 보이게 합니다.

방법론

본 논문은 TFM 교사 모델을 CPU 준비가 된 GBDT 또는 MLP 학생 모델로 변환하는 누출 인식 증류 파이프라인을 제안합니다. 핵심 혁신은 컨텍스트 누출을 방지하기 위한 층화 아웃 - 오브 - 폴드 (OOF) 라벨링의 사용입니다.

1. OOF 라벨링: 훈련 데이터셋 $D$를 $K=5$개의 층화된 폴드로 분할합니다. 각 폴드 $k$에 대해, 교사 모델은 $D \setminus D_k$에 적합시키고 홀드아웃 세트 $D_k$에 대해서만 확률을 예측합니다. 이를 통해 교사 모델이 컨텍스트에서 쿼리 예제를 절대 보지 못하도록 하여 증류에 필요한 소프트 레이블 신호 (클래스 간 구조) 를 보존합니다.
2. 다중 교사 평균화: 다중 교사 설정에서 TabPFN, TabICL, LimiX 등 서로 다른 TFM 계열의 OOF 예측값은 목표값으로 사용되기 전에 동일한 가중치로 평균화됩니다.
3. 학생 목표 함수: 학생 모델 (XGBoost, CatBoost, LightGBM 또는 MLP) 은 힌톤 스타일의 혼합 손실 함수를 사용하여 훈련됩니다:
   $$\mathcal{L} = \alpha \sum_i w_i T_i^2 \text{KL}(\hat{p}_i^{T_i} \parallel q_i^{T_i}) + (1-\alpha) \sum_i w_i \ell_{\text{CE}}(y_i, q_i)$$
   여기서 $\alpha=0.7$이며, $T_i$는 교사 엔트로피에 따른 샘플별 온도 스케일링이고, $w_i$는 과도하게 확신하거나 거의 무작위적인 예측의 가중치를 낮추는 신뢰도 가중치입니다. 트리 기반 학생 모델의 경우, KL 항은 소프트 레이블 로짓에 대한 클래스별 MSE 회귀로 단순화됩니다.
4. 교사 선택: 본 논문은 홀드아웃 샘플에서의 단독 AUC 로 교사 모델을 순위 매긴 것이 결과 학생 모델의 순위와 정확히 일치한다는 것을 발견했습니다. 따라서 최고의 교사를 선택하는 데 복잡한 아키텍처 탐색이 필요하지 않으며, 단순히 가장 높은 단독 AUC 를 가진 교사를 선택하면 됩니다.

주요 기여

1. 필수적인 OOF 라벨링: 저자들은 ICL 기반 TFM 에게 OOF 라벨링이 선택 사항이 아님을 확인했습니다. 이를 수행하지 않으면 소프트 레이블 신호가 붕괴되어 증류가 하드 레이블 기준선보다 열등한 학생 모델을 생성합니다.
2. 성능 대 지연 시간: TabICLv2 를 XGBoost 로 증류하면 단일 CPU 코어에서 1.9ms 의 지연 시간으로 0.882 의 매크로 평균 AUC 를 달성하여 교사의 AUC 의 96.5% 를 유지합니다. 이는 교사 대비 38 배에서 860 배의 속도 향상을 의미하며, 153 개 데이터셋 중 51% 에서 조정된 CatBoost 기준선보다 통계적으로 우수한 성능을 보입니다 (Wilcoxon $p=0.0008$).
3. 예측 가능한 이점: 증류의 이점은 데이터 차원에 따라 매우 예측 가능합니다:
   • 저차원 데이터 (특성 수 $\le$ 21): 증류는 CatBoost 대비 유의미한 이점 (+0.011 AUC) 을 제공합니다.
   • 고차원 데이터 (특성 수 $>$ 21): 이점은 미미하며 (+0.001 AUC), 교사 모델 자체가 CatBoost 보다 성능이 낮다면 증류는 성능을 더욱 저하시킬 수 있습니다.
4. 다중 교사 역학:
   • 트리 학생 모델: 다중 교사 평균화는 통계적으로 감지 가능하지만 실질적으로 미미한 이점 (+0.0006 AUC) 을 제공합니다. 단일 강력한 교사만으로도 충분합니다.
   • MLP 학생 모델: 다중 교사 평균화는 효과적인 레이블 스무딩으로 작용하여 단일 교사 훈련 대비 유의미한 이점 (+0.006 AUC, $p=0.003$) 을 제공합니다.
5. 오픈소스 파이프라인: 전체 파이프라인은 TabTune 라이브러리의 일부로 구현되어 공개되었습니다.

결과

평가에는 TALENT, OpenML-CC18, TabZilla, TabArena 의 153 개 분류 데이터셋이 포함되었습니다.

• 정확도: 가장 좋은 증류 구성 (TabICLv2 $\to$ XGBoost) 은 0.882 AUC 를 달성하여 교사의 성능을 96.5% 유지합니다.
• 지연 시간: 증류된 트리 학생 모델은 CPU 에서 1.94.0ms 에 실행되는 반면, 교사는 GPU 에서 1511,275ms 가 소요됩니다.
• 기준선: 증류된 모델은 153 개 데이터셋 중 51% 에서 조정된 CatBoost 기준선을 능가합니다. 반면, 더 약한 교사인 OrionMSP 는 학생 모델이 22.6% 만 승리합니다.
• 애블레이션: 저차원 데이터의 MLP 학생 모델의 경우, OOF 를 사용한 하드 레이블 ($\alpha=0$) 이 전체 힌톤 소프트 레이블 파이프라인보다 0.034 AUC 더 우수한 성능을 보였는데, 이는 깨끗한 저차원 데이터의 경우 복잡한 소프트 레이블 메커니즘은 선택 사항일 수 있지만 OOF 라벨링은 여전히 필수적임을 시사합니다.

의의

본 논문은 표 형식 파운데이션 모델의 "배포 문제"가 정확도 장점을 희생하지 않고 해결될 수 있음을 보여줍니다. ICL 교사 모델에 내재된 특정 누출 문제를 해결함으로써, 저자들은 TFM 의 "클래스 간 구조"를 경량 CPU 네이티브 GBDT 로 증류할 수 있음을 입증했습니다.

이 접근법의 의의는 예측 가능성과 우아한 실패에 있습니다:

• TFM 이 강력하다면, 증류는 그 강점을 빠른 학생 모델로 전달합니다.
• TFM 이 약하거나 (또는 데이터가 고차원인 경우), 증류는 인위적으로 성능을 높이지 않습니다. 대신 깔끔하게 실패하여 종종 표준 GBDT 보다 낮은 성능을 보임으로써 교사가 좋은 후보가 아님을 신호합니다.
• 이 방법은 학생 모델에 대한 복잡한 하이퍼파라미터 튜닝을 요구하지 않으며, 주요 결정은 단순히 가장 성능이 좋은 교사를 선택하는 것입니다.

저자들은 저차원 구조화된 데이터의 경우, 신호 붕괴를 피하기 위해 OOF 라벨링 프로토콜을 엄격히 준수한다면 추론 비용의 1% 수준에서 교사 모델에 근접한 정확도를 달성할 수 있다고 결론지었습니다.