PaperlayerView exact revision on arXiv

Source-linked AI summary

LaViDa-R1: Advancing Reasoning for Unified Multimodal Diffusion Language Models

Shufan Li, Yuchen Zhu, Jiuxiang Gu, Kangning Liu, Zhe Lin, Yongxin Chen, Molei Tao, Aditya Grover, Jason Kuen

arXiv:2602.14147v2cs.CV

TL;DR

Reasoning dLLMs have lacked broad multimodal task coverage and faced training instability. LaViDa-R1 unifies SFT and multi-task RL with answer-forcing, tree search, and complementary likelihood estimation, achieving strong performance across multimodal benchmarks. Its reported gains include a +0.10 ImgEdit boost, while remaining behind state-of-the-art autoregressive multimodal models.

  • Problem

    Existing reasoning dLLMs focus on limited tasks and dataset-specific finetuning, while RL can collapse, restrict exploration, and lose training signal on difficult prompts.

  • Method

    LaViDa-R1 unifies diverse multimodal tasks through post-training that integrates SFT and multi-task RL with answer-forcing, tree search, and complementary likelihood estimation.

  • Results

    LaViDa-R1 achieves strong reasoning performance across multiple benchmarks, including MathVerse, ChartQA, Lisa-Grounding, and ImgEdit, with a +0.10 ImgEdit boost.

  • Takeaways & Limitations

    The unified framework supports reasoning across visual math, visual understanding, grounding, language-only tasks, and image editing.

  • Takeaways & Limitations

    A considerable performance gap remains between multimodal dLLMs and state-of-the-art autoregressive multimodal models such as Qwen3-VL.

Abstract

from arXiv · show

Diffusion language models (dLLMs) recently emerged as a promising alternative to auto-regressive LLMs. The latest works further extended it to multimodal understanding and generation tasks. In this work, we propose LaViDa-R1, a multimodal, general-purpose reasoning dLLM. Unlike existing works that build reasoning dLLMs through task-specific reinforcement learning, LaViDa-R1 incorporates diverse multimodal understanding and generation tasks in a unified manner. In particular, LaViDa-R1 is built with a novel unified post-training framework that seamlessly integrates supervised finetuning (SFT) and multi-task reinforcement learning (RL). It employs several novel training techniques, including answer-forcing, tree search, and complementary likelihood estimation, to enhance effectiveness and scalability. Extensive experiments demonstrate LaViDa-R1's strong performance on a wide range of multimodal tasks, including visual math reasoning, reason-intensive grounding, and image editing.

1. Introduction

LaViDa-R1 addresses the limited task coverage and training challenges of reasoning dLLMs with a unified multimodal post-training framework. It combines SFT, RL, and guided generation techniques, achieving strong results across diverse benchmarks.

  • Motivation: Reasoning dLLM research has largely focused on limited tasks and dataset-specific finetuning, leaving general-purpose multimodal reasoning underexplored.The identified gap includes tasks such as image editing and reason-intensive object grounding, alongside RL collapse and weak training signals for difficult prompts.
  • Approach: LaViDa-R1 unifies visual and language tasks, including mathematical reasoning, visual question answering, reason-intensive grounding, and image editing.
  • Approach: Its post-training objective integrates supervised finetuning and reinforcement learning while replacing KL regularization with SFT regularization.The stated design aims to support exploration beyond the pretrained model distribution while preventing collapse.
  • Approach: Answer-forcing and tree search generate higher-quality training samples when ground-truth answers are available or unavailable, respectively.Complementary likelihood estimation addresses missing-signal and imbalanced-gradient problems in dLLM likelihood estimation.
  • Results: LaViDa-R1 achieves strong reasoning performance across MathVerse, ChartQA, Lisa-Grounding, and ImgEdit benchmarks.

2. Background and Related Works

