The Design Space of Tri-Modal Masked Diffusion Models¶

会议: NeurIPS 2026
arXiv: 2602.21472
OpenReview: https://openreview.net/forum?id=NBUCuy2uaq
领域: 多模态VLM
关键词: 掩码扩散模型, 三模态预训练, SDE重参数化, 多模态Scaling Law, 反向掩码

一句话总结¶

本文首次提出并从零预训练了覆盖文本、图像与音频的统一离散掩码扩散模型（Tri-Modal MDM），基于SDE重参数化消除了最优批大小调节，并给出了计算最优缩放定律与各模态专属采样超参数。

研究背景与动机¶

序列建模的核心经验表明，在全上下文可见的任务场景中，双向信息交互往往显著优于单向因果建模。然而，当前主流的多模态与大语言模型几乎全由因果自回归（Autoregressive, AR）架构主导。因果约束虽然带来了逐词似然分解与高效 KV 缓存推理的便利，但在处理证据分散、需要跨多模态双向填补与条件合成的任务时，单向严格从左至右的因式分解天然缺乏全局语境感知与任意位置迭代精化的灵活性。离散掩码扩散模型（Masked Diffusion Models, MDM）通过引入统一的破损（掩码）与逆向迭代去噪过程，重新确立了全局双向注意力的核心优势。

尽管双向离散扩散在纯文本生成上已逐步逼近自回归模型表现，但在多模态领域，绝大多数现有工作仍停留在对已有单模态基础模型做下游微调或蒸馏复用的阶段（如基于 LLaDA 微调视文双模态，或基于 Qwen 等 AR 骨干蒸馏）。这种方案不仅无法从根本上探索多模态联合表征在预训练计算预算下的演化规律，更将模型受限于双模态组合。一旦将数据域扩展至文本、图像与音频三模态联合空间，离散词表膨胀、跨模态梯度方差异构、不同模态对噪声掩码敏感度差异巨大，导致超参数调节与缩放行为变得极其复杂脆弱。

现有研究缺乏在大规模从零预训练体制下对多模态 MDM 设计空间的系统性探索。本文的核心 idea 是：构建统一离散词表与单流 Transformer 骨干的三模态掩码扩散模型（Tri-Modal MDM），通过随机微分方程（SDE）重参数化彻底解耦物理批大小与梯度优化方差（消除 \(B_{\text{opt}}\)），进而推导出以非加和幂律为核心的三模态计算最优 Scaling Law，并揭示各模态在推理采样阶段的异构最优解。

方法详解¶

整体框架¶

本文构建了一个端到端的原生三模态离散掩码扩散模型，涵盖文本（Text）、图像（Image）与音频（Audio）。模型采用统一的离散 Token 序列表示与单一 Transformer 骨干网络，无需为特定模态配置独立编码头或外部适配器，直接支持文本生图像（T2I）、图像描述（Captioning）、文本生语音（TTS）及语音识别（ASR）等多向条件推理任务。

整体流程包括：统一词表与任务引导构建、连续时间掩码破损过程、双向 Transformer 联合去噪以及反向迭代离散展开。以下是核心阶段与处理管线：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["多模态输入流<br/>纯文本/图文对/音文对"] --> B["统一离散词元化与任务引导封装"]
    B --> C["连续时间独立掩码破损过程"]
    C --> D["双向Transformer主干去噪<br/>Cut-Cross-Entropy与Z-Loss约束"]
    D --> E["SDE重参数化与超参数缩放<br/>解耦物理批大小与梯度方差"]
    E --> F["模态感知反向迭代重构<br/>自适应步数/CFG/温度采样"]
    F --> G["多模态生成产物<br/>文本/256x256图像/语音"]

关键设计¶

1. 统一离散词元化与任务引导封装：消除模态专用头

