跳转至

Diffusion Model as a Generalized Segmentation Learner

会议: ECCV 2026
论文: ECCV 2026 Poster
项目主页: https://wang-haoxiao.github.io/DiGSeg/
领域: 语义分割 / 扩散模型 / 医学图像
关键词: 扩散模型, 通用分割学习器, 开放词表分割, 潜在去噪, 跨域迁移

一句话总结

DiGSeg 将预训练扩散模型重塑为通用分割学习器,通过将图像与掩码编码至潜在空间、多尺度注入 CLIP 文本特征并结合退火多分辨率噪声调度与 DDIM-trailing 单步高效去噪,实现免后处理的高精度语义分割与跨域泛化。

研究背景与动机

现代图像分割系统在封闭集语义分割、实例分割以及遥感、农业和医学影像等专用垂直领域均取得了显著进展。然而,当前的分割生态呈现出严重的碎片化割裂状态:不同任务与领域往往依赖完全不同的定制架构、标签空间与训练流水线,一旦从封闭词表向开放词表迁移,或从自然图像拓展到航空遥感和医疗图像,系统通常需要重构大量组件。这一瓶颈促使研究者思考:能否构建一个单一、统一的通用分割框架,在跨任务、跨词表和跨视觉领域之间保持稳健的预测能力?

近期研究发现,在大规模互联网图文数据上预训练的文本到图像扩散模型,其去噪轨迹中蕴含着丰富的空间对应关系与语义先验。然而,早期试图发掘这一潜力的尝试(如 DiffSeg、DiffCut、DiffuMask、Seg4Diff 等)大多停留在事后提取扩散 U-Net 内部的交叉注意力或自注意力特征图。这些注意力图往往噪声严重、分辨率受限且跨层一致性差,导致生成的掩码碎片化严重,必须依赖复杂的启发式图割或聚类后处理;更关键的是,这些方案大多仅适配单一任务,未能将扩散模型转变为通用的密集预测引擎。

本文打破了“将扩散模型当作冻结特征提取器或事后聚类器”的传统范式,提出通过轻量微调将扩散模型的隐式先验转化为显式的分割条件接口。核心 idea:将图像与真实掩码共同映射到预训练 VAE 潜在空间,以图像 latent 和多尺度 CLIP 文本嵌入为条件微调去噪 U-Net,结合退火多分辨率噪声与 DDIM-trailing 调度,使扩散模型直接作为端到端的通用分割学习器运行。

方法详解

整体框架

DiGSeg 将分割任务形式化为条件潜在扩散去噪生成过程 \(p_\theta(y|x)\),其中 \(x\) 为输入的 RGB 图像,\(y\) 为对应的分割掩码。整个框架包含三大核心组件:视觉潜在通路(Visual Latent Pathway)、CLIP 对齐文本条件模块(CLIP-Aligned Text Conditioning Module)以及分割一致去噪网络(Segmentation-Consistent Denoising U-Net)。

在训练阶段,输入图像 \(x\) 与单通道掩码 \(y\)(复制为 3 通道)分别通过冻结的 Stable Diffusion VAE 编码器压缩至紧凑潜在空间 \(z_x\) 与 \(z_y\)。对分割 latent \(z_y\) 施加退火多分辨率噪声得到扰动 latent \(z_y^t\),输入去噪 U-Net \(\epsilon_\theta\);同时,文本类别或提示经冻结 CLIP 文本编码器编码后,在多个分辨率层级注入 cross-attention,引导 U-Net 预测噪声并以 MSE 目标进行优化。在推理阶段,从高斯白噪声 \(z_y^T\) 开始,在图像 latent \(z_x\) 和文本 embedding 引导下执行 DDIM-trailing 逆向去噪,最终解码器 \(\mathcal{D}\) 将纯净 latent \(z_y^0\) 重建为连续型预测图,经全局阈值 \(\tau\) 离散化输出最终高质量分割掩码。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入图像 x 与标签掩码 y"] --> B["视觉潜在编码<br/>冻结 VAE 压缩至 z_x 与 z_y"]
    B --> C["退火多分辨率加噪<br/>多尺度高斯扰动注入 z_y 得到 z_y^t"]
    C --> D["多尺度文本条件注入<br/>冻结 CLIP 嵌入过 Cross-Attention"]
    D --> E["分割一致 U-Net 去噪<br/>预测噪声并执行 DDIM-trailing 逆向更新"]
    E --> F["轻量测试时集成与解码<br/>多种子 latent 平均后过 VAE Decoder 解码"]
    F --> G["固定阈值离散化<br/>输出最终分割掩码"]

