Beyond ViT Tokens: Masked-Diffusion Pretrained Convolutional Pathology Foundation Model¶
会议: NeurIPS 2026
arXiv: 2605.08276
论文: https://arxiv.org/abs/2605.08276
领域: 医学图像
关键词: 计算病理学, 基础模型, 掩码扩散模型, 卷积神经网络, 细胞级密集预测
一句话总结¶
针对主流 ViT 病理基础模型切块分词破坏空间连续性与局部细胞形态的问题,本文提出掩码扩散卷积基础模型 CMD,在像素空间通过结构化掩码去噪预训练纯卷积 ConvNeXt-UNet 并结合冻结病理大模型条件注入,实现对细胞级密集预测任务的高精度、强泛化与极低标注开销表征迁移。
研究背景与动机¶
细胞级密集预测(如细胞核分割、炎性细胞检测以及组织微环境分析)是计算病理学的核心基石。然而,病理切片包含微小且形态异质性极高的组织结构,细胞边界模糊纤细,加之不同组织来源、染色批次与扫描仪引入的严重域偏移,使得密集人工像素级标注成本极其昂贵。
近年来以 UNI、Virchow、Prov-GigaPath 为代表的病理基础模型大多依赖 Vision Transformer(ViT)架构,并在切片级弱监督多示例学习中取得了卓越成就。但这类 ViT 模型依赖固定的 patch 分词机制,切块操作不可避免地割裂了空间连续结构,丢失了 patch 内部对于像素与实例级识别至关重要的微观形态、染色质纹理及细胞间边界细节。与此同时,具备归纳偏置与多尺度特性的现代卷积网络(如 ConvNeXt)却长期缺乏像掩码图像建模(MIM)那样可扩展的自监督预训练范式。
受生成模型通过高保真图像重构学习视觉表征的启发,本文探索将生成式建模引入卷积病理基础模型的设计空间中。与其通过高斯噪声扩散合成图像,不如将时间步放松为可控的结构化遮挡毁损难度。核心 idea:构建基于纯卷积 ConvNeXt-UNet 的像素级掩码扩散预训练模型(CMD),以结构化掩码去噪取代高斯噪声,并通过自适应归一化注入冻结病理大模型的语义先验,学习保留微观局部形态与多尺度上下文的细胞级密集表征。
方法详解¶
整体框架¶
CMD 的核心目标是学习可复用的生成式病理密集表征,而非训练特定任务的端到端分割网络。预训练阶段,模型接收经时间步控制遮挡比例的病理图像块,以纯卷积 ConvNeXt-UNet 结合冻结病理基础模型的全局特征条件在像素空间进行遮挡重建。下游评估阶段,生成采样与重建头被完全舍弃,预训练骨干网络被冻结作为密集特征提取器,仅训练轻量分割头以极少标注量完成细胞核分割等密集预测。
整个前向与迁移流程包含自监督掩码扩散预训练、多尺度密集特征抽取与下游轻量预测适配三个阶段:
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["未标注病理切片图像 x₀<br/>(256×256 像素空间)"] --> B["1. 结构化掩码扩散前向过程<br/>时间步 t 映射掩码率 r_t = t/(T+1) 掩码补丁生成 x_t"]
A --> C["冻结病理大模型编码器<br/>(提取全局语义先验 z_pfm)"]
B --> D["2. 纯卷积 ConvNeXt-UNet 骨干网络<br/>5 级下采样-瓶颈-5 级对称上采样 + 密集跳跃连接"]
C --> E["时间步与全局语义先验注入<br/>c = φ_t(t) + φ_z(z_pfm) 过 AdaLN 调制各卷积块"]
E --> D
D --> F["3. 像素级多目标重建优化<br/>联合 L1 像素误差与 SSIM 结构相似度损失"]
F -->|预训练完成 / 冻结骨干网络| G["4. 冻结多尺度特征抽取<br/>抽取解码器各层特征图映射至下游轻量分割头"]
关键设计¶
1. 结构化掩码扩散目标:以渐进组织遮挡替代高斯噪声毁损 传统 DDPM 的高斯加噪过程破坏了病理图像的高频形态结构,不利于引导网络感知组织微观拓扑。CMD 将扩散过程重构为时间步可控的结构化重构任务,采样离散时间步 \(t \in [1, T]\)(其中 \(T=1000\)),将其映射为遮挡比例 \(r_t = \frac{t}{T+1}\)。输入图像划分为大小为 \(8 \times 8\) 的不重叠补丁,通过二值掩码矩阵 \(m_t \in \{0, 1\}^{H \times W \times C}\) 生成毁损图像 \(x_t = m_t \odot x_0\)。这一设计将时间步 \(t\) 的物理含义从连续噪声方差尺度转变为结构遮挡难度,迫使网络必须依据未遮挡的局部细胞与间质上下文推断隐匿的微观解剖结构。
2. 纯卷积 ConvNeXt-UNet 骨干:兼顾微观局部偏置与多尺度上下文 为了克服 ViT 分词造成的细胞边界破坏,CMD 采用对称的 5 阶段下采样与 5 阶段上采样 ConvNeXt-UNet 拓扑。特征分辨率沿 \(256^2 \to 128^2 \to 64^2 \to 32^2 \to 16^2 \to 8^2\) 递减再对称恢复,并在 \(8 \times 8\) 瓶颈层设置 6 个堆叠 ConvNeXt 块用于聚合宏观组织上下文。每个 ConvNeXt 块采用 \(7 \times 7\) 深度可分离卷积、LayerNorm、反向瓶颈 MLP(扩展比为 3)与全局响应归一化(GRN),并通过 Pixel-unshuffle 下采样与 Pixel-shuffle 上采样避免空间信息损失。编码器各阶段的多尺度特征图直接通过通道拼接跳跃连接至对应解码器,使得极微小的细胞核轮廓与染色质边界得以原样传递。
3. 时间步与病理大模型语义自适应调制:AdaLN 双重特征融合注入 单纯依赖局部重构易陷入浅层纹理拼合,缺乏高层病理类型与微环境先验。CMD 引入冻结的病理大模型(如 H-Optimus-0 mini 或 UNI)提取整图级全局嵌入 \(z_{\text{pfm}}\),将其与扩散时间步编码 \(\phi_t(t)\) 线性投影到同一空间并相加融合: $\(c = \phi_t(t) + \phi_z(z_{\text{pfm}})\)$ 随后,在主干网络的每一个 ConvNeXt 块中采用自适应层归一化(AdaLN-Zero)对特征进行仿射调制: $\(\operatorname{AdaLN}(u, c) = \gamma(c) \odot \operatorname{LN}(u) + \beta(c)\)$ 其中 \(\gamma(c)\) 与 \(\beta(c)\) 由 \(c\) 线性映射生成。该机制使得卷积主干既能感知当前的遮挡毁损程度,又能受到全局组织病理学分类语义的约束,实现微观形态建模与宏观病理语义的双向对齐。
4. 像素空间与联合结构损失优化:保留高频微观结构 不同于潜在扩散模型(LDM)在 VAE 压缩潜空间中去噪,病理切片中诸如细胞膜间隙、核仁等极细微结构会在 VAE 降采样下被严重平滑。CMD 坚持在原生像素空间展开掩码扩散,联合计算像素级 \(L_1\) 误差与结构相似度(SSIM)损失: $\(\mathcal{L}_{\text{CMD}} = \mathbb{E}_{x_0, t} \left[ \lambda_1 \|x_0 - \hat{x}_0\|_1 + \lambda_s \frac{1 - \operatorname{SSIM}(x_0, \hat{x}_0)}{2} \right]\)$ 该联合目标既保证了重构结果对染色质吸光强度的像素保真度,又利用 SSIM 局部协方差统计强化了细胞交界边缘的高频对比度感知。
损失函数 / 训练策略¶
预训练在 100 万张来自 HistAI 集合的 \(512 \times 512\) 未标注病理图像块上展开。输入经混合裁剪策略(80% 概率随机裁剪至 \(256 \times 256\),20% 概率全图缩放至 \(256 \times 256\))以兼顾多分辨率感受野。优化器选用 AdamW,初始学习率 \(3 \times 10^{-5}\),无权重衰减,采用 BF16 混合精度与 EMA 衰减率 0.9999。模型经过 80K 步迭代完成收敛。下游微调时冻结预训练骨干网络,仅微调包含 7.15M 参数的轻量多尺度融合分割头(或 21.21M 参数的四层卷积解码分割头),以余弦退火学习率调度与联合 Cross-Entropy + Dice 损失完成训练。
实验关键数据¶
主实验¶
在 CPM-15、CPM-17 和 TNBC 三大公认细胞核分割基准上,冻结骨干网络的 CMD-L 与代表性 ViT 基础模型、端到端全监督分割网络进行全面对比(结果含 95% 置信区间):
| 方法 | 骨干状态 | 可训练参数 (M) | CPM-15 Dice | CPM-15 Precision | CPM-17 Dice | CPM-17 Precision | TNBC Dice | TNBC Precision |
|---|---|---|---|---|---|---|---|---|
| CONCHv1.5 | 冻结 ViT | 7.06 | 0.748 | 0.740 | 0.858 | 0.863 | 0.826 | 0.806 |
| Prov-Gigapath | 冻结 ViT | 9.42 | 0.759 | 0.753 | 0.871 | 0.863 | 0.828 | 0.809 |
| Hibou | 冻结 ViT | 7.06 | 0.794 | 0.781 | 0.875 | 0.872 | 0.841 | 0.823 |
| H-Optimus-0 | 冻结 ViT | 9.42 | 0.781 | 0.795 | 0.881 | 0.869 | 0.843 | 0.811 |
| UNI2-h | 冻结 ViT | 9.42 | 0.767 | 0.761 | 0.873 | 0.864 | 0.833 | 0.804 |
| nnUNet | 全监督端到端 | 12.19 | 0.791 | 0.751 | 0.875 | 0.842 | 0.803 | 0.754 |
| TransUNet | 全监督端到端 | 105.28 | 0.809 | 0.770 | 0.897 | 0.881 | 0.843 | 0.825 |
| ViT-Adapter | 全监督端到端 | 133.52 | 0.876 | 0.889 | 0.872 | 0.877 | 0.854 | 0.838 |
| CMD-L (轻量头 7.15M) | 冻结 ConvNeXt | 7.15 | 0.864 | 0.849 | 0.865 | 0.901 | 0.854 | 0.851 |
| CMD-L (SOTA头 21.21M) | 冻结 ConvNeXt | 21.21 | 0.889 | 0.885 | 0.880 | 0.883 | 0.854 | 0.830 |
消融实验¶
在相同无监督数据集(55K 病理图像)与线性探测评估协议下,对骨干网络选型、表示空间及大模型条件注入进行的系统消融:
| 实验组别 | 消融配置 | CPM-15 Dice / BF1 | CPM-17 Dice / BF1 | Kumar Dice / BF1 | TNBC Dice / BF1 | 说明 |
|---|---|---|---|---|---|---|
| 骨干架构 | DiT | 44.39 / 19.65 | 56.33 / 23.24 | 48.80 / 11.12 | 55.11 / 21.12 | 分词破坏边界,边界 F1 (BF1) 极低 |
| Attention U-Net | 53.68 / 19.31 | 56.04 / 23.19 | 50.37 / 20.80 | 54.05 / 16.39 | 经典结构局部性好但容量受限 | |
| ConvNeXt-UNet | 50.19 / 33.66 | 55.20 / 35.04 | 51.61 / 34.38 | 53.23 / 25.09 | 边界敏感性 (BF1) 显著占优 | |
| 空间选择 | VAE Latent (256²) | 50.19 / 33.66 | 55.20 / 35.04 | 51.61 / 34.38 | 53.23 / 25.09 | 潜空间压缩过滤细胞高频形态 |
| High-res VAE (512²) | 47.96 / 11.52 | 63.80 / 30.49 | 53.09 / 33.31 | 58.24 / 22.43 | 提升输入分辨率仍无法弥补潜空间损失 | |
| Pixel-space (256²) | 63.25 / 56.03 | 76.74 / 73.45 | 62.76 / 62.25 | 73.51 / 75.63 | 像素原生空间极大幅度提升微观细节捕捉 | |
| 语义注入 | CMD w/o PFM | 55.17 / 43.52 | 71.29 / 64.36 | 60.21 / 68.02 | 60.24 / 45.90 | 缺少大模型先验指导,语义辨识度下降 |
| CMD w/ PFM (UNI) | 63.25 / 56.03 | 76.74 / 73.45 | 62.76 / 62.25 | 73.51 / 75.63 | 全局语义与局部掩码重建形成绝佳互补 |
少样本与扩展性关键发现¶
- 极限制度少样本适配优势显著:在 1-shot 极度匮乏设定下(CPM-17 与 TNBC),冻结的 ViT 大模型(UNI2-h、H-Optimus-1)Dice 仅有 0.46~0.49,而冻结 CMD-L(仅 7M 参数头)分别取得 0.769 与 0.714 的 Dice,大幅超越完全端到端微调的 ViT-Adapter(364M 全量参数下仅 0.714 与 0.498)。证明其表征直接嵌入了微观形态连续性,无需大量样本矫正局部拓扑。
- 分辨率不敏感特性:现有 ViT 模型必须依赖 384~512 的较高输入分辨率,而 CMD 统一在 \(256 \times 256\) 输入下即达到或超越前者的高分辨率表现,大幅节省下游推理与内存成本。
- 优异的算力扩展规律:预训练步数从 10K 扩展至 80K,CPM-15 性能稳步上升(0.851 升至 0.864),未见自监督过拟合;骨干网络从 CMD-B(217M)扩展至 CMD-L(602M),在 CPM-15、CPM-17 和 TNBC 上的 Dice 分别从 0.852/0.829/0.836 提升至 0.865/0.863/0.850。
亮点与洞察¶
- 范式创新:为纯卷积架构开辟类似 MIM 的生成式自监督预训练路径。不同于传统对比学习或端到端分割,将扩散模型的时间步诠释为空间掩码遮挡比率,使 ConvNeXt 得以进行像素级自监督学习并捕捉密集细胞解剖拓扑。
- 架构解耦:局部形态与全局语义的分工协作。由冻结病理大模型通过 AdaLN 负责输入全局宏观组织先验,全卷积扩散网络专注重构局部连续细胞几何与边缘,巧妙化解了单一架构在全局上下文与局部连续性之间的天然权衡。
- 即插即用迁移范式:验证了密集预测领域无需对数百兆巨型大模型进行昂贵的端到端全量微调,冻结密集表征配合微型预测头(7M)即可匹敌端到端 SOTA 模型。
局限与展望¶
- 预训练计算成本与数据规模仍有扩展空间:当前预训练数据为 100 万 patch 规模,相较于数千万乃至亿级 WSI patch 的 ViT 病理大模型仍属中等规模,且像素级扩散多步训练开销较大。
- 全切片(WSI)层级整合尚存空白:当前实验集中在 patch 级细胞核与组织密集预测,如何将 CMD 提取的细胞图谱表征与切片级临床预后/诊断弱监督 MIL 模型有机联通,有待进一步探索。
- 未来方向:探索多模态文本报告条件引导下的掩码扩散预训练,以及在超高分辨率千兆像素全切片分析中的层级化部署。
相关工作与启发¶
- vs UNI / Virchow / Prov-GigaPath:传统 PFM 专注于基于 ViT 的切片全局特征提取,通过固定的 patch 切割损失了微观连续轮廓;CMD 则重塑纯卷积网络的局部偏置,专注服务细胞级密集预测。
- vs TransUNet / nnUNet / ViT-Adapter:全监督模型依赖大量密集真值标注且跨域迁移泛化弱;CMD 通过百万级无标注病理切片自监督掩码扩散预训练,实现了在极少量标注(1-shot/5-shot)下以固定特征超越全监督网络的表现。
- vs 标准潜在扩散模型 (LDM):LDM 依赖 VAE 潜在编码压缩,极易抹杀纳米/微米级的细胞膜染色边界;CMD 实验证明原生像素空间扩散对于细粒度病理形态学至关重要。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙将掩码扩散机制与纯卷积 UNet 及病理大模型条件调制结合,开拓了纯卷积病理密集表征学习新范式。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 5 大数据集消融、主干网络全维度对比、1/5/10-shot 极限少样本评测及步数/参数扩展性分析。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑缜密,动机推导自洽,消融实验针对架构痛点层层剖析。
- 价值: ⭐⭐⭐⭐⭐ 为计算病理学由全局分类向精细化细胞与空间生物学表征演进提供了极具启发性的技术路线。