为打破模态割裂，本文将异构连续信号彻底离散化，构建非重叠并集统一词表 \(\mathcal{V}=\mathcal{V}_{\text{text}}\sqcup\mathcal{V}_{\text{audio}}\sqcup\mathcal{V}_{\text{image}}\)，总词表大小达 117,698。其中，文本采用 Tiktoken（100,281 词元）；图像采用 SBER-MoVQGAN 压缩为 \(256\times 256\) 分辨率下的 1024 个离散 Patch 词元（词表 16,387）；音频采用 Higgs Audio Tokenizer v2 并选用 4 个码本进行解码重构，压缩因子达 960（词表 1,027）。为了界定模态边界并支持多任务调度，引入模态专属边界元 \(\mathcal{V}_{\text{spec}}=\{\text{BOS}_m, \text{EOS}_m, \text{MASK}_m\}\) 与任务引导元 \(\mathcal{V}_{\text{task}}=\{\text{TASK}_{\text{text}}, \text{TASK}_{\text{image-text}}, \text{TASK}_{\text{audio-text}}\}\)。所有输入被拼接为统一序列并在末尾使用 \(\text{PAD}_{\text{text}}\) 补齐至固定上下文长度 \(L^\star=3256\)（纯文本则采用无缝拼接 Packing）。任务引导元在扩散训练与推理中永远不被掩码，从而保证任务条件始终清晰可见。

2. 连续时间伯努利掩码与无偏重要性加权去噪

扩散前向过程定义在连续时间区间 \(t\in[0, 1]\) 上。对于属于模态 \(m(i)\) 的位置 \(i\)，其以独立伯努利概率 \(\beta_t\) 被替换为该模态专有的 \(\text{MASK}_{m(i)}\) 标记。累积未掩码概率为 \(\bar{\alpha}_t\)。由于后验转移概率在非掩码状态下直接为恒等，仅在掩码状态发生信息破坏，模型训练目标等价于最大化变分下界（ELBO）。为校正随机下采样带来的梯度偏差，损失函数引入关键的反比时间重要性权重 \(w(t)=1/t\)：

\[\mathcal{L}(\theta)=\mathbb{E}_{s\sim\mathcal{D}, t\sim\mathrm{U}(\epsilon,1)}\left[\frac{1}{t\cdot |\mathcal{I}_{t}|}\sum_{i\in\mathcal{I}_{t}}\ell_{i}(\theta,s)\right] + \lambda\,\mathcal{L}_{z}\]

其中 \(\mathcal{I}_t\) 是当前被掩码且非 Padding 的 token 索引集合，\(\ell_i(\theta, s)\) 采用高效的 Cut-Cross-Entropy 计算以规避大词表下的显存爆炸，\(\mathcal{L}_z\) 为 z-loss 正则项以稳定对数幅度。因子 \(1/t\) 正好抵消了伯努利采样期望值 \(\mathbb{E}[\mathbf{1}_{i\in\mathcal{I}_t}]=t\)，使得所有时间步在期望上对每个 token 的重建误差给予无偏等权贡献。

3. SDE 重参数化消除最优批大小：解耦物理算力与优化动态

针对传统自回归与扩散模型训练中普遍存在的“最优批大小”（\(B_{\text{opt}}\)）调优负担，本文将 AdamW 优化过程形式化为随机微分方程（SDE）的离散化表征。在非 SDE 体系下，批大小同时影响优化轨迹长度（SDE Horizon）与梯度噪声尺度（SDE Drift），形成必须精心平衡的 U 形性能曲线。本文通过 SDE 缩放因子 \(\kappa\) 将物理批大小 \(B\) 与学习率、动量参数及 \(\epsilon\) 动态绑定：

\[\kappa=\left(\frac{D^{\text{base}}}{D}\right)^{\gamma}\left(\frac{B}{B^{\text{base}}}\right), \quad \text{lr}=\text{lr}^{\text{base}}\sqrt{\kappa}, \quad \beta_1=(\beta_1^{\text{base}})^{\kappa}, \quad \beta_2=(\beta_2^{\text{base}})^{\kappa}, \quad \epsilon=\frac{\epsilon^{\text{base}}}{\sqrt{\kappa}}\]

在临界批大小 \(B_{\text{crit}}\)（或临界步数 \(S_{\text{crit}}\)）以内，SDE 重参数化保证了最终验证损失与物理批大小完全无关。研究发现临界迭代步数 \(S_{\text{crit}}\) 与模型参数量 \(N\) 无关，仅随数据量 \(D\) 呈亚线性增长。进一步在 Drift-Horizon 权衡实验中，通过调优控制参数 \(\gamma\)，测定出介于保持噪声不变（\(\gamma=0\)）与保持时间步不变（\(\gamma=1\)）之间的最优权衡系数 \(\gamma^*\approx 0.44\)。

4. 渐进数据高效的多模态计算最优 Scaling Law

