Pocket Foundation Models: TFM을 CPU 기반 Gradient-Boosted Trees로 증류하기
요약
고성능 정형 데이터 파운데이션 모델(TFM)을 CPU 기반의 Gradient-Boosted Trees(XGBoost, CatBoost)로 증류하여 추론 속도를 획기적으로 개선하는 방법을 제안합니다. 층화된 Out-of-Fold(OOF) 레이블링을 통해 인컨텍스트 학습 모델의 레이블 유출 문제를 해결함으로써, CPU 환경에서도 교사 모델 성능의 96.5%를 유지하며 최대 860배의 속도 향상을 달성했습니다.
핵심 포인트
- TFM을 XGBoost/CatBoost로 증류하여 GPU 대비 CPU에서 밀리초(ms) 단위의 초고속 추론 가능
- Stratified Out-of-Fold(OOF) 기법을 사용하여 증류 과정에서의 레이블 유출 및 소프트 타겟 붕괴 방지
- TabICLv2를 XGBoost로 증류 시, 교사 모델 성능의 96.5%를 유지하며 38~860배의 속도 향상 기록
- 증류 성능은 저차원 데이터(피처 21개 미만)에서 더 효과적이며, 고차원 작업에서는 주의가 필요함
- 전체 파이프라인은 TabTune 라이브러리를 통해 오픈 소스로 제공됨
사기 탐지 스코어러(fraud scorer)는 2ms 이내에 응답해야 합니다. 최고의 정형 데이터 파운데이션 모델 (Tabular Foundation Models, TFMs)은 GPU에서 151~1,275ms가 소요됩니다. 우리는 TFM을 CPU에서 네이티브로 실행되는 XGBoost 또는 CatBoost 학생 모델로 오프라인 증류 (distilling)함으로써 이 격차를 해소합니다. 핵심적인 장애물은 인컨텍스트 학습 (In-Context Learning, ICL) 교사 모델의 특이성에 있습니다. 이들은 자신의 학습 세트를 스코어링할 때 레이블을 유출(leak)하며, 이로 인해 소프트 타겟 (soft targets)이 원-핫 벡터 (one-hot vectors)에 가깝게 붕괴되어 증류할 클래스 간 구조가 남지 않게 됩니다. 층화된 아웃-오브-폴드 (Stratified Out-of-Fold, OOF) 교사 레이블링이 이를 방지합니다. TALENT, OpenML-CC18, TabZilla, TabArena에서 추출한 153개의 분류 데이터셋에 걸쳐, TabICLv2를 XGBoost로 증류했을 때 CPU에서 1.9ms 만에 0.882의 macro-mean AUC (교사 AUC의 96.5%)를 달립니디. 이는 교사-학생 쌍에 따라 38배에서 860배의 속도 향상을 의미하며, 튜닝된 CatBoost 베이스라인 대비 통계적으로 유의미한 우위(Wilcoxon p = 0.0008; 승률 51%)를 보입니다. 네 가지 추가 발견 사항은 다음과 같습니다: 교사의 순위(rank)가 학생의 순위로 정확히 전이됩니다; 성능 향상은 저차원 데이터(< 21개 피처: CatBoost 대비 +0.011, > 21개 피처: +0.001)에 집중됩니다; 멀티 교사 평균화 (multi-teacher averaging)는 MLP 학생 모델에는 도움이 되지만 (+0.006, p = 0.003), 트리 학생 모델에는 0.001 미만의 효과만을 더합니다; 그리고 교사 모델 자체가 CatBoost에 뒤처지는 고차원 작업에서는 증류가 상황을 개선하기보다 악화시킵니다. 전체 파이프라인은 TabTune 라이브러리의 일부로 오픈 소스화되었습니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기