Anchors improve control and reasoning in diffusion language models
Anchors Enhance Controllability and Reasoning of Diffusion Language Models
Machine Learning
Summary
Diffusion language models can generate text while following certain rules, but controlling them well is hard. The authors designed a way to create and place small 'anchor' pieces of text in the right spots, which helps guide the model to produce better and more accurate answers. They teach a special system to pick good anchors using feedback about how well the final text meets the rules. This method helps the models do better on programming and math problems.
What this means in practice
- •For software developers: Generate code snippets that follow exact instructions with improved accuracy by guiding diffusion language models using anchor tokens.
- •For data scientists: Create precise mathematical proofs and reasoning outputs by controlling language models with learned anchors to ensure correct answers and constraints.
- •For virtual assistant developers: Improve response accuracy in assistants by injecting control tokens that steer text generation toward user constraints and correctness.$Commercial implications: Enables more reliable AI assistants that meet user demands with higher accuracy, forming a basis for commercial assistant products.
Authors
Feiyu Qu, Zheng Huang, Andrea Zanette
Abstract
Diffusion language models (DLMs) offer new opportunities for controllable text generation. However, their controllability, defined as the ability to satisfy task constraints while preserving their original capabilities, remains underexplored. DLMs' bidirectional attention allows control tokens to guide generation from arbitrary positions in the output sequence, but effective guidance depends on the compatibility between their content and placement. To this end, we propose AnchorDLM, a framework in which an anchor generator produces an anchor consisting of a short output fragment and its position for each query. The anchor is placed as control tokens at the specified position to initialize the output sequence of a frozen DLM and guide subsequent denoising. We train the anchor generator with reinforcement learning (RL) to jointly optimize anchor content and placement, using rewards computed from the DLM's generated outputs based on answer correctness and constraint satisfaction. Moreover, through the lens of discrete flow matching (DFM), we interpret anchor placement as helping to reshape the initial velocity field that guides generation. Experiments on code generation and mathematical reasoning across five DLM backbones show that AnchorDLM improves controllability and further yields gains in reasoning performance across diverse benchmarks.