Scaling Multi-Teacher Distillation for Digital Pathology¶

会议: NeurIPS 2026
arXiv: 2610.12574
OpenReview: https://openreview.net/forum?id=tz5T3eSFDA
领域: 医学图像
关键词: 多教师蒸馏, 数字病理学, 基础模型, 专家混合模型, 结构化注意力掩码

一句话总结¶

针对传统多教师蒸馏在教师数量扩展时表征崩溃的瓶颈，提出可学习教师专属 token、结构化注意力掩码、MoE 尾部层与跨样本 InfoNCE 对比损失，将 10 个 SOTA 病理大模型蒸馏为紧凑高效的 OMNI 编码器系列（7M-330M），计算开销降低达 10 倍且全面超越教师模型。

研究背景与动机¶

近年来数字病理学领域涌现出大量基于大规模自监督预训练的视觉基础模型（如 UNI2-h、Virchow2、H-optimus-1、Prov-GigaPath 等）。这些模型基于不同的预训练数据分布、掩码自编码或对比学习目标构建，各自在特定器官组织、癌症亚型或染色协议上具备互补表征优势。然而，现有病理基础模型参数量庞大（最高达 1.1B）且计算复杂度极高（单切片推理达数百 GFLOPs），极大地限制了全视野数字病理切片（WSI）高通量分析在真实临床工作流中的实时部署。

多教师知识蒸馏（Multi-Teacher Distillation, MTD）理论上能够将多个异构教师的互补特长融合进紧凑的学生网络。但深入研究表明，现有的 MTD 范式存在显著的扩展性失效（scaling failure）：当简单引入更多教师模型时，基于单共享全局 CLS token 压缩重构的方法在下游任务上的性能不仅不再提升，反而出现明显衰退。排查证实这种失效并非源于优化不稳定或梯度方向冲突，而是根本性的表征瓶颈——单个低维潜在表征强行兼顾过多异构几何流形时，引发了不可调和的信息挤压与重构崩塌。

本文的切入角度是打破单共享潜在空间的传统定势，将教师知识的统一提取转变为解耦的专属表征学习。核心 idea：通过引入可学习的教师专属 token 与结构化注意力掩码，令各教师在 patch 级别共享视觉计算、在 token 聚合阶段完全阻断跨教师干扰，并结合尾部 MoE 层与跨样本 InfoNCE 对比损失，实现可向多教师平滑扩展的高效病理基础模型 OMNI。

方法详解¶

整体框架¶

OMNI 的核心目标是将 \(M\) 个各异的冻结教师模型 \(\mathcal{T} = \{t_{\psi_i}\}_{i=1}^M\) 蒸馏压缩进一个紧凑轻量的 ViT 学生编码器 \(s_\theta\) 中。输入组织病理图像切片（tile）首先被切分成 \(N_p\) 个不重叠的 patch 并线性映射为特征序列 \(\mathbf{Z}^{\mathrm{patch}}\)。与标准 ViT 仅添加单个共享全局 CLS token 不同，模型针对 \(M\) 个教师分别引入 \(M\) 个可学习的专属教师 token，与 patch token 共同构成长度为 \(M + N_p\) 的序列。

在整个 ViT 主干中，模型引入特制的结构化注意力掩码，确保 patch 之间自由交互、各教师 token 能够感知全部 patch，但教师 token 之间严格隔离互不窥探；在 Transformer 骨干尾部的最后若干层，标准 FFN 被替换为基于输入自适应分配的稀疏混合专家（MoE）层，为不同教师对应的表征提供额外的非线性拟合容量而保持推理计算量恒定。最终提取出的 \(M\) 个教师专属特征通过多任务投影头联合监督，在下游评估时仅需直接进行平均池化（mean-pooling）即可获得鲁棒且互补的综合特征向量。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入病理切片图像<br/>H×W×C 切块为 Np 个 Patch"] --> B["前置输入序列构建<br/>拼接 M 个独立可学习教师 Token"]
    B --> C["教师专属注意力掩码<br/>Patch 全局交互，教师 Token 间严格屏蔽"]
    C --> D["尾部混合专家层<br/>Top-2 路由自适应分配前馈计算容量"]
    D --> E["双路对齐与对比学习目标<br/>点对点余弦/L1重构 + 批次 InfoNCE 对比损失"]
    E --> F["下游多教师表征汇聚<br/>均值池化聚合为紧凑全局切片向量"]

