三模态掩码扩散模型的设计空间 The Design Space of Tri-Modal Masked Diffusion Models
首个统一文本、图像、音频的三模态掩码扩散模型，通过SDE参数化消除最优批量大小搜索
前置知识
掩码扩散模型（Masked Diffusion Models）
掩码扩散模型是一种离散扩散模型，通过逐步将输入序列中的token替换为特殊的[MASK]标记来进行前向扩散过程。在反向过程中，模型学习从完全被掩码的序列重建原始数据。与连续扩散模型不同，MDM直接在离散token空间操作，避免了连续离散化的问题。训练目标是最大化变分下界（ELBO），损失函数为$L(\theta) = E_{s \sim D, t \sim U(\epsilon,1)} \left[ \frac{w(t)}{|I_t|} \sum_{i \in I_t} \ell_i(\theta, s) \right]$，其中$I_t$是被掩码位置的集合，$w(t)=1/t$是权重函数。
本文的核心就是将MDM从单模态扩展到三模态，理解MDM的基本原理是理解论文贡献的基础
扩散模型（Diffusion Models）
扩散模型是一类生成模型，通过定义前向扩散过程逐步向数据添加噪声，然后学习反向去噪过程来生成新样本。前向过程将数据分布$q_0:=p_{data}$转换为简单的先验分布$q_T$（通常是高斯噪声或完全掩码状态）。反向过程则从$q_T$开始，通过迭代去噪重建数据。扩散模型在图像生成领域取得了巨大成功，如DDPM、Stable Diffusion等。
理解扩散模型的基本框架有助于把握MDM与传统连续扩散模型的区别和联系
变分自编码器/向量量化（VQ-VAE/VQ-GAN）
VQ-VAE和VQ-GAN是用于将连续数据（如图像、音频）离散化的神经网络架构。它们通过学习码本（codebook）将高维输入映射到离散token序列，实现数据的压缩表示。VQ-GAN在VQ-VAE基础上引入了对抗训练和感知损失，提高了重建质量。在本文中，图像使用SBER-MoVQGAN进行tokenization，音频使用Higgs Audio v2，文本使用Tiktoken。
多模态MDM的基础是将不同模态的数据统一转换为离散token，理解tokenization过程对理解模型架构至关重要
缩放定律（Scaling Laws）
缩放定律描述了模型性能（如损失函数）如何随模型参数量$N$和训练数据量$D$变化的幂律关系。经典的缩放定律公式为$L(N,D) = E + A/N^a + B/D^b$，其中$E$是不可压缩误差，$a$和$b$是缩放指数。这些定律为计算最优的模型规模和数据量提供指导，帮助在给定计算预算下做出最优的资源分配决策。
本文的核心贡献之一是建立了三模态MDM的缩放定律，这对于指导大规模多模态MDM的训练至关重要
SDE参数化（SDE Parameterization）
SDE参数化是一种将随机优化器（如AdamW）的超参数重新参数化的方法，将优化过程视为随机微分方程（SDE）的离散化。通过这种重新参数化，可以使训练损失对批量大小不敏感，只要批量大小不超过临界值$B_{crit}$。关键思想是AdamW的超参数（学习率、动量等）与批量大小存在冗余关系，可以通过调整超参数来补偿批量大小的变化。
本文通过SDE参数化消除了对最优批量大小$B_{opt}$的搜索需求，这是提高训练效率的关键技术贡献
研究动机
当前多模态生成模型主要存在三个核心问题。首先，大多数多模态MDM工作集中在双模态（文本-图像）设置，如MMaDA，缺乏对三模态（文本、图像、音频）的统一建模。现有方法要么从预训练的单模态模型微调（如LLaDA），要么依赖自回归模型的蒸馏（如Dream），没有从零开始预训练原生的三模态MDM。其次，训练扩散模型时面临最优批量大小$B_{opt}$的搜索问题。Bergsma等人（2025）的研究表明，存在一个最优批量大小$B_{opt}$使得在固定token预算$D$下损失最小化，但搜索这个$B_{opt}$需要大量计算资源。第三，缺乏对多模态MDM缩放行为的系统性研究，不清楚模型规模$N$和数据量$D$如何影响三模态MDM的性能，也没有为这类模型提供计算最优的训练指导。
本文的目标是本文的具体目标是：（1）构建首个从零预训练的三模态MDM，统一处理文本、图像和音频三种模态；（2）通过SDE参数化消除对最优批量大小$B_{opt}$的搜索需求，使训练损失在批量大小$B \leq B_{crit}$时保持不变；（3）建立三模态MDM的缩放定律，提供计算最优的训练指导，确定最优token数$D^*(N)$与模型规模$N$的关系；（4）系统研究模态混合比例、噪声调度、推理超参数等设计选择对模型性能的影响。最终目标是为社区提供一个经过验证的三模态MDM预训练方案，以及实用的设计空间指南。
与已有工作不同的是，本文的独特切入角度在于系统性地研究三模态MDM的'设计空间'（Design Space），而不仅仅是提出一个新模型。与现有工作相比，本文抓住了几个被忽视的关键点：（1）从零预训练而非微调：大多数现有工作（如MMaDA、Dream）都是在预训练模型基础上微调，但本文认为预训练阶段才是塑造潜在空间和决定最终性能的关键阶段；（2）批量大小与优化的解耦：通过SDE参数化将物理批量大小（由计算资源决定）与逻辑批量大小（用于平衡梯度方差）解耦，这是一个被忽视但极其重要的工程问题；（3）模态依赖的推理设计：发现不同模态在推理时需要不同的最优超参数（CFG、温度、采样步数），这挑战了'一刀切'的推理策略；（4）缩放定律的模态特异性：三模态MDM的缩放行为与自回归模型显著不同，需要专门的缩放定律指导。
核心方法
本文的方法可以类比为'建造一座多语言翻译塔'。想象有一座塔，底层是三种语言（文本、图像、音频）的混合文本，顶层是完全被掩码的'未知语言'。训练过程就是学习如何从顶层逐步'翻译'回底层，而推理过程则是从完全未知的状态开始，逐步'破译'出有意义的内容。技术路线分为四个关键部分：（1）构建统一的多语言词典：将三种模态的token合并为统一词汇表$V = V_{text} \sqcup V_{audio} \sqcup V_{image}$，并添加模态特定的边界标记（BOS/EOS/MASK）；（2）设计掩码扩散过程：对输入序列进行独立的伯努利掩码，每个位置以概率$\beta_t$被替换为对应的MASK token；（3）训练双向Transformer：使用标准的双向Transformer架构，通过预测被掩码位置的原始token来学习；（4）迭代推理生成：从完全掩码的序列开始，按照预定义的线性调度逐步揭示token。
本文的核心创新点是通过SDE参数化消除对最优批量大小$B_{opt}$的搜索需求，并首次建立了三模态MDM的缩放定律。与已有方法的本质区别体现在三个方面：首先，与MMaDA、Dream等依赖预训练模型的方法不同，本文从零开始预训练三模态MDM，确保在预训练阶段（主要计算消耗所在）塑造最优的潜在空间。其次，通过SDE参数化将物理批量大小（由GPU饱和度、FLOP效率、墙钟时间决定）与逻辑批量大小（用于平衡随机优化中的梯度方差）解耦。这使得训练损失在批量大小$B \leq B_{crit}$时保持不变，消除了$B_{opt}$搜索的需要。第三，建立了三模态MDM的缩放定律$L = E + A/N^{a/b} + B/D^b$，发现$a \approx 0.14$，$b \approx 0.17$，最优token数$D^*(N) = 7754 \cdot N^{0.84}$，表明三模态MDM在数据效率上优于自回归模型。
方法步骤详情
方法分为训练和推理两个主要阶段，具体步骤如下： **训练阶段**： 1. **序列构建**：对于文本-图像对，构建序列$s = [TASK_{image\text{-}text}, BOS_{image}, x^1_{image}, ..., x^{L_{image}}_{image}, EOS_{image}, BOS_{text}, x^1_{text}, ..., x^{L_{text}}_{text}, EOS_{text}]$，其中$x^i_m$是模态$m$的第$i$个token。文本序列被填充到最大序列长度$L^*$。 2. **前向掩码**：对每个位置$i$独立应用伯努利掩码，以概率$\beta_t$将其替换为对应的MASK token：$s^t_i = MASK_{m(i)}$（概率$\beta_t$）或$s^t_i = s_i$（概率$1-\beta_t$）。任务token永不被掩码。 3. **模型预测**：双向Transformer接收被掩码的序列$s_t$，输出每个位置的logits $h = f_\theta(s_t)$。 4. **损失计算**：仅在被掩码位置$I_t = \{i | s^t_i = MASK_{m(i)}\}$上计算损失：$\ell_i(\theta, s) = -\log p_\theta(s_i | s_t) = -\log \frac{\exp h_{i,s_i}}{\sum_{v \in V} \exp h_{i,v}}$，使用cut-cross-entropy避免内存爆炸。 5. **SDE参数化**：通过重新参数化AdamW超参数（学习率、动量$\beta_1, \beta_2$、$\epsilon$），使训练损失对批量大小不敏感。 **推理阶段**： 1. **初始化**：从完全掩码的序列$s_K$开始，包含任务token和边界token。 2. **迭代去噪**：对于$k = 1, ..., K$：模型预测logits $h_k = f_\theta(s_{k-1})$，从模态约束的分布中采样候选token，根据线性调度揭示子集被掩码位置。 3. **输出**：当所有位置都被揭示时，生成完成。
技术新颖性
本文的技术新颖性主要体现在五个方面：（1）首个三模态MDM：与MMaDA（双模态）、LLaDA（单模态）不同，本文首次将文本、图像、音频三种模态统一在一个MDM框架中，使用单一Transformer骨干和统一词汇表，无需模态特定的头或适配器。（2）SDE参数化在多模态MDM中的首次应用：将Mlodozeniec等人（2025）的SDE参数化引入多模态MDM领域，消除了$B_{opt}$搜索需求，这是该领域的首创。（3）三模态MDM的缩放定律：首次建立了三模态MDM的缩放定律，发现$L(N,D) = E + A/N^{a/b} + B/D^b$，其中$a \approx 0.14$，$b \approx 0.17$，最优token数$D^*(N) = 7754 \cdot N^{0.84}$。这与自回归模型的Chinchilla缩放定律（$D \propto 20N$）有显著差异。（4）模态依赖的推理设计：系统发现不同模态需要不同的最优推理超参数。例如，文本生成需要较低的CFG（~3.0），图像生成需要较高的CFG（~6.0），音频生成在CFG、温度、top-p上存在权衡。（5）反掩码（Anti-masking）训练技巧：通过为每个样本生成两个互补的掩码版本来减少梯度方差，在不增加计算成本的情况下提升性能。在7B文本模型上，反掩码在大多数任务上带来1-5%的一致性提升。
实验结果
本文的核心发现涵盖六个方面，每个都有具体的实验数据支撑： **1. 统一3B模型的强大性能**：预训练的3B三模态MDM在文本、图像、音频三个模态上都展示了竞争力。在图像生成方面，FID-Inception为10.41（训练集评估种子），FID-DINOv2为112.12；在CC12M上FID-Inception为10.06，FID-DINOv2为107.61。GenEval评估显示单物体准确率93.12%，双物体63.38%，计数33.44%，颜色64.89%，位置11.50%，颜色属性27.00%，总体得分48.89。在文本基准上，LM Harness评估显示OpenBookQA 37.40，TruthfulQA MC2 40.76，BBH 24.97，MMLU 41.57，WinoGrande 64.69，ARC-Easy 72.52，HellaSwag 65.88，LogiQA2 30.66，PIQA 71.55，ARC-Challenge 43.09。在音频生成方面，训练集FAD为0.218，WER为0.124，PQ 6.89，CU 6.20，CE 5.45，PC 1.89；LibriSpeech-PC上FAD为0.368，WER为0.164，PQ 6.01，CU 5.34，CE 5.07，PC 3.01。 **2. SDE参数化消除$B_{opt}$**：在320M模型、13B token的实验中，SDE参数化在批量大小$B \leq B_{crit}$时保证恒定损失。例如，在$10^5$到$10^7$ token的批量大小范围内，SDE参数化的损失保持稳定，而非SDE参数化则呈现典型的U型曲线。临界迭代次数$S_{crit}$与模型规模$N$无关，这与Bergsma等人（2025）的发现一致，但本文首次在SDE参数化框架下验证。 **3. 三模态MDM的缩放定律**：通过训练262个不同规模的模型，建立了缩放定律$L = E + A/N^{a/b} + B/D^b$，其中$a \approx 0.14$，$b \approx 0.17$，$R^2$得分99.3%，MRE为0.5%。最优token数$D^*(N) = 7754 \cdot N^{0.84}$，表明一个3B模型需要至少480B token，远高于Chinchilla对自回归模型的60B token建议。最优计算分配为$N^*(C) \propto C^{0.55}$，$D^*(C) \propto C^{0.45}$。 **4. 模态混合比例**：在320M模型、13B token的实验中，发现模态间在当前规模下不存在协同效应。平均损失仅取决于该模态的混合比例，与其他模态比例无关。因此，1/3:1/3:1/3的默认混合比例是合理的。 **5. 推理超参数敏感性**：文本生成最优CFG为~3.0，图像生成最优CFG为~6.0，音频生成在CFG、温度、top-p上存在权衡。生成步数增加带来收益递减，文本和图像生成在~1024步后改善有限。 **6. 反掩码训练技巧**：在7B文本模型上，反掩码在大多数任务上带来一致性提升。在多模态1.3B模型上，反掩码将FID-Inception从18.69降至17.81（训练集），从26.77降至21.04（CC12M）；FAD从0.24降至0.22（训练集），从0.79降至0.55（LibriSpeech）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 图像生成（FID-Inception） | FID-Inception ↓ | 10.41（训练集）/ 10.06（CC12M） | MMaDA 8B（双模态） | 首次实现三模态，性能接近或超过双模态基线 |
| 文本基准（MMLU） | MMLU准确率 | 41.57% | LLaDA 8B（单模态） | 在更小模型规模（3B vs 8B）下达到可比性能 |
| 文本基准（PIQA） | PIQA准确率 | 71.55% | LLaDA 8B（单模态） | 在3B规模下表现竞争力 |
| 音频生成（FAD） | FAD ↓ | 0.218（训练集）/ 0.368（LibriSpeech-PC） | 无直接基线 | 首个三模态MDM的音频生成能力 |
| 音频生成（WER） | WER ↓ | 0.124（训练集）/ 0.164（LibriSpeech-PC） | 无直接基线 | 低词错误率表明高质量语音生成 |
| 缩放定律拟合 | R²得分 | 99.3% | 无 | 高拟合度验证了缩放定律的有效性 |
局限与改进
本文存在以下主要局限性： **1. 缺乏与其他多模态模型的直接对比**：虽然论文展示了3B模型的性能，但没有与同期的多模态模型（如Flamingo、BLIP-2、LLaVA等）进行直接基准比较。这使得难以准确评估本文方法相对于现有技术的实际优势。论文提到Quokka的计算最优曲线在20B规模（约2T token）处与本文交叉，但没有提供直接的性能对比。 **2. 推理效率问题未充分解决**：论文承认'naive implementations still exhibit substantial latency compared to autoregressive baselines and further work is required to improve sampling efficiency in MDM'。尽管展示了推理超参数的消融研究，但没有提供具体的推理延迟或吞吐量数据，也没有与自回归模型的推理效率进行对比。 **3. 模态间协同效应有限**：在320M模型规模的模态混合比例实验中，没有观察到模态间的协同效应——平均损失仅取决于该模态的混合比例。这表明在当前规模下，模态间可能存在竞争而非协同，限制了多模态学习的潜在优势。 **4. 数据集规模和多样性**：训练数据包括2M小时音频（来自网络爬取并经Whisper转录），以及CC3M、CC12M、COYO等图像数据集。虽然论文声称进行了PII过滤和质量过滤，但没有提供具体的数据质量评估或多样性分析。此外，数据混合比例（1/3:1/3:1/3）是凭经验选择的，缺乏理论指导。 **5. 模型规模限制**：最大的预训练模型是3B参数，这对于展示三模态MDM的可行性是足够的，但可能不足以充分发挥多模态学习的潜力。论文的缩放定律表明更大的模型可能带来显著改进，但没有提供更大规模模型的实验结果。
独立分析的弱点
基于对论文的深入分析，我识别出以下几个主要弱点及相应的改进方向： **1. 缺乏模态间交互机制的深入探索**：当前架构使用单一Transformer骨干处理所有模态，但模态间的交互仅通过注意力机制隐式实现。在多模态混合比例实验中没有观察到协同效应，这可能是因为模型没有显式地学习模态间的对齐和融合。改进方向：可以引入跨模态注意力层、模态特定的适配器模块，或者设计显式的模态对齐损失函数，促进模态间的知识共享。 **2. 推理效率优化不足**：论文承认推理延迟问题，但没有提出具体的解决方案。MDM的迭代推理本质上比自回归解码更慢，因为需要多次前向传播。改进方向：可以研究自适应掩码调度（根据序列的'难度'动态调整揭示顺序）、并行采样技术、知识蒸馏到更快的模型，或者探索混合架构（部分自回归+部分扩散）。 **3. 缩放定律的局限性**：建立的缩放定律$L = E + A/N^{a/b} + B/D^b$中，$E$（不可压缩误差）的解释不够清晰。此外，缩放定律是在特定数据混合和模态比例下建立的，可能不适用于其他数据分布。改进方向：需要研究$E$的组成，分析不同模态对不可压缩误差的贡献；探索数据质量、模态比例对缩放指数$a$和$b$的影响；建立更通用的缩放定律框架。 **4. 评估基准的完整性**：虽然论文使用了多种评估指标（FID、FAD、WER、LM Harness等），但缺乏统一的多模态评估框架。不同模态使用不同的评估协议，使得跨模态比较困难。改进方向：可以开发统一的多模态评估基准，或者至少提供标准化的评估流程，使不同研究之间能够公平比较。 **5. 反掩码的理论理解不足**：反掩码被证明有效，但论文没有提供充分的理论解释。为什么互补的掩码模式能够减少梯度方差？在什么条件下反掩码最有效？改进方向：需要建立反掩码的理论框架，分析其对梯度方差、收敛速度、泛化性能的影响，并探索更一般的'互补训练'范式。
未来方向
基于本文的贡献和发现，未来研究可以在以下几个方向展开： **1. 更大规模的三模态MDM**：本文的缩放定律表明3B模型需要至少480B token，且$D^*(N) = 7754 \cdot N^{0.84}$。未来工作可以训练10B、30B甚至更大规模的三模态MDM，验证缩放定律的预测，并探索更大规模下是否出现模态协同效应。特别地，论文提到与Quokka的计算最优曲线在20B规模交叉，这为研究不同缩放行为提供了具体目标。 **2. 推理效率的根本性改进**：当前MDM的推理效率远低于自回归模型。未来可以研究：（1）自适应采样策略，根据序列的局部'难度'动态调整揭示顺序和步数；（2）并行解码技术，利用MDM的双向性质同时生成多个token；（3）混合架构设计，将MDM与自回归组件结合，平衡质量和效率。 **3. 新模态的引入**：本文已经建立了三模态（文本、图像、音频）的框架，但可以进一步扩展到更多模态，如视频、3D点云、触觉数据等。关键挑战是如何设计统一的tokenization方案和模态边界标记，以及如何处理不同模态的序列长度差异。 **4. 模态协同机制的理论研究**：本文没有观察到模态间的协同效应，但这可能是因为模型规模或数据量不足。未来工作可以：（1）研究更大规模下是否出现协同效应；（2）设计显式的跨模态对齐机制；（3）分析不同任务（如视觉问答、语音翻译）对模态协同的需求。 **5. 缩放定律的理论深化**：当前的缩放定律是经验性的，缺乏理论解释。未来可以：（1）从信息论角度分析$E$（不可压缩误差）的组成；（2）研究数据质量、模态多样性对缩放指数的影响；（3）建立更通用的多模态缩放理论框架。 **6. 下游应用和微调**：本文主要关注预训练，但展示了模型在文本到图像、文本到语音等任务上的能力。未来可以研究：（1）如何高效微调预训练的三模态MDM适应特定任务；（2）少样本学习能力；（3）在现实世界应用中的潜力和局限性。
复现评估
本文的复现性评估如下： **开源情况**：论文没有明确提到代码或模型的开源计划。这是一个显著的局限性，因为复现三模态MDM需要大量的工程工作和计算资源。不过，论文提供了详细的超参数表格（Table 5），包括模型架构、训练配置、数据混合等信息，为复现提供了重要参考。 **数据集**：论文使用的数据集包括：（1）文本：Nemotron-CC、DCLM、The Pile子集（Wikipedia、HackerNews、Ubuntu IRC、Arxiv、DM-mathematics、Openwebtext）、StackOverflow数据、Qwen-32B推理轨迹等；（2）音频：2M小时网络爬取音频，经Whisper转录并过滤；（3）图像：CC3M、CC12M、COYO及其他许可数据集。这些数据中，CC3M、CC12M、COYO是公开可用的，但音频数据和一些许可数据集可能难以获取。 **算力需求**：根据论文的缩放定律，训练一个3B模型需要480B token。假设每个token需要约6 FLOPs（对于Transformer），总计算量约为$2.88 \times 10^{21}$ FLOPs。使用1000个A100 GPU（约312 TFLOPS FP16），需要约$9.2 \times 10^6$ GPU秒，即约106天。这表明复现需要相当大的计算资源，但并非不可及。 **复现难度**：中等偏高。主要挑战包括：（1）SDE参数化的实现需要对AdamW优化器进行修改；（2）cut-cross-entropy的实现需要高效的内存管理；（3）多模态数据的预处理和tokenization需要特定的编码器（SBER-MoVQGAN、Higgs Audio v2、Tiktoken）；（4）分布式训练的工程复杂性。不过，论文提供了详细的超参数和架构描述，如果有足够的工程能力和计算资源，复现是可行的。 **建议**：对于希望复现或基于本文工作的研究者，建议：（1）首先在小规模（如320M模型）上验证SDE参数化和缩放定律；（2）使用公开可用的数据集（CC12M、LibriSpeech等）作为起点；（3）逐步增加模态数量，从双模态开始；（4）重点关注SDE参数化和反掩码这两个相对容易实现的技术贡献。
论文图表
