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 라이브러리를 통해 오픈 소스로 제공됨