Learnability-Informed Fine-Tuning of Diffusion Language Models

We aim to improve the reasoning capabilities of diffusion language models (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, thereby aligning 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.

Diffusion language models (DLMs) are a new kind of large language model that work like image generators such as Stable Diffusion, but for words. Instead of writing one word at a time the way ChatGPT does, DLMs start from a blank sentence and gradually fill it in, somewhat like solving a crossword. Teaching DLMs to reason, for instance to solve math problems, usually involves additional training on examples of good answers. This recipe works well for ChatGPT-style models but surprisingly fails for DLMs, sometimes even making them worse. We found the cause: the recipe ignores when each word can actually be learned. When most of the sentence is still blank, only common words like "the" can be guessed; rare or technical words are anyone's guess. When the sentence is nearly complete, predicting common words is too easy to teach the model anything new. Our method, LIFT, matches the training to this reality: it focuses on easy words early, when the sentence is mostly blank, and on hard words later, when there is enough context to make them learnable. Across six reasoning tests, LIFT beats existing methods, with up to a 3× improvement on a math competition benchmark.