The background presents discrete diffusion models, reinforcement-learning approaches for reasoning, and alternative alignment methods. LaViDa-R1 extends these foundations toward broader multimodal tasks by combining RL with best-of-N distillation.

  • Discrete diffusion models: Discrete diffusion language models progressively unmask tokens and use bidirectional context, offering an alternative to autoregressive generation.Inference begins from a fully masked sequence and iteratively applies the learned reverse process until a clean sequence is obtained.
  • Reinforcement learning: GRPO-style reinforcement learning samples response groups, computes per-sample rewards and advantages, and optimizes the policy using these signals.The cited formulation includes a reverse-KL estimator and normalized reward-based advantages.
  • Reinforcement learning: A unified objective can represent SFT, RL, and DPO by changing the coefficient β and the definition of the advantage term.
  • RL for multimodal tasks: Existing dLLM reinforcement-learning work has mostly targeted language-only tasks with task-specific training, while multimodal coverage remains limited.LaViDa-R1 extends RL to reason-intensive object grounding and image editing.
  • Non-online-RL methods: Best-of-N distillation minimizes the divergence between the model policy and the best response distribution, and LaViDa-R1 combines this objective with standard RL.

3. Method

LaViDa-R1 unifies supervised finetuning, reinforcement learning, self-distillation, guided rollout generation, and complementary masking within one post-training framework for multimodal diffusion models.

  • 3.1. Unified Post-training: The unified policy-gradient formulation combines objectives by concatenating prompt-response pairs with their corresponding sample weights.Pairs may come from offline datasets or online generation, after which the policy computes sequence log-likelihoods for optimization.
  • 3.1. Unified Post-training: LaViDa-R1 combines SFT, online GRPO, and online self-distillation, using SFT regularization instead of KL regularization.The final design sets γ = 0.5; SFT supports exploration while preventing collapse and removes the need for a reference model.
  • 3.2. Guided Rollout Generation: Guided generation addresses vanishing RL signals by producing high-quality rollouts when all sampled responses receive low rewards.Answer-forcing applies when ground-truth answers are available, while tree search targets tasks such as image editing without ground-truth answers.
  • 3.2.1. ANSWER FORCING: Answer-forcing injects a ground-truth answer into a partially masked sequence and uses diffusion unmasking to generate an intermediate reasoning trace.The synthetic trace is conditioned on the inserted answer and then added to the training group.
  • 3.2.2. TREE SEARCH: Tree search selects the highest-reward sample, resumes generation from its early diffusion state, and repeats this process to expand the effective group size to N^k.The method generates N new samples from a partially masked state rather than fully noised sequences at each iteration.
  • 3.3. Complementary-Masking Likelihood Estimator: The complementary-masking estimator uses two complementary masked views to account for every token while reducing the training-inference gap.It uses t1 sampled uniformly from [0, 1] and t2 = 1−t, contrasting with independent masking or always-masking estimators.

4. Experiments

Experiments evaluate LaViDa-R1 across visual reasoning, VQA, language-only understanding, image editing, and reason-intensive grounding. The results show broad improvements, including substantial gains from unified post-training on image editing and grounding.

  • Experimental Setup: The two-stage training procedure first applies SFT to reasoning data, then unified post-training mixes SFT and RL data with SFT, RL, and self-distillation losses.RL data cover math reasoning, visual question answering, reason-intensive object grounding, and image editing.
  • Image Understanding Results: LaViDa-R1 is evaluated across visual reasoning, VQA, and language-only benchmarks using accuracy, with improvements reported across all tasks.The largest gains are observed on the language-only GSM8K and Math500 datasets.
  • Image Understanding Results: LaViDa-R1 shows consistent improvements on additional multimodal benchmarks that are less reason-intensive, including MMMU, MMBench, and MME.
  • Image Editing Results: +0.10 is the ImgEdit improvement from LaViDa-R1, compared with +0.01 from SFT, whose gain is described as negligible.ImgEdit scores measure visual quality and prompt compliance through a GPT-4 judge model.
  • Reason-Intensive Grounding: +22.1 mIoU is the additional gain from unified post-training over SFT on the Lisa-Grounding benchmark.SFT first improves performance by +10.8 mIoU over the LaViDa-O baseline; the benchmark reports P@0.5 and mIoU.

