基于注意力的多实例学习的极限学习机用于全切片图像分类.
0. TL; DR
这篇论文将极限学习机(Extreme Learning Machine, ELM)与现代的Attention MIL框架进行了结合,大幅降低训练复杂度和参数量。注意力网络中用于将实例特征映射到高维空间的变换矩阵$V$在随机初始化后就被“冻结”,不再参与梯度更新。唯一需要训练的只有最后用于生成注意力分数的权重向量$w$。
Attention-based ELM模型在将需要训练的参数量减少了5倍的情况下,其平均AUC与完全训练的深度MIL模型相比差距在1.5%以内。
1. 背景介绍
在计算病理学领域,全切片图像(WSI)的自动分析是核心任务。由于WSI尺寸巨大且缺乏像素级标注,多实例学习(MIL)成为了解决这一问题的方法。其中基于注意力机制的深度MIL(Attention-based deep MIL)凭借其优异的性能和一定的可解释性,成为了事实上的标准范式。
其核心逻辑是:通过一个可学习的注意力网络,为WSI中的每个图块(instance)分配一个权重,然后对图块的特征进行加权聚合,得到整张WSI的表示,用于最终的诊断。这种机制使得模型能够聚焦于那些含有诊断信息的关键区域,尤其是在处理罕见细胞检测等任务时显得至关重要。
注意力模块本身通常是一个包含非线性激活和多个可学习参数矩阵的小型神经网络。这意味着整个模型,包括特征提取器、注意力网络和最终的分类器,都需要通过端到端的反向传播进行训练,这在处理数百万图块时,依然是一个计算密集型的过程。
极限学习机(Extreme Learning Machine, ELM)对单隐层前馈网络的训练提出了一个颠覆性的观点:隐藏层的参数(输入权重和偏置)完全不需要训练,随机初始化后固定即可。唯一需要学习的是从隐藏层到输出层的权重,而这可以通过一次简单的线性回归(最小二乘法)解析求解。ELM的优势是显而易见的:极快的训练速度和极少的训练参数。
本文旨在将ELM与Attention MIL框架相结合,在保留注意力机制强大能力的同时,大幅简化其训练过程,降低其计算负担。
2. Extreme MIL
2.1 Attention-based MIL
标准的Attention-based Deep MIL模型,给定一个包(WSI)中的$n$个实例特征${h_1, \dots, h_n}$,计算注意力权重$a_k$:
\[a_k = \frac{\exp(w^T \tanh(V h_k))}{\sum_{j=1}^n \exp(w^T \tanh(V h_j))}\]其中$V$是一个可学习的矩阵,用于将原始特征$h_k$映射到一个更高维的、非线性的空间。$w$是一个可学习的向量,用于从这个高维空间中计算出一个标量的“注意力分数”。tanh是非线性激活函数。
然后聚合经过$V$和$tanh$变换后的高维特征,生成包表示$z$:
\[z = \sum_{k=1}^n a_k \tanh(V h_k)\]最后将包表示$z$送入一个前馈神经网络分类器。
整个模型,包括$V$、$w$和分类器的参数,都是通过反向传播进行端到端训练的。
2.2 Attention-based Extreme MIL
作者将ELM的思想应用于上述的注意力模块。Attention-based Extreme MIL的设置与Deep MIL几乎完全相同,但有一个根本性的区别:在注意力权重的计算公式中,用于特征变换的矩阵$V$,在随机初始化后,其参数被完全“冻结”,在整个训练过程中不进行任何更新。唯一的可训练参数在注意力模块中的权重向量$w$。

矩阵$V$通常是注意力模块中参数量最大的部分。例如如果输入特征是1000维,隐藏层是512维,则$V$就有超过50万个参数。现在这50万个参数都不需要训练了。此外由于$V$被固定,反向传播的链条在这一步被“截断”,使得优化过程更简单、更快速。
ELM的理论已经证明,即使隐藏层参数是随机的,只要隐藏层节点足够多,单隐藏层全连接层依然是通用函数逼近器。作者将这个思想迁移过来,认为随机的$V$同样能提供一个足够丰富的高维特征空间,足以让可学习的$w$从中挑选出有用的信息来计算注意力。
3. 实验分析
作者在一个真实的血液病理学数据集BloodMNIST上进行了一系列严谨的控制变量实验,任务是检测罕见的幼红细胞(erythroblasts)。这是一个典型的低证据比率MIL问题。

所有的Attention-based MIL模型,其性能都显著优于简单的逻辑回归基线(平均AUC提升15-25%)。这证明了注意力机制对于罕见细胞检测这类任务是必不可少的。

Attention-based Extreme MIL在训练参数量减少了5倍的情况下,其性能与完全训练的Deep MIL相比差距极小。它意味着,注意力网络中那个最庞大的参数矩阵$V$,可能真的不需要精细的梯度下降训练。一个好的随机投影,已经足以创建一个足够丰富的特征空间。
