구조화된 의료 데이터를 위한 Tabular Foundation Models 증류
요약
Tabular Foundation Models(TFMs)의 높은 추론 비용과 인프라 요구 사항을 해결하기 위해 지식 증류(Knowledge Distillation)를 활용한 경량화 연구를 수행했습니다. 층화된 Out-of-fold 교사 라벨링 방식을 통해 컨텍스트 누수 문제를 해결하였으며, 실험 결과 학생 모델이 교사 모델 성능의 90% 이상을 유지하면서도 CPU 환경에서 최대 26배 빠른 속도를 기록했습니다.
핵심 포인트
- 지식 증류를 통해 TFM의 성능을 경량화된 모델로 성공적으로 전이 가능함
- 층화된 Out-of-fold 교사 라벨링을 사용하여 In-context TFM의 컨텍스트 누수 문제를 방지함
- 증류된 모델은 의료 데이터 분석에 필수적인 보정(Calibration) 및 공정성(Fairness)을 유지함
- CPU 환경에서 기존 TFM 대비 최소 26배 이상의 추론 속도 향상을 달성함
- 멀티 교사 평균화가 항상 단일 교사 모델보다 우수한 성능을 보장하지는 않음을 발견함
Tabular Foundation Models (TFMs)는 의료 데이터셋에서 강력한 성능을 달성하지만, 추론 비용(Inference cost)과 인프라 요구 사항이 실제 사용을 제한합니다. 본 연구에서는 지식 증류 (Knowledge Distillation)를 통해 TFM의 예측 동작을 경량화된 Tabular 모델로 전이할 수 있는지 연구합니다. In-context TFM은 추론 시점에 훈련 데이터셋에 의존하므로, 단순한 증류 방식은 컨텍스트 누수 (Context leakage)를 유발할 수 있습니다. 우리는 이를 층화된 Out-of-fold 교사 라벨링 (Stratified out-of-fold teacher labeling)을 통해 해결합니다. 19개의 의료 데이터셋, 6개의 TFM 교사 (Teacher) 모델, 4개의 학생 (Student) 모델군, 그리고 여러 멀티 교사 앙상블 (Multi-teacher ensembles)을 대상으로 실험한 결과, 증류된 학생 모델은 교사 모델 AUC의 최소 90%를 유지하며 일부 사례에서는 교사 모델을 능가하는 성능을 보였습니다. 또한, CPU에서 최소 26배 더 빠르게 작동하면서도 의료 애플리케이션에 필수적인 보정 (Calibration) 및 공정성 (Fairness)을 유지했습니다. 나아가, 멀티 교사 평균화 (Multi-teacher averaging)가 최상의 단일 교사 모델보다 일관되게 성능을 향상시키지는 않는다는 점을 발견했습니다. 따라서 누수 방지형 증류 (Leakage-aware distillation)는 추론 제약이 있는 의료 환경에 TFM 수준의 예측 성능을 도입할 수 있는 실행 가능한 경로입니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 arXiv cs.LG의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기