跳转至

MMDiff: Extending Diffusion Transformers for Multi-Modal Generation

会议: ECCV 2026
论文: ECCV 2026 原文
领域: 3D视觉
关键词: 扩散Transformer, 多模态生成, 特征聚合, 密集感知预测, 概念注意力

一句话总结

针对扩散 Transformer 在生成过程中丢弃丰富感知表征的问题,MMDiff 在完全冻结 DiT 主干的前提下,通过跨去噪轨迹的多时间步自适应特征聚合、单向概念注意力提取以及与判别表征的互补融合,仅需轻量级解码器即可在生成图像的同时高精度预测语义分割、显著性与深度图。

研究背景与动机

图像与视频生成模型在工业界与下游学术研究中,往往不仅需要渲染高质量视觉像素,还需要同步获取场景的几何与语义信息(例如深度估计、语义分割与显著性边界),以支撑自动驾驶标注、3D 重建或实体机器人规划。然而,现有的前沿扩散模型(特别是 FLUX、Stable Diffusion 3 等 Diffusion Transformer,即 DiT)在数十个去噪步和多层网络中处理成千上万个 token,一旦最终图像渲染完毕,所有中间计算的高维感知表征便被直接丢弃。下游系统不得不引入一整套外部视觉感知模型对生成的像素重新做逆向提取,这不仅造成巨大计算浪费,还容易受到重提取模型的误差干扰。

以往尝试利用生成先验提取多模态标注的方法(如 DatasetDM、DiffuMask)多建立在早期的 U-Net 扩散架构之上,且普遍仅依赖单步特征或最终去噪潜在变量。U-Net 依靠层级空间下采样来凝聚语义,而现代 DiT 架构采用扁平结构,所有图块 token 在所有层中均保持相同的空间分辨率。在没有空间层级压缩的情况下,DiT 的语义信息不再集中于特定尺度的深层特征,而是高度弥散于整个去噪时间步轨迹之中:早去噪阶段的特征熵极高,主导全局语义与物体布局轮廓;中间时间步精细化物体边界;晚去噪阶段则保留局部微观纹理与高频细节。因此,简单沿用传统 U-Net 的单时间步(如 \(t=0\))或特征均值抽取,会损失大量关键感知信息,导致密集预测精度严重坍塌。

面对这一结构特性,本文的核心切入点是直接针对 DiT 的去噪动力学轨迹进行时间轴特征重构,使冻结的生成底座无需微调即可同时具备多模态生成能力。核心 idea:通过稀疏层采样与自适应多时间步特征融合网络捕捉去噪全过程的互补感知信息,并结合互不干扰的单向概念注意力机制提取任务空间先验,以极低参数量的轻量解码头将冻结 DiT 转变为高质量多模态生成器。

方法详解

整体框架

MMDiff 的输入为文本提示词(或对真实图像加入扩散噪声后的潜变量),生成主干采用完全冻结的 12B 参数 FLUX.1-dev DiT 模型。在去噪生成过程中,系统在特定的多个去噪时间步与多个代表性层中抽取出中间图块特征,送入自适应特征聚合模块,同时引入独立的单向概念注意力流与可选的判别式视觉表征(如 DINO-v3),最后通过参数量仅约 36M 的轻量级任务专属解码头输出目标模态(语义分割掩码、深度图或显著性图)。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入文本提示词 / 真实图像加噪潜变量"] --> B["多时间步稀疏特征提取<br/>采样FLUX单流层与去噪时间步"]
    B --> C["自适应时空特征聚合<br/>3层Transformer预测动态权重 + CBAM细化"]
    C --> D["概念驱动单向注意力引导<br/>解耦概念Token输出点积相似度图"]
    D --> E["判别表征互补注入<br/>拼接DINO-v3对比学习特征"]
    E --> F["轻量任务专属解码器<br/>DeepLabV3+ / DPT输出多模态预测"]

关键设计