5. Ablation Studies

The ablations examine answer forcing, tree search, likelihood estimation, self-distillation, and the unified loss. Results indicate that guided sampling and SFT-regularized objectives improve stability or performance under the tested settings.

  • Answer Forcing: 10% answer forcing has the best overall performance, whereas high injection ratios can cause collapse from misleading correctness rewards.Answer-forced samples may receive high correctness rewards despite ill-formed reasoning traces, producing misleading learning signals.
  • Tree Search: Steps [0,8] are a good tree-search choice, while [0,8,16,32] yields almost identical performance because later branching adds little.Later diffusion steps introduce less uncertainty and therefore contribute less to the search.
  • Likelihood Estimator: The likelihood-estimator ablation compares random masking, full masking, two i.i.d. samples, and the proposed estimator.The supplied passages identify the compared setups but do not report their individual numerical results.
  • Self-Distillation Loss: Combining the self-distillation and other loss functions yields better performance than using standard GRPO alone.The combined loss assigns greater importance to the best-generated samples and supports the flexibility of the unified paradigm.
  • Unified Loss: The unified loss with SFT regularization is more stable and yields higher reward than standard online GRPO, which diverges even with β = 0.1 KL regularization.The comparison is based on average reward per sample during training.

A.2. Unified paradigm for post-training

The appendix frames several post-training objectives as weighted policy-gradient objectives and describes guided rollout procedures for answer-forcing and tree search. This unification supports combining objectives while tailoring sampling to reward availability.

  • Unified Objective: GRPO, Online DPO, SFT, and Best-of-N self-distillation share gradients with weighted policy-gradient objectives under suitable advantage weights.The appendix provides derivations showing how objective-specific gradients can be represented in a common form.
  • SFT and Self-Distillation: SFT assigns constant advantage 1 to all samples, while Best-of-N self-distillation assigns advantage 1 only to the highest-reward sequence.Best-of-N therefore corresponds to SFT on the self-generated best sequence.
  • Answer Forcing: Answer forcing adds a ground-truth-answer sample for tasks with verifiable rewards such as mathematical correctness or object-grounding IoU.The procedure initializes an answer-prefilled masked sample and evaluates it alongside policy-generated samples.
  • Answer Forcing: The implementation generates N + 1 samples in parallel, evaluates rewards, and conditionally retains the extra answer-forced sample based on a threshold.The threshold is applied to the maximum reward among the non-answer-forced samples.
  • Tree Search: Tree search branches from saved trajectories at selected restart timesteps, using the best-reward trajectory to initialize additional samples.The method is designed for real-valued rewards and is not applicable to 0-1 rewards when all rewards are zero.

A.5. Likelihood Estimator

The likelihood estimator uses complementary masking to cover tokens while balancing their importance. The design modifies vanilla complementary masking and is evaluated alongside established estimator variants.

  • Compared Estimators: The appendix compares d1, UniGRPO, two-sample UniGRPO, vanilla complementary masking, and the proposed estimator.Figure 6 visualizes the different likelihood-estimation choices.
  • Compared Estimators: d1 uses a fully masked sequence and MC = 1, whereas UniGRPO samples a mask ratio uniformly and computes likelihood only over masked tokens.The two approaches differ in token coverage and masking strategy.
  • Proposed Estimator: The proposed estimator uses two complementary mask samples to achieve all-token coverage while balancing token importance uniformly.Complementary masking couples the two samples so masked positions collectively cover the sequence.
  • Proposed Estimator: The estimator replaces the vanilla timestep weighting with w(t) = 1, a modification reported to have a profound impact on model performance.The paper attributes this design choice to the proposed likelihood-estimation recipe.
  • Training Context: The training setup uses two stages: SFT alone in Stage 1, followed by unified online-sampling and offline-data training in Stage 2.Stage 2 retains Stage 1 datasets and adds reinforcement-learning datasets spanning reasoning and multimodal tasks.