关键设计¶

1. 教师专属 Token 与结构化注意力掩码：切断潜在空间瓶颈与跨教师互扰

针对单个共享 CLS 向量无法承载多个教师异构流形导致的重构崩塌问题，模型在输入序列显式维护可学习教师 token 矩阵 \(\mathbf{C} \in \mathbb{R}^{M \times D}\) 及对应的位置编码 \(\mathbf{E}^{\mathrm{teach}} \in \mathbb{R}^{M \times D}\)，使得每个教师 token 能够独立、专注地从底层 patch 中提炼归属于该教师视角的视觉抽象。

为了防止不同教师 token 在深层自注意力计算中过早同质化或相互干扰，模型在所有注意力层应用分块结构化注意力掩码 \(\mathbf{A} \in \{0, 1\}^{(M + N_p) \times (M + N_p)}\)：

\[\mathbf{A} = \begin{pmatrix} \mathbf{I}_M & \mathbf{1}_{M \times N_p} \\ \mathbf{1}_{N_p \times M} & \mathbf{1}_{N_p \times N_p} \end{pmatrix}\]

其中 \(\mathbf{I}_M\) 强制令任意教师 token 仅能关注自身，严格屏蔽与其他教师 token 的横向注意力；\(\mathbf{1}_{M \times N_p}\) 与 \(\mathbf{1}_{N_p \times M}\) 允许教师 token 与全部 patch token 之间双向注意交互，确保空间上下文的充分聚合；\(\mathbf{1}_{N_p \times N_p}\) 则保留底层 patch token 之间的标准自注意力。这种机制实现了“底层视觉特征联合计算共享、高层语义聚合完全解耦独立”，从根本上保障了表征在面对更多教师时的扩展能力。

2. 尾部混合专家层：兼顾参数表征容量与常数级推理吞吐

随着蒸馏教师数量从 3 个扩展到 10 个，学生网络在顶层需要拟合截然不同的特征映射。若单纯加宽或加深整个 ViT 模型，推理延时将显著上升，违背了轻量化的初衷。为此，OMNI 将学生网络最后若干个 Transformer Block（小型/微型为最后 3 层，大型为最后 5 层）中的前馈网络（FFN）替换为含 \(E=5\) 个专家的 MoE 结构，每次计算采用 Top-2 门控路由机制。

路由器根据当前 token 的隐藏状态，自适应地选择两个最匹配的 FFN 专家进行激活计算。为了避免门控网络坍塌到单一常用专家，模型引入了 Switch Transformer 的负载均衡正则项：

\[\mathcal{L}_{\mathrm{MoE}} = E \cdot \sum_{e=1}^{E} f_e \cdot \bar{p}_e\]

其中 \(f_e\) 为批次内路由至专家 \(e\) 的 token 比例，\(\bar{p}_e\) 为专家 \(e\) 的平均门控概率。MoE 层的加入使得学生模型在不增加活跃计算 FLOPs 的前提下，拥有更充裕的参数空间以协调 10 个异构教师的非线性分布转换。

3. 实例判别式对比正则化：抑制跨样本表征退化

常规蒸馏目标仅采用逐样本的点对点回归（如余弦相似度与平滑 L1 损失），这在多目标联合强监督下容易导致学生输出向均值漂移，缺乏区分批次内不同组织切片细微病理差别的全局对比结构。为此，设计为每个教师专属 token 分支增加跨批次的对称 InfoNCE 对比损失。

对于批大小为 \(B\) 的训练样本，经投影和 \(\ell_2\) 归一化后的学生教师专属向量 \(\tilde{\mathbf{z}}_i^{(b)}\) 与真实教师 CLS 向量 \(\tilde{\mathbf{y}}_i^{(b)}\) 构成两两相似度矩阵 \([\mathbf{S}_i]_{ab} = \frac{\tilde{\mathbf{z}}_i^{(a)\top} \tilde{\mathbf{y}}_i^{(b)}}{\tau}\)（温度超参 \(\tau=0.1\)）。对比损失将对角线样本作为正样本对，对同批次其他切片进行对比排斥：

