QG-MIL: 一种用于医学影像中跨域多实例学习的门控 Transformer 聚合器

TL;DR QG-MIL 提出门控 Transformer 聚合器，通过 RMSNorm、QK 归一化、注意力门控和 SwiGLU 前馈模块，解决医学影像 MIL 的注意力集中问题，实现更分散的实例权重与稳定训练，跨病理与血液学六项基准平均提升 +6.1 宏 F1。

问题

问题背景

在医学影像分析中，多实例学习（MIL） 已成为处理弱监督场景（仅有切片/病人级标签）的核心范式，尤其适用于全切片病理（WSI） 和细胞级血液学。基于注意力机制的聚合器因其可解释性（通过实例权重指示关键区域）而被广泛采用，但它们频繁出现注意力集中（attention concentration） 现象，导致权重塌缩到极少数实例上。

现有方法局限

主流注意力 MIL 聚合器（如 ABMIL、TransMIL、CLAM）训练时容易将几乎全部注意力分配给个别实例，引发：

• 过自信预测：模型仅依赖单一区域做出决策，忽略全局证据，导致预测分布过度尖锐；
• 不稳定预测：对输入扰动极度敏感，注意力模式在不同训练轮次间剧烈波动；
• 泛化受损：在跨机构或跨尺度数据上性能显著下降。 现有缓解方案（如辅助损失、实例掩码、多阶段正则化）引入了额外超参数和训练复杂度，且经常损害原始分类精度，难以作为通用模块嵌入不同架构。

为什么此问题重要且困难

• 标签成本极高：医学图像实例级标注需要昂贵而又稀缺的专家人力，MIL 是释放海量无注释数据价值的关键桥梁，但注意力集中从根本上削弱了这一设定——模型习得的“解释性”权重与真实临床证据脱节，会动摇 AI 辅助诊断的可信度。
• 跨域泛化需求迫切：从全切片病理到细胞级血液学，实例数量相差多个数量级（数千 vs 数百万），单一聚合策略很难同时避免注意力集中并保持高准确率，这对模型架构的鲁棒性提出了极大挑战。
• 工程落地限制：实际部署中希望聚合器即插即用，不依赖额外的实例级干扰或复杂的训练流程。然而现有分散注意力的方法往往以牺牲简洁性为代价。

行业类比

类似推荐系统中，模型若过度关注少数热门物品会导致“信息茧房”，医学 MIL 的注意力集中同样会造成诊断盲区——忽略散布在不同实例中的微弱但关键信号，最终损害整体决策的稳健性。

核心洞察

• 通过归一化与门控的协同组合，不依赖辅助损失就瓦解注意力集中：QG-MIL 将 RMSNorm 预归一化、每头 QK 归一化和细粒度注意力输出门控融合进 Transformer 聚合器，这三点合力让实例权重自然分散，无需引入额外的熵正则化、掩码或多阶段训练。相较以往 MIL 方法需要精心设计辅助目标或牺牲训练效率来对抗注意力坍塌，该设计从架构内部就稳定了梯度流，是一种更根本且成本更低的解法。
• 跨域一致性来自全组件协同，而非单点优化：消融实验表明，单个组件（如仅 SwiGLU FFN）在个别数据集上能达到或接近完整模型的效果，但只有四者联合才在所有六个跨尺度、跨成像模态的基准上低方差地稳定领先。这说明 QG-MIL 并非为某一特定数据分布调优，而是提供了一种对实例数量、染色差异等域偏移鲁棒的聚合范式，这对实际临床多场景部署极具价值。

方法

输入与动机
QG-MIL 面向包级弱监督分类任务：每个包包含若干实例（如全切片病理图像的图块或血细胞图像的单细胞），仅提供包级标签。包内实例先经预训练特征提取器（如 ResNet、ViT）转换为固定维度的特征向量。注意力聚合器的目标是学习实例重要性权重并生成包级表示，但标准注意力机制易出现注意力集中——少数实例垄断权重，导致预测过度自信且不稳定。

关键模块：门控 Transformer 聚合器
QG-MIL 的聚合器由四个协同的架构组件构成，直接嵌入 Transformer 块：