在 CompleteP 结合 SDE 重参数化架构下，本文训练了 262 个不同规格的模型（参数量排除嵌入层后计算），覆盖 \(5\times 10^{18}\) 至 \(1\times 10^{22}\) FLOPs 的 24 条等算力曲线。研究发现 Hoffmann 等人（Chinchilla）提出的加和形式无法准确拟合三模态扩散曲面，而 Kaplan 形式（引入不可压缩固有熵底项 \(E\)）表现出极高拟合优度（\(R^2=99.3\%\), \(\text{MRE}=0.5\%\)）：

\[L(N, D) = E + \left(A N^{-a/b} + B D^{-1}\right)^{b}, \quad a\approx 0.14, \; b\approx 0.17\]

由此推导出三模态 MDM 的计算最优前沿为：

\[D^{\star}(N) \approx 7754 \cdot N^{0.84}\]

这一指数 \(\alpha=a/b=0.84<1\) 揭示了关键理论特性：随着模型规模增大，单位参数所需的计算最优 token 量 \(D^\star/N \propto N^{-0.16}\) 反而逐步递减，展现出离散扩散模型在渐进尺度上优于自回归模型的数据利用效率。然而在 3B 等实用规模下，由于首项常数较大，3B 扩散模型在计算最优路径上需要约 480B token，远超 Chinchilla 对 AR 语言模型推荐的 60B token。

5. 补正梯度方差的反向掩码与模态异构采样配置

由于离散掩码扩散训练面临极高的采样方差，本文探索了反向掩码（Anti-masking）技术：在相同批次内成对输入反相掩码样本（即同一序列分别应用掩码模式及其布尔非操作）。在相同总计算量（Token Horizon）严格对齐下，反向掩码显著压低了批梯度方差。在推理采样阶段，本文确立了高度分化的模态最优超参：图像生成偏好较高的分类器自由引导（CFG 6.0~9.0）、中等温度（\(T\approx 1.0\)）和多采样步数（1024 步），且多项式（Polynomial）掩码调度显著优于余弦与线性调度；音频生成（TTS）则严格受制于保真度与文本对齐的张力，偏好较低的引导尺度（CFG 1.0~3.0，过高 CFG 会严重恶化 FAD 音频保真度）、轻微升温（\(T=1.0\sim 1.2\)）与 512~1000 步采样。

实验关键数据¶

主实验¶

基于上述统一架构，团队从零预训练了统一三模态 3B 模型（24 层、隐藏维度 3072、24 头、序列长度 3256，总计处理 6.4T tokens，其中含 3.4T 文本 tokens、1B 图像样本与 1B 音频样本）。以下展示其在跨模态基准上的表现及与不同训练策略的对齐验证：

模态任务	评测数据集 / 基准	核心评估指标	3B 统一预训练模型得分	基线 / 对比参考（Base）	增益 / 评价
图像生成 (T2I)	CC12M (256x256)	FID-Inception \(\downarrow\)	10.06	26.77 (1.3B Base)	显著领先更小规模预训练模型
图像生成 (T2I)	CC12M (256x256)	FID-DINOv2 \(\downarrow\)	107.61	395.73 (1.3B Base)	语义结构与视觉质量大幅跃升
视觉组合生成	GenEval Suite	Overall Accuracy (%) \(\uparrow\)	48.89	-	单物体93.12%, 双物体63.38%, 颜色64.89%
语言理解推理	LM-Harness (MMLU)	Accuracy (%) \(\uparrow\)	41.57	30.20 (7B Base text)	强双向去噪能力展现良好通识水准
语言理解推理	LM-Harness (ARC-Easy)	Acc-norm (%) \(\uparrow\)	72.52	63.72 (7B Base text)	常识推理基础坚实
语音合成 (TTS)	LibriSpeech-PC	FAD \(\downarrow\)	0.368	0.79 (1.3B Base)	生成音频分布贴近自然语音
语音合成 (TTS)	LibriSpeech-PC	WER \(\downarrow\)	0.164	-	词错误率表现稳定，发音可懂度高
语音合成 (TTS)	LibriSpeech-PC	Production Quality (PQ) \(\uparrow\)	6.01	-	Audiobox 美学评分达高实用度

消融实验¶

表 1：多模态反向掩码（Anti-masking）在严格算力对齐下的消融对比（1.3B 模型）