关键设计

1. 视觉潜在通路与通道复制:无缝适配预训练 VAE 空间 传统直接在像素空间训练扩散分割模型会带来极其庞大的计算开销,而直接抽取中间特征又无法直接产生锐利的边界。DiGSeg 利用开源预训练 Stable Diffusion 的 VAE 编码器 \(\mathcal{E}\) 与解码器 \(\mathcal{D}\) 构建潜在表示空间。针对单通道分割掩码 \(y \in \mathbb{R}^{H \times W}\) 无法直接送入 3 通道 RGB VAE 的维度不兼容问题,本文未引入额外的跨域投影层或修改编码器架构,而是直接将单通道掩码在通道维度上复制为 3 通道伪彩色图像送入 \(\mathcal{E}\): $\(z_x = \mathcal{E}(x), \quad z_y = \mathcal{E}(\text{Concat}(y, y, y))\)$ 这种设计在完全冻结 VAE 参数的前提下,完美保持了潜在空间的连续先验与感知保真度,使去噪 U-Net 能够直接在具有高度结构化语义的隐空间内建立图像模式与空间掩码之间的密集对应。

2. 多尺度 CLIP 文本条件注入:构建开放词表统一语义接口 为了摆脱传统语义分割模型受限于封闭分类头(固定类别的卷积核或分类权重)的局限,DiGSeg 建立了全流程的语言引导机制。给定目标类别标签或任意自然语言查询短语,利用预训练冻结的 CLIP 文本编码器提取密集词向量 \(t_{\text{clip}}\)。不同于仅在瓶颈层或输入端拼接提示的做法,文本向量在去噪 U-Net 的各个下采样与上采样尺度中通过交叉注意力机制持续注入。该多尺度注入使得浅层特征能够捕获粗粒度类别的整体布局,而深层特征能专注于精细局部的边界语义对齐,使同一套权重无需架构修改即可无缝处理开放词表查询。

3. 退火多分辨率噪声调度:强化跨尺度几何与语义一致性 标准扩散模型通常采用各向同性的单分辨率高斯白噪声,这容易使模型在去噪早期忽略大尺度的轮廓结构、而在后期难以恢复极细粒度的边界。DiGSeg 引入退火多分辨率噪声机制(Annealed Multi-Scale Noise):在加噪时组合多个下采样与上采样尺度的空间高斯分量。在扩散初始高噪声时间步,高频扰动成分占据主导,迫使网络重点关注局部的微观轮廓一致性;随着去噪深入至低时间步,低频结构成分主导演化,引导网络强化全局语义与轮廓拓扑。训练目标采用标准 MSE 准则: $\(\mathcal{L}_{\text{seg}} = \mathbb{E}_{t, \epsilon, x, y} \left[ \| \epsilon - \epsilon_\theta(z_y^t, z_x, t, t_{\text{clip}}) \|_2^2 \right]\)$ 该策略显著提升了复杂交界处的掩码连续性,从根本上消除了早期注意力抽取方法中常见的斑状空洞。

4. DDIM-trailing 采样与测试时隐空间集成:兼顾极速与鲁棒性 扩散密集预测的核心痛点在于多步迭代导致的推理延迟。DiGSeg 采用后移时间步(trailing timesteps)采样策略构建快速 DDIM 推理,避免了传统 leading timesteps 导致的训练与推理阶段分布失配,使得模型在仅需 1 步去噪迭代时便能达到高度可用的分割质量。为进一步抑制扩散固有的随机方差,DiGSeg 在推理时支持轻量级的多随机种子测试时集成(Test-Time Ensembling):并行执行若干次前向去噪轨迹,在送入解码器之前直接在潜在特征空间对预测的 \(z_y^0\) 进行均值聚合,再经 \(\mathcal{D}(z_y^0)\) 重建并以全局固定阈值 \(\tau = 0.7\) 完成掩码二值化,有效消除毛刺且无需任何形态学后处理。

损失函数 / 训练策略

