CODE :
https://github.com/Gen-Verse/MMaDA
Abstract
我们介绍了一种新型的多模态扩散基础模型MMaDA,它被设计用于在文本推理、多模态理解和文本到图像生成等不同领域实现卓越的性能。该方法的特点是三个关键创新:(i) MMaDA采用统一的扩散架构,具有共享的概率公式和模态不可知设计,消除了对模态特定组件的需求。这种体系结构确保了跨不同数据类型的无缝集成和处理。(ii)我们实施了一个混合的长思维链(CoT)微调策略,该策略策划了跨模式的统一的CoT格式。通过在文本和视觉域之间对齐推理过程,该策略促进了最终强化学习(RL)阶段的冷启动训练,从而从一开始就增强了模型处理复杂任务的能力。(iii)我们提出了UniGRPO,这是一种专门为扩散基础模型量身定制的统一的基于策略梯度的RL算法。利用多样化的奖励模型,UniGRPO统一了推理和生成任务的后训练,确保了一致的性能改进。实验结果表明,MMaDA-8B作为统一的多模态基础模型具有较强的泛化能力。它在文本推理方面超过了LLaMA-3-7B和Qwen2-7B等强大的模型,在多模态理解方面优于Show-o和SEED-X,在文本到图像生成方面优于SDXL和Janus。这些成就突出了MMaDA在统一扩散架构中弥合预训练和后训练之间差距的有效性,为未来的研究和发展提供了一个全面的框架。
Introduction
大语言模型(LLM)在从文本生成(如ChatGPT)到复杂推理(如DeepSeek - R1)等各种任务中取得了最先进的性能,彻底改变了自然语言处理(NLP)领域。受其成功的启发,研究界将大语言模型扩展到多模态领域,催生了多模态大语言模型(MLLM)或视觉语言模型(VLM),如GPT -4和Gemini。这些模型旨在为跨异构模态(文本、图像等)的理解和生成提供一个统一的框架。
早期的多模态方法将语言模型与扩散模型相结合,分别处理离散(如文本)和连续(如图像)模态。随后的自回归(AR)方法通过训练一个带有下一个标记预测功能的单一Transformer来简化架构,在单一模型中统一了离散和连续生成。另一类工作在共享架构中利用特定模态的训练目标:例如,Show - o和Transfusion分别结合自回归和扩散建模来处理文本和视觉语义。
尽管近期的进展已经探索了基于扩散的架构用于全局上下文建模和并行生成 [20,21],但现有的统一多模态基础模型主要集中在模型架构设计和预训练策略上,在探索训练后方法方面存在关键空白,尤其是在非自回归设置中。
为了填补这一空白,我们系统地研究了统一多模态扩散基础模型的设计空间,引入了一个新颖的框架,推动了架构和训练范式的发展(表1中进行了全面比较)。这项工作弥合了统一多模态扩散模型预训练和训练后之间的差距,为这一新兴领域的未来研究提供了一个整体框架。
表1 不同统一多模态基础模型族采用的具体设计选择,包括它们的核心损失函数。下一个标记预测损失定义为
L
N
T
P
=
E
x
i
[
−
log
P
θ
(
x
i
∣
x
<
i
)
]
L_{NTP} = E_{x_i} [-\log P_{\theta}(x_i | x_{<i})]
LNTP=Exi[−logPθ(xi∣x<i)],表示在其前面的上下文
x
<
i
x_{<i}
x<i条件下生成下一个标记
x
i
x_i
xi 的标准负对数似然。连续扩散模型的训练目标由
L
D
i
f
f
-
c
o
n
t
=
E
t
,
x
0
∼
q
(
x
0
)
,
ϵ
∼
N
(
0
,
I
)
,
c
∥
ϵ
−
ϵ
θ
(
x
t
,
t
,
c
)
∥
2
L_{Diff\text{-}cont} = E_{t,x_0\sim q(x_0),\epsilon\sim N(0,I),c} \|\epsilon - \epsilon_{\theta}(x_t, t, c)\|^2
LDiff-cont=Et,x0∼q(x0),ϵ∼N(0,I),c∥ϵ−ϵθ(xt,t,c)∥2给出,其中
x
t
x_t
xt表示在时间步
t
t
t时原始数据
x
0
x_0
x0 的加噪版本,
c
c
c是一个可选的条件信号。模型学习预测添加到数据中的噪声
ϵ
\epsilon
ϵ,从而通过定义为
x
t
−
1
=
F
θ
(
x
t
,
t
)
x_{t -1} = F_{\theta}(x_t, t)
xt−1=Fθ(xt,t) 的迭代去噪过程实现
x
0
x_0
x0 的重建。离散扩散(掩码标记预测)损失由
L
D
i
f
f
-
d
i
s
c
=
E
z
i
∗
[
−
1
∣
M
∣
∑
i
∈
M
log
p
θ
(
z
i
∗
∣
z
(
M
)
masked
)
]
L_{Diff\text{-}disc} = E_{z^*_i} \left[ -\frac{1}{|M|} \sum_{i\in M} \log p_{\theta}(z^*_i |z_{(M)}^{\text{masked}}) \right]
LDiff-disc=Ezi∗[−∣M∣1∑i∈Mlogpθ(zi∗∣z(M)masked)]给出,它衡量在掩码序列
z
(
M
)
masked
z_{(M)}^{\text{masked}}
z(M)masked提供的可见上下文条件下,正确预测被集合
M
M
M掩码位置处的原始离散标记
z
i
∗
z^*_i
zi∗ 的平均负对数似然。
贡献:
• 统一扩散基础架构:我们提出了MMaDA,这是一类基于扩散的模型,通过共享的概率公式和模态无关的架构,将传统生成器扩展为通用任务求解器。这种设计在消除特定模态组件的同时,还能在各项任务中保持出色的性能。
• 混合长思维链后训练:我们引入了混合长思维链(CoT)微调来实现冷启动训练。通过为各项任务精心设计统一的CoT格式,我们使不同模态(如文本和视觉)之间的推理过程保持一致,促进跨模态协同,并在最终输出生成前学习中间推理。
• 统一强化学习(UniGRPO):我们开发了一种专为多模态生成量身定制的以扩散为中心的统一强化学习算法(UniGRPO)。这种方法利用多样化的奖励建模来增强模型进行复杂推理的能力,并在生成过程中保持事实一致性。
• 最先进的性能:MMaDA在三项关键任务(文本推理、多模态理解和文本到图像生成)中实现了卓越且均衡的性能。值得注意的是,在准确性、效率和任务适应性方面,它的表现优于自回归和基于扩散的基线模型。
Related Work
Multimodal Large Language Models for Multimodal Understanding 近期,诸如Gemini -2.0、o1 - preview和DeepSeek - R1等大语言模型(LLM)的发展推动了多模态大语言模型(MLLM)的演进。在该领域的早期探索中,包括LLaVA、MiniGPT -4和InstructBLIP等模型展示了令人瞩目的多模态理解能力。这些研究通过将预训练的特定模态编码器(如CLIP)的特征投影到大语言模型的输入空间,推动了大语言模型在多模态场景中的融合,从而促进了在统一变换器内的多模态理解和推理。在多模态大语言模型方面,人们在视觉编码器、对齐适配器和精心策划的数据集等方面做出了诸多努力,并且大多数都遵循自回归生成范式,该范式已被证明在大语言模型的文本生成中是有效的。然而,它们通常无法同时进行文本推理和多模态推理。在这项工作中,MMaDA开发了扩散基础模型来填补这一空白。
Diffusion Models and Autoregressive Models for Visual Generation 用于视觉生成的扩散模型和自回归模型大量的扩散模型(如相关文献)在视觉生成方面取得了显著成功。除了连续空间上典型的去噪扩散过程外,一系列框架,如 D3PM 和 VQDiffusion,采用离散扩散建模进行视觉生成。具体来说,使用预训练的图像分词器将图像表示为离散标记序列。在训练阶段,模型经过优化以恢复这些被随机掩码的部分标记的原始值。Transformer系列在自然语言处理任务的自回归建模中展现出强大能力。许多方法尝试将自回归建模应用于视觉生成,通过对视觉细节内的语义依赖进行建模。例如,LlamaGen采用 Llama架构并改进码本设计,以提高离散分词器在类别条件图像生成中的性能。VAR 通过设计多尺度图像分词器,用“下一尺度预测”范式取代了“下一个标记预测”范式。然而,现有的自回归方法在视觉生成能力方面仍落后于扩散方法。在这项工作中,MMaDA训练扩散模型对文本和视觉内容进行建模,并通过自回归(AR)或半自回归(Semi - AR)采样进行高效推理。
Unified Vision-Language Foundation Models 最近,众多研究聚焦于开发在理解和生成方面均表现出色的统一多模态基础模型。像 SEED - X 和 DreamLLM等方法,以及其他一些方法,将所有模态表示为一系列标记,并采用统一的Transformer架构对整个系统进行端到端训练。例如,Emu3 使用多模态标记化序列的混合从零开始训练单个Transformer,仅通过下一个标记预测进行优化。虽然这些统一的自回归模型前景可观,但它们在视觉生成任务上可能存在困难。Transfusion 和 Show - o分别采用自回归建模进行文本生成和扩散建模进行视觉生成。然而,它们主要关注预训练策略。现有的统一多模态基础模型仍缺乏对有效后训练设计的探索。
MMaDA: Multimodal Large Diffusion Language Models
Pretraining with Unified Diffusion Architecture and Objective
Data Tokenization 为了建立一个能够同时处理文本和视觉数据的统一建模框架,我们在两种模态上采用了一致的离散标记化策略。这种设计使模型能够在单一的建模目标下运行,即预测离散的掩码标记。对于文本标记化,我们使用了来自 LLaDA [22] 的标记器,它是我们 MMaDA模型的主干。对于图像标记化,我们利用了从 Show-o [20]采用的预训练图像量化器,该量化器基于 MAGVIT -2 [24]架构,可将原始图像像素转换为离散语义标记序列。给定一个尺寸为 (H × W) 的输入图像,编码器会生成一个尺寸为 (H\f ×W\f) 的标记图,其中 (f)表示下采样因子。在这个实现中,我们采用的下采样因子 (f =16),码本大小为8192。这种配置将一个512 × 512像素的图像转换为一个包含 (32×32 =1024)个离散标记的序列。转换后的离散图像标记用于理解和生成建模任务。
LLaDA 、Show-o
Unified Probabilistic Formulation for Pretraining 近期的统一多模态框架旨在将多个建模目标(如自回归生成和基于扩散的去噪)整合到一个单一架构中,以进行联合理解和生成任务 [20,21](详见附录 A.1中的预备知识)。然而,这些方法通常会引入复杂的混合机制,从而影响模型的效率和一致性。相比之下,我们提出了一个简化的框架,它不仅降低了架构的复杂性,还引入了一个统一的扩散目标,以在共享概率公式下对视觉和文本模态进行建模。通过对齐各模态的噪声损坏和语义恢复过程,我们能够在预训练期间实现更有效的跨模态交互,促进异构数据源的无缝集成。
具体而言,我们将MMaDA构建为一个掩码标记预测器(用于图像和文本标记),这是一个参数化模型
p
θ
(
⋅
∣
x
t
)
p_{\theta}(\cdot|x_t)
pθ(⋅∣xt),它以
x
t
x_t
xt作为输入并同时预测所有掩码标记。该模型使用仅在掩码图像/文本标记上计算的统一交叉熵损失进行训练:
L
u
n
i
f
y
(
θ
)
=
−
E
t
,
x
0
,
x
t
[
1
t
∑
i
=
1
X
I
[
x
i
t
=
[
M
A
S
K
]
]
log
p
θ
(
x
i
0
∣
x
t
)
]
L_{unify}(θ)= -\mathbb{E}_{t,\mathbf{x}_0,\mathbf{x}_t}\left[\frac{1}{t}\sum_{i =1}^{X} \mathbb{I}[\mathbf{x}_i^t = [MASK]] \log p_{\theta}(\mathbf{x}_i^0|\mathbf{x}_t)\right]
Lunify(θ)=−Et,x0,xt[t1∑i=1XI[xit=[MASK]]logpθ(xi0∣xt)]
其中
x
0
x_0
x0是真实值,时间步
t
t
t是从
[
0
,
1
]
[0,1]
[0,1]中均匀采样得到的,
x
t
x_t
xt是通过对
x
0
x_0
x0应用前向扩散过程得到的。
I
[
⋅
]
I[\cdot]
I[⋅]表示指示函数,用于确保仅对掩码标记计算损失。具体的预训练任务在第4节中详细介绍。
Post-Training with Mixed Long-CoT Finetuning.
Cold Start Long-CoT Data Curation 我们研究思维链(CoT)机制 [26]如何能够增强我们统一多模态扩散框架的后训练效果,并观察其在促进跨模态协同方面的有效性。为此,我们精心整理了一个涵盖三种核心任务的长思维链轨迹的精简数据集,这三种核心任务分别为:文本推理、多模态推理和文本到图像生成。该数据集通过以下原则使我们预训练的MMaDA模型能够进行稳定的后训练:
- 统一思维链格式:视觉语言基础模型面临的一个关键挑战是不同任务输出格式的异构性(例如,文本生成与图像生成)。我们提出了一种与任务无关的思维链格式:
|<特殊令牌>| <推理过程> |<特殊令牌>| <结果>。
“<推理过程>”对最终输出之前的逐步推理轨迹进行编码。这种统一结构弥合了特定模态输出之间的差距,并促进了任务间的知识转移。例如,通过使语义逻辑与视觉合成保持一致,增强的文本推理能力可直接提高生成图像的真实感。 - 多样性、复杂性和准确性:我们利用开源大语言模型和视觉语言模型(LLM/VLM)在各个任务中生成多样化的推理轨迹(如图2所示)。为确保质量,我们采用最先进的模型作为验证器,过滤掉不准确或浅显的推理,仅选择高质量、长篇幅的思维链样本。与以往专注于通用理解和生成的统一模型不同,我们的MMaDA专门设计用于:
- 推理密集型任务(例如,数学问题求解)
- 以及具备世界知识感知的文本到图像生成,其中事实一致性至关重要。
Mixed Long-CoT Finetuning 利用我们统一的扩散架构和概率公式,我们开发了一种混合任务长思维链微调策略,以跨异构任务联合优化模型。这种方法不仅增强了特定任务的能力,还为后续的强化学习(RL)阶段创建了强大的初始化。训练过程遵循以下步骤:
- Prompt Preservation and Token Masking: 我们保留原始提示 p 0 p_0 p0,并独立地对结果 ( x 0 ) (x_0) (x0)中的标记进行掩码处理,记为 r t r_t rt。
- Joint Input and Loss Computation: 将拼接后的输入 [ p 0 , r t ] [p_0, r_t] [p0,rt]输入到我们预训练的掩码预测器中以计算损失。这使得模型能够利用提示和受损结果的上下文信息来重构掩码区域 ( r 0 ) (r_0) (r0)。
目标函数定义为:
L
Mixed-SFT
=
−
E
t
,
p
0
,
r
0
,
r
t
[
1
t
∑
i
=
1
L
′
I
[
r
i
,
t
=
[
MASK
]
]
log
p
θ
(
r
i
,
0
∣
p
0
,
r
t
)
]
L_{\text{Mixed-SFT}} = -\mathbb{E}_{t,p_0,r_0,r_t} \left[ \frac{1}{t} \sum_{i =1}^{L'} \mathbb{I}[r_{i,t} = [\text{MASK}]] \log p_{\theta}(r_{i,0}|p_0, r_t) \right]
LMixed-SFT=−Et,p0,r0,rt[t1∑i=1L′I[ri,t=[MASK]]logpθ(ri,0∣p0,rt)]
其中
L
′
L'
L′表示序列长度。这里,
[
p
0
,
r
0
]
[p_0, r_0]
[p0,r0] 和
[
p
0
,
r
t
]
[p_0, r_t]
[p0,rt]分别对应干净数据
x
0
x_0
x0及其加噪后的数据
x
t
x_t
xt。这种公式化表达确保模型在学习恢复被掩码的标记时,能与原始提示和特定任务的推理逻辑保持一致。
Post-Training with Unified Reinforcement Learning
Unified GRPO for Diffusion Foundation Models
通过我们的混合长思维链微调,MMaDA展现出了在生成最终输出之前生成统一且连贯推理链的能力。为了进一步提升其在知识密集型任务以及复杂推理/生成场景中的性能,我们提出了UniGRPO,这是一种专为扩散基础模型量身定制的、基于策略梯度的新型强化学习算法。该方法实现了一个以扩散为中心的强化学习训练框架,能够统一跨多种模态和推理范式的特定任务目标。该方法主要由两部分组成:基于扩散的强化学习的统一数学公式,以及使策略梯度与特定任务奖励相匹配的多样化奖励建模。
Challenges in Adapting Autoregressive GRPO to Diffusion Models 原始的广义响应概率优化(GRPO)[27]依赖于计算词元级别的对数似然
π
θ
(
o
i
,
t
∣
q
,
o
i
,
<
t
)
\pi_{\theta}(o_{i,t}|q, o_{i,<t})
πθ(oi,t∣q,oi,<t)以及序列级别的概率
π
θ
\pi_{\theta}
πθ 和
π
r
e
f
\pi_{ref}
πref(附录 A.2中有初步介绍)。在自回归(AR)大语言模型中,这些指标可通过生成的链式法则高效推导得出。然而,扩散模型带来了三个关键挑战:
- Local Masking Dependency: 在扩散过程中,词元级别的对数似然 log π θ ( o i , t ∣ q , o i , < t ) \log \pi_{\theta}(o_{i,t}|q, o_{i,<t}) logπθ(oi,t∣q,oi,<t)仅在掩码区域内有效,这与自回归(AR)模型不同,在自回归模型中所有词元都是有效的。
- Mask Ratio Sensitivity: :由于扩散动态依赖于掩码模式,因此必须为响应片段采样一个均匀的掩码比例,以近似策略分布 π θ \pi_{\theta} πθ。
- Non-Autoregressive Sequence-Level Likelihoods: 由于扩散模型中不存在自回归链式法则,序列级对数似然不能直接由词元级概率累加得到。
先前的方法采用了次优策略来解决这些问题。LLaDA [22]对众多掩码比率(例如128个样本)进行蒙特卡罗采样,这会给在线策略强化学习带来高昂的计算成本。d1 [28]固定掩码比率并随机化问题掩码,这降低了噪声多样性,并且忽略了扩散的多步去噪特性。
Unified Formulation for Diffusion GRPO 为了克服这些局限性,我们引入了 UniGRPO,这是一种为扩散架构设计的计算高效的近似算法。给定针对查询 q q q 的一批响应 { o i } i = 1 G \{o_i\}_{i =1}^G {oi}i=1G,在梯度更新期间每个响应都保持固定,以确保策略评估的稳定性。我们强调 UniGRPO 的三个关键点如下:
- Structured Noising Strategy: 对于每个 o i o_i oi,我们从 [ 0 , 1 ] [0,1] [0,1]中均匀采样一个掩码率 p i p_i pi,并通过将标记替换为 [ MASK ] [\text{MASK}] [MASK]来构建一个扰动版本 o ~ i , p \tilde{o}_{i,p} o~i,p。 p i p_i pi 的随机种子在不同的梯度步长中变化。此策略旨在保持随机性,同时确保模型能够经历扩散去噪过程的各个阶段,从几乎完全被掩码的答案到几乎完全去噪的答案。通过这样做,UniGRPO可以从多步去噪信息中学习,这与扩散模型的传统训练方法一致,并能充分利用其多步生成能力。
- Efficient Log-Likelihood Approximation:我们将扰动分布下的每个标记的期望对数似然定义为:
π ′ θ ( o i , t ∣ q , o ~ , p i ) = E p i ∼ [ 0 , 1 ] [ I [ o i , t , p = [ MASK ] ] log p θ ( o i , t , p ∣ q ) ] \pi'\theta(o_{i,t} | q, \tilde{o}, p_i) = \mathbb{E}_{p_i\sim[0,1]} [\mathbb{I}[o_{i,t,p} = [\text{MASK}]] \log p_{\theta}(o_{i,t,p}|q)] π′θ(oi,t∣q,o~,pi)=Epi∼[0,1][I[oi,t,p=[MASK]]logpθ(oi,t,p∣q)]
然后,序列级别的对数似然通过对掩码标记求平均来近似:
π θ ′ = 1 M ∑ o i , t ∈ M log p θ ( o i , t ∣ q ) \pi'_{\theta} = \frac{1}{M} \sum_{o_{i,t} \in M} \log p_{\theta}(o_{i,t}|q) πθ′=M1∑oi,t∈Mlogpθ(oi,t∣q)
其中 M M M表示掩码标记的数量。 - Policy Gradient Objective::每个标记的奖励计算为当前策略似然与旧策略似然之比:
r
i
,
t
′
(
θ
)
=
π
θ
′
(
o
i
,
t
∣
q
,
o
~
,
p
i
)
π
old
′
(
o
i
,
t
∣
q
,
o
~
,
p
i
)
r'_{i,t}(\theta) = \frac{\pi'_{\theta}(o_{i,t}|q,\tilde{o},p_i)}{\pi'_{\text{old}}(o_{i,t}|q,\tilde{o},p_i)}
ri,t′(θ)=πold′(oi,t∣q,o~,pi)πθ′(oi,t∣q,o~,pi)最终的 UniGRPO目标整合了裁剪后的替代奖励和 KL正则化:
J U n i G R P O ( θ ) = E ( q , a ) ∼ D , { o i } i = 1 G ∼ π θ old ( ⋅ ∣ q ) , { p i ∈ [ 0 , 1 ] } i = 1 G [ 1 G ∑ i = 1 G 1 ∣ o i ∣ ∑ t = 1 ∣ o i ∣ ( min ( r i , t ′ ( θ ) A ^ i , t , clip ( r i , t ′ ( θ ) , 1 − ε , 1 + ε ) A ^ i , t ) − β D KL ( π θ ′ ∣ ∣ π ref ′ ) ) ] J_{UniGRPO}(θ) = \mathbb{E}_{(q,a)\sim D,\{o_i\}_{i =1}^G\sim\pi_{\theta_{\text{old}}}(\cdot|q),\{p_i\in[0,1]\}_{i =1}^G}\left[\frac{1}{G}\sum_{i =1}^G\frac{1}{|o_i|}\sum_{t =1}^{|o_i|}(\min\left(r_{i,t}'(\theta)\hat{A}_{i,t},\text{clip}\left(r_{i,t}'(\theta),1 - \varepsilon,1 + \varepsilon\right)\hat{A}_{i,t}\right) - \beta D_{\text{KL}}(\pi_{\theta}'||\pi_{\text{ref}}')\right)] JUniGRPO(θ)=E(q,a)∼D,{oi}i=1G∼πθold(⋅∣q),{pi∈[0,1]}i=1G[G1∑i=1G∣oi∣1∑t=1∣oi∣(min(ri,t′(θ)A^i,t,clip(ri,t′(θ),1−ε,1+ε)A^i,t)−βDKL(πθ′∣∣πref′))]
其中 A ^ i , t \hat{A}_{i,t} A^i,t表示优势估计值, ε \varepsilon ε控制裁剪范围, β \beta β平衡 KL散度惩罚项。
通过这种设计,UniGRPO捕捉到了扩散模型中关键的多步去噪动态过程。通过让模型在不同的掩码条件下预测答案,同时保留输入的自然结构,它避免了计算效率低下(如LLaDA)和预测过于简化(如d1)的问题。UniGRPO的训练过程在算法1中概述。有关UniGRPO算法的更多详细信息,请参考附录B和第5.2节。
Diversified Reward Modeling
我们进一步将UniGRPO的优化目标(公式(5))简化如下:
UniGRPO的通用优化目标:
J
UniGRPO
(
θ
)
=
E
o
∼
π
θ
(
⋅
∣
q
)
[
F
(
R
Uni
(
o
)
)
−
β
P
(
o
)
]
J_{\text{UniGRPO}}(\theta) = \mathbb{E}_{o\sim\pi_{\theta}(\cdot|q)}[F(R_{\text{Uni}}(o)) - \beta P(o)]
JUniGRPO(θ)=Eo∼πθ(⋅∣q)[F(RUni(o))−βP(o)]
其中
R
U
n
i
(
o
)
R_{Uni}(o)
RUni(o)表示从模型生成的回复
o
o
o中获得的奖励,
P
(
⋅
)
P(\cdot)
P(⋅)是惩罚项,它表示如公式 (5)中所指定的 KL散度。这是一个统一的基于规则的奖励系统,其中
R
U
n
i
(
⋅
)
R_{Uni}(\cdot)
RUni(⋅)可以用针对不同任务的多样化奖励来实例化。为了满足不同任务的多样化需求,我们在统一公式 (6)下定义了一系列奖励,为每个任务分支提供量身定制的强化学习优化方向。我们主要采用三种类型的奖励:
• Textual Reasoning Rewards: 我们在GSM8K [29]数据集的训练分割上应用UniGRPO,并定义一个综合奖励。这包括答对时给予2.0的正确性奖励,以及如果回复符合我们预定义的格式 “…”则给予0.5的格式奖励。
• Multimodal Reasoning Rewards: 对于诸如GeoQA [30]和CLEVR [31]等数学任务,我们采用与文本推理中相同的正确性和格式奖励。此外,对于基于描述的任务,我们进一步引入一个CLIP奖励,其值为 0.1 \cdot \text{CLIP}(\text{image}, \text{text}),其中原始用于衡量文本 -图像对齐程度的CLIP分数乘以0.1以平衡其影响。
• Text-to-Image Generation Rewards: 对于图像生成任务,我们纳入相同的CLIP奖励来评估文本 -图像语义对齐程度,同时还有一个反映人类偏好分数的图像奖励。这两个奖励都乘以0.1,以确保在优化过程中贡献均衡。
Flexible Sampling Strategies at Inference Time
Semi-Autoregressive Sampling for Text Generation 在文本生成方面,我们采用了 LLaDA [22]中引入的半自回归去噪策略,该策略将自回归解码与基于扩散的去噪相结合。具体而言,输出序列被划分为多个块,并从左到右依次生成。在每个块内,为所有被掩码的位置计算对数概率,然后随机或根据置信度分数选择一部分标记进行去噪。掩码调度遵循线性调度,与 LLaDA一致。去噪过程会重复给定的步数。
在我们的评估中,我们将总序列长度设为
N
=
1024
N=1024
N=1024,并执行
N
/
2
=
512
N/2=512
N/2=512 步去噪。序列被分成每块64个令牌。在每一步中,我们不考虑位置,对当前块内置信度最低的2个令牌进行去掩码操作。一旦一个块中的所有令牌都被去噪,该过程就进入下一个块。下面提供了定性比较。如图所示,半自回归去噪策略倾向于生成更复杂和详细的描述,而非自回归的固定长度生成通常会产生非常短的响应。这一观察结果与LLaDA的报告发现一致。对于经过指令微调的模型,考虑到训练过程中包含了大量的|EOS|
令牌,如果不分块直接应用最低置信度重掩码策略,会导致生成句子中|EOS|
令牌的频率异常高。
Parallel Non-Autoregressive Sampling for Image Generation 在图像生成方面,我们采用低置信度重新掩码策略,并遵循余弦噪声调度,这与 MAGVIT -2 [24]中的设置一致。与文本生成不同,我们不采用半自回归方法;相反,整个输出序列被视为一个单一的生成块。在评估过程中,我们生成长度为1024 的序列,对应于512×512分辨率的图像。去噪过程包含50个时间步长,并且我们应用无分类器引导,引导系数设置为3.5。
Experiments
Experimental Setup
Datasets 为了训练MMaDA,我们针对相应的训练阶段使用了多种不同的数据集,具体如下:
- 基础语言和多模态数据:为了实现基本的文本生成能力,我们采用了RefinedWeb [37]数据集。对于多模态理解和生成任务,我们纳入了广泛使用的开源图像 -文本数据集 [38 -42]。
- 指令调优数据:为了增强遵循指令的能力,我们使用Alpaca [43]进行文本指令调优,使用LLaVA -1.5 [32]进行视觉指令调优。
- **推理数据:**对于混合长思维链(Mixed Long - CoT)微调,我们精心整理了一组多样化的推理数据集。对于文本数学和逻辑推理,我们采用了来自ReasonFlux [44]、LIMO [45]、s1k [46]、OpenThoughts [47]和AceMath - Instruct [48]的数据集。对于多模态推理,我们使用LMMR1 [49]模型在GeoQA [30]和CLEVR [31]上生成响应,并保留回答正确的实例。此外,对于具备世界知识感知的图像生成,我们使用GPT -4.1合成了涵盖科学、文化和地标等方面的事实性物品 -描述对,并将其格式化为统一的思维链风格轨迹。
- 强化学习数据:对于UniGRPO训练,我们采用了推理部分使用的原始数学和逻辑数据集 [30,31,50]。
Evaluation and Baselines 我们使用特定任务的指标和基线在三个不同的任务上评估我们的MMaDA:
- 多模态理解:遵循LLaVA [32],我们在POPE、MME、Flickr30k、VQAv2、GQA和MMMU上进行评估,并与仅用于理解的模型 [32–36]以及统一模型 [8,9,11,12,16,19,20]进行比较。
- 图像生成:我们使用测试集中的50K个提示来评估生成质量,计算CLIP分数 [51] 和ImageReward [52]以评估文本对齐和人类偏好对齐。我们采用GenEval [53]进行通用评估,采用WISE [54]评估基于世界知识的生成,并与特定的生成模型 [55–58] 和统一基线 [8,9,11–13,16,19,20,59]进行比较。
- 文本生成:我们在MMLU、GSM8K和相关基准测试上评估指令遵循和推理性能,并与LLaMA2 -7B、Qwen2 -7B和LLaDA -8B进行比较。
Implementation Details 我们使用LLaDA -8B - Instruct的预训练权重[22]初始化MMaDA,并使用Show - o的预训练权重初始化图像分词器。我们分三个阶段进行联合训练:Stage1: 初始模型使用基础语言和多模态数据进行训练,包括用于文本生成的RefinedWeb、用于类别条件图像生成的ImageNet -1k,以及用于图像字幕生成的额外图像 -文本数据集,持续训练200K步。随后再进行400K步的训练,此阶段用更多样化的图像 -文本对替代ImageNet。Stage2: 然后使用指令调优数据和推理数据对模型进行联合训练,持续50,000步。Stage3: 最后一个阶段是使用强化学习数据进行UniGRPO训练,持续50,000步。训练在64块A100(80GB)GPU上进行,全局批量大小为1,280.采用AdamW优化器,初始学习率为 5 × 10 − 5 5\times10^{-5} 5×10−5,并使用余弦学习率调度器。
Multimodal Understanding
表2报告了我们的方法在标准基准测试(包括POPE、MME、Flickr30k、VQAv2、GQA和MMMU)上的多模态理解性能。对于MMaDA输出中包含推理过程的情况,我们将最终答案作为预测结果。与LLaVA - v1.5、InstructBLIP和Qwen - VL - Chat等专门的仅用于理解的模型相比,尽管我们的模型是在统一目标下进行训练的,但在大多数基准测试中都取得了相当或更优的结果。与其他统一模型(如SEED - X、DreamLLM、Janus、Emu3和Show - o)相比,我们的方法在多个基准测试中始终表现更优,尤其得益于所提出的混合长思维链微调(Mixed Long - CoT Finetuning)和统一广义奖励策略优化(UniGRPO)强化学习阶段。值得注意的是,这是首次证明基于扩散的多模态大语言模型(MLLM)具有强大的理解能力,凸显了我们的统一架构在衔接生成和理解任务方面的潜力。定性结果见附录C和附录D。
Text-to-Image Generation
表3展示了在文本到图像生成基准测试上的评估结果。我们的模型在仅生成模型和统一模型中,在 CLIP分数和 ImageReward方面均取得了最高性能,这归因于 UniGRPO训练阶段,该阶段的奖励与这些指标明确对齐。此外,由于理解分支进行了大量推理训练,我们的方法在 GenEval上展现出了卓越的组合性和对象计数能力。值得注意的是,在用于评估世界知识感知生成能力的 WISE [54]文化基准测试中,我们的模型显著优于先前的方法,这是因为它进行了基于文本推理的联合训练,而现有的统一模型通常缺乏这一点。定性结果见附录 C 和附录 D。
Textual Reasoning
表4详细列出了MMaDA在一系列基准测试中的语言建模性能,这些基准测试涵盖了如MMLU、ARC - C和TruthfulQA等通用任务,以及包括GSM8K、MATH和GPQA在内的数学任务。尽管MMaDA仅在有限的特定任务标记和纯开源文本数据上进行训练,但在MMLU、ARC - C测试中,它与Qwen2 -7B和LLaMA3 -8B等强大的基线模型表现相当,并且在数学基准测试中始终优于LLaDA -8B。值得注意的是,MMaDA(我们的模型)开创了用于文本生成、多模态推理和图像生成的统一扩散模型的联合训练——这是一种在以往统一架构中很少探索的多任务配置。这些结果强调了基于扩散的模型作为通用大语言模型的可行性,并表明通过增强文本数据和模型扩展,未来可能会有更出色的表现。定性结果见附录C。
Observations, Analysis and Conclusion
Ablation Study
我们展示了我们的多模态多任务自适应(MMaDA)模型在不同训练阶段的定量消融实验结果,这些阶段包括混合长思维链(Mixed Long-CoT)微调阶段和统一梯度投影优化(UniGRPO)阶段。所有结果均按照第3节中的采样过程生成。如表5所示,在第一阶段之后,我们的模型仍落后于大多数基线模型。在第二阶段,混合长思维链微调显著增强了模型的推理能力,特别是在数学和几何领域。在第三阶段,统一梯度投影优化进一步提升了模型性能,使模型在包括数学推理、几何问题求解以及图像生成基准(如CLIP分数和图像奖励分数)等各种任务上取得了与最先进方法相当的结果。这些结果表明,统一梯度投影优化有效地提升了模型的理解/推理能力和生成能力。
Design Choices of UniGRPO
Effect of General Masking Strategy 为了评估我们提出的掩码策略的影响,我们首先在强化学习框架内与d1 [28]进行了对比分析。鉴于大规模消融研究的计算成本很高,我们在GSM8K数据集上进行了这些实验,使用了8块A100 GPU。原始的d1方法和我们的UniGRPO方法都应用于该数据集,均从我们的MMaDA的同一预训练检查点开始。我们在图3中展示了训练期间的奖励趋势。如图所示,我们的方法在训练期间始终能获得更高的奖励值,这与我们的理论分析相符。与d1不同,UniGRPO不对问题进行掩码处理,而是对答案进行部分掩码,而非完全掩码。这使得输入序列保留了部分噪声,促使模型在多个去噪时间步上进行学习。因此,这能更好地利用扩散模型的内在特性,提高整体学习能力。
Effect of Uniformly Random masking 在各轮迭代中,我们没有采用完全随机的掩码策略,而是对答案部分采用了均匀随机掩码策略。具体来说,我们首先随机采样一个起始时间步,然后在整个扩散时间步(在我们的实验中设置为1000)中均匀生成其余的去噪时间步。例如,若随机选择的起始时间步为100,且总共有5轮训练迭代,那么其余时间步将均匀分布,分别设置为300、500、700 和900。这种设计在保留随机性的同时,确保了对扩散过程更一致的覆盖。我们在图4中展示了这种结构化掩码策略所产生的训练奖励趋势。如图所示,完全随机选择时间步的基线方法在训练期间往往会引入不稳定性,导致奖励波动更频繁,并且需要更多的步骤才能收敛。相比之下,我们的均匀间隔采样策略在对数似然估计中能有效近似蒙特卡罗平均的行为,从而提高了稳定性并加快了收敛速度。
Synergy Across Various Tasks
在联合训练过程中,我们观察到文本生成、多模态理解和图像生成这三个任务类别之间存在明显的协同效应。如图6所示,在第二阶段(训练步数120K -200K),所有关键性能指标均呈现出持续提升,这体现了我们统一训练框架的互利性。这种协同效应在质量上也很明显:如图5所示,模型的响应(包括文本和视觉方面)变得越来越复杂且连贯。具体而言,文本输出的信息更加丰富,逻辑结构更加清晰,而视觉理解则能产生更精确、更贴合实际的描述。因此,对于相同的提示,生成的图像变得更加准确、详细,并且更符合给定的指令,这证明了联合优化在增强跨模态对齐和组合性方面的有效性。
Sampling Efficiency
我们认为采样效率是扩散模型相对于自回归(AR)方法的一个关键优势。与按顺序生成标记的 AR模型不同,扩散模型允许在每个去噪步骤中并行生成标记,从而大幅减少所需的前向传播次数。为了量化这一优势,我们评估了 MMaDA 在不同去噪步数下的性能。在我们的设置中,生成过程从1024个 [MASK]标记开始,最多允许1024个去噪步骤,这对应于一张512×512分辨率的图像。如表6所示,即使只进行15步或50步去噪,图像生成仍能保持良好的性能。对于文本和多模态任务,只需使用全部步数的四分之一或一半就能得到连贯的输出。这些结果凸显了基于扩散的语言模型在效率方面的潜力,并表明未来采样技术的进步或高阶求解器的应用可能会进一步提升它们的速度和质量。
Task Extension
基于扩散的模型的一个显著优势是,它们无需额外的微调就能自然地进行图像修复和外推。这是因为这些任务可以被表述为掩码标记预测问题,而这本质上已融入扩散模型的训练目标中。虽然像 Show-o等先前的工作仅在图像生成的背景下展示了这一特性,但 MMaDA 将其进一步扩展到多模态理解和文本生成。如图7所示,我们的模型支持跨三种模态的图像修复:预测文本序列中缺失的片段;在给定图像和部分输入的情况下完成视觉问答的答案;根据不完整的视觉提示进行图像修复。这些示例展示了我们统一的扩散架构在各种生成和推理任务中的灵活性和泛化能力。
Conclusion
这项工作引入了一个统一的扩散基础模型,即 MMaDA,它在单一概率框架内整合了文本推理、多模态理解和生成。据我们所知,MMaDA是首个系统探索基于扩散的基础模型设计空间的模型,并提出了新颖的后训练策略。在各种视觉 -语言任务上进行的大量实验表明,MMaDA与专门模型相当,甚至表现更优,凸显了扩散模型作为多模态智能下一代基础范式的潜力。然而,由于当前模型规模(80亿参数)的限制,MMaDA也存在一定的局限性,我们将采用更大的模型规模以提升性能。