PaperlayerView exact revision on arXiv

Source-linked AI summary

LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Models

Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, Chongxuan Li

arXiv:2505.19223v2cs.LG

TL;DR

MDM preference alignment is hindered by bias and variance from ELBO-based likelihood estimation, while broader alignment tasks remain underexplored. The paper introduces VRPO, which combines theoretically grounded, unbiased variance-reduction strategies, and applies it to LLaDA to produce LLaDA 1.5, which consistently outperforms LLaDA across mathematics, coding, and alignment benchmarks.

  • Problem

    MDM alignment remains relatively underexplored, and replacing intractable DPO log-likelihoods with ELBO estimates introduces bias and variance into preference optimization.

  • Method

    VRPO bounds preference-optimization errors through score-estimator variance and reduces that variance using larger sampling budgets, optimal timestep allocation, and antithetic sampling.

  • Results

    LLaDA 1.5 consistently outperforms LLaDA across mathematics, coding, and alignment benchmarks and remains highly competitive mathematically against strong MDMs and ARMs.

  • Takeaways & Limitations

    VRPO provides a theoretically analyzed and empirically validated approach for improving large-scale MDM alignment, with stronger general capabilities demonstrated by LLaDA 1.5.

  • Takeaways & Limitations

    The models may still generate discriminatory, biased, or otherwise harmful content despite curated and filtered preference data and ongoing safeguards.

Abstract

from arXiv · show

While Masked Diffusion Models (MDMs), such as LLaDA, present a promising paradigm for language modeling, there has been relatively little effort in aligning these models with human preferences via reinforcement learning. The challenge primarily arises from the high variance in Evidence Lower Bound (ELBO)-based likelihood estimates required for preference optimization. To address this issue, we propose Variance-Reduced Preference Optimization (VRPO), a framework that formally analyzes the variance of ELBO estimators and derives bounds on both the bias and variance of preference optimization gradients. Building on this theoretical foundation, we introduce unbiased variance reduction strategies, including optimal Monte Carlo budget allocation and antithetic sampling, that significantly improve the performance of MDM alignment. We demonstrate the effectiveness of VRPO by applying it to LLaDA, and the resulting model, LLaDA 1.5, outperforms its SFT-only predecessor consistently and significantly across mathematical (GSM8K +4.7), code (HumanEval +3.0, MBPP +1.8), and alignment benchmarks (IFEval +4.0, Arena-Hard +4.3). Furthermore, LLaDA 1.5 demonstrates a highly competitive mathematical performance compared to strong language MDMs and ARMs. Project page: https://ml-gsai.github.io/LLaDA-1.5-Demo/.

1 INTRODUCTION

MDMs have made strong language-modeling progress, but aligning them with human preferences remains underexplored because ELBO-based likelihood estimates introduce bias and variance. VRPO addresses this challenge with unbiased variance reduction and improves LLaDA across broad benchmarks.

  • MDMs have achieved significant language-modeling progress, including performance comparable or superior to ARMs at small scale.
  • Preference alignment for MDMs remains underexplored, especially beyond specialized reasoning and code-generation tasks.
  • Exact DPO log-likelihoods are intractable for diffusion models, so ELBO substitution creates a preference score from four ELBO terms with nested expectations.
  • ELBO estimation adds bias and variance to the preference-optimization loss and gradient, with both governed by the preference-score estimator's variance.
  • VRPO reduces score-estimator variance through larger ELBO sampling budgets, optimal timestep allocation, and antithetic sampling.
  • LLaDA 1.5 improves LLaDA consistently on mathematics, coding, and alignment tasks while maintaining highly competitive mathematical performance against MDMs and ARMs.The model was produced by applying VRPO to LLaDA 8B Instruct with 350k preference pairs.

2 PRELIMINARIES

This section introduces alignment methods, masked diffusion models, and the ELBO-based likelihood estimation used when exact MDM likelihoods are intractable. The resulting estimator is unbiased but can have substantial variance under computational limits.

  • Traditional alignment uses reward modeling followed by reinforcement learning, while DPO avoids explicit reward-model training.
  • Masked diffusion models progressively mask tokens in a forward process and learn to denoise them through a reverse process.
  • MDM exact log-likelihoods are often approximated with an ELBO involving the forward diffusion process and per-step mask-prediction loss.
  • The ELBO estimator averages mask-prediction losses over n = n_t × n_yt masked samples and is unbiased by linearity of expectations.
  • Computational constraints typically limit n, making estimator variance a central challenge for MDM preference optimization.
  • LLaDA is an 8B-parameter MDM pretrained on 2.3 trillion tokens and fine-tuned on 4.5 million SFT pairs.

3 METHOD

