ICLR 2026会议论文《ASMIL:用于全切片成像的注意力稳定多示例学习》研究报告
本文是一篇发表于顶级机器学习会议ICLR 2026的研究论文,题为《ASMIL:用于全切片成像的注意力稳定多示例学习》。论文由多伦多大学、斯坦福大学和北海道大学的研究人员合作完成,主要作者包括林峰·叶(Linfeng Ye)、沙扬·莫哈杰尔·哈米迪(Shayan Mohajer Hamidi)、池志翔(Zhixiang Chi)等。本研究报告旨在向中文研究社区详细介绍这项工作的背景、方法、结果与价值。
一、 研究背景与动机
本研究的核心科学领域是计算病理学,特别是基于全切片图像(Whole Slide Image, WSI)的弱监督学习。WSI通过数字扫描技术将玻璃病理切片转换为分辨率高达十亿像素级别的数字图像,是癌症诊断和分型的关键工具。然而,其极端的大小和极度的稀疏性(例如,肿瘤区域可能仅占整张切片的极小部分)带来了巨大挑战。在临床实践中,逐像素或逐切片标注成本高昂,因此大多数数据集仅提供弱监督的切片级标签。
多示例学习(Multiple Instance Learning, MIL)框架自然适用于此场景。它将一张WSI视为一个“包”(bag),其中包含成千上万个图像块(即“实例”)。模型仅使用包级标签进行训练,目标是学习从实例特征到包级预测的映射。近年来,基于注意力的MIL(Attention-based MIL)方法,如ABMIL、TransMIL等,因其能够通过注意力权重突出关键区域并增强模型可解释性,已成为WSI分析的事实标准。
尽管如此,作者团队发现并系统性地指出了现有基于注意力的MIL方法存在的三个主要问题: 1. (PI)不稳定的注意力动态:这是本研究首次发现和深入分析的新问题。在WSI的弱监督训练中,同一张切片在不同训练轮次(epoch)上的注意力分布会发生剧烈振荡,而非收敛到一个稳定的模式。这种不稳定性通过连续轮次间注意力分布的Jensen-Shannon散度(JSD)进行量化,如图1所示(TransMIL示例),其导致训练过程不稳定并损害最终性能。 2. (PII)注意力过度集中:已有研究指出的问题,即模型倾向于将所有注意力分配给极少数图像块,而忽略其他潜在有用信息,这损害了模型的泛化能力和可解释性。作者将其归因于softmax函数的指数敏感性。 3. (PIII)过拟合:病理数据集中训练样本数量通常有限(仅几百张),而模型参数众多,容易导致过拟合。
本研究的核心目标,就是提出一个统一的框架,以同时克服以上三个挑战,提升基于注意力的MIL在WSI分析中的稳定性、泛化能力和可解释性。
二、 研究方法与流程
为实现上述目标,作者提出了注意力稳定多示例学习(Attention-Stabilized Multiple Instance Learning, ASMIL)框架。其核心工作流程如图2所示,主要由三个创新组件构成:
通过锚定模型(Anchor Model)稳定注意力:这是解决注意力不稳定问题的关键。ASMIL维护一个与在线模型(Online Model)注意力模块结构完全相同、但参数不通过反向传播更新的“锚定模型”。锚定模型的参数通过在线模型参数的指数移动平均(Exponential Moving Average, EMA)进行更新(公式4)。在每次训练中,在线模型和锚定模型接收相同的输入特征令牌(tokens)。在线模型使用标准softmax计算注意力分布α,而锚定模型的注意力则用于提供一个稳定、平滑的参考分布。为了将这种稳定性传递给在线模型,ASMIL引入了一项稳定损失(公式6),即在线注意力分布α与锚定注意力分布之间的Kullback-Leibler(KL)散度。该损失的梯度将推动在线模型的注意力向更稳定的锚定注意力对齐(公式8)。
在锚定模型中使用归一化Sigmoid函数(Normalized Sigmoid Function, NSF)防止注意力过度集中:为解决softmax导致的注意力过度集中问题,作者提出在锚定分支中用NSF替换softmax(公式5)。作者通过理论证明(定理1),NSF具备“选择性平坦化”特性:它能在高置信度的实例(“高权重实例”)之间提供更均衡的权重分配,同时有效抑制低置信度实例(“低权重实例”)的权重。而单个温度参数的softmax无法同时满足对高权重实例的均衡化目标和对低权重实例的抑制目标。如图3所示,在Camelyon-16数据集上,相较于softmax产生的极度集中的注意力热图,NSF产生了更分散、更能覆盖整个肿瘤区域的注意力。作者发现,将NSF直接用于在线模型会导致梯度消失,因此将其巧妙地应用于锚定模型,使其作为指导在线模型的稳定先验,而不影响其学习动态。
通过令牌随机丢弃(Token Random Dropping)缓解过拟合:这是一种专门为ASMIL设计的正则化策略。在训练阶段,ASMIL会以一定概率b随机丢弃一部分可学习的特征令牌(feat tokens),然后将剩余的令牌与一个可训练的[CLS]令牌一起输入第二个Transformer编码器,以生成最终的包级表示用于分类。在推理阶段,则保留所有令牌(b=0)。这种方法防止了特征令牌之间的过度共适应,鼓励模型不依赖于固定的令牌子集,从而提升了泛化能力。消融实验表明,丢弃率b≈0.5时效果最佳。
ASMIL的整体训练目标是分类损失(交叉熵)与注意力稳定损失(KL散度)的加权和(公式9)。在训练过程中,仅在线模型通过梯度下降进行更新(公式10),锚定模型通过EMA同步,并且其梯度被阻止回传。重要的是,在推理阶段,锚定模型被完全丢弃,因此不会增加任何额外的计算开销。
三、 实验结果与发现
研究在三个公开的WSI数据集上进行了全面的实验评估:Camelyon-16(淋巴结转移检测)、Camelyon-17(多中心转移检测)、以及BRACS(乳腺癌亚型分类)。
1. 亚型分类性能: 如表1所示,ASMIL在几乎所有数据集和骨干网络(ResNet-18和ViT-SSL)上都取得了最先进的(state-of-the-art)或极具竞争力的性能。例如,在使用ViT-SSL骨干的BRACS数据集上,ASMIL的F1分数达到0.781,AUC达到0.914,相较之前的最佳方法有显著提升(F1提升3.9%)。在具有稀疏肿瘤区域挑战的Camelyon-16和Camelyon-17数据集上,ASMIL的优势更为明显,分别带来了3.3%的F1分数提升和高达6.49%的F1分数提升。
2. 锚定模型与NSF的通用性提升: 作者将锚定模型(+Anchor)及锚定模型+NSF(+Anchor+NSF)作为插件模块,应用到四种主流的注意力MIL方法(ABMIL, CLAM-SB, TransMIL, DSMIL)上。如表2所示,这两种模块的加入几乎总能一致性地提升基线模型的性能。其中,ABMIL在BRACS数据集上应用“+Anchor+NSF”后,F1分数提升了10.73%。这证明了ASMIL框架中核心组件的普适价值。
3. 肿瘤定位能力: 除了分类性能,作者在Camelyon-16数据集上定量和定性地评估了模型的定位能力。定性分析(图4)显示,与基线方法(如TransMIL, DTFD-MIL, CAMIL)相比,ASMIL生成的注意力热图能够更一致、更全面地突出肿瘤区域,避免了过度集中于少数点的问题。这归功于NSF对注意力过度集中的缓解。定量指标(病灶级自由响应ROC曲线下面积FROC、Dice系数、特异性)也支持ASMIL在定位任务上达到了最优性能。
4. 消融分析与验证: 如表3所示,对ASMIL三个核心组件(Anchor, NSF, Random Drop)的组合消融实验证实了每个组件的有效性。完整模型表现最佳,移除任何一个组件都会导致性能下降,其中锚定模型的影响最大。此外,t-SNE可视化(图5, 图6)表明,使用锚定模型的ASMIL和TransMIL在训练过程中能够更快地形成分离良好的包级特征簇,证明了锚定模型在稳定训练动态、加速收敛方面的作用。
四、 研究结论与意义
本研究的核心贡献在于首次系统性地识别并解决了基于注意力的MIL在WSI分析中的“不稳定注意力动态”问题,并提出了统一的ASMIL框架来协同应对注意力不稳定、过度集中和过拟合三大挑战。
科学价值:
应用价值:
五、 研究亮点
ASMIL代表了WSI弱监督学习领域的一项重要进展,它通过稳定训练过程、优化注意力机制和加强正则化,为构建更准确、更可靠、更可解释的计算病理学模型提供了有力的新工具。论文代码已开源,保证了研究的可复现性。