LaViDa-R1：推进统一多模态扩散语言模型的推理能力 LaViDa-R1: Advancing Reasoning for Unified Multimodal Diffusion Language Models
提出统一后训练框架，通过SFT与多任务RL结合，大幅提升多模态扩散语言模型的推理能力
前置知识
扩散语言模型（dLLM）
扩散语言模型是一种基于离散扩散过程的语言模型，与自回归模型不同，它从完全掩码的序列开始，通过多步扩散过程逐步解掩码，最终生成完整的输出序列。dLLM具有双向上下文建模能力，可以在生成过程中同时考虑前后文信息，并且能够并行解码多个token，提供更快的推理速度。形式化地，给定长度为L的序列$y_0$和条件提示$x$，前向掩码扩散过程$q(y_t|y_s)$在时间区间$[0,1]$上逐步掩码token，其中$t=1$时序列完全掩码为$[M, M, \cdots, M]$，$t=0$时序列无掩码。dLLM策略模型$\pi_\theta$训练用于建模反向过程$p(y_s|y_t, x)$。
本文的核心是提升dLLM的推理能力，理解dLLM的基本工作原理是理解本文方法的前提。dLLM的掩码-解掩码机制使得answer-forcing等创新技术成为可能。
GRPO（Group Relative Policy Optimization）
GRPO是一种高效的强化学习算法，通过组内相对优势计算来优化策略模型。给定一组N个响应$y_1, \ldots, y_N$，GRPO计算每个样本的优势值$A_i^{GRPO} = \frac{r_i - \text{mean}(r_1, \ldots, r_N)}{\text{std}(r_1, \ldots, r_N)}$，其中$r_i$是每个样本的奖励。GRPO的优势在于不需要单独的价值网络，通过组内归一化直接估计优势函数，计算效率高且训练稳定。该算法已被证明在纯在线策略设置下可以简化为策略梯度目标。
LaViDa-R1的统一后训练框架基于GRPO策略梯度形式构建，理解GRPO的原理是理解本文如何将SFT、RL和自蒸馏统一在同一个目标函数下的关键。
ELBO（证据下界）
ELBO是扩散语言模型中用于估计序列对数概率$\log \pi_\theta(y|x)$的替代目标函数。由于dLLM的似然函数难以直接计算，通常通过ELBO进行蒙特卡洛估计：$\log \pi_\theta(y|x) = \mathbb{E}_{t,y_t}\left[ w(t) \sum_{k \in \mathcal{M}(y_t)} \log \pi_\theta(y[k]|y_t, x) \right]$，其中$w(t)$是权重函数，$\mathcal{M}(y_t)$是被掩码位置的集合。不同的MC采样策略和权重函数选择会产生不同的估计方法，直接影响训练的稳定性和效率。
本文提出的互补掩码似然估计器直接改进了ELBO的计算方式，解决了现有方法中token梯度缺失和不平衡的问题，是理解本文技术贡献的核心概念。
文本填充（Text Infilling）
文本填充是dLLM的一项独特能力，指在给定部分已知token的条件下，模型能够生成中间部分的token来填补空缺。与自回归模型只能从左到右生成不同，dLLM可以利用双向上下文信息，在已知开头和结尾的情况下，生成合理的中间内容。这种能力使得answer-forcing技术成为可能：通过在序列末尾注入正确答案，模型可以反向生成通向该答案的推理过程。
answer-forcing是LaViDa-R1的核心创新之一，它依赖于dLLM的文本填充能力来构造高质量的推理轨迹，这是自回归模型无法实现的特性。
研究动机
扩散语言模型（dLLM）作为自回归模型的有前景替代方案，近年来在语言和多模态任务中展现出强大能力，具有更快的推理速度、双向上下文和统一的视觉-文本生成范式等优势。然而，将推理能力引入dLLM面临四个关键挑战：第一，现有工作主要聚焦于有限的任务集（如数学推理），且通常需要针对特定数据集进行微调，将强化学习扩展到支持图像编辑和目标定位等多样化多模态任务的研究仍处于空白状态。第二，使用强化学习训练dLLM容易导致模型崩溃，即使存在KL散度正则化器也难以避免，而且KL正则化器会阻碍模型在训练过程中的探索能力，降低任务性能。第三，对于复杂或困难的任务，模型可能无法生成高质量的样本，导致低质量的训练信号，甚至由于收益递减而产生零训练信号。第四，与自回归模型可以精确高效地评估序列似然不同，dLLM的序列似然计算是难以处理的，通常通过蒙特卡洛方法估计，这会产生缺失和不平衡的token梯度问题，影响训练稳定性。
本文的目标是本文旨在提出LaViDa-R1，一种构建强性能多模态推理dLLM的完整训练方案。具体目标包括：建立统一的后训练框架，无缝集成监督微调（SFT）和多任务强化学习（RL），支持数学推理、视觉问答、图像编辑和目标定位等多种任务；开发引导式rollout生成算法（answer-forcing和tree search），在训练过程中高效采样高质量输出，解决训练信号消失问题；设计稳定且高效的互补掩码似然估计器，改进现有蒙特卡洛方法的不足。最终实现一个通用的多模态推理dLLM，在多个基准测试上取得显著提升。
与已有工作不同的是，LaViDa-R1的独特切入角度在于三个方面：首先，它提出了统一的策略梯度目标函数框架，将SFT、在线GRPO和自蒸馏损失整合在同一个目标函数$\mathcal{J}_{Unified}(\theta) = \frac{1}{N} \sum_{i=1}^{N} A_i \log \pi_\theta(y_i|x_i)$中，通过不同的权重$A_i$和数据来源控制不同训练目标，这是首次在dLLM中实现这种统一。其次，创新性地利用dLLM的双向生成和文本填充能力，提出answer-forcing机制，在有正确答案时通过注入答案反向生成推理轨迹，这是自回归模型无法实现的特性。第三，提出互补掩码似然估计器，使用两个互补的掩码样本并采用$w(t)=1$的均匀权重，解决了token梯度缺失和不平衡问题，相比现有方法更加稳定高效。
核心方法
LaViDa-R1的训练框架包含三个主要组件：统一后训练策略梯度目标、引导式rollout生成算法和互补掩码似然估计器。整体思路是先通过统一的目标函数框架将SFT和RL训练整合，然后针对RL训练中的信号消失问题设计两种引导生成策略（有答案时的answer-forcing和无答案时的tree search），最后通过改进的似然估计器确保训练的稳定性。技术路线从LaViDa-O基础模型出发，第一阶段进行推理数据的监督微调，第二阶段使用混合SFT、RL和自蒸馏损失的统一后训练，在数学推理、视觉问答、目标定位和图像编辑四类任务上进行多任务、多奖励、跨模态的强化学习。
LaViDa-R1的核心创新在于观察到SFT、在线GRPO和自蒸馏损失都可以写成统一的策略梯度形式$\mathcal{J}_{Unified}(\theta) = \frac{1}{N} \sum_{i=1}^{N} A_i \log \pi_\theta(y_i|x_i)$，区别仅在于数据来源和权重$A_i$的选择。这一洞察带来三个重要优势：首先，可以通过简单拼接训练批次来组合不同目标，无需额外采样；其次，使用SFT正则化替代KL散度正则化，允许模型在不被次优参考模型约束的情况下充分探索动作空间，同时防止崩溃；第三，从计算效率角度，移除参考模型的需求显著降低了RL训练成本，因为不再需要将参考模型加载到GPU内存或在单独的服务器上托管。此外，引入自蒸馏目标放大了组中最佳样本的训练信号，产生更强的学习效果。
方法步骤详情
LaViDa-R1的训练过程分为两个阶段：第一阶段是推理数据的监督微调（SFT），使用链式思维（CoT）数据对基础模型LaViDa-O进行微调；第二阶段是统一后训练，同时优化SFT、在线GRPO和自蒸馏损失，其中$\gamma=0.5$用于平衡GRPO和自蒸馏。具体步骤包括：（1）数据引擎提供提示-响应对$(y_i, x_i)$和对应的权重$A_i$，可以来自离线数据集或在线生成；（2）对于在线生成，使用策略模型$\pi_\theta$采样N个响应；（3）计算每个响应的奖励（数学任务用正确性奖励，目标定位用IoU奖励，图像编辑用EditScore奖励模型）；（4）如果所有响应奖励较低且有正确答案，使用answer-forcing生成合成推理轨迹；如果没有正确答案，使用tree search扩展高质量轨迹；（5）计算组内优势值，对GRPO和自蒸馏使用加权平均$A_i^{aggri} = \gamma A_i^{distill} + (1-\gamma) A_i^{GRPO}$；（6）使用策略模型计算每个序列的对数似然$\log \pi_\theta(y_i|x_i)$，通过互补掩码似然估计器进行稳定估计；（7）优化统一目标函数更新模型参数。
技术新颖性
LaViDa-R1的技术新颖性体现在多个层面：首先，统一后训练框架首次将SFT、在线GRPO和自蒸馏损失整合在同一个策略梯度目标中，通过简单的权重聚合实现多目标优化，避免了传统方法中需要为不同目标单独采样rollout的计算开销。其次，answer-forcing机制创造性地利用了dLLM特有的文本填充能力，在ground-truth答案可用时通过在完全掩码序列末尾注入答案token，让模型反向生成通向答案的推理轨迹，这是自回归模型无法实现的特性。第三，tree search算法通过在早期扩散步骤进行分支，从高奖励轨迹的中间状态生成新样本，有效扩大了有效组大小。第四，互补掩码似然估计器使用两个互补掩码样本（$y_{t_1}$和$y_{t_2}$，其中$y_{t_2}[i] = y[i]$当$y_{t_1}[i] = M$，否则$y_{t_2}[i] = M$）并采用$w(t)=1$的均匀权重，相比使用$w(t)=1/t$的方法避免了因掩码比率差异导致的token权重不平衡问题，确保了训练的稳定性。
实验结果
LaViDa-R1在多个基准测试上取得了显著且一致的性能提升。在视觉数学推理方面，MathVista从基础模型的56.9提升至60.0（+3.1），MathVerse从36.9提升至38.7（+1.8），展现出对复杂视觉数学问题的更强推理能力。在视觉问答任务中，ChartQA从80.0提升至81.7（+1.7），AI2D从76.7提升至78.9（+2.2），MMMU-Pro从31.2提升至32.8（+1.6），证明了模型在图表理解和科学图像理解上的进步。最令人印象深刻的是语言任务上的提升，GSM8K从47.4跃升至81.5（+34.1），MATH-500从23.4提升至38.6（+15.2），这表明基础模型的预训练数据以视觉为中心，语言能力存在较大提升空间，统一后训练有效弥补了这一缺陷。在图像编辑方面，LaViDa-R1在ImgEdit基准上取得3.90的整体分数（基础模型3.71，+0.19），值得注意的是SFT仅带来+0.01的微小提升，表明监督学习已达到性能饱和，而RL框架成功驱动了超越监督分布的探索。在推理密集型目标定位任务上，Lisa-Grounding数据集的mIoU从26.1大幅提升至60.0（+33.9），P@0.5从29.2提升至66.7（+37.5），证明了统一RL框架在需要复杂视觉推理的定位任务上的强大效果。消融实验进一步验证了各组件的有效性：10%的answer-forcing注入率取得最佳性能，过高注入率会导致崩溃；tree search的步骤[0,8]是较好的选择；互补掩码似然估计器优于其他MC方法；$\gamma=0.5$的自蒸馏权重平衡了GRPO和自蒸馏的优势。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 视觉数学推理 - MathVista | Accuracy | 60.0 | 56.9 (LaViDa-O) | +3.1 |
| 视觉数学推理 - MathVerse | Accuracy | 38.7 | 36.9 (LaViDa-O) | +1.8 |
| 视觉问答 - ChartQA | Accuracy | 81.7 | 80.0 (LaViDa-O) | +1.7 |
| 视觉问答 - AI2D | Accuracy | 78.9 | 76.7 (LaViDa-O) | +2.2 |
| 视觉问答 - MMMU-Pro | Accuracy | 32.8 | 31.2 (LaViDa-O) | +1.6 |
| 语言推理 - GSM8K | Accuracy | 81.5 | 47.4 (LaViDa-O) | +34.1 |
| 语言推理 - MATH-500 | Accuracy | 38.6 | 23.4 (LaViDa-O) | +15.2 |
| 图像编辑 - ImgEdit | Overall Score | 3.90 | 3.71 (LaViDa-O) | +0.19 |
| 目标定位 - Lisa-Grounding | mIoU | 60.0 | 26.1 (LaViDa-O) | +33.9 |
| 目标定位 - Lisa-Grounding | P@0.5 | 66.7 | 29.2 (LaViDa-O) | +37.5 |
局限与改进
论文在多个方面存在局限性。首先，模型规模方面，LaViDa-R1基于LaViDa-O构建，而LaViDa-O本身是一个相对较小的模型（约7-8B参数级别），与GPT-4o等商业模型相比仍有差距，特别是在ImgEdit基准上GPT-4o取得4.20的整体分数，而LaViDa-R1仅为3.90。其次，计算效率方面，虽然统一框架移除了参考模型的需求，但tree search和多步扩散过程本身仍然带来较高的计算开销，特别是tree search需要多次采样和奖励评估。第三，训练数据方面，RL训练依赖于高质量的奖励信号，对于图像编辑任务使用GPT-4作为评判模型可能存在偏差，且奖励模型的质量直接影响最终性能。第四，任务覆盖方面，虽然论文展示了在多个任务上的改进，但主要集中在视觉理解和简单生成任务，对于更复杂的多步推理、长文本生成等场景的效果尚未验证。第五，模型的泛化能力需要进一步评估，当前实验主要在标准基准上进行，对于实际应用场景中的分布外数据、领域特定任务等的适应性有待考察。此外，消融实验表明answer-forcing的注入率需要仔细调优（10%最优），过高会导致崩溃，这增加了超参数调优的复杂性。
独立分析的弱点
LaViDa-R1存在几个值得关注的弱点：首先，answer-forcing机制的鲁棒性不足，消融实验显示100%注入率会导致模型崩溃（MathVista仅4.1%），即使10%的注入率也可能引入噪声信号，因为强制注入的样本即使推理轨迹质量不佳也会获得高正确性奖励，可能产生误导性的学习信号。改进方向可以探索自适应注入策略，根据模型当前能力和任务难度动态调整注入率。其次，tree search的有效组大小扩展有限，从16×2到16×4的提升几乎可以忽略（3.90 vs 3.87），说明简单增加分支可能遇到收益递减问题，需要探索更智能的分支策略，如基于不确定性或多样性的分支选择。第三，互补掩码似然估计器虽然解决了梯度不平衡问题，但使用两个固定掩码样本的估计方差可能仍然较高，可以探索重要性采样或控制变量等方差 reduction 技术。第四，统一框架中$\gamma=0.5$的权重选择是基于经验的，不同任务可能需要不同的平衡策略，可以引入任务自适应的权重调整机制。第五，对于图像编辑任务，RL训练的提升相对较小（+0.19），可能是因为EditScore奖励模型本身的局限性，需要探索更精细的奖励设计。
未来方向
基于LaViDa-R1的成果，未来研究可以在多个方向展开：首先，模型规模扩展方面，将统一后训练框架应用到更大规模的dLLM上（如100B+参数），探索规模效应如何影响推理能力和多模态性能。其次，任务扩展方面，将框架扩展到更多样的任务，如视频理解与生成、3D场景理解、多轮对话推理等，验证统一后训练的通用性。第三，训练效率优化方面，探索更高效的RL算法，如减少采样次数的方法、课程学习策略、或基于模型能力的自适应难度调整。第四，奖励模型改进方面，开发更精细的多模态奖励模型，特别是对于图像编辑任务，可以结合人类偏好学习来设计更好的奖励信号。第五，推理能力深化方面，探索更复杂的推理模式，如多步推理、自我纠正、工具使用等，进一步提升模型在复杂任务上的表现。第六，部署优化方面，研究如何将推理增强的dLLM高效部署到实际应用中，包括推理加速、模型压缩、边缘设备适配等。最后，理论分析方面，深入研究统一后训练框架的理论基础，理解SFT正则化与KL正则化的本质区别，以及自蒸馏目标如何影响模型的探索-利用平衡。
复现评估
论文的复现性存在一定挑战。首先，基础模型LaViDa-O的开源情况需要确认，如果基础模型不可获取，则完整复现LaViDa-R1的训练过程将非常困难。其次，训练数据方面，论文使用了多个数据集进行多任务训练，包括数学推理、VQA、目标定位和图像编辑数据，这些数据的获取和预处理需要大量工作。第三，计算资源方面，统一后训练涉及在线采样、奖励计算和模型更新，需要较大的GPU内存和计算时间，特别是tree search和多步扩散过程的计算开销较高。第四，超参数方面，论文涉及多个关键超参数（answer-forcing注入率、tree search步骤、自蒸馏权重$\gamma$、扩散步数等），需要仔细调优才能达到报告的性能。第五，奖励模型方面，图像编辑任务使用了EditScore奖励模型，该模型的可用性和版本需要确认。总体而言，如果基础模型和主要训练数据开源，结合论文提供的详细方法描述，中等规模的研究团队应该能够在合理时间内复现核心结果，但完整的多任务训练可能需要较大的计算投入。
论文图表