1. 多时间步稀疏特征提取:消除空间与时间冗余 DiT 每一层包含上千个 token 且通道维度高达 3072,若对全层全时间步进行暴力拼接会导致显存爆炸,而简单平均则会抹杀不同去噪阶段的信息互补性。实验分析表明,FLUX 单流层的激活幅度随深度单调递增且相邻层高度相关,而特征熵在去噪轨迹中差异巨大。为此,方法设计了稀疏采样策略,在单流块中选取具有代表性的深浅层集合 \(\mathcal{L} = \{4, 15, 26, 37\}\),并在去噪时间步中选取覆盖粗到细尺度的四个代表点 \(\mathcal{T} = \{0.78, 0.52, 0.26, 0.001\}\)。针对 DiT 激活值偏大且跨层极差显著的问题,特征首先经由线性投影层降维至通道数 \(C'\) 并施加层归一化(LayerNorm),保障数值稳定。

2. 自适应时空特征聚合:动态内容自适应赋权 平滑均匀的物体内部区域更依赖早去噪时间步提供的全局语义一致性,而复杂的物体边界与边缘轮廓则必须依靠后期去噪步的高频几何信息。为了让融合权重随像素空间位置动态调节,MMDiff 构建了一个由 3 层 Transformer(8 头自注意力、隐层维度 768、带有时间位置编码与 LayerScale 初始化)组成的聚合网络。该网络为每个像素位置预测一组归一化的时间步权重向量 \(\alpha \in \mathbb{R}^{|\mathcal{T}|}\)(满足 \(\sum_{t} \alpha_t = 1\))。加权求和后的特征进一步通过卷积注意力模块(CBAM)强化通道与空间响应,并保留来自最纯净步(\(t=0\))的残差跳跃连接:

\[F_{\text{fused}} = \text{CBAM}\left(\sum_{t \in \mathcal{T}} \alpha_t \cdot F_t^{\text{proj}}\right) + F_0^{\text{proj}}\]

这一机制使得空间各处自适应兼顾高层语意与微观边界,在 VOC 分割实验上相较于单步提取带来高达 28.7% 的 mIoU 爆发式增益。

3. 概念驱动单向注意力引导:无损提取精确空间先验 传统的文本-图像交叉注意力会受到定冠词、介词和标点等无关词汇的稀释,且多模态 DiT 的提示词嵌入在层间不断演化,难以在不干扰原始生成图面貌的前提下抽取任意概念掩码。MMDiff 引入解耦的概念 token 流,将预设的概念词(如分割中的“object/background”、深度估计中的“near/far”、显著性检测中的“salient/contour”)输入文本编码器获得初始向量,在多模态块中执行单向交叉注意力更新:

\[o_c = \text{softmax}(q_c k_{xc}^T) v_{xc}\]

其中概念 query 仅向拼接后的图像与概念键值对 \(k_{xc}, v_{xc}\) 寻址,而图像 token 严禁反向向概念寻址,从而完全不干涉生成图像的渲染轨迹。随后在多模态输出空间计算图块向量 \(o_x\) 与概念向量 \(o_c\) 的点积相似度 \(\phi(o_x, o_c) = \text{softmax}(o_x o_c^T)\) 并跨层平均,为密集预测头提供极度清晰且可解释的边界先验。

4. 生成-判别表征互补与轻量解码:极低代价解锁多模态输出 扩散生成的表征擅长捕获物体全局语义结构与生成组合逻辑,而对比学习模型(如 DINO-v3)则长于刻画局部边缘判别特征。MMDiff 将融合后的 DiT 特征、概念引导图与冻结的 DINO-v3-ViT-B/16 特征进行通道拼接,构建全面的多源特征池。后端采用专用的轻量感知头:语义分割采用带有 ASPP 模块的 DeepLabV3+ 解码器,深度估计与显著性检测采用分层重组的 DPT 解码器。整个解码器系统参数量仅约 36M,远小于 12B 的生成底座,既保护了生成主干的通用能力,又赋予其近乎零边际成本的多任务输出能力。

损失函数 / 训练策略

训练时 FLUX 生成主干全程冻结,仅利用少量带标注的真实数据集微调轻量解码器与聚合网络。采用 AdamW 优化器,基础学习率设为 \(10^{-4}\),任务专属解码头使用差分学习率(基础速率的一半 \(5 \times 10^{-5}\)),并引入衰减系数 \(\beta = 0.999\) 的指数移动平均(EMA)以保证收敛鲁棒性。

实验关键数据

主实验

在真实图像测试集评估特征质量时,MMDiff 与各类代表性扩散特征提取器及判别编码器对比,展现出优异的感知精度。

方法 编码器冻结 PASCAL VOC 2012 (mIoU ↑) DUTS 显著性 (\(S_m\) ↑ / \(F_m\) ↑ / MAE ↓) NYU Depth v2 (AbsRel ↓ / RMSE ↓ / \(\delta_1\) ↑)
DatasetDM (U-Net) 41.19% 0.845 / 0.909 / 0.077 0.1536 / 0.612 / 0.772
Diffusion Hyperfeatures (U-Net) 67.57% 0.892 / 0.910 / 0.021 0.1348 / 0.503 / 0.844
VPD (U-Net 微调) 82.36% 0.912 / 0.916 / 0.024 0.1244 / 0.402 / 0.867
MMDiff (纯DiT特征) 78.90% 0.918 / 0.929 / 0.020 0.1175 / 0.370 / 0.951
DINO-v3 (判别模型) 83.09% 0.920 / 0.922 / 0.021 0.1288 / 0.389 / 0.945
VPD + DINO-v3 83.77% 0.918 / 0.914 / 0.020 0.1223 / 0.387 / 0.946
MMDiff + DINO-v3 (全模型) 84.95% 0.934 / 0.947 / 0.018 0.1164 / 0.365 / 0.952

消融实验

消融实验系统验证了各个组件对最终密集感知表现的增益,尤其是多时间步聚合的必要性。

配置 VOC'12 mIoU (%) ↑ DUTS \(S_m\) DUTS MAE ↓ NYU AbsRel ↓ NYU \(\delta_1\) (%) ↑ 说明
FLUX 双流块 (单步 \(t=0\)) 40.52 0.884 0.033 0.622 74.0 仅利用前期双流文本-图像交互特征
FLUX 单流块 (单步 \(t=0\)) 50.20 0.890 0.032 0.478 78.0 切换为统一的单流注意力层表征
+ 传统交叉注意力图 58.31 0.869 0.043 0.455 80.0 文本 Prompt 平均交叉注意力
+ 概念驱动单向注意力 63.77 0.903 0.024 0.431 82.0 引入解耦概念 Token 引导 (+13.57%)
多时间步聚合 + 概念注意力 78.90 0.918 0.020 0.118 95.1 加入 4 时间步融合 (+28.70% mIoU)
完整模型 (无自适应权重聚合) 78.21 0.918 0.017 0.130 81.4 移除 Transformer 像素级权重与 CBAM
完整模型 (4T + 概念 + DINO-v3) 84.95 0.934 0.018 0.116 95.2 全配置结合判别互补表征

关键发现

  • 多时间步聚合是 DiT 感知跃迁的核心引擎:由单时间步 \(t=0\) 的 50.20% 到多时间步融合的 78.90%,暴增 28.7% mIoU,证明扁平 DiT 的高维概念并非在末端汇集,而是时间展开的。
  • 概念注意力优于全词平均注意力:解耦的 concept tokens 相比常规全 prompt 交叉注意力,mIoU 从 58.31% 跃升至 63.77%,消除了无用停用词带来的空间注意力弥散。
  • 纯合成数据训练超越传统 U-Net SOTA:在合成数据预训练实验中,仅依靠 MMDiff 生成的图像和自动提取的标注,训练下游模型即在 PASCAL VOC 上获得 78.9% mIoU,大幅超越同类方案(DatasetDM 65.9%,Dataset Diffusion 68.2%)。
  • 可直接扩展至 3D 视频生成领域:在视频扩散模型 Wan-2.1 与 RealEstate10K 的 3D 高斯泼溅(3DGS)重建实验中,将传统的去噪潜在变量替换为 MMDiff 聚合特征,PSNR 从 21.37 提升至 21.61,SSIM 从 0.705 提升至 0.747。

亮点与洞察

  • 算力再利用的绿色范式:将原本在图像生成过程中被当成副产物丢弃的高维去噪轨迹激活值回收利用,避免了用外部大模型做二次逆向推理的重复浪费。
  • 单向注意力防污染设计:让概念 token 只单向读取生成画面的 token,而阻断反向流动,既获得了高度可解释的特定概念分割指引,又完全保持了原有生成图像像素的保真度与独立性。
  • 打破生成与判别表征壁垒:首次证实完全冻结的纯 DiT 去噪特征能够在密集几何与语义任务上抗衡甚至超越微调的判别模型,且两者在通道层面的融合产生明显正向互补。

局限与展望

  • 推理期内存开销:由于需要缓存 4 个时间步在 4 个关键层的特征图,相比单步推理增加了约 4 倍的瞬时显存占用,在极端大分辨率生成时面临显存瓶颈。
  • 开集概念人工指定:概念注意力引导中的概念 token 目前需要在推理时由任务专家手工设定(例如“object”、“contour”、“near”),尚未实现端到端的自适应概念挖掘。
  • 未来方向:可进一步将时间步聚合机制深度内化为常驻低秩适配器(LoRA),并探索在实时交互式视频编辑与世界模型物理场预测中的实时在线感知。

相关工作与启发

  • vs DatasetDM / DiffuMask: 后两者基于 SD U-Net,特征提取受限于 U-Net 粗糙的空间下采样分辨率,且只取最后时间步;MMDiff 基于扁平 DiT 架构,开创性地进行了全去噪轨迹的时空自适应融合。
  • vs VPD: VPD 需要在下游数据集上对昂贵的扩散网络主干进行有监督微调;MMDiff 全程冻结 12B FLUX 主干,仅需训练 36M 极轻量解码头,泛化与保留原有生成先验能力更强。
  • vs DINO-v3: DINO-v3 依赖自监督对比学习擅长判别细节;MMDiff 证明生成式 DiT 的多步轨迹表征与之正交,两者拼接后在多项密集任务中达到超越单一模型的最高水平。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 深入剖析 DiT 时序去噪分布特性,提出自适应多时间步特征融合与单向概念注意力解耦设计。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖语义分割、显著性检测、深度估计三项核心任务,设立多模态一致性、合成数据驱动与真实图像评估三大协议。
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑缜密,特征演化分析透彻,理论机制与实验验证高度统一。
  • 价值: ⭐⭐⭐⭐⭐ 为无标注合成数据自生成与多模态联合生成落地提供了极具实用价值的范式。