Distilling Tabular Foundation Models for Structured Health Data

Motivation

• Known: TabPFNv2.5, TabICLv2, TabDPT, LimiX, Orion-MSP 등 TFM은 소규모 tabular 데이터셋에서 튜닝된 gradient-boosted trees와 동등하거나 더 나은 성능을 보이며, knowledge distillation은 hard label보다 student의 일반화 성능을 높이는 것으로 알려져 있다.
• Gap: TFM은 추론 시 GPU 기반 인프라와 training set에 대한 in-context conditioning을 요구하여 CPU 기반·air-gapped 병원 환경에 배포하기 어렵고, 이러한 in-context 특성으로 인해 naive distillation을 적용하면 teacher가 이미 관측한 샘플에 대한 soft label이 overconfident해지는 context(teacher-identity) leakage 문제가 발생하는데, 이를 다룬 연구가 부재하다.
• Why: 병원 시스템은 CPU 전용, air-gapped, data-residency 제약 하에서 calibrated probability와 공정성을 보장하는 예측 모델이 필요하므로, TFM 수준의 예측 품질을 경량 모델에 안전하게 전이하는 방법은 임상 의사결정 지원 시스템의 실질적 배포 가능성을 크게 높인다.
• Approach: Hinton loss 기반 knowledge distillation을 사용하되, teacher가 in-context로 training set을 조건화하는 문제를 k=5 stratified cross-validation 기반 out-of-fold teacher labeling으로 해결하여 leakage-free soft label을 생성하고, 이를 tree-based 및 MLP student에 adaptive temperature와 confidence weighting을 적용해 학습한다.

Achievement

Figure 2. Inference latency, log scale. Teachers (red) sit at 200–4000 ms on GPU. Distilled students sit at 4–11 ms on C

1. 정확도 유지: distilled tree-based student(LightGBM, CatBoost, XGBoost)는 teacher AUC의 최소 98% 이상을 유지하며 일부 데이터셋에서는 teacher를 능가한다.
2. 속도 개선: 모든 distilled student가 CPU 상에서 teacher 대비 최소 26배 이상 빠르며(MLP는 최대 약 49배), 모델 크기도 1MB 미만으로 경량화되었다.
3. 보정 및 공정성: tree student는 teacher의 calibration과 fairness 특성을 대체로 보존하는 반면, MLP student는 post-hoc calibration이 필요하고 fairness profile이 불안정하다.
4. multi-teacher ensemble의 한계: 여러 teacher의 soft label을 균등 평균하는 multi-teacher distillation은 대체로 최고 단일 teacher 대비 일관된 성능 향상을 주지 못한다.

How

• 19개 헬스케어 데이터셋(16 binary, 3 multiclass; 심장학, 종양학, 신장학, 피부과, 중환자실 등)을 대상으로 소규모 데이터 영역(≤10K rows)에서 실험 수행
• 6개 TFM teacher(TabPFNv2.5, TabPFNv2.6, TabICLv2, TabDPT, LimiX, Orion-MSP v1.5)와 4개 student family(LightGBM, CatBoost, XGBoost, MLP) 조합, 다중 teacher ensemble 포함
• teacher는 k=5 stratified cross-validation으로 fold별로 D\Dk에 학습되고 Dk에 대해서만 soft label을 생성하여 context leakage 방지 (Figure 1의 OOF distillation pipeline)
• student 학습은 Hinton mixed loss(α=0.7, KL divergence + cross-entropy)를 사용하며, teacher entropy 기반 adaptive per-sample temperature(Ti∈[1,5])와 confidence weight(wi, moderate-entropy 예측에 peak)를 적용
• LogisticRegression, 기본 설정 XGBoost/LightGBM(zero-imputed input)을 baseline으로 사용, 평가지표는 average AUC, retention(AUC(Student)/AUC(Teacher)×100), latency/memory, calibration, fairness이며 TabTune 프레임워크로 전처리·평가 일관성 확보

Originality

• ICL 기반 TFM teacher 특유의 context leakage(teacher-identity leakage) 문제를 최초로 명시적으로 정의하고, stacking 관행을 응용한 stratified out-of-fold teacher labeling으로 이를 해결
• 19개 헬스케어 데이터셋, 6개 최신 TFM teacher, 4개 student family, multi-teacher ensemble을 아우르는 대규모·체계적인 벤치마크 설계
• 단순 정확도뿐 아니라 calibration과 fairness까지 distillation 관점에서 종합적으로 분석한 점이 헬스케어 응용에 특화된 독창적 기여

Limitation & Further Study

• multi-teacher averaging이 균등 가중치 방식에 국한되어 있어, teacher 간 불일치를 고려한 가중 앙상블이나 선택적 앙상블 기법에 대한 탐구가 부족함
• MLP student의 calibration 및 fairness 불안정성에 대한 근본 원인 분석이나 개선 방안(예: distillation-aware calibration)이 충분히 다뮈지지 않음
• 데이터셋이 모두 10K rows 이하의 소규모 영역으로 제한되어 있어, 대규모 헬스케어 데이터셋에 대한 일반화 가능성이 검증되지 않음
• teacher 선택에 따른 student 성능 편차의 원인(예: teacher 간 아키텍처 차이)에 대한 심층 분석이 부족하여 실제 배포 시 teacher 선택 가이드라인 제시가 제한적임

Evaluation

총평: TFM의 in-context 특성으로 인한 leakage 문제를 명확히 짚고 이를 실용적인 out-of-fold labeling으로 해결한 견고한 실증 연구로, 헬스케어 배포 맥락에서의 실질적 가치가 크지만 multi-teacher 전략과 MLP calibration 개선에 대한 추가 연구가 필요하다.