Distilling Tabular Foundation Models for Structured Health Data
作者: Aditya Tanna, Nassim Bouarour, Mohamed Bouadi, Vinay Kumar Sankarapu, Pratinav Seth
分类: cs.LG, cs.AI
发布日期: 2026-05-18
💡 一句话要点
提出面向结构化健康数据的表格基础模型蒸馏方法,实现轻量化部署。
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)
关键词: 表格数据 知识蒸馏 健康数据 基础模型 模型压缩 医疗应用 分层留一法
📋 核心要点
- 表格基础模型在医疗领域表现优异,但高昂的推理成本限制了其应用。
- 论文提出一种知识蒸馏方法,将表格基础模型的知识迁移到轻量级模型,降低推理成本。
- 实验表明,蒸馏后的模型在保持性能的同时,显著提升了推理速度,并保证了校准性和公平性。
📝 摘要(中文)
表格基础模型(TFMs)在健康数据集上表现出色,但其推理成本和基础设施要求限制了实际应用。本文研究了是否可以通过知识蒸馏将TFMs的预测能力转移到轻量级表格模型。由于上下文TFMs在推理时依赖训练集,直接蒸馏可能导致上下文泄露;本文通过分层留一法教师标记来解决这个问题。在19个医疗保健数据集、6个TFM教师模型、4个学生模型族以及多个多教师集成模型上,实验表明蒸馏后的学生模型保留了至少90%的教师模型AUC,在某些情况下甚至优于教师模型,同时在CPU上运行速度至少快26倍,并保持了对健康应用至关重要的校准性和公平性。此外,多教师平均并没有始终优于最佳单教师模型。因此,感知泄露的蒸馏是使TFM质量的预测应用于推理受限的健康环境的可行途径。
🔬 方法详解
问题定义:表格基础模型(TFMs)在医疗健康领域展现了强大的预测能力,但其模型规模庞大,推理计算成本高昂,对硬件资源要求较高,难以在资源受限的医疗环境中部署和应用。直接应用TFMs会带来较高的经济成本和时间成本,限制了其在实际医疗场景中的应用。
核心思路:论文的核心思路是通过知识蒸馏技术,将大型TFMs的知识迁移到轻量级的学生模型中。学生模型结构简单,推理速度快,易于部署。通过让学生模型学习教师模型的预测结果,使其在保持较高预测性能的同时,显著降低推理成本。为了解决直接蒸馏可能导致的上下文泄露问题,论文提出了分层留一法教师标记策略。
技术框架:整体框架包含教师模型训练和学生模型蒸馏两个阶段。首先,训练多个表格基础模型作为教师模型。然后,使用分层留一法策略,为每个样本生成教师模型的预测标签。最后,使用这些标签训练轻量级的学生模型。学生模型学习教师模型的预测分布,从而获得与教师模型相似的预测能力。该框架支持使用不同的教师模型和学生模型。
关键创新:论文的关键创新在于提出了分层留一法教师标记策略,以解决直接蒸馏可能导致的上下文泄露问题。传统的知识蒸馏方法直接使用教师模型在整个训练集上的预测结果作为学生模型的训练目标,这可能导致学生模型学习到训练集的上下文信息,从而影响其泛化能力。分层留一法通过将训练集划分为多个fold,并使用在其他fold上训练的教师模型为当前fold的样本生成标签,从而避免了上下文泄露。
关键设计:论文使用了多种表格基础模型作为教师模型,包括AutoGluon、TabTransformer等。学生模型则选择了包括Logistic Regression、MLP、Random Forest和GBDT在内的多种轻量级模型。损失函数方面,使用了交叉熵损失函数来衡量学生模型和教师模型预测结果之间的差异。分层留一法中,将数据集划分为5个fold。实验中,对比了单教师蒸馏和多教师蒸馏的效果,并评估了蒸馏后模型的性能、推理速度、校准性和公平性。
🖼️ 关键图片
📊 实验亮点
实验结果表明,通过知识蒸馏,学生模型能够保留至少90%的教师模型AUC,在某些情况下甚至超越教师模型。同时,学生模型的推理速度在CPU上至少提升了26倍。此外,蒸馏后的模型保持了良好的校准性和公平性,这对于医疗应用至关重要。多教师平均并没有始终优于最佳单教师模型。
🎯 应用场景
该研究成果可应用于医疗健康领域的多种场景,例如疾病诊断、风险预测、患者分层等。通过将大型表格基础模型的知识迁移到轻量级模型,可以在资源受限的环境中部署高性能的预测模型,为临床决策提供支持,提高医疗效率和质量。该方法还有助于降低医疗成本,并促进个性化医疗的发展。
📄 摘要(原文)
Tabular foundation models (TFMs) achieve strong performance on health datasets, but their inference cost and infrastructure requirements limit practical use. We study whether their predictive behavior can be transferred to lightweight tabular models through knowledge distillation. Since in-context TFMs condition on the training set at inference time, naive distillation can introduce context leakage; we address this with stratified out-of-fold teacher labeling. Across $19$ healthcare datasets, $6$ TFM teachers, $4$ student families, and several multi-teacher ensembles, we find that distilled students retain at least $90\%$ of teacher AUC, outperforming teachers in some cases, while running at least $26\times$ faster on CPU and preserving calibration and fairness critical for health applications. Moreover, multi-teacher averaging does not consistently improve over the best single teacher. Leakage-aware distillation is thus a viable route for bringing TFM-quality predictions into inference-constrained health settings.