论文122：Agent aggregator with mask denoise mechanism for histopathology WSI analysis (2024, ACM MM)

1 要点

题目：用于WSI分析的掩码去噪代理聚合器 (Agent aggregator with mask denoise mechanism for histopathology whole slide image analysis)

代码：

研究目的：
WSI分辨率巨大且缺乏细粒度注释，使其分类和分析任务面临挑战。此外，现有基于注意力机制的方法在处理WSI时存在计算复杂度高和无法捕捉实例间信息的问题。因此，本文所提出的AMD-MIL方法旨在通过动态调整特征表示，并引入掩码去噪机制来改进注意力分配，从而提高对WSI的检测能力和模型的可解释性。

关键技术：

1. 代理聚合器：
   引入代理Token，作为查询和键之间的中间变量，用于计算实例的重要性；
2. 掩码去噪机制：
   利用从代理聚合值映射的掩码和去噪矩阵，动态地掩盖低贡献表示并消除噪声；
3. 可训练代理Token：
   将代理令牌从通过均值池化得到的非训练参数转换为可训练的矩阵，以更有效地映射特征。

数据集：

1. Camelyon16
2. Camelyon16
3. TCGA-KIDNEY
4. TCGA-LUNG

引用：

@inproceedings{Ling:2024:agent:19,
author		=	{Ling, Xitong and Ouyang, Minxi and Wang, Yizhi and Chen, Xinrui and Yan, Renao and Cheng, Junru and Guan, Tian and Liu, Xiaoping and Tian, Sufang and He, Yonghong},
title		=	{Agent aggregator with mask denoise mechanism for histopathology whole slide image analysis},
booktitle	=	{{ACM MM}}，
years		=	{2024},
pages		=	{1--9},
url			=	{https://openreview.net/pdf?id=qQTr32f832}
}

2 方法

2.1 MIL和特征提取

在多示例学习 (MIL)中，每张WSI被概念化为一个有标记的包，其中包含的各个切片被视为具有不确定标签的实例。以二分类为例，每个WSI$X$被划分为多个区块$\{(x_1,y_1),\dots ,(x_N,y_N)\}$，其中$N$是实例数量。在MIL框架下，包标签$Y$与实例标签$y_i$之间的相关性如下：
$$Y=\{\begin{array}{ll}1,& \text{iff }{\sum }_{i=1}^N{y}_i>0;\\ 0,& \text{else},\end{array}\text{\hspace{1em}(1)}$$在训练期间，实例$y_i$的标签不可用。此外，WSI分类的目标是开发一个分类器$\mathcal{M}(X)$，用于预测包标签$\hat{Y}$：
$$\hat{Y}\leftarrow \mathcal{M}(X):=h(g(f(X)))),\text{\hspace{1em}(2)}$$其中$f$、$g$，以及$h$分别代表特征提取器、特征聚合器，以及MIL分类器。

特征聚合器被认为是总结特征的最重要部分，它可以聚合不同区块的特征。注意力机制可以识别WSI中区块的重要性，并且在特征聚合器中广泛使用：

1. 基于注意力机制的特征聚合器：
   $$G=\sum _{i=1}^N{a}_i{h}_i=\sum _{i=1}^N{a}_{i}f(x_i)\in {\mathbb{R}}^D,\text{\hspace{1em}(3)}$$其中$G$是包表示、$h_i\in {\mathbb{R}}^D$是通过特征提取器$f$提取的区块$x_i$的特征、$a_i$可训练注意力系数，以及$D$是向量$G$和$h_i$的维度；
2. 基于自注意力的特征聚合器：
   $$Q=H{W}_Q,K=H{W}_K,V=H{W}_v,\text{\hspace{1em}(4)}$$$$O=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_q}}\right)V=SV\text{\hspace{1em}(5)}$$其中$W_q$、$W_K$，以及$W_v$分别代表的可训练参数矩阵、$H$表示区块特征的集合、$O$是所有特征的整合，以及$d_q$是查询向量的维度。

2.2 注意力聚合器

在公式5，$\text{Sim}(Q,K)$的计算复杂度为$O(N^2)$。由于$N$通常包含数千个元素，这显著延长了预期的计算时间。线性注意力可以减少计算时间，但以牺牲信息为代价。为了缓解这一问题，TransMIL采用了Nyström近似来处理公式5中的自注意力，即构建矩阵$\tilde{Q}$和$\tilde{K}$，并计算每个片段 (segment) 的均值如下：
$$\tilde{Q}=[{\tilde{q}}_1;\dots ;{\tilde{q}}_m],{\tilde{q}}_j=\frac{1}{m}\sum _{i=(j-1)\times l+1}^{(j-1)\times l+m}{q}_i,\forall j=1,\dots ,m\text{\hspace{1em}(6)}$$$$\tilde{K}=[{\tilde{k}}_1;\dots ;{\tilde{k}}_m],{\tilde{k}}_j=\frac{1}{m}\sum _{i=(j-1)\times l+1}^{(j-1)\times l+m}{k}_i,\forall j=1,\dots ,m\text{\hspace{1em}(7)}$$其中，$\tilde{Q},\tilde{K}\in {\mathbb{R}}^{m\times D}$。然后，公式5中的$S$可以被近似为：
$$\hat{S}=\text{softmax}\left(\frac{Q{\tilde{K}}^T}{\sqrt{d_q}}\right)Z^{\ast }\text{softmax}\left(\frac{\tilde{Q}{K}^T}{\sqrt{d_q}}\right),\text{\hspace{1em}(8)}$$其中，$Z^{\ast }$表示$z(\tilde{Q},\tilde{K},Z)=0$的近似解，其仅需线性数量的迭代来收敛。

