Don't Let It Fade: Preserving Edits in Diffusion Language Models via Token Timestep Allocation

TL;DR

This paper introduces Token Timestep Allocation (TTA), a novel method that improves the controllability and fluency of diffusion language models by explicitly ordering token updates, thereby reducing update forgetting and enhancing refinement.

Contribution

The paper proposes TTA, a new inference-time technique for token ordering in diffusion language models, which significantly improves control and quality of generated text.

Findings

TTA increases sentiment control accuracy by over 20%.

TTA nearly halves perplexity compared to baseline methods.

TTA reduces toxicity levels in detoxification tasks.

Abstract

While diffusion language models (DLMs) enable fine-grained refinement, their practical controllability remains fragile. We identify and formally characterize a central failure mode called update forgetting, in which uniform and context agnostic updates induce token level fluctuations across timesteps, erasing earlier semantic edits and disrupting the cumulative refinement process, thereby degrading fluency and coherence. As this failure originates in uniform and context agnostic updates, effective control demands explicit token ordering. We propose Token Timestep Allocation (TTA), which realizes soft and semantic token ordering via per token timestep schedules: critical tokens are frozen early, while uncertain tokens receive continued refinement. This timestep based ordering can be instantiated as either a fixed policy or an adaptive policy driven by task signals, thereby supporting a broad spectrum of refinement strategies. Because it operates purely at inference time, it applies uniformly across various DLMs and naturally extends to diverse supervision sources. Empirically, TTA improves controllability and fluency: on sentiment control, it yields more than 20 percent higher accuracy and nearly halves perplexity using less than one fifth the steps; in detoxification, it lowers maximum toxicity (12.2 versus 14.5) and perplexity (26.0 versus 32.0). Together, these results demonstrate that softened ordering via timestep allocation is the critical lever for mitigating update forgetting and achieving stable and controllable diffusion text generation.