Discrete Tilt Matching
Abstract
Masked diffusion large language models (dLLMs) are a promising alternative to autoregressive generation. While reinforcement learning (RL) algorithms have been adapted to be compatible with dLLMs for fine-tuning them, their reliance on the computation of the marginal likelihood to evaluate policy objectives is intractable. To overcome this, we exploit a dynamical relation between the unmasking posterior of the base model and that which targets the reward-tilted distribution to derive Discrete Tilt Matching (DTM), an algorithm that avoids intractable likelihood evaluation entirely. DTM can be phrased as a cross-entropy loss that only requires forward evaluation of rewards and whose variance can be adaptively controlled, improving training stability. We motivate DTM on maze planning tasks, and show that fine-tuning LLaDA-8B-Instruct with DTM achieves higher accuracy at lower compute costs than prior RL-based fine-tuning methods across the Sudoku, Countdown, and MATH500 benchmarks.
Lay Summary
Most large language models write answers one word after another. Masked diffusion language models work differently: they start with masks and gradually unveil them into clean words, sometimes in many places at once. This flexibility could be useful for tasks that require planning or structure, such as puzzles. But it also makes these models harder to improve using rewards, because the same final answer can be reached through many different revealing orders. Scoring the whole answer in the usual way becomes unreliable or too expensive. We introduce Discrete Tilt Matching (DTM), a method for improving these models without needing that hard whole-answer score. Instead of judging only the finished text, DTM teaches the model to make better small decisions along the way: which missing piece to fill next and what to put there, guided by rewards for good final answers. We also add simple tools that make training more stable. In experiments, DTM improves an 8-billion-parameter diffusion language model on structured reasoning tasks such as Sudoku and Countdown, while staying competitive on standard math tests. This helps make diffusion-style language models more practical to train toward useful goals.