1. RMSNorm 预归一化：在多头注意力之前对输入特征施加 RMSNorm，替代常规的 LayerNorm，稳定训练初期的梯度流动。
2. 逐头 QK 归一化：对每个注意力头内的查询-键点积结果进行归一化，控制注意力分数幅度，避免因部分特征向量范数过大而导致注意力尖峰。此项无需额外可学习参数。
3. 细粒度注意力输出门控：在注意力加权求和后、进入前馈网络前，引入一个可学习的门控向量，对每个特征通道进行元素级缩放。这相当于对注意力输出做软掩码，抑制噪声通道并保留有用信号。
4. SwiGLU 风格前馈模块：将传统 FFN 中的两层线性变换与 ReLU 替换为带有 Swish 激活的门控线性单元（即 SwiGLU 结构），增强非线性表达能力且训练更鲁棒。

四个组件共同作用，使注意力质量分布更均匀（通过注意力叠加图和注意力质量分析验证），同时避免引入辅助损失、掩码或多阶段正则化。整个模块可堆叠多层，但实验表明单层即可达到最优。

输出与训练
最终聚合的包表示送入分类头得到预测。训练仅使用包级交叉熵损失，无需实例级监督。

与同类工作的差异
不同于 ABMIL、DSMIL 等依赖温度系数、辅助损失或两阶段训练来对抗注意力集中的方法，QG-MIL 仅凭架构设计便实现稳定的注意力分布和跨领域一致性，在六个涵盖组织病理与血液细胞的数据集上平均 macro F1 提升 6.1 点，且方差更低。

实验

实验设计与评估范式

QG-MIL 在 六个基准数据集 上进行了全面评估，覆盖 全切片病理学 与 细胞级血液学 两种本质不同的 MIL 尺度（数十万实例 vs 数百实例）。

• 对比基线包括当前主流的注意力 MIL 聚合器（如 ABMIL、TransMIL、DSMIL 等）以及近期基于 Transformer 的变体。
• 消融实验系统性地移除或替换四个核心组件（RMSNorm 预归一化、逐头 QK 归一化、细粒度注意力输出门控、SwiGLU 式前馈模块），以验证各组件贡献。
• 评估指标主要是 macro F1，并辅以注意力权重分布的度量（如注意力质量分析、注意力叠加可视化）。

关键发现

1. 性能提升：最佳 QG-MIL 变体在全部六个数据集中均超越基线，平均 macro F1 提高 +6.1 点，且在跨尺度场景下表现稳定。
2. 注意力分布：定性（注意力热力图）与定量（注意力质量指数）分析均证实 QG-MIL 的实例权重分布更为均匀，有效缓解了传统注意力聚合中的注意力集中现象。
3. 组件协同：消融实验表明，尽管单个组件可能在特定数据集上匹配完整模型的表现，但四项设计协同使用时提供了最强的跨域一致性和最小的性能方差。

与基线的深度对比

传统注意力 MIL 聚合器（如 ABMIL）倾向于将绝大部分权重分配给少数关键实例，导致过自信、不稳定的预测，且常依赖多阶段正则化或辅助损失来缓解。QG-MIL 则通过门控 Transformer 架构在不引入额外损失、掩码或正则化的情况下实现了天然的注意力分散。与 TransMIL 等 Transformer 基线相比，QG-MIL 的门控机制和归一化策略显式地抑制了早期层的注意力尖峰，使训练更稳定，收敛更一致。这种设计哲学对于临床部署极为关键——它意味着模型在不同医院、不同染色条件下能保持校准的置信度与可解释性，而无需针对每个场景手工调参。综合来看，QG-MIL 证明了通过精心架构协同而非外部约束，即可在领域无关的 MIL 任务中达成鲁棒且可泛化的注意力分布。

行业影响

落地场景

