论文

ADMIL:用于病理学选择性基础模型推理的注意力蒸馏多实例学习

ADMIL: Attention-Distilled Multiple Instance Learning for Selective Foundation Model Inference in Pathology

模型推理推理加速

摘要

使用病理学基础模型嵌入的基于注意力的多实例学习(ABMIL)对于幻灯片级任务是有效的,但详尽的推理需要将大型图像编码器应用于每个前景图块,尽管随后的注意力分布通常集中在一小部分信息区域上。我们引入了 ADMIL(注意力蒸馏多实例学习),这是一种选择性计算框架,可将 ABMIL 教师的注意力提取到轻量级图块选择模型 PriorNet 中。使用 EfficientNet 架构,PriorNet 从具有 KL 散度的原始图块像素中学习教师注意力分布;在推理时,它会对前景池进行评分,选择前 K 个图块,并在选定的 ABMIL 学生预测幻灯片标签之前仅在该子集上调用昂贵的基础模型。在 BRACS、PANDA 和 CAMELYON16 中,ADMIL 分别在 K=4、8 和 128 个图块上与全教师标题性能相匹配,避免了 > 98% 的基础模型 (Virchow2) 图块嵌入和模型推理失败。随机和教师注意的预言机控制表明,这个结果取决于与任务相关的选择,而不是仅仅取决于图块计数的减少。定量和定性分析表明,PriorNet 以高保真度恢复了教师的图块排序,同时关注与任务相关的形态区域。 ADMIL 表明,几乎所有昂贵的图块编码都可以在不牺牲幻灯片级性能的情况下删除,从而为在延迟和计算成本是关键考虑因素的临床环境中更有效地部署提供了潜在的途径。