B.2. Additional Ablation Studies

The appendix reports additional ablation studies to verify LaViDa-R1’s design choices.

  • Additional ablation studies are reported to verify the design of LaViDa-R1.

B.3. Timestep Weighting

The timestep-weighting ablation compares complementary masking with w = 1 and w = 1/t, finding that w = 1 performs better, especially for image editing.

  • w = 1 works better than w = 1/t under complementary masking.The comparison is reported in Table 7 and identified as a timestep-weighting ablation.
  • w = 1/t significantly degrades performance on image-editing tasks.The passage attributes this to visual-generation samples containing many more image than text tokens, which exacerbates per-token loss imbalance.
  • Image-editing samples contain 256 text tokens and 4096 image tokens, compared with 512 tokens for understanding-only tasks.This token-count difference is given as a possible explanation for the weighting behavior.

B.5. Qualitative Results

The paper evaluates additional loss combinations and presents qualitative outputs showing LaViDa-R1's reasoning capabilities across object grounding, image editing, and visual math tasks.

  • Adding online DPO, online DPO-smooth, or online SLiC to the GRPO+SFT baseline performs worse.These alternatives can be introduced by modifying the weight Ai in Equation 5 without substantially changing the training pipeline.
  • Qualitative results cover object grounding, image editing, and visual math reasoning.The outputs are shown in Figures 8, 9, and 10.
  • The qualitative results demonstrate strong reasoning capabilities across diverse multimodal tasks.The paper uses these outputs to support the model's behavior across the three evaluated task types.

C. Limitation

The paper identifies limitations in reasoning quality, inference efficiency, task coverage, and reward modeling for text-to-image reasoning.

  • LaViDa-R1 remains behind state-of-the-art autoregressive MLLMs in reasoning performance.The authors specifically identify a considerable gap relative to Qwen3-VL despite improvement over LaViDa-O.
  • dLLM-specific inference optimization is incomplete, creating a throughput bottleneck during online sampling.The training process uses eager Python execution because low-level frameworks such as vLLM have not been fully adapted to dLLMs.
  • More tasks remain to be explored despite the expanded scope of multimodal dLLM reinforcement learning.The authors leave further unification of task coverage as an open area.
  • Current reward models cannot reliably distinguish reasoning-grounded text-to-image generations.The authors report that VLM-based models may rely on surface-level concepts rather than the compositional and historical reasoning required by prompts.
  • Effective and scalable reward modeling for reasoning-driven text-to-image generation remains an open challenge.The paper leaves reinforcement-learning fine-tuning for text-to-image reasoning to future work.

D. Additional Discussions for Related Works

The discussion situates LaViDa-R1 among dLLM reinforcement learning and unified SFT-RL work, while distinguishing answer forcing from related inpainting-based reasoning methods.

  • Reinforcement Learning with dLLMs: Prior dLLM reinforcement-learning work includes d1, wd1, and DMPO, which develop alternative probability estimation or policy distribution-matching approaches.These methods address reinforcement learning for dLLMs through different advantage or likelihood-estimation designs.
  • Unified SFT and RL: LaViDa-R1 applies unified SFT-RL to unified multimodal models, whereas prior unified strategies largely focused on text-only LLMs.Its loss formulation is intended to stabilize reinforcement learning rather than make single-stage SFT or RL obsolete.
  • Connection between Answer Forcing and IGPO: Answer Forcing uses dLLM inpainting while requiring only the final answer, unlike IGPO, which requires high-quality ground-truth reasoning traces.The distinction concerns the supervision needed during online sampling.
  • The paper reports qualitative results for reason-intensive object grounding, reason-based image editing, and visual math problems.These results correspond to Figures 8, 9, and 10.
  • Text-to-Image Generation: The text-to-image example asks for an image representing the light source that replaced candles in early twentieth-century homes.The prompt is illustrated in Figure 11 and appears as a generation input.
Loading 2602.14147v2…