模型基于 Stable Diffusion v2 基础架构搭建,在训练过程中完全冻结 VAE 编码器/解码器和 CLIP 文本编码器,仅对去噪 U-Net 参数进行全量微调。训练采用标准 DDPM 加噪流程(\(T=1000\) 步),使用 AdamW 优化器,学习率设为 \(1 \times 10^{-4}\),批大小为 16,输入图像分辨率统一调整为 \(512 \times 512\),并配以 \([0.8, 1.2]\) 的随机缩放抖动与水平翻转。在 8 张 NVIDIA A100 GPU 上,整个训练过程仅需约 1 天(约 30,000 次迭代)即可收敛。

实验关键数据

主实验

在开放词表分割与封闭集语义分割主流基准上,DiGSeg 相比于过往专有模型及基于注意力的扩散分割基线展现了全面的性能优势。

开放词表语义分割性能对比(mIoU, %):

方法 VLM 骨干 视觉骨干 训练集 A-847 PC-459 A-150 PC-59 Cityscapes
ODISE CLIP ViT-L/14 Stable Diffusion COCO-Panoptic 11.0 13.8 28.7 55.3 -
OVSeg CLIP ViT-L/14 Swin-B COCO-Stuff 9.0 12.4 29.6 55.7 -
SAN CLIP ViT-L/14 Side Adapter COCO-Stuff 12.4 15.7 29.9 51.8 -
CAT-Seg CLIP ViT-L/14 - COCO-Stuff 16.0 23.8 31.5 62.0 -
Seg4Diff CLIP ViT-L/14 Stable Diffusion COCO-Stuff - - 35.2 51.2 26.0
ESC-Net CLIP ViT-L/14 - COCO-Stuff 18.1 27.0 41.8 65.6 -
DiGSeg (本文) CLIP ViT-L/14 Stable Diffusion COCO-Stuff 19.9 29.2 43.2 68.4 38.5
提升 vs 次优 - - - +1.8 +2.2 +1.4 +2.8 +0.6
ESC-Net CLIP ViT-B/16 - COCO-Stuff 13.3 21.1 35.6 59.0 -
DiGSeg (本文) CLIP ViT-B/16 Stable Diffusion COCO-Stuff 17.5 23.1 37.2 62.7 36.5
提升 vs 次优 - - - +3.3 +2.0 +1.6 +3.7 +1.3

标准语义分割基准性能对比(mIoU, %):

方法 COCO 输入尺寸 COCO mIoU ADE20K 输入尺寸 ADE20K mIoU
SegFormer-B5 \(512^2\) 46.7 \(640^2\) 51.8
Mask2Former-Swin-L - - \(640^2\) 57.3
OneFormer - - \(640^2\) 57.0
LDMSeg - - \(512^2\) 52.2
SegMAN \(512^2\) 48.2 \(512^2\) 53.2
EoMT \(512^2\) 48.7 \(512^2\) 57.1
DiGSeg (本文) \(512^2\) 50.8 (+2.1) \(512^2\) 58.6 (+1.3)

在遥感 DeepGlobe 道路分割基准上,未经任何遥感架构定制的 DiGSeg 斩获 65.78% IoU、79.93% Precision、78.92% Recall 与 78.79% F1,显著超越了专为卫星影像设计的经典基线。

消融实验

训练噪声调度消融对比(mIoU, %):

噪声配置 COCO mIoU ADE20K mIoU
标准高斯噪声(Standard Gaussian) 48.9 56.7
仅退火噪声(w/ Annealed noise) 49.2 57.1
仅多分辨率噪声(w/ Multi-resolution noise) 49.7 57.6
多分辨率 + 退火噪声(Multi-res. + Ann. 本文) 50.8 58.6

采样步数与集成策略推理权衡(E-Step, ADE20K & COCO):

采样配置 (E-Step) COCO mIoU ADE20K mIoU 推理帧率 (FPS)
\(1 \times 1\)(1步单次前向) 48.2 56.8 11.27
\(1 \times 2\)(2步单次前向) 48.5 57.1 10.61
\(4 \times 2\)(2步4次集成) 50.5 58.5 5.82
\(8 \times 2\)(2步8次集成,默认) 50.8 58.6 3.15
\(20 \times 50\)(无 trailing 标准 DDIM) 50.9 58.8 0.12

