Distilling foundation models for robust and efficient models in digital pathology

📄 arXiv: 2501.16239v3 📥 PDF

作者: Alexandre Filiot, Nicolas Dop, Oussama Tchita, Auriane Riou, Rémy Dubois, Thomas Peeters, Daria Valter, Marin Scalbert, Charlie Saillard, Geneviève Robin, Antoine Olivier

分类: cs.CV

发布日期: 2025-01-27 (更新: 2025-02-28)

备注: Preprint

💡 一句话要点

提出H0-mini模型，通过知识蒸馏提升数字病理学中模型的鲁棒性和效率。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 数字病理学 知识蒸馏 基础模型 模型压缩 鲁棒性 轻量级模型 病理图像分析

📋 核心要点

  1. 现有数字病理学基础模型依赖大规模数据集和模型尺寸，导致计算成本高昂和推理速度慢。
  2. 论文提出知识蒸馏方法，将大型基础模型提炼为小型模型H0-mini，显著降低参数量和推理成本。
  3. 实验表明，H0-mini在多个基准测试中表现出色，且在染色和扫描条件变化下具有更强的鲁棒性。

📝 摘要（中文）

近年来，数字病理学领域的基础模型（FM）严重依赖于扩展预训练数据集和模型规模，从而产生大型且强大的模型。虽然这提高了各种下游任务的性能，但也增加了计算成本和推理时间。本文探讨了将大型基础模型提炼成较小模型的方法，从而将参数数量减少几个数量级。通过利用蒸馏技术，我们提炼的模型H0-mini在显著降低推理成本的同时，实现了与大型FM几乎相当的性能。它在多个公共基准上进行了评估，在HEST基准上获得第3名，在EVA基准上获得第5名。此外，在PLISM数据集上进行的鲁棒性分析表明，我们提炼的模型在染色和扫描条件变化方面达到了出色的鲁棒性，显著优于其他最先进的模型。这为设计用于数字病理学的轻量级且鲁棒的模型开辟了新的视角，而不会影响性能。

🔬 方法详解

问题定义：数字病理学领域的大型基础模型虽然性能优异，但计算成本高昂，推理速度慢，难以在资源受限的环境中部署。现有方法缺乏在保持性能的同时降低模型复杂度的有效手段。

核心思路：利用知识蒸馏技术，将大型、高性能的基础模型（教师模型）的知识迁移到小型模型（学生模型）H0-mini。学生模型在学习教师模型输出的同时，也学习原始数据，从而在参数量大幅减少的情况下，尽可能地保持甚至超越教师模型的性能。

技术框架：整体框架包含两个阶段：首先，使用大型基础模型作为教师模型，对病理图像进行特征提取和预测。然后，利用教师模型的输出作为软标签，结合原始病理图像数据，训练小型学生模型H0-mini。学生模型的训练目标是同时模仿教师模型的预测结果和学习原始数据的特征表示。

关键创新：该论文的关键创新在于成功地将知识蒸馏技术应用于数字病理学领域的大型基础模型，并设计了一种有效的蒸馏策略，使得小型模型H0-mini在保持高性能的同时，显著降低了计算成本和推理时间。此外，该模型在鲁棒性方面表现出色，能够有效应对染色和扫描条件的变化。

关键设计：论文中使用了交叉熵损失函数来衡量学生模型和教师模型预测结果之间的差异，并结合KL散度损失来进一步约束学生模型的输出分布与教师模型保持一致。此外，还使用了数据增强技术来提高模型的泛化能力和鲁棒性。具体的网络结构细节（如层数、每层神经元数量等）可能在论文正文中详细描述，但摘要中未提及。

🖼️ 关键图片

fig_0

📊 实验亮点

H0-mini模型在HEST基准测试中获得第3名，在EVA基准测试中获得第5名，证明了其在性能上与大型基础模型具有可比性。更重要的是，在PLISM数据集上的鲁棒性分析表明，H0-mini模型在染色和扫描条件变化方面显著优于其他最先进的模型，表明其具有更强的泛化能力和实用价值。

🎯 应用场景

该研究成果可应用于数字病理学辅助诊断、病理图像分析和远程医疗等领域。轻量级且鲁棒的模型H0-mini更易于部署在资源受限的环境中，例如移动设备或边缘计算平台，从而实现快速、准确的病理图像分析，辅助医生进行诊断，提高诊断效率和准确性，并促进远程医疗的发展。

📄 摘要（原文）

In recent years, the advent of foundation models (FM) for digital pathology has relied heavily on scaling the pre-training datasets and the model size, yielding large and powerful models. While it resulted in improving the performance on diverse downstream tasks, it also introduced increased computational cost and inference time. In this work, we explore the distillation of a large foundation model into a smaller one, reducing the number of parameters by several orders of magnitude. Leveraging distillation techniques, our distilled model, H0-mini, achieves nearly comparable performance to large FMs at a significantly reduced inference cost. It is evaluated on several public benchmarks, achieving 3rd place on the HEST benchmark and 5th place on the EVA benchmark. Additionally, a robustness analysis conducted on the PLISM dataset demonstrates that our distilled model reaches excellent robustness to variations in staining and scanning conditions, significantly outperforming other state-of-the art models. This opens new perspectives to design lightweight and robust models for digital pathology, without compromising on performance.