The method replaces intractable diffusion-model likelihoods in DPO with ELBO estimates, then analyzes how score-estimator variance creates loss bias and variance. VRPO reduces this variance through larger budgets, optimal allocation, and shared sampling, with extensions beyond DPO.

  • 3.1 SUBSTITUTING LIKELIHOODS WITH ELBOS IN DPO: ELBO substitution makes DPO feasible for MDMs but introduces nested expectations over diffusion time and masked data.The resulting preference score combines four ELBOs, each estimated through Monte Carlo sampling.
  • 3.1 SUBSTITUTING LIKELIHOODS WITH ELBOS IN DPO: Nonlinear log σ(·) transforms an unbiased score estimate into a loss estimate with additional bias, while stochastic sampling adds loss and gradient variance.The target is log σ(E[ŝθ]), whereas the estimated objective averages log σ(ŝθ).
  • 3.2 VARIANCE-REDUCED PREFERENCE OPTIMIZATION: Theorem 1 bounds empirical DPO loss bias and variance using the variance of the preference-score estimator.The analysis relies on the Lipschitz continuity of log σ(·) and unbiasedness of the score estimator.
  • 3.2 VARIANCE-REDUCED PREFERENCE OPTIMIZATION: VRPO increases ELBO sampling, allocates n samples across timesteps with nt = n and nyt = 1, and shares samples between policy and reference ELBOs.The latter two techniques reuse or redistribute samples without increasing the total sample count; sharing samples lowers score variance under the stated correlation condition.
  • 3.3 EXTENSION TO OTHER ALIGNMENT METHODS: The techniques reduce score-estimator variance without bias and extend to alignment methods involving ELBO estimation or differences between correlated ELBOs.The paper specifically discusses applicability beyond DPO, including likelihood and likelihood-ratio estimation in PPO- and GRPO-style objectives.
  • 3.2 VARIANCE-REDUCED PREFERENCE OPTIMIZATION: VRPO reduces score-estimator variance by decreasing individual ELBO variance and increasing correlation between paired ELBO estimates.These are the two variance-reduction principles underlying the proposed techniques.

4 EXPERIMENTS

Experiments evaluate VRPO on LLaDA across mathematics, coding, and alignment benchmarks, comparing it with LLaDA Instruct and naive DPO. Ablations show that variance reduction generally improves estimator stability, optimization, and task performance.

  • Experimental Setup: LLaDA 1.5 trains LLaDA 8B Instruct for one epoch on 350K preference pairs using VRPO.The evaluation covers mathematics, scientific reasoning, code generation, and alignment tasks.
  • Benchmark Results: LLaDA 1.5 with VRPO consistently outperforms LLaDA Instruct and naive DPO across all reported benchmarks.The benchmark comparison includes the fully variance-reduced configuration against both baselines.
  • Ablation Experiments: Lower preference-score estimator variance generally coincides with lower loss and gradient variance and better task performance.This empirical pattern supports controlling optimization errors through the score-estimator variance.
  • Ablation Experiments: Increasing the sampling budget from 1 to 8 reduces Vˆsθ from 44.0 to 1.0 and improves GSM8K accuracy from 80.1 to 83.3.The result validates the predicted benefit of increasing the ELBO sampling budget.
  • Ablation Experiments: Under fixed sampling budgets, optimal timestep allocation generally lowers variance and improves results compared with repeated masked samples per timestep.Antithetic sampling also notably decreases variance, although variance reductions do not always yield substantial downstream gains.

5 RELATED WORK

Related work situates this paper at the intersection of masked diffusion language modeling, MDM alignment, and Monte Carlo variance reduction. Prior alignment efforts adapt existing methods, while VRPO connects ELBO estimation challenges with principled variance-control techniques.

  • Masked Diffusion Models: Masked diffusion models build on discrete diffusion advances and can match autoregressive models in perplexity.Simplified masked-diffusion objectives have also been proposed for efficient training.
  • Alignment of MDMs: Recent MDM alignment studies include policy-gradient methods, continuous-time DPO variants, and other adaptations validated on specialized or small-scale tasks.These approaches represent emerging efforts to align diffusion language models.
  • Variance Reduction Techniques: VRPO relates to Monte Carlo variance reduction, doubly stochastic optimization, and variational inference.Its antithetic sampling adapts antithetic variates to couple correlated ELBO terms.

6 CONCLUSION

The paper analyzes bias and variance in ELBO-based preference optimization for MDMs and proposes VRPO as a theoretically grounded variance-reduction framework. Applied to LLaDA, VRPO yields LLaDA 1.5 with stronger general capabilities and supports further work on MDM alignment.

  • Conclusion: VRPO addresses high bias and variance in ELBO-based likelihood estimation for aligning MDMs with human preferences.The framework combines theoretical guarantees with empirical validation and offers insights beyond specific architectures or datasets.
  • Conclusion: LLaDA 1.5 demonstrates stronger general capabilities than LLaDA, especially in mathematics, coding, and alignment.Ablations further examine how VRPO components affect variance reduction, optimization stability, and efficiency.
  • Conclusion: The paper discusses extending its variance-reduction techniques to broader reinforcement-learning-based alignment algorithms.It presents the work as guidance for future research on MDM alignment and diffusion-based language models.

