Beyond ViT Tokens: Masked-Diffusion Pretrained Convolutional Pathology Foundation Model for Cell-Level Dense Prediction¶
会议: NeurIPS 2026
arXiv: 2605.08276
OpenReview: https://openreview.net/forum?id=u0riwBjIls
论文: https://arxiv.org/abs/2605.08276
领域: self_supervised
关键词: 计算病理学, 掩码扩散模型, 卷积基础模型, 细胞级密集预测, 自监督表征学习
一句话总结¶
针对现有基于 ViT 的病理大模型因固定分块破坏空间连续性与细胞细粒度形态细节的瓶颈,提出结合掩码扩散预训练与 ConvNeXt-UNet 架构的卷积病理基础模型 CMD,在保持骨干冻结且仅微调极少下游参数下显著超越 ViT 基线与端到端分割网络。
研究背景与动机¶
细胞级密集预测(如细胞核分割、炎性细胞检测以及组织微环境解析)是计算病理学落地的关键核心任务。与常规自然图像语义分割不同,病理组织切片结构高度复杂,细胞核形态多样、尺度极小且细胞间隙极其狭窄,边界模糊多变;与此同时,多中心采集带来的染色差异、扫描仪色散以及组织源异质性造成强烈的域偏移,而像素级的密集标注成本又异常高昂,迫切需要具备强泛化能力的自监督病理基础模型提供结构化的下游表征。
然而,现阶段主流的病理基础模型几乎全盘基于 Vision Transformer(ViT)架构,采用对比学习或掩码图像建模进行自监督预训练。这类方法依赖固定的图像 Patch Token 化处理,在处理全局滑片级(WSI)多示例分类时表现优异,但在细胞级密集预测中暴露了天然缺陷:刚性的分块硬性切断了连续的细胞组织学结构,使得 Token 内部的细微形态、染色质纹理和极小间隙的边缘线索发生不可逆丢失,在密集预测时往往只能捕捉粗粒度的组织区域,导致邻近细胞粘连或漏检。卷积神经网络所固有的归纳偏置(局部性、平移不变性与空间连续性)本是密集预测的理想载体,但长久以来缺乏一种能与 ViT 掩码自编码器相媲美的大规模生成式可扩展预训练范式。
本文的切入视角在于:既然扩散模型能够通过高保真去噪重建捕获细腻的局部纹理分布,完全可以将扩散过程重构为时间步可控的结构化掩码去噪自编码范式,并在纯卷积架构与像素空间中完成。核心 idea:提出基于掩码扩散机制与纯卷积 ConvNeXt-UNet 的病理基础模型 CMD,在像素空间直接以时间步控制结构遮挡比例实施去噪重构,结合冻结病理大模型的语义先验调制,构建出保持精细细胞形态与连续轮廓的密集表征提取器。
方法详解¶
整体框架¶
CMD 框架将自监督表征学习与下游密集预测任务完全解耦:首先在海量无标注病理切片上,利用部分遮挡的图像块与冻结病理模型的全局语义先验进行基于像素空间的掩码扩散去噪预训练;预训练完成后彻底舍弃耗时的反向迭代采样过程,将预训练模型直接作为冻结的特征编码器,从 ConvNeXt-UNet 的多尺度解码器块中提取富含空间细节的稠密特征,仅训练轻量级的分割头即可高效迁移至下游密集预测任务。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入病理图像切片<br/>256×256 像素空间"] --> B["1. 结构化掩码扩散预训练目标<br/>时间步 t 映射遮挡率 rt 随机掩码"]
C["冻结病理大模型特征<br/>z_pfm 提供全局语义"] --> D["3. 冻结全局病理先验注入<br/>自适应层归一化 AdaLN 调制"]
B --> E["2. 纯卷积 ConvNeXt-UNet 骨干<br/>多尺度编解码与跳跃连接重构"]
D --> E
E --> F["重构损失监督<br/>L1 像素误差 + SSIM 结构相似度"]
E -.->|预训练收敛后冻结骨干| G["4. 冻结稠密特征提取与轻量微调<br/>固定 t=50 抽取多尺度 Decoder 特征"]
G --> H["轻量级/SOTA 分割头<br/>细胞核分割与密集预测输出"]
关键设计¶
1. 结构化掩码扩散预训练目标:以空间掩码替代高斯噪声建立形态去噪机制 传统 DDPM 在正向过程中添加连续高斯噪声,虽然能生成高保真样本,但难以对空间组织学形态建立针对性的局部结构推理能力。CMD 将扩散模型重新表述为以时间步条件化的结构恢复任务:对于无标注图像切片 \(x_0 \in \mathbb{R}^{H \times W \times C}\),均匀采样时间步 \(t \in \{1, \ldots, T\}\),并将遮挡比率定义为 \(r_t = \frac{t}{T+1}\)。通过在 \(8 \times 8\) 的 Patch 网格上生成随机二值掩码 \(m_t\),得到遮挡图像 \(x_t = m_t \odot x_0\)。此时时间步 \(t\) 不再控制噪声方差,而是控制结构丢失的严重程度,强迫网络学习从可见组织上下文中根据病理形态学规律修复缺失的细胞核与基质纹理。
2. 纯卷积 ConvNeXt-UNet 骨干:在像素空间无损维持显微局部性与多尺度轮廓 针对 ViT 固定 Patch 切分所带来的边界断裂问题,CMD 构建了完全对称的 ConvNeXt-UNet 架构。输入输出均在原生像素空间(\(256 \times 256\))中进行,彻底规避了 VAE 潜在空间压缩对细胞微小边缘、染色质颗粒度等高频信号的丢失。模型采用 5 级下采样阶段、低分辨率(\(8 \times 8\))瓶颈层与 5 级对称上采样解码器,各阶段利用 \(7 \times 7\) 深度可分离卷积、LayerNorm、GELU 和全局响应归一化(GRN)模块构建 ConvNeXt 块,并借助跳跃连接跨层复用细粒度边缘线索,天然保留了微观组织学所需的局部归纳偏置。
3. 冻结全局病理先验注入:融合时间步与大模型特征的自适应调制 尽管纯卷积网络具有极佳的局部几何建模能力,但在局部遮挡极重时可能缺乏全局宏观组织语义。CMD 引入冻结的预训练病理大模型(如 UNI 或 H-Optimus)抽取整图级特征向量 \(z_{\text{pfm}}\),与时间步正弦嵌入 \(\phi_t(t)\) 线性变换后融合为条件向量 \(c = \phi_t(t) + \phi_z(z_{\text{pfm}})\)。该条件向量通过自适应层归一化(AdaLN-Zero)动态缩放和偏置每个 ConvNeXt 块的归一化激活: $\(\operatorname{AdaLN}(u,c) = \gamma(c) \odot \operatorname{LN}(u) + \beta(c)\)$ 这种设计使得卷积骨干无需耗费参数去重学高层语义,而是专注于在全局上下文指引下精细重构细胞微观结构。
4. 冻结稠密特征提取与轻量微调:单步前向推理驱动高效下游细胞分割 在完成自监督预训练后,CMD 不执行生成模型昂贵的百步迭代去噪推理,而是将其作为静态稠密特征提取器。给定测试切片,固定选择浅层去噪时间步(实验中设定为 \(t=50\)),直接执行一次单步前向传播,并抽取 ConvNeXt-UNet 解码器中多层中间特征。这些特征具有高分辨率的空间对齐性,通过轻量级多尺度融合头(参数量仅约 7.15M)或多通道卷积头即可完成精准的细胞核实例与语义分割,不仅参数效率极高,且完全解耦了基础表征与下游微调。
损失函数 / 训练策略¶
CMD 的自监督预训练联合优化像素级绝对误差与结构一致性,损失函数定义为: $\(\mathcal{L}_{\text{CMD}} = \mathbb{E}_{x_0, t} \left[ \lambda_1 \|x_0 - \hat{x}_0\|_1 + \lambda_s \mathcal{L}_{\text{SSIM}}(x_0, \hat{x}_0) \right]\)$ 其中结构相似性损失 \(\mathcal{L}_{\text{SSIM}}(x_0, \hat{x}_0) = \frac{1 - \text{SSIM}(x_0, \hat{x}_0)}{2}\),通过局部均值、方差与协方差保证重构细胞边缘和微观对比度的保真度。在 100 万张 HistAI 无标注病理切片上,采用 AdamW 优化器(学习率 \(3 \times 10^{-5}\),BF16 混合精度,EMA 衰减率 0.9999)训练 80k 步。下游任务保持 CMD 骨干完全冻结,仅以余弦退火调度和交叉熵结合 Dice 损失训练下游分割头。
实验关键数据¶
主实验¶
在三项代表性细胞级密集预测基准(CPM-15、CPM-17 和 TNBC)上,将冻结的 CMD-L 表征与多种前沿冻结 ViT 病理大模型及全量微调端到端分割网络进行对比(表中指标汇报 Dice / Precision,带 95% 置信区间):
| 模型体系 | 模型名称 | 冻结参数 (M) | 分割头参数 (M) | 输入分辨率 | CPM-15 (Dice / Prec) | CPM-17 (Dice / Prec) | TNBC (Dice / Prec) |
|---|---|---|---|---|---|---|---|
| ViT 病理大模型 | Prov-Gigapath | 1100 | 9.42 | 392 / 490 / 504 | 0.759 / 0.753 | 0.871 / 0.863 | 0.828 / 0.809 |
| ViT 病理大模型 | Hibou | 304 | 7.06 | 384 / 496 / 512 | 0.794 / 0.781 | 0.875 / 0.872 | 0.841 / 0.823 |
| ViT 病理大模型 | H-Optimus-0 | 1100 | 9.42 | 392 / 490 / 504 | 0.781 / 0.795 | 0.881 / 0.869 | 0.843 / 0.811 |
| ViT 病理大模型 | UNI | 303 | 7.06 | 384 / 496 / 512 | 0.738 / 0.726 | 0.867 / 0.854 | 0.820 / 0.786 |
| 端到端分割基线 | TransUNet | - | 105.28 | 384 / 496 / 512 | 0.809 / 0.770 | 0.897 / 0.881 | 0.843 / 0.825 |
| 端到端分割基线 | ViT-Adapter | - | 133.52 | 384 / 496 / 512 | 0.876 / 0.889 | 0.872 / 0.877 | 0.854 / 0.838 |
| 纯卷积基础模型 (Ours) | CMD-L (轻量头) | 602 | 7.15 | 统一 256 | 0.864 / 0.849 | 0.865 / 0.901 | 0.854 / 0.851 |
| 纯卷积基础模型 (Ours) | CMD-L (SOTA头) | 602 | 21.21 | 统一 256 | 0.889 / 0.885 | 0.880 / 0.883 | 0.854 / 0.830 |
在统一的 256×256 分辨率且冻结骨干条件下,CMD-L 仅凭轻量头就在各项数据集上显著超过了输入分辨率高达 512 的 ViT 大模型,并使用 21.21M 参数的分割头全面超越了参数量高达 133M 的端到端全量优化网络 ViT-Adapter。
消融实验¶
为验证核心组件的合理性,论文在 55k 数据集上对骨干架构、预训练空间及大模型条件注入进行了系统消融(以 Linear Probe 评估 Dice 与边界 F1 分数 BF1):
| 消融维度 | 具体配置 | CPM-15 (Dice / BF1) | CPM-17 (Dice / BF1) | Kumar (Dice / BF1) | TNBC (Dice / BF1) |
|---|---|---|---|---|---|
| 骨干架构 (统一 VAE+PFM) | DiT | 44.39 / 19.65 | 56.33 / 23.24 | 48.80 / 11.12 | 55.11 / 21.12 |
| 骨干架构 | 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 |
| 表征空间 (统一 ConvNeXt) | VAE 潜在空间 (256²) | 50.19 / 33.66 | 55.20 / 35.04 | 51.61 / 34.38 | 53.23 / 25.09 |
| 表征空间 | 高分辨率 VAE (512²) | 47.96 / 11.52 | 63.80 / 30.49 | 53.09 / 33.31 | 58.24 / 22.43 |
| 表征空间 | 像素空间 (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 |
关键发现¶
- 像素空间建模带来质的飞跃:从 VAE 潜在空间转为直接像素空间后,边界 F1(BF1)分数在 CPM-15 和 CPM-17 上分别暴涨了 +22.37 和 +38.41。这证明潜在扩散模型所依赖的 VAE 下采样是丢失微观细胞边界的主要根源,密集病理预测必须扎根于像素级重构。
- 极小样本下的惊人迁移能力:在 1-shot 极度匮乏设定下,冻结的 CMD-L 在 CPM-17 和 TNBC 上分别取得 0.769 和 0.714 的 Dice,而强大的 ViT 大模型(如 H-Optimus-1)仅能达到 0.468 和 0.486。CMD 仅凭 7M 微调参数就展现出强大的形态学少样本泛化力。
- 分辨率不敏感与计算扩展性:CMD 统一使用 \(256 \times 256\) 即可在性能上匹敌使用 496 或 512 分辨率的 ViT 模型;同时,随着预训练步数从 10k 扩增至 80k、参数量从 CMD-B(217M)扩展至 CMD-L(602M),下游 Dice 指标保持单调稳定增长。
亮点与洞察¶
- 将扩散机制转为可控掩码去噪自编码器:用结构遮挡率取代随机高斯噪声,将生成扩散模型改造为高表现力的自监督特征预训练范式,摆脱了生成扩散繁琐多步采样的负担。
- 打破 ViT 在病理基础模型中的绝对垄断:在密集预测这一细粒度领域证明了现代纯卷积(ConvNeXt-UNet)的先天归纳偏置优势,为细胞级分析提供了一条高质量、低计算负担的技术路径。
- AdaLN 条件注入巧妙融合语义与几何:冻结已有 ViT 病理大模型充当“导师”提供高阶病理语义,由卷积骨干专注学习高频形态重构,实现了先验复用与微观表征学习的解耦双赢。
局限与展望¶
- 预训练数据来源单一与规模受限:目前预训练基于 HistAI 约 100 万张图像,相较于行业顶级前沿模型的数千万切片仍有扩充空间,未来可继续验证超大规模数据下的扩展定律。
- 下游微调的时间步超参数依赖:特征抽取默认选取固定时间步 \(t=50\),尽管实验证实该步长能很好平衡语义与底层纹理,但缺乏对不同下游任务的最优自适应选择机制。
- 跨模态与全切片诊断串联待探索:目前主要聚焦于局部细胞与细胞核级密集预测,后续可探索结合病理报告等文本信息的跨模态引导,并向 WSI 级别肿瘤微环境级联建模延展。
相关工作与启发¶
- vs ViT-based PFMs (Prov-GigaPath / UNI / H-Optimus):主流 ViT 大模型基于 Patch Token 化,注重全局滑片级语义,但在细胞微观边缘上缺乏连续性且易断裂;CMD 采用纯卷积与像素级掩码扩散,在冻结下游微调时细胞轮廓完整度与边界精度显著领先。
- vs 端到端分割网络 (nnUNet / TransUNet / ViT-Adapter):传统方法依赖大量全监督密集标注进行端到端全量参数优化,且跨数据集泛化脆弱;CMD 采用自监督大规模无标注预训练,下游仅微调 7M 参数即可在少样本下取得压倒性优势。
- vs 潜在扩散表示学习 (Latent Diffusion):潜在空间扩散受制于 VAE 压缩重建的模糊化效应;CMD 证实了对于显微细胞预测,直接在像素空间执行多尺度掩码扩散重构是保持细胞核尖锐边缘的决定性因素。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙重构掩码扩散为时间步可控自编码器,成功开拓纯卷积病理基础模型路线。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 5 大骨干消融、多分辨率对比、少样本微调及特征聚类可视化,证据详实严谨。
- 写作质量: ⭐⭐⭐⭐⭐ 逻辑结构清晰严密,设计动机层层递进,消融对比切中痛点。
- 价值: ⭐⭐⭐⭐⭐ 为计算病理学细胞级密集分析确立了高效鲁棒的预训练基准,极具落地与参考意义。