模型配置	图像 CC12M FID (Inception) \(\downarrow\)	图像 CC12M FID (DINOv2) \(\downarrow\)	音频 Train FAD \(\downarrow\)	音频 LibriSpeech FAD \(\downarrow\)	结论分析
标准掩码训练 (Base)	26.77	395.73	0.24	0.79	双周期常规掩码基线
反向掩码训练 (Anti-mask)	21.04	361.00	0.22	0.55	方差骤降，尤其在音频生成上取得显著泛化提升（FAD降幅超30%）

表 2：7B 纯文本扩散模型在 LM-Harness 上的反向掩码消融对比

模型配置	MMLU (%) \(\uparrow\)	ARC-Challenge (%) \(\uparrow\)	HellaSwag (%) \(\uparrow\)	BBH (%) \(\uparrow\)	Winogrande (%) \(\uparrow\)
标准掩码训练 (Base)	30.20 \(\pm\) 0.39	31.66 \(\pm\) 1.36	52.20 \(\pm\) 0.50	21.52 \(\pm\) 0.46	52.01 \(\pm\) 1.40
反向掩码训练 (Anti-mask)	32.86 \(\pm\) 0.39	34.90 \(\pm\) 1.39	55.15 \(\pm\) 0.50	27.05 \(\pm\) 0.50	54.06 \(\pm\) 1.40

关键发现¶

模态间缺乏协同但不存在恶性互斥：在对文本、图像-文本、音频-文本混合比例进行的 15 组系统消融中，在 320M/13B tokens 规模下，某一模态的验证损失只取决于自身在混合中的权重占比，不受另外两模态相对比例的交叉影响。模态间主要竞争网络容量与计算配额，默认等比例（各 33%）分配在实操中稳妥有效。
SDE 重参数化彻底消除批大小超参搜索：在临界步数 \(S_{\text{crit}}\) 之前，不同物理批大小在相同的总 Token 预算下达成几乎一致的验证损失曲线，消除了传统训练中对 \(B_{\text{opt}}\) 的穷举寻找。
推理采样的模态分裂：图像生成对高 CFG 强引导展现出高度耐受并受益于此（最佳 CFG 7~9）；但音频 TTS 生成存在极强的 trade-off——提升 CFG 能够降低文本转写 WER，但当 CFG 超出 3.0 后会急剧劣化音频信号的自然度与 FAD。

亮点与洞察¶

SDE 视角下的批大小与超参联合转移：巧妙借助连续随机微分方程映射，将 AdamW 的超参数与物理批大小建立动态绑定。这一机制成功消除了昂贵的批大小网格搜索，使工程部署可在保证 FLOP 效率的前提下自由调整硬件并行度。
颠覆 Chinchilla 常规的多模态数据缩放规律：揭示了离散扩散模型的计算最优曲线呈现次线性增长趋势（\(D^\star \propto N^{0.84}\)），这意味着超大规模下的扩散模型渐进数据利用率高于因果语言模型（\(D \propto N^{1.0}\)），为大规模多模态扩散基座提供了关键量化基准。
分模块超参数调优（Per-module Tuning）的显著增益：附录证明针对嵌入层、注意力投影和 MLP 门控赋予非均匀的学习率与阻尼系数，可在相同损失目标下实现约 1.81 倍的 Token 吞吐节省。

局限与展望¶

采样效率与延迟瓶颈：当前离散掩码模型在推理阶段通常需要数百至 1024 步反向迭代去噪，虽然支持高质量无序生成与全局协调，但单步缺乏因果模型的流式 KV 缓存机制，实际推理延迟显著高于自回归基线。
中小规模下未见模态协同效应：在 320M 参数规模下的混合比例消融中，未观察到跨模态促进的“涌现协同”，模态之间仍呈现直接争夺参数容量的现象，未来需要验证在更大参数量（如 >10B）下多模态预训练是否能自发诱导跨模态泛化。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次提出真正统一的从零预训练三模态离散掩码扩散架构，并在理论上引入 SDE 消除最优批大小。
实验充分度: ⭐⭐⭐⭐⭐ 包含 262 次 Scaling Law 训练、跨模态系统性采样消融、批大小不变性验证与反向掩码评估。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严谨，从连续时间 CTMC、ELBO 加权到 SDE 离散化分析层层递进。
价值: ⭐⭐⭐⭐⭐ 为探索非自回归多模态基础模型设计空间提供了奠基性的基准方案与开源缩放范式。