ETHICS STATEMENT

The paper aligns MDMs with human preferences to improve helpfulness while acknowledging that harmful outputs may persist. It describes data filtering and continued safeguard evaluation as mitigation measures.

  • The paper focuses on aligning MDMs with human preferences to improve helpfulness.
  • Models may still generate discriminatory, biased, or otherwise harmful content.
  • The authors filtered preference data to remove harmful material where feasible and will continue evaluating and refining safeguards.

B.2.3 TIGHTNESS ANALYSIS (THEOREM 2, THEOREM 3)

The tightness analysis connects preference-optimization error to the variance of the estimated preference score. It also shows that VRPO’s unbiased sampling strategies reduce ELBO and gradient-estimator variance under stated assumptions.

  • Theorem 2: Theorem 2 states that DPO loss bias scales proportionally to the square root of preference-score estimator variance under a first-order Taylor expansion.
  • Theorem 3: Theorem 3 states that DPO loss variance scales proportionally to preference-score estimator variance under a first-order Taylor expansion.
  • VRPO unbiasedness: The preference-score estimator remains unbiased under increased sampling budgets, optimal allocation, and antithetic sampling.
  • Variance reduction: Given n = n_t × n_yt masked samples, ELBO variance is Θ(1/n) and is minimized at n_t = n, n_yt = 1.
  • Variance reduction: Sharing Monte Carlo samples between model and reference policies yields lower preference-score variance when their correlation is positive.
  • Gradient analysis: Under bounded per-step loss gradients, the preference-score gradient is bounded and its optimization error can be analyzed through score and gradient variance.

C DETAILS OF FIGURE 2

Figure 2 uses Gaussian inputs with varying variance to examine how the nonlinear log-sigmoid transformation affects empirical variance and bias. Both quantities increase monotonically with input variance.

  • Figure 2 samples 1,000 points from zero-mean Gaussian distributions across ten input-variance levels from 0.1 to 1.0.
  • Panel (a): Panel (a) compares the Gaussian input density, the nonlinear log σ(x) function, the transformed-sample histogram, and the reference log σ(E[X]) value.
  • Panel (b): Panel (b) plots input variance against the empirical variance and absolute bias of log σ(X).
  • Panel (b): Both transformed variance and bias grow monotonically with V[X], supporting variance reduction as a way to mitigate both effects.

D.3 TRAINING

LLaDA 1.5 is trained with 350K preference pairs using a one-epoch VRPO setup. Variance reduction produces smoother, less variable training losses alongside faster decline and lower final loss.

  • 350K preference pairs are filtered, deduplicated, reward-ranked, and partly supplemented with state-of-the-art LLM responses.
  • Training runs for one epoch with batch size 64, AdamW, β = 0.2 DPO loss, and a 0.05-weighted MDM SFT loss.
  • Variance reduction strategies produce smoother training-loss trajectories with substantially lower variability.
  • Variance-reduced configurations also show faster loss decrease and lower final loss.

D.4 EVALUATION

The evaluation examines how sampling strategies and variance reduction affect LLaDA 1.5, including training stability, benchmark performance, and robustness across random seeds.

  • Variance reduction: Variance-reduction configurations produce smoother loss trajectories, faster loss decreases, and lower final losses than configurations without them.Figure 5 compares antithetic sampling, timestep budgets, and masked-sample budgets; curves are exponentially smoothed for visualization.
  • Inference details: Setting the |EOS| confidence score to zero improves LLaDA performance by reducing excessive |EOS| generation and truncated outputs.The issue is attributed to |EOS| padding during SFT.
  • Sampling strategies: LLaDA 1.5 is evaluated with diffusion, semi-autoregressive, and low-confidence-remasking strategies across representative benchmarks.The study reports best configurations after tuning answer and block lengths for the different sampling methods.
  • Training randomness: Across three VRPO runs, LLaDA 1.5 consistently outperforms LLaDA with higher mean scores and small standard deviations across benchmarks.For most tasks, its 95% confidence intervals lie entirely above the corresponding LLaDA means; judge-based benchmarks were omitted.

E CASE STUDY

The case study compares LLaDA 8B Instruct with LLaDA 1.5 on alignment and mathematical examples, using separate tables for each task category.

  • Alignment cases: Tables 7 and 8 present alignment cases comparing LLaDA 8B Instruct with LLaDA 1.5.
  • Math cases: Tables 9 and 10 present mathematical cases comparing LLaDA with LLaDA 1.5.
Loading 2505.19223v2…