\[\mathcal{L}_i^{\mathrm{ctr}} = \frac{1}{2} \left[ -\frac{1}{B} \sum_{b=1}^{B} \log \frac{e^{[\mathbf{S}_i]_{bb}}}{\sum_{b'=1}^{B} e^{[\mathbf{S}_i]_{bb'}}} - \frac{1}{B} \sum_{b=1}^{B} \log \frac{e^{[\mathbf{S}_i]_{bb}}}{\sum_{a=1}^{B} e^{[\mathbf{S}_i]_{ab}}} \right]\]

该判别信号为点对点距离损失提供了互补的正交几何约束，显著增强了特征空间的全局判别力与泛化边界。

损失函数 / 训练策略¶

模型综合采用 token 级重构损失、密集 patch 级重构损失、InfoNCE 对比损失以及 MoE 负载均衡损失共同端到端训练。针对第 \(i\) 个教师的总蒸馏目标定义为：

\[\mathcal{L}_i = \mathcal{L}^{\mathrm{cls}}_{\mathrm{MTD}, i} + \mathcal{L}^{\mathrm{patch}}_{\mathrm{MTD}, i} + \lambda_{\mathrm{ctr}} \mathcal{L}_i^{\mathrm{ctr}}\]

其中 \(\mathcal{L}^{\mathrm{cls}}_{\mathrm{MTD}, i}\) 与 \(\mathcal{L}^{\mathrm{patch}}_{\mathrm{MTD}, i}\) 均采用复合重构度量 \(\Phi(\mathbf{x}, \mathbf{y}) = [1 - \cos(\mathbf{x}, \mathbf{y})] + \lambda_{L_1} \text{smooth}_{L_1}(\mathbf{x} - \mathbf{y})\)，损失权重设为 \(\lambda_{L_1}=1.0\)、\(\lambda_{\mathrm{ctr}}=0.1\)。

总体联合优化目标为所有 \(M=10\) 个教师的平均损失与 MoE 正则化项之和：

\[\mathcal{L} = \frac{1}{M} \sum_{i=1}^{M} \mathcal{L}_i + \alpha_{\mathrm{MoE}} \mathcal{L}_{\mathrm{MoE}}\]

其中 \(\alpha_{\mathrm{MoE}} = 0.05\)。所有 OMNI 模型从头随机初始化，使用 Adam 优化器在源自 TCGA 的 10,000 张 WSI 抽取的 1000 万张切片上训练 700k 步，基础学习率设为 \(1 \times 10^{-4}\)，批大小为 128。

实验关键数据¶

主实验¶

在涵盖 13 个多器官病理切片数据集的 THUNDER 基准（切片级）以及 7 个 CPTAC 队列共 25 项分子/亚型预测任务的 PathoBench 基准（全片级）上进行全面线性探测（Linear Probing）评估。

模型	参数量	嵌入维度	GFLOPs	切片级 F1 (%) (13 个数据集)	全片级 Bacc (%) (25 个任务)	综合得分 (Mean)	综合平均排名 (↓)
H-optimus-1	1.1B	1536	296.0	86.3	62.5	74.4	5.4
Prov-GigaPath	1.1B	1536	223.4	84.2	61.0	72.6	10.0
UNI2-h	681M	1536	180.4	87.4	61.5	74.5	6.8
Virchow2	631M	1280	164.6	85.9	61.5	73.7	7.4
CONCH 1.5	307M	1024	79.5	84.3	61.5	72.9	10.2
GPFM (蒸馏基线)	303M	1024	81.1	83.2	60.9	72.0	9.3
Vanilla MTD (ViT-B)	86M	768	22.0	85.4	61.8	73.6	7.3
OMNI-T (本文)	7.3M	192	1.39	83.9	63.2	73.6	7.5
OMNI-S (本文)	28.7M	384	5.48	85.3	62.9	74.1	5.8
OMNI-B (本文)	114M	768	21.8	87.1	63.2	75.2	4.2
OMNI-L (本文)	330M	1024	67.5	87.8	63.5	75.7	2.3

消融实验¶

在 ViT-Base 架构（OMNI-B）上，系统验证结构化注意力掩码、尾部 MoE 模块以及 InfoNCE 对比损失对最终性能的独立影响。

