Scaling Whole-Slide Pathology Foundation Model Pretraining with Billions Off-the-Shelf Tokens¶
会议: ECCV 2026
论文: ECCV 官方链接
领域: 医学图像
关键词: 计算病理、全切片图像 (WSI)、基础模型、自监督学习、Token 裁剪与保留
一句话总结¶
针对计算病理中仅用 tile 级 [CLS] token 导致空间细粒度病理特征丢失且全空间 token 计算过重的两难困境,提出免训练贪心多样性保留模块 TokRet 与块稀疏长上下文预训练框架 PathTokScale,在仅需现有模型 10% 切片数量下将训练规模扩展至 20 亿 token 级别,在肿瘤分级与生存分析等多项任务中超越现有顶尖病理大模型。
研究背景与动机¶
病理全切片图像(Whole Slide Image, WSI)分析在临床疾病诊断、肿瘤分级、预后生存分析及基因突变预测中占据核心地位。然而,单张 WSI 拥有高达数十亿像素(十亿级,Giga-pixel)的极高分辨率,且标注成本高昂,促使研究界广泛采用弱监督多示例学习(MIL)与自监督基础模型。现有主流方法(如 UNI、Virchow、TITAN、GigaPath 等)通常利用在 patch 上预训练的视觉 Transformer(ViT)提取局部特征,并为了计算可行性,统一截取 tile 级别的任务不可知 [CLS] token 作为该区域的全局表征。
然而,过度依赖 [CLS] token 形成了病理基础模型严重的表征瓶颈。研究表明,单纯增加 ViT 规模(如扩展到 ViT-Large/Giant)在此类架构中只能带来微弱收益,且即使将 [CLS] 与平均池化的空间 token 简单拼接,也无法挽回细粒度局部病理结构与异质性线索的丢失。病理诊断往往高度依赖于离散分布、微小而关键的细胞形态变化。若直接保留 tile 内全部空间 token(单张切片可达数十万至上百万 token),Transformer 的二次方计算复杂度又在工程上不可承受。深入分析表明,局部瓦片(tile)内的空间 token 存在大量信息重叠与冗余,这说明无需全量保留,只需挑选最富代表性的子集即可还原切片的全局特征分布。
基于此切入点,本文跳出“盲目堆叠切片数据量或盲目扩大参数量”的传统范式,转向有效 token 维度的可扩展性。核心 idea:将瓦片内空间 token 保留建模为组合优化中的最大最小多样性问题,设计免额外训练的贪心多样性选择模块 TokRet,并结合基于块路由的稀疏长上下文 Transformer 与 RoPE 位置编码构建 PathTokScale 框架,以十亿级离架 token 完成切片级自监督预训练。
方法详解¶
整体框架¶
PathTokScale 的整体处理流程涵盖图像分层预处理、免训练空间 token 贪心保留、块稀疏长上下文 Transformer 编码与区域级 DINOv2 自监督预训练四大阶段。首先将 WSI 划分为非重叠的宏观区域(如 \(7168 \times 7168\) 像素)以及更小的瓦片(\(448 \times 448\) 像素),并通过切片 ViT 将瓦片拆分为 \(16 \times 16\) 像素的小 patch 提取嵌入;随后通过 TokRet 在瓦片内保留代表性多样化 token(约 10% 比例);最后利用带有块级路由(Block Router)的高效长序列 Transformer 进行全局交互与多视图自监督蒸馏。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["WSI 切片输入<br/>分层划分为区域与瓦片"] --> B["瓦片级 ViT 编码<br/>提取全部空间 patch tokens"]
B --> C["TokRet 贪心多样性保留<br/>最大最小距离迭代选取代表 token"]
C --> D["块稀疏路由 Transformer<br/>粗粒度块打分 + Top-M 细粒度注意力"]
D --> E["DINOv2 多视图自蒸馏<br/>两路数据增强视图对齐预训练"]
E --> F["下游切片级任务微调<br/>肿瘤分级 / 生存预测 / 基因突变"]
关键设计¶
1. TokRet 贪心多样性保留:以最大最小距离免训练压缩瓦片空间冗余 为了兼顾细粒度表征丰富性与推理吞吐,TokRet 放弃了易受启发式打分漂移影响的自注意力幅值截断,而是将空间 token 筛选形式化为组合优化中的最大最小多样性问题(Max-Min Diversity Problem, MMDP)。给定一个瓦片内的候选 token 集合 \(T\)(基数 \(m\)),目标是选出一个大小为 \(\tilde{m}\)(约保留 10%)的子集 \(\tilde{T}\),最大化所选 token 之间的最小余弦距离:
其中 \(d(\gamma, \omega)\) 为两 token 间的余弦距离。为避免组合爆炸,TokRet 采用贪心策略:首先一次性预计算候选集合内所有 token 之间的余弦距离矩阵;随后选取与其它所有 token 平均距离最大的单个 token 作为初始锚点;在后续每一步迭代中,计算剩余候选 token 到当前已选集合 \(\tilde{T}\) 中所有 token 的最小距离,并选取该最小距离最大的候选 token 加入 \(\tilde{T}\),直至达到目标数量 \(\tilde{m}\)。这种机制确保了所保留的 token 在特征空间中均匀分散在原始分布边界,有效捕获稀疏且关键的病理异质性。
2. 块路由稀疏长上下文编码:打破数十万 token 的二次方计算瓶颈 在经过 TokRet 筛选后,一张完整切片所包含的空间 token 仍可达 10 万至 20 万量级,常规自注意力机制依然无法承载。PathTokScale 借鉴 Mixture-of-Block-Attention(MoBA)思想,引入块级稀疏路由机制。将保留的 token 序列划分为大小为 512 的非重叠块 \(\{B_1, \dots, B_K\}\)。对于任意 query token \(q_i\),首先与各个 key 块的均值向量进行粗粒度相似度打分:
系统根据打分选出相关度最高的前 \(M\) 个块(默认 \(M=4\)),仅在这 \(M\) 个块内部执行细粒度的 Softmax 局部自注意力计算,将二次方复杂度降至近似线性复杂度。同时,引入二维旋转位置编码(2D RoPE),同一瓦片内部的不同空间 token 共享瓦片的网格坐标,由 ViT 本身捕获 patch 间相对空间关系,使模型在预训练时仅需以 \(7168 \times 7168\) 区域进行训练,在下游任务中能无缝外推至超大尺寸整张 WSI。
3. PathTokScale 自监督预训练:十亿级离架 Token 的切片级自蒸馏 在区域和切片级别,模型采用 DINOv2 框架构建自监督学习目标。通过对选取的区域 token 序列施加随机裁剪、水平翻转及特征空间高斯噪声等增强操作,构建两个增强视图 \(T_r^{(1)}\) 和 \(T_r^{(2)}\),分别输入 student 网络和基于指数移动平均(EMA)更新的 momentum teacher 网络,通过对齐两者的输出概率分布完成自蒸馏。由于 TokRet 保留了高信息量 token,该框架基于 HistAI 公开数据集及约 1 万张切片,在无需专有超大规模数据集的前提下,将有效训练 token 总量推向约 20 亿(2 Billion)级别,比传统基于切片级特征的研究提升了 20 倍。
实验关键数据¶
主实验¶
在切片级分类(肿瘤分级 BRACS、生物标志物预测 BCNB ER、基因突变预测 TCGA-LUAD TP53)与生存预测任务中,PathTokScale 全面超越传统 MIL、现代长序列 Mamba 以及大参数量病理基础模型。
表 1:切片级分类任务宏观性能对比(摘自原论文 Table 1)
| 方法 | BRACS (F1) | BRACS (AUC) | BCNB ER (F1) | BCNB ER (AUC) | TCGA-LUAD TP53 (F1) | TCGA-LUAD TP53 (AUC) |
|---|---|---|---|---|---|---|
| ABMIL (UNI-1) | 0.692±0.03 | 0.875±0.02 | 0.754±0.08 | 0.857±0.07 | 0.646±0.04 | 0.701±0.07 |
| TransMIL (UNI-1) | 0.592±0.03 | 0.859±0.02 | 0.690±0.06 | 0.810±0.08 | 0.608±0.04 | 0.674±0.08 |
| GMMamba (UNI-1) | 0.670±0.04 | 0.893±0.03 | 0.784±0.08 | 0.873±0.07 | 0.678±0.04 | 0.734±0.05 |
| PathVQ (UNI-1) | 0.730±0.02 | 0.902±0.01 | 0.782±0.07 | 0.872±0.05 | 0.684±0.03 | 0.733±0.06 |
| GigaPath | 0.677±0.03 | 0.862±0.03 | 0.791±0.07 | 0.875±0.06 | 0.670±0.04 | 0.738±0.06 |
| TITAN | 0.696±0.04 | 0.891±0.01 | 0.813±0.06 | 0.905±0.05 | 0.702±0.03 | 0.765±0.06 |
| UNI-2 (ABMIL) | 0.698±0.03 | 0.887±0.02 | 0.770±0.07 | 0.861±0.06 | 0.661±0.03 | 0.729±0.06 |
| Ours + UNI-1 | 0.779±0.03 | 0.914±0.03 | 0.801±0.05 | 0.890±0.06 | 0.714±0.03 | 0.763±0.04 |
| Ours + Conch-1.5 | 0.767±0.05 | 0.910±0.05 | 0.809±0.07 | 0.898±0.05 | 0.725±0.05 | 0.771±0.03 |
表 2:多癌种患者疾病特异性生存预测(C-Index,摘自原论文 Table 2)
| 类别 / 模型 | BRCA | CRC | BLCA | UCEC | KIRC | 平均 C-Index |
|---|---|---|---|---|---|---|
| ABMIL (UNI-1) | 0.633±0.06 | 0.612±0.08 | 0.540±0.07 | 0.671±0.08 | 0.691±0.08 | 0.629 |
| TransMIL (UNI-1) | 0.612±0.07 | 0.684±0.06 | 0.595±0.06 | 0.695±0.08 | 0.671±0.10 | 0.651 |
| PANTHER (UNI-1) | 0.758±0.06 | 0.665±0.10 | 0.612±0.07 | 0.757±0.10 | 0.716±0.10 | 0.702 |
| CHIEF | 0.737±0.04 | 0.680±0.08 | 0.599±0.02 | 0.758±0.10 | 0.736±0.06 | 0.702 |
| GigaPath (170k+ slides) | 0.687±0.08 | 0.628±0.08 | 0.589±0.05 | 0.779±0.10 | 0.751±0.07 | 0.687 |
| TITAN (330k+ slides) | 0.713±0.04 | 0.710±0.11 | 0.657±0.05 | 0.789±0.09 | 0.774±0.06 | 0.729 |
| Ours + UNI-1 | 0.712±0.07 | 0.699±0.09 | 0.635±0.04 | 0.769±0.06 | 0.767±0.09 | 0.716 |
| Ours + Conch-1.5 | 0.717±0.04 | 0.719±0.07 | 0.655±0.05 | 0.784±0.08 | 0.778±0.07 | 0.731 |
消融实验与关键发现¶
论文针对 Token 选择机制、保留数量 \(K\) 及预训练 Token 规模进行了细致消融: - Token 选择策略对比(图 3a):在 BRACS 分类任务上,对比了基础方案 [CLS](准确率约 0.70)、[CLS]+均值池化(约 0.70)、全局均值池化(\(4 \times 4\) 为 0.68,\(7 \times 7\) 为 0.70)、基于自注意力幅值的 Attention Top-16 / Top-64(分别为 0.72 和 0.74),而 TokRet Top-20 达到 0.76,TokRet Top-32 达到约 0.77。证明最大最小多样性筛选在极少 token(20 个)下的表征覆盖率远胜均值平均和注意力启发式。 - Token 保留数量敏感性(图 4):当每个瓦片保留的 token 极少(\(K \le 4\) 时,如 Top-2 降至 0.64,Top-4 降至 0.66)时,性能甚至落后于 [CLS],因为过度稀疏打乱了全局语义完整性;但当 \(K \ge 8\) 后性能急剧提升并在 \(K=20 \sim 32\) 达到稳定饱和区。 - 预训练 Token 扩展律(图 3b):将预训练 token 数量从 0.1B(纯 CLS 起步,BRACS 性能约 0.70)逐步扩展到 0.8B(约 0.74)、1.6B(约 0.77)及 2.0B(达到近 0.79),下游准确率呈现清晰单调上升趋势,确立了病理切片领域在 token 维度的 Scaling Law。
亮点与洞察¶
- 将 Token 压缩转为组合多样性问题:巧妙将切片 token 筛选映射为 Max-Min Diversity Problem,无需任何神经网络反向传播和前置训练,仅依靠特征距离矩阵即可从原始瓦片中选出覆盖几何边界的代表性子集。
- 以 Token 扩展打破数据量扩展壁垒:以往基础模型(如 GigaPath、TITAN)消耗 10 万至 30 余万张切片堆砌预训练数据,而本文证明仅仅借助公开展出的 ~3 万张切片,通过深度挖掘瓦片内未被使用的空间 token,就能在 10% 切片规模下取得同等或更优的表征能力。
- 长序列稀疏注意力与 RoPE 的无缝移植:将大语言模型中的块路由注意力(MoBA)和 2D 旋转位置编码引入病理切片建模,使轻量 6 层 Transformer 具备了高效吞吐 20 万 token 切片的工程能力。
局限与展望¶
- 瓦片内坐标细粒度丢失:为了降低位置编码复杂度,目前同一瓦片内的所有保留 token 均被赋予了相同的网格 \((x, y)\) 坐标,依赖 ViT 隐层特征来隐式编码瓦片内细微相对几何位置,未来可探索更精细的层次化坐标编码。
- 全局预计算矩阵的内存开销:当单张切片组织面积极大且需要批量处理海量瓦片时,全局成对距离矩阵与贪心循环在极端瓦片数下会产生轻微的 CPU/GPU 内存峰值。
- 多模态扩展空间:当前自监督预训练仅在纯视觉表征下完成,未来若能将保留的多样化高价值 token 与临床病理诊断报告文本进行对齐,有望构建更通用的病理多模态大模型。
相关工作与启发¶
- vs GigaPath / TITAN: GigaPath 与 TITAN 依赖海量私有数据(分别超过 17 万和 33 万张切片)通过 MAE / iBOT 在瓦片级 [CLS] token 上做切片级长上下文预训练。本文证明限制其效能的不是切片总数而是 [CLS] 本身的表征损失;PathTokScale 仅用其约 10% 切片数量(3 万张切片),通过释放空间 token 并扩展至 20 亿 token 规模即打平甚至反超了其各项下游表现。
- vs PathVQ: PathVQ 尝试通过自编码器离散量化(VQ-VAE)压缩空间 token 维度以保留空间信息,但自编码器需要额外训练,且易受重构语义损失影响。TokRet 采用免训练的贪心多样性筛选,保留原始高保真特征并彻底剔除冗余,不仅推理效率更高,在各项下游指标上也显著优于 PathVQ。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ [将组合优化最大最小多样性策略引入病理瓦片空间 token 筛选,摆脱了启发式注意力与繁重调优]
- 实验充分度: ⭐⭐⭐⭐⭐ [覆盖肿瘤分类、生物标志物预测、基因突变与 5 个癌种生存分析,附带完整的 token 数量消融与扩展律曲线]
- 写作质量: ⭐⭐⭐⭐⭐ [逻辑严密,对计算病理基础模型 Scaling 瓶颈的分析直击要害,图表精美]
- 价值: ⭐⭐⭐⭐⭐ [为学术界在有限切片资源下开展高质量病理基础模型预训练提供了切实可行的高效范式]