QG-MIL 可直接嵌入需要弱监督实例级聚合的医疗影像产品中，尤其适用于全切片病理图像（WSI）分析和血涂片细胞学分类。其稳定且分散的注意力机制可降低因注意力集中导致的假阳性 / 假阴性，适用于癌症筛查 AI 辅助诊断系统、血液病自动化分型平台 以及远程病理会诊中的智能预标注模块。此外，该聚合器的设计不依赖特定尺度，也可迁移至工业缺陷检测、视频异常帧检测等泛 MIL 场景，为需要从 bag 中定位关键实例的产品提供更可靠的注意力分布。

商业价值

从降本角度，QG-MIL 减少了因预测不稳定带来的复审成本，降低了对大规模实例级标注数据的依赖，直接节约标注支出。从增收角度，更一致的预测结果可提升医疗 AI 产品的诊断置信度，增加医院或第三方检验中心的付费意愿；同时，其在多个基准上平均 +6.1 宏观 F1 的提升直接转化为更强的产品竞争力。从体验角度，更分散的注意力热图使病理学家能更信任模型给出的标注区域，减少人工校验时间，提高诊断效率。

与现有产品/工作流的接口

QG-MIL 作为 MIL 聚合器，可无缝替换现有 Attention-based MIL 流程中的聚合模块（如 ABMIL、CLAM-SB），无需改动特征提取网络（如预训练的 ResNet、ViT）和下游任务头。其输入为实例特征嵌入，输出为 bag 级预测和实例权重，与主流病理分析栈（如 CLAM、Slideflow、QuPath 的 AI 扩展）完全兼容。开箱即用的配置化实现和明确的训练策略（无辅助损失、无掩码）使其易于通过持续集成 / 持续部署（CI/CD）管道被集成到现有 MLOps 流水线中，仅需替换模型定义文件即可快速验证。

具体落地 Use Case

• 数字病理云平台：在提供远程诊断服务的 SaaS 平台中，QG-MIL 可替代现有 MIL 聚合器，用于前列腺癌、乳腺癌等全切片图像的分级任务。其稳定的注意力分布可提升解释性报告的质量，并在多中心数据上保持一致的性能，降低因设备或染色差异带来的方差。
• 全自动血液分析流水线：在高通量实验室的血液分析仪中，QG-MIL 可处理大量单个细胞图像（实例数可达数千），为急性白血病分型等任务提供更均衡的实例重要性评分，减少因少数异常细胞导致的注意力集中误判，提高稀有类型检测的召回率。

局限

• **跨域泛化证据不足**：论文在两类医学影像任务（全切片病理和细胞血液学）上验证，但均属于生物医学图像中的MIL场景。对于其他常见MIL应用（如药物活性预测、视频异常检测）是否同样有效缺乏实验支撑。虽然标题强调“Domain-Agnostic”，但评估仅覆盖两种MIL尺度，未涉及不同模态或目标分布更不同的任务，通用性断言仍需更多领域证据。消融研究表明，部分组件在某些数据集上可达到与完整模型相当的性能，暗示组合优势并非在所有场景下都显著，实际部署时可能需要根据任务简化架构。
• **计算效率未量化**：QG-MIL引入了RMSNorm预归一化、逐头QK归一化、细粒度注意力输出门控和SwiGLU式前馈模块，这些操作会增加每个聚合步骤的计算开销。对于包含数万实例的全切片图像，额外的归一化和门控可能显著拖慢推理速度并消耗更多GPU内存。然而论文未提供任何训练吞吐量、推理延迟或显存占用的测量，无法评估其在大规模部署中的实际可行性。与基于普通注意力的AB-MIL或CLAM等轻量基线相比，复杂度的增加是否有对等性能回报，缺少定量性价比分析。
• **注意力均匀性可能损害稀有类别检测**：该方法的核心优势是将注意力更均匀地分布在实例上，以避免过度集中导致的预测不稳定。但在某些临床应用中，对少数关键实例（如微小转移灶、稀有形态异常）的强聚焦是必要的，均匀化可能削弱对稀疏信号的响应，降低对稀有类别的召回率。论文未测试含极端类别不平衡或罕见子类型的场景，注意力质量分析和性能改进仅基于宏观F1，缺少针对小类的细粒度评估。若应用于此类任务，可能需要额外的稀疏性引导，否则会牺牲检测敏感度。