难点：

• 主要在于理解片段，因为公式6–7中的$m$和$l$未解释清楚，这需要提前了解transmil才能了解；
• 简单来说，给定的WSI，其有$N$个实例，所有的实例可以按顺序被分为$L$个片段，这里的$L$是为了和本文的$l$对应，原文并没有这个符号；
• 每个片段的大小为$m$，因此公式6的含义即为：在给定的第$l$个片段内，如何通过索引来计算新的${\tilde{q}}_j$，公式7同理；

在MIL任务中，Nyström注意力会因为采样机制而过滤掉具有重要特征的切片。此外，$N$的差异会导致局部下采样时的整体不平衡。因此，本文考虑具有线性时间复杂度的代理注意力方法：
$$O=\sigma (Q{A}^T)\sigma (A{K}^T)V\text{\hspace{1em}(9)}$$其中$\sigma (\cdot )$是Softmax函数、$A\in {\mathbb{R}}^{n\times D}$是池化$Q$得到的代理矩阵 (这里的使用的平均池化)，以及$n$是代理的维度，其是一个超参数。由于代理是非训练的，并且注意力分数的分布可能不是最优的，因此有必要建立一个能够动态调整注意力分数分布以增强模型性能和灵活性的自适应代理。

2.3 代理掩码去噪机制

公式5和公式9构成了本文的基底，如图1。图2则展示了所提出方法的总体框架。在预处理特征输入到模型之前，一个类别Token与这些特征直接拼接，这一点和transmil类型，得到特征矩阵$D\in {\mathbb{R}}^{D\times (N+1)}$，其中$D$是特征的维度，以及$(N+1)$表示包括嵌入的类别Toekn在内的区块数量。

3.3.1 可训练代理

公式9中的矩阵$A$最初是在$Q$的基础上，使用均值池化获得，即$A=pooling(Q)\in {\mathbb{R}}^{n\times D}$，这限制了$Q$中信息的完整性。对此，本文将$A$定义为一个可训练的矩阵：

1. 基于矩阵$A\in {\mathbb{R}}^{n\times D}$，可以计算中间矩阵$Q_A=Q{A}^T\in {\mathbb{R}}^{(N+1)\times n}$和$K_A = AK^T \in \mathbb{R}^{n \times (N+1)} )；
2. 使用通用注意力策略，中间变量$V_A$计算为：
   $$V_A=\sigma (K_A)V=\sigma (A{K}^T)V\in {\mathbb{R}}^{n\times D}.\text{\hspace{1em}(10)}$$
   难点：

• $n$：可训练矩阵$A$的节点参数，具体可以参见算法1，其实就是在网络初始化时，生成一个$B\times n\times D$的矩阵。

3.3.2 掩码代理

在MIL任务中，WSI中的大多数区块对预测的贡献不大，因此使用可训练的阈值生成可学习的掩码：
$$\tau =\sigma (p(W_{\tau }{V}_A^T)),\text{\hspace{1em}(11)}$$其中$W_{\tau }\in {\mathbb{R}}^{1\times D}$、$p$是一个类似于均值池化的可调节聚合函数，以及$\tau$阈值。进一步，每个特征的重要性将被计算，以优化隐藏空间中的重要特征。然而，特征选择会有丢失信息的风险。为了平衡重要信息选择和聚合原始特征之间的特性，提出了一个新的模块：
$$V_{M{D}_{ij}}=V_{A_{ij}}\mathbb{I}(M_{ij}>\tau )+D{N}_{ij},\text{\hspace{1em}(12)}$$其中$M=W_M{V}_A$是阈值矩阵，用于获得每个特征的重要性、$DN=W_{DN}{V}_A$是用于聚合信息的去噪矩阵，以及$W_M$和$W_{DN}$是可学习的参数。

3.3.3 代理可视化

代理注意力无法像通常注意力机制那样为每个实例生成一个注意力分数，这使得实验中的区块预测可视化变得困难。对此，提出了一种注意力分数的可视化策略：
$$At{t}_i=\sum _{j=1}^n{Q}_{A_{0,j}}{K}_{A_{j,i+1}}\text{\hspace{1em}(13)}$$其中$At{t}_i$就可以作为特征$h_i$的注意力分数。

3.3.4 AMD

引入了一个名为掩码去噪机制的新框架，如图2所示，其包括基于学习的代理注意力机制、表示细化。以及特征聚合。算法过程如算法1，且该模块可以表示为：
$$O=\sigma (Q{A}^T)V_{MD},\text{\hspace{1em}(14)}$$其中$md$表示掩码去噪机制，以及$V_{MD}$表示基于公式12计算出的矩阵。

由于阈值选择方法的差异，另外两种特征阈值选择策略如下：

• Mean-AMD：均值选择，选择所有特征的平均值作为阈值；
• CNN-AMD。CNN选择，通过组卷积方法减少不同组的特征，并使用组间平均值作为阈值。