配置	InfoNCE 损失	MoE 尾部层	结构化注意力掩码	切片级 F1 (%)	全片级 Bacc (%)	综合均值 (%)	说明
OMNI-B (完整模型)	✓	✓	✓	87.1	63.2	75.2	完整方案
w/o InfoNCE 对比损失	✗	✓	✓	86.6 (-0.5)	61.1 (-2.1)	73.9 (-1.3)	去除跨样本判别信号，全片性能跌幅显著
w/o MoE 尾部层	✓	✗	✓	86.6 (-0.5)	61.2 (-2.0)	73.9 (-1.3)	缺少专长路由，多教师表征容量受限
w/o 结构化注意力掩码	✓	✓	✗	86.2 (-0.9)	61.6 (-1.6)	73.9 (-1.3)	教师 token 全局混杂，表征隔离被破坏
Vanilla MTD (单 CLS 基线)	✗	✗	✗	85.4 (-1.7)	61.8 (-1.4)	73.6 (-1.6)	传统 MTD 范式，表征瓶颈严重

关键发现¶

扩展法则发生质变：Vanilla MTD 随教师数量从 1 增加至 10 时，重构余弦相似度从 0.90 剧烈坍塌至 0.83，下游微调 F1 不升反降；而 OMNI 在 10 个教师下重构质量仍维持在 0.87 以上，性能随教师数量增加保持单调稳健上升，且无论教师按强到弱或弱到强排列，扩展收益均持续成立。
极致的计算性价比：参数量仅 7.3M、FLOPs 仅 1.39G 的 OMNI-T 在全片级任务上即可达到 63.2% Bacc，超越了 1.1B 参数的 Prov-GigaPath 与 H-optimus-1；OMNI-L 以约 \(1/4\) 的 FLOPs 全面击败 SOTA 教师模型 UNI2-h。
卓越的对抗鲁棒性与校准度：在 PGD 对抗攻击下，OMNI-L 的性能衰减仅 16.1%（SOTA 教师模型普遍衰减 29.4%~58.1%），期望校准误差（ECE）低至 2.8%，证明多教师互补特征解耦融合后具备更高的预测置信度与几何平滑度。

亮点与洞察¶

解耦式教师专属 Token 机制：抛弃了传统 MTD 依赖单个共享全局瓶颈向量的惯性假设，通过引入并行专属 Token 成功解除了异构潜在空间融合的容量挤压。
简洁高效的注意力隔离设计：仅需引入极轻巧的块状阻断掩码即可实现各教师特征提取分支的解耦，无需构建多个独立的昂贵学生骨干或复杂级联解码头。
下游即插即用的无参数均值融合：通过 CKA 相似度与 UMAP 分析证实，学习到的教师表征在潜在空间中既相互解耦又平滑可插值，直接均值池化即可取得强于复杂加权聚合器的综合表征。

局限与展望¶

下游聚合策略仍有探索空间：目前主要采用无参数的均值池化进行简单整合；虽然在 BRACS 等任务上初步验证了可学习加权平均能带来额外收益，但跨器官自适应的动态 token 融合机制尚未全面铺开。
密集 Patch 表征尚未专属化：当前 patch 级别的重构监督仍沿用传统统一特征投影方式，尚未在 patch 级别针对性构建不同教师偏好的多尺度局部表征。
临床前泛化验证仍需深化：尽管在 39 个公开多中心任务上验证充分，但在真实医院环境下罕见突变、复杂混合组织学切片等极具挑战性的前瞻性临床验证仍有待落地。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次揭示并解决了多教师蒸馏在多教师扩展时的表征崩塌难题，提出的掩码与专属 Token 架构简洁优雅。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖 10 个前沿病理基础模型教师、39 项切片及全片级任务，并详尽探究了鲁棒性、校准度与专家路由行为。
写作质量: ⭐⭐⭐⭐⭐ 论证脉络极度清晰，由现象驱动到假说验证与架构设计，实验呈现与图表分析严谨扎实。
价值: ⭐⭐⭐⭐⭐ 开源的 OMNI 系列极大地推进了边缘及临床算力受限环境下高通量计算病理分析的实用性。