关键发现

  • 噪声调度机制协同增效:单纯引入多分辨率噪声在 COCO 上带来 +0.8% mIoU 提升,而退火调度与多分辨率噪声叠加时产生更强共振,使总体性能大幅跃升至 50.8%(提升 +1.9%),证实了跨尺度渐进去噪对稠密掩码重建的决定性价值。
  • 高效采样与集成收益收敛:得益于 DDIM-trailing 设计,仅需 1 步去噪就能获得 56.8% 的 ADE20K mIoU,推理速度突破 11 FPS;增加测试时集成至 8 次(\(8 \times 2\))可进一步稳步增益约 1.8%~2.6%,超过 10 次后精度提升进入平缓期。
  • 卓越的小样本数据效率:在仅使用 ADE20K 训练集 50% 数据的情况下,DiGSeg 达到了与全量数据几乎一致的分割水平;即使缩减至 25% 监督数据,依然保持极高稳健性,展现了预训练扩散视觉先验在低监督场景下的抗过拟合优势。

亮点与洞察

  • 将扩散模型由先验特征库重塑为端到端生成执行器:过去做法普遍将扩散网络视为冻结的特征提取工具或提取脆弱的注意力热图,DiGSeg 证明直接在隐空间通过逆向条件去噪优化掩码潜变量,能彻底摒除复杂的图割与 CRF 等后处理。
  • DDIM-trailing 规避跨域密集预测分布漂移:敏锐捕捉到传统 DDIM 初始时间步偏置导致的密集预测域失配问题,改用 trailing 调度后使单步去噪成为现实,有效打破了扩散模型应用于密集感知任务的推理速度壁垒。
  • 零结构定制实现全领域通用迁移:同一套模型与训练配方在无需修改骨干或引入额外检测头的情况下,直接跨越自然图像、遥感俯拍(DeepGlobe)、作物表型(Phenobench)乃至自动驾驶(BDD100K),展现出强大的多域通用表征能力。

局限与展望

  • 极端专有医学领域的表征局限:在微观病理及高度特异的医疗影像分割中(如 BraTS、REFUGE-2),DiGSeg 的表现相对平缓。作者与评测指出这主要源于通用 CLIP 文本编码器与自然图像扩散底模缺乏细粒度医学概念的语义先验,未来需引入 Med-CLIP 或领域自适应提示微调。
  • 多步集成开销与部署瓶颈:尽管单步推理已达 11 FPS,但追求极致精度仍依赖 \(8 \times 2\) 这一多分支测试时集成方案,在边缘端或实时视频流交互任务中依然面临算力与显存压力。
  • 离散掩码阈值机制的自适应性有待增强:当前依赖全局固定阈值 \(\tau = 0.7\) 完成连续 logit 截断,面对极大尺度跨度或小目标稀疏类别时存在边界欠拟合风险,探索自适应动态阈值或直接建模离散概率分布是重要的改进方向。

相关工作与启发

  • vs ODISE / Seg4Diff: ODISE 和 Seg4Diff 仍依赖预先提取扩散内部特征或依赖 Mask2Former 提议头进行分类匹配;DiGSeg 摒弃了多阶段提议匹配,直接利用去噪扩散生成隐变量掩码,结构更简练且全基准超越。
  • vs Marigold / DiffusionDepth: 虽同属利用预训练扩散先验重塑密集预测的探索,Marigold 面向单通道连续深度估计且无文本引导;DiGSeg 创新性地构建了多尺度 CLIP 文本通路,不仅处理连续分布,更能应对开放词表的多类别离散语义绑定。

评分

  • 新颖性: ⭐⭐⭐⭐☆(将去噪扩散显式作为开放词表与通用分割学习器,退火多分辨率加噪与 DDIM-trailing 结合合理)
  • 实验充分度: ⭐⭐⭐⭐⭐(横跨封闭语义、开放词表、遥感、农业与多域自适应,消融完备翔实)
  • 写作质量: ⭐⭐⭐⭐⭐(结构条理清晰,图表逻辑自洽,理论推导与工程设计结合紧密)
  • 价值: ⭐⭐⭐⭐☆(为视觉生成先验赋能密集感知理解提供了清晰可复现的范例,对通用感知大模型构建有重要借鉴意义)