Timestep placement boosts quality and speed in diffusion language models
A Step in the Right Place: Understanding and Optimizing Timestep Placement in Continuous Diffusion Language Models
Computation and Language
Summary
Generating text quickly with continuous diffusion language models often requires a careful balance between quality and speed. The authors found that changing where the model takes steps during the text generation process can improve the final output. They created a method called StepTune that finds better step placements, making the model produce good text using fewer steps. This means the model can generate text faster without losing quality, as shown in tasks like translation and summarization.
What this means in practice
- •For machine translation teams: Improve translation speed and quality by optimizing timesteps in diffusion-based language models with fixed model parameters.
- •For automatic summarization developers: Reduce latency and maintain summary quality using calibrated step placement for faster conditional text generation.
Authors
Shutian Zheng, Zhihong Yu, Linye Wei, Zeyang Zhang, Meng Li
Abstract
Continuous diffusion language models represent an emerging paradigm for language modeling. They generate text in parallel by iteratively refining continuous representations and mapping them back to tokens through a decoder. However, realizing the efficiency advantage of parallel generation requires maintaining generation quality with only a few denoising steps. Even when the number of steps is fixed, their placement along the denoising trajectory remains flexible. In this paper, we demonstrate that changing timestep placement alone can substantially alter the final latent representations and thereby affect the decoded token distributions. Building on these findings, we introduce StepTune, an algorithm that searches for timestep placements that improve generation quality. The search is conducted on a small calibration set and guided by a language-aware objective, with the model and sampling budget fixed. Once calibrated, the timestep placement is fixed and reused across inputs. Experiments demonstrate consistent improvements in unconditional generation perplexity and conditional task metrics across all evaluated sampling budgets. In particular, StepTune enables 16-step sampling with a frozen model to outperform its native 64-step baseline on both XSum summarization and WMT14 translation and reduces end-to-end latency by approximately $3.8\times$. These results establish timestep placement as a practical means of improving the quality-efficiency trade-off in continuous diffusion language models.