Trajectory Alignment of Diffusion Language Models

Diffusion based sequence models are increasingly being used as fast data generators for many applications. Once pre-trained, reinforcement learning (RL) based methods are then used to align these models with a suitable reward function. However applying RL based optimization to these models is made difficult by one of two problems: a) propagating the signal through the denoising process and b) lack of efficient likelihood for diffusion models. We solve this by proposing a new framework for aligning the diffusion model output with a given reward function. Our method maps the standard KL-regularized RL optimization to a trajectory-level Bellman consistency problem over the reverse diffusion process. We show that the corresponding objective has a simplified form which does not need a learned value model. The trajectory Bellman identity leaves only a prompt-dependent normalizer. We handle this term with two grouped objectives: a soft reward normalizer and a centered objective with GRPO-style reward. Experiments across multiple domains evaluate these variants for reward alignment of discrete diffusion models.