Seg4Diff: Unveiling Open-Vocabulary Segmentation in Text-to-Image Diffusion Transformers¶

会议: NeurIPS 2025
arXiv: 2509.18096
代码: GitHub
领域: 图像分割
关键词: 扩散模型, 开放词汇分割, MM-DiT, 注意力分析, 语义对齐

一句话总结¶

通过系统分析多模态扩散Transformer（MM-DiT）的联合注意力机制，发现特定层（"语义定位专家层"）天然具备高质量语义分割能力，并提出轻量微调方法MAGNET同时提升分割与生成性能。

研究背景与动机¶

文本到图像扩散模型通过交叉注意力机制将语言概念隐式地定位到图像区域，已有工作表明U-Net架构的扩散模型的交叉注意力图可以用于零样本语义分割。然而U-Net产生的注意力图通常噪声较多且空间碎片化，限制了分割质量。

近年来，扩散Transformer（DiT）架构逐步取代U-Net，特别是多模态扩散Transformer（MM-DiT，如Stable Diffusion 3）引入了联合自注意力——将图像和文本token拼接后统一做自注意力，实现更强的跨模态交互。但对于MM-DiT内部注意力如何贡献于图像生成的理解仍然有限，尤其缺乏对其语义定位能力的深入分析。

本文的核心动机有三个方面： 1. MM-DiT的联合注意力机制与传统U-Net的交叉注意力本质不同，需要专门分析 2. 如果MM-DiT内部存在语义定位能力，能否直接用于开放词汇分割？ 3. 能否通过强化这种能力同时提升分割和生成质量？

方法详解¶

整体框架¶

Seg4Diff是一个系统性分析和利用MM-DiT语义定位能力的框架，分为三个递进阶段：(1) 分析MM-DiT联合注意力的内部交互模式；(2) 基于发现的语义定位专家层构建零样本分割方案；(3) 提出MAGNET轻量微调策略增强分割与生成。

关键设计¶

语义对齐的涌现分析（Emergent Semantic Alignment）: 作者首先将MM-DiT的联合注意力分解为四种交互类型：图像到图像（I2I）、图像到文本（I2T）、文本到图像（T2I）、文本到文本（T2T）。通过定量分析发现I2T注意力得分不成比例地高于I2I（尽管I2T区域面积约为I2I的1/40），说明I2T主导了整体注意力预算。进一步通过PCA可视化和Value投影的L2范数分析，发现特定层（尤其是第9个MM-DiT块）对文本token的Value范数显著高于图像token，表明文本信息主要在这些层注入到语义对齐的图像区域。通过高斯模糊扰动实验因果验证了这些层对图像-文本对齐的关键作用。
语义分组的涌现（Emergent Semantic Grouping）: 基于上述分析，作者提出利用I2T注意力图进行开放词汇分割。具体做法是：将输入图像通过VAE编码为潜空间表示 \(x_{\text{img}}\)，在中间时间步加噪保留空间结构；文本提示由类别名拼接而成，编码为 \(x_{\text{text}}\)。在语义定位专家层提取I2T注意力图，对所有头取平均得到 \(\bar{A}_{I2T} = \frac{1}{H}\sum_{h=1}^{H} A_{I2T}^h\)，然后reshape为每个文本token的mask logit \(M^{(j)} \in \mathbb{R}^{h \times w}\)。对属于同一类别的多个文本token的注意力图取平均，最终对每个像素做argmax得到分割预测。实验表明第9层表现最佳。此外，<pad> token在无条件生成中也能自发地将图像分解为有意义的语义区域，可用于无监督分割。
MAGNET轻量微调（Mask Alignment for Segmentation and Generation）: 在语义定位专家层应用LoRA微调（rank=16），优化两个互补损失：流匹配损失 \(\mathcal{L}_{\text{FM}}\) 监督扩散过程，以及mask损失 \(\mathcal{L}_{\text{mask}}\) 强化I2T注意力的语义分组能力。mask损失通过二分图匹配将 \(l \cdot H\) 个I2T注意力图与ground-truth mask一一配对，然后计算focal loss和dice loss的加权和：\(\mathcal{L}_{\text{mask}} = \lambda_{\text{focal}} \mathcal{L}_{\text{focal}} + \lambda_{\text{dice}} \mathcal{L}_{\text{dice}}\)。总损失为 \(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{FM}} + \lambda_{\text{mask}} \mathcal{L}_{\text{mask}}\)。

损失函数 / 训练策略¶

训练使用SA-1B或COCO的10k图像，CogVLM生成的文本描述。图像分辨率1024×1024，AdamW优化器（lr=\(1 \times 10^{-5}\)），两张A6000 GPU训练，有效batch size为16。对SA-1B使用头级注意力图匹配细粒度mask，对COCO使用token级注意力图匹配更粗粒度的mask。

实验关键数据¶

主实验¶

开放词汇语义分割（mIoU）:

方法	架构	VOC20	COCO-Obj	PC59	ADE
DiffSegmenter	SD1.5	66.4	40.0	45.9	24.2
iSeg	SD1.5	82.9	57.3	39.2	24.2
ProxyCLIP	CLIP-H/14	83.3	49.8	39.6	24.2
CorrCLIP	CLIP-H/14	91.8	52.7	47.9	28.8
Seg4Diff	SD3	89.2	62.0	49.0	34.2
Seg4Diff+MAGNET	SD3+COCO	89.8	62.9	51.2	35.2

无监督分割（mIoU）:

方法	VOC21	PC59	Object	Stuff-27	ADE
DiffSeg	49.8	48.8	23.2	44.2	37.7
DiffCut	62.0	54.1	32.0	46.1	42.4
Seg4Diff	54.9	52.6	38.5	49.7	44.9
Seg4Diff+MAGNET	56.1	53.5	38.8	53.5	45.4

消融实验¶

配置	CLIPScore	说明
Baseline (SD3)	27.14	无mask alignment
+MAGNET (SA-1B)	27.24	图像生成质量也提升
+MAGNET (COCO)	27.28	COCO训练效果略优

T2I-CompBench++评测也证实MAGNET在属性绑定、数值概念等方面优于baseline。

关键发现¶

语义定位是MM-DiT的涌现属性，集中在特定层（第9层为"语义定位专家层"）
注意力头呈现多粒度语义分组：不同头关注目标的不同部位（如熊的耳朵和腿），汇总后形成完整mask
<pad> token在无条件生成中也能发现有意义的语义区域
强化语义分组不仅提升分割，还改善生成质量

亮点与洞察¶

将扩散模型"分析→发现→利用→增强"的研究范式做到极致，逻辑链条完整
发现MM-DiT第9层作为语义定位专家层的洞察非常有价值，为统一生成与感知提供了新视角
MAGNET仅需10k图像和LoRA微调即可同时提升两个任务，实用性强
<pad> token的语义分组发现为无监督分割提供了全新途径

局限与展望¶

当前仅分析了SD3和少数MM-DiT变体，可扩展到更多DiT架构
开放词汇分割仍依赖已知类别名作为文本提示，真正的开放场景需要结合类别发现
分割精度受限于潜空间分辨率（远低于像素级），高分辨率应用需额外上采样
MAGNET的mask损失引入了分割标注依赖，如何完全无监督增强值得探索

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次系统分析MM-DiT的语义定位能力并提出利用方案
实验充分度: ⭐⭐⭐⭐ 多数据集多任务评测全面，消融实验充分
写作质量: ⭐⭐⭐⭐⭐ 逻辑递进清晰，分析-发现-利用的叙事结构优美
价值: ⭐⭐⭐⭐ 为统一生成与感知模型指明了方向