Diffusion language models improved to better follow specific text instructions

In-Place Instruction Following in Diffusion Language Models

Computation and LanguageArtificial Intelligence

Summary

When computers write text, it's important that they follow instructions exactly where given. The authors studied a way to make diffusion language models better at following these in-place instructions, which tell the model what to do at specific spots in the text. They created a new test called IIF-Bench to check how well the models follow different types of instructions. They then developed a method called GRAFT that helps the models pay more attention to the instructions during writing, improving their accuracy without hurting their overall writing skills.

Diffusion language modelsIn-place promptingInstruction followingIterative denoisingAttention mechanismFine-tuningPreference optimizationBenchmark evaluationConstraint satisfaction

Authors

Zheng Nie, Zherui Li, Jiaming Zhang, Kun Wang, Zhenhong Zhou, Yufei Guo

Abstract

Diffusion Large Language Models (dLLMs) generate text via bidirectional iterative denoising, naturally supporting user-specified constraints anchored at arbitrary output positions, a paradigm known as In-place Prompting (IPP). We formalize this as the In-place Instruction Following (IIF) task and construct IIF-Bench, a hierarchical benchmark spanning literal, style, and discourse-function constraints, paired with a rubric-based local-global evaluation protocol. An inference-time attention-bias probe suggests that vanilla dLLMs often under-prioritize constraint spans during denoising. We then propose GRAFT, an IPP-oriented post-training framework combining constraint-aware SFT and preference optimization. On four representative dLLMs, GRAFT raises the average IIF score from 57.75 to 73.10 (+15.35 points), with absolute gains of 15.91 and 15.57 points on literal and discourse-function constraints, while preserving general generation ability.