0. TL; DR

作者提出了一个简单而高效的max-pooling新范式FocusMIL，引入变分编码器，不再学习确定的实例表示，而是将每个实例编码为一个高斯分布，并从中采样得到潜在表示。采用Mini-batch训练，改变了传统MIL逐包训练（batch size=1）的方式，在一个批次中同时处理多个包。这使得max-pooling在每一轮更新时，都能看到来自不同包的、恶性程度各异的“最强阳性实例”，从而学习到更泛化的决策边界。

在Camelyon16/17数据集上，简单的FocusMIL的分类性能与复杂得多的SOTA注意力方法相当甚至更优。在更关键的图块级别分类和肿瘤定位任务上，FocusMIL显著超越了所有注意力方法。在一个专门设计的、含有“毒药样本”的半合成数据集中，注意力方法集体“中毒”失效，而FocusMIL几乎不受影响，强有力地证明了其因果鲁棒性。

1. 背景介绍

在过去的几年里，基于注意力机制的MIL方法（如ABMIL, TransMIL, DSMIL等）席卷了WSI分类领域。它们的核心思想是通过一个可学习的网络，为每个实例（patch）分配一个注意力权重，然后对实例特征进行加权聚合，得到包（slide）的表示。

作者假设，一个patch的特征由三类潜在因素构成：

$Z_c$ (Causal): 因果因素。真正决定细胞是否癌变的内在特征，如细胞核形态、异型性等。
$Z_s$ (Style): 风格因素。与标签无关的实例自身属性，如细胞的位置、旋转、大小等。
$Z_B$ (Bag-context): 包/背景因素。整个WSI共享的、与单个细胞因果无关的混淆信息，如H&E染色深浅、组织整体形态、扫描仪伪影等。

标准MIL假设指出“一个包为正，当且仅当它至少包含一个正实例”，要求模型的决策只与最关键的阳性证据有关，而与包的$Z_B$信息无关。注意力方法违反了这一假设，因为它聚合了所有实例，决策受到了非关键实例（包括所有阴性实例）的影响。而Max-pooling方法天然遵循了这一假设，因为它只选择得分最高的那个实例，其决策完全独立于其他实例。

2. FocusMIL

作者首先通过实验分析了传统Max-pooling方法（如mi-Net）失败的原因。mi-Net在训练集上可以达到近乎完美的slide-level准确率，但在验证集上却像随机猜测一样。同时其patch-level的预测能力也几乎为零。

传统MIL训练通常采用batch size=1的模式，即每次只处理一个包。Max-pooling机制使得在每次反向传播时，只有一个实例的梯度被更新。当模型遇到难以学习的阳性实例时，最简单的方法就是直接记住这个实例的特征，而不是学习其内在的、可泛化的模式。这导致了严重的过拟合。

作者提出了FocusMIL，在mi-Net的基础上引入了两个简单但至关重要的改进：采用变分编码器学习实例表示（注入随机性）和采用Mini-batch训练（引入多样性），缓解了max-pooling的过拟合问题，使其能够专注于学习真正的因果知识。

⚪ 变分编码器

FocusMIL不再让编码器直接输出一个确定的实例特征向量$x$，而是让它输出一个多元高斯分布的参数（均值$μ$和方差$σ²$）。然后从这个分布中采样出一个潜在表示$z$，用于后续的分类。

\[z = \mu + \sigma \odot \epsilon, \quad \text{where } \epsilon \sim \mathcal{N}(0, I)\]

引入KL散度损失，约束编码器产生的分布向标准正态分布$N(0, I)$看齐。

\[\mathcal{L} = -\log p(Y | z^*) + \beta \cdot \text{KL}[q(z|x) \| p(z)]\]

其中，$z^{*}$是经过max-pooling后选出的那个最关键实例的潜在表示。

每次从分布中采样，都为实例表示引入了随机性。模型无法再通过简单地记住一个固定的特征向量来作弊，它被迫去学习一个分布区域，这个区域代表了某一类实例的本质特征。KL散度正则化使得潜在空间更加平滑和结构化，有助于模型学习更有意义的表示。

⚪ Mini-batch训练

FocusMIL改变传统MIL逐包训练的方式，在每个mini-batch中同时处理多个（$n > 1$）包。

max-pooling机制会在每个包中选出一个“最阳性”的实例。当一个batch包含多个阳性包时，这些被选出的实例的恶性程度是不同的——有些来自高度恶性的肿瘤（易分类），有些来自低度恶性的肿瘤（难分类）。模型在一次梯度更新中，必须同时考虑这些“易”和“难”的样本，从而被迫学习一个能够覆盖所有恶性程度的、更宽容、更泛化的决策边界。

当一个batch中的不同WSI来自不同医院、采用不同染色工艺时，它们的背景信息$Z_B$是不同的。如果模型试图依赖$Z_B$（如颜色）来预测，它会在这个多样化的batch上得到很高的损失。为了最小化损失，模型需要忽略掉不稳定的$Z_B$，转而专注于在所有包中都保持不变的因果特征$Z_c$（如细胞形态）。

3. 实验分析

3.1 Camelyon16：真实世界WSI分类

使用ResNet-18和更强大的CTransPath作为特征提取器，对比FocusMIL与多种SOTA注意力方法。

在更考验模型真实理解能力的patch级别分类上，所有基于max-pooling的方法（mi-Net, CausalMIL, FocusMIL）的性能都远超所有基于注意力的方法。在slide级别分类上，FocusMIL的性能与最复杂的注意力方法相当甚至更优。

这个结果表明，尽管注意力方法在slide-level的AUC上刷出了很高的分数，但它们很可能是在用“虚假关联”作弊，其对单个patch的理解是错误的。而FocusMIL虽然简单，却学到了更本质的东西，因此在更细粒度的任务上表现出色。

3.2 Camelyon17：分布外（OOD）泛化挑战

训练集和测试集来自不同的医院，存在明显的领域漂移（domain shift）。无论使用何种特征，FocusMIL都取得了SOTA级的性能。传统的mi-Net和CausalMIL在使用简单特征时，出现了严重的过拟合。这说明简单max-pooling仍需FocusMIL的变分编码器和mini-batch策略来稳定。而注意力方法虽然没有完全失效，但性能普遍不如FocusMIL。FocusMIL学会了忽略医院、设备、染色等带来的背景偏差$Z_B$，专注于不变的因果特征$Z_c$。

3.3 半合成数据集：“毒药”测试

在训练集的正常slide中人为加入一种“毒药”信号（如特定的绿色噪声）；而在测试集的肿瘤slide中加入同样的“毒药”信号。一个遵守MIL假设的模型，应该学会忽略这个只在正常slide中出现的毒药，因为它与阳性概念无关。而一个违反MIL假设的模型，则可能学会“有毒药=阴性”这条虚假规则。

所有被测试的注意力方法（ABMIL, DSMIL, DTFD-MIL等），在测试集上的AUC都暴跌至接近0。这明确地表明，它们确实学会了“有毒药=阴性”的错误规则，当在测试集看到“有毒药的肿瘤slide”时，会坚定地将其预测为阴性。所有基于max-pooling的方法（mi-Net, CausalMIL, FocusMIL）的性能几乎不受影响。实验证明了注意力方法在聚合信息时，确实会利用包内所有实例（包括阴性实例和背景）的上下文信息来做决策，从而陷入虚假关联的陷阱。而max-pooling的机制则使其天然地对这类背景混淆信息“免疫”。