Learnability-Informed Fine-Tuning of Diffusion Language Models

Abstract: We aim to improve the reasoning capabilities of diffusion LLMs (DLMs). While SFT is a popular post-training recipe for autoregressive models, its use in DLMs faces challenges and can even hurt performance, though the underlying causes remain understudied. Our analysis reveals that vanilla SFT overlooks learnability, namely what and when tokens are learned. Specifically, rare tokens are difficult to learn when most of the input is masked, whereas it is straightforward and thus of little value to learn common tokens when most of the input is unmasked. Motivated by our analysis, we propose LIFT, an efficient SFT-based post-training algorithm for DLMs. LIFT learns easy tokens when most of the input is masked and hard tokens when more context is available, thus aligning the training with the information available at different diffusion time steps. Our results show that LIFT outperforms existing SFT baselines across six reasoning benchmarks, achieving up to a 3x relative gain on AIME'24 and AIME'25. Our code is publicly available at https://github.com/divelab/LIFT.

• The paper introduces LIFT, a learnability-informed fine-tuning method that dynamically selects training tokens based on frequency and diffusion timestep.
• It achieves up to a 3x improvement in pass@16 accuracy on reasoning tasks while offering compute-efficient post-training, outperforming RL-based baselines with 1000x lower GPU-hours.
• Empirical studies show that jointly modeling token frequency and timing enhances generalization across datasets and architectures in Diffusion Language Models.

Learnability-Informed Fine-Tuning for Diffusion LLMs

Introduction

The emergence of Diffusion LLMs (DLMs) has shifted the paradigm for large-scale language modeling from autoregressive decoding to parallel token generation, drastically improving inference throughput. While supervised fine-tuning (SFT) has been foundational in boosting reasoning and performance for autoregressive LLMs (ARLMs), its application to DLMs has remained underexplored and oftentimes detrimental. In "Learnability-Informed Fine-Tuning of Diffusion LLMs" (2605.22939), the authors rigorously analyze the underpinnings of SFT dynamics in DLMs and introduce LIFT, a principled, learnability-informed post-training algorithm targeting reasoning tasks.

Motivation & Analysis of SFT Challenges in DLMs

The paper identifies two orthogonal but interacting dimensions governing effective SFT in DLMs:

• What tokens are learned? Empirical analysis across 0.5B tokens aggregated from various reasoning datasets reveals that rare, low-frequency tokens are significantly harder for DLMs to predict, especially when masked.
• When are tokens learned? The difficulty of predicting tokens varies with the diffusion timestep. When a large fraction of input tokens is masked (late timesteps), model confidence on rare tokens collapses due to insufficient context, rendering them effectively unlearnable. Conversely, frequent tokens remain trivial to predict regardless of context.

Prior SFT baselines either attempt to address token frequency or diffusion timestep in isolation but fail to exploit their joint effect. The authors demonstrate, through detailed statistical analysis, that the interaction between "what" and "when" is critical to maximize the learning signal and optimize the loss landscape for DLMs.

The LIFT Algorithm: Learnability-Informed Fine-Tuning

LIFT operationalizes the above insights into an adaptive SFT framework that dynamically selects training targets based on token difficulty and available context. The methodology works as follows:

• At each training step, a secondary masking ratio is sampled to create a more corrupted input.
• Token-wise confidence scores under the model are computed (i.e., the likelihood assigned to the ground truth token given the masked input).
• Supervision is restricted to a subset of masked tokens, dynamically partitioned along the diffusion timestep:
  • Late timesteps (high masking): Train on easy, frequent tokens (top-K by confidence).
  • Early timesteps (low masking): Train on difficult, rare tokens (bottom-K by confidence).
  • Intermediate timesteps: Standard random SFT masking.
• This selection schedule is governed by a hyperparameter $H \geq 2$, controlling the transition between regimes.

A compute-efficient variant, LIFT-A, eliminates one forward pass by approximating the loss at the secondary timestep, trading some accuracy for substantial efficiency.

Numerical Results and Empirical Claims

LIFT demonstrates robust and consistent gains across a suite of reasoning tasks (GSM8K, MATH, Countdown, Sudoku, AIME'24, AIME'25) on LLaDA-8B-Instruct and LLaDA-1.5 models:

• AIME'24/AIME'25: LIFT3 achieves up to a 3x improvement in pass@16 accuracy compared to vanilla SFT, notably outperforming all baseline methods (SFT, GIFT, CART).
• Pareto Frontier: LIFT matches or surpasses RL-based approaches such as d1 [Zhao et al., 2025] at roughly 1000x lower GPU-hours, formally establishing a new compute-efficient Pareto frontier for DLM post-training.
• Robustness: Performance gains generalize across various fine-tuning datasets, generation lengths, and extend to other DLM architectures such as Dream-7B.
• Ablations: Empirical studies confirm that time-independent variants (randomized token selection) or focusing solely on difficulty without diffusion timestep scheduling perform inconsistently and degrade on complex benchmarks. Only the joint modeling of "what" and "when" delivers uniformly strong results.

Practical & Theoretical Implications

The implications of this work are multifaceted:

• Efficiency of Post-Training: LIFT radically reduces the compute burden for high-fidelity post-training, granting DLMs reasoning capacities comparable to RL-based methods but at orders of magnitude less cost.
• Token Curriculum for DLMs: The learnability-informed masking schedule constitutes an adaptive, token-level curriculum. Training focuses on tokens when they are most learnable, maximizing convergence rates and preventing wastage of learning iterations on effectively unlearnable tokens.
• Generalization: The approach is robust across model architectures, fine-tuning datasets, task domains (math reasoning, code generation), and diverse generation configurations, supporting broad adoption.
• Open Questions: Efficient RL algorithms for DLMs remain a challenge due to sequence-level likelihoods. LIFT closes much of the performance gap, possibly diminishing the practical relevance of RL for DLM post-training under current technical limitations.

Future Prospects

Potential future research directions include:

• Integration of curriculum learning and adaptive sampling strategies with LIFT, leveraging both token-level confidence and sequence-level curriculum advancements.
• Extensions of LIFT to multi-modal DLMs, further scaling reasoning abilities in broader generative domains.
• Theoretical analysis of the convergence properties and generalization bounds of learnability-informed fine-tuning schedules.

Conclusion

The learnability-informed fine-tuning paradigm embodied by LIFT systematically addresses the dual dimensions of token difficulty and context availability in DLMs. Through thorough empirical analysis and algorithmic innovation, the paper demonstrates that joint modeling of "what" and "when" during SFT unlocks significant performance and efficiency gains for reasoning and generative tasks, setting new standards for post-training in diffusion language modeling (2605.22939).