Research map / Image generation / Fast diffusion sampling
Don’t Let It Fade: Preserving Edits in Diffusion Language Models via Token Timestep Allocation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Token Timestep Allocation (TTA), an inference-time soft token-ordering method that assigns per-token timesteps (fixed or gradient-adaptive) to mitigate update-forgetting in classifier-guided diffusion LMs, plus progressive step reduction for fewer steps.
Problem
Classifier-guided diffusion language models suffer token-level fluctuations across timesteps that overwrite earlier semantic edits (update-forgetting), degrading fluency and control while requiring many denoising steps.
Approach
Assign each token its own timestep via a token timestep function, freezing critical tokens early (smaller timesteps, chosen by classifier-gradient magnitude or linear position) so guided edits persist, and progressively fine-tune the model for fewer steps.
Key results
- Detoxification: TTA (T=200) reaches max toxicity 26.0 and PPL 40.6, vs. 32.3 and 58.8 for TESS (T=1000).
- Sentiment control: TTA (T=200) reaches 94.7% accuracy and PPL 20.5, vs. 71.1% and 31.7 for TESS.
- Lexical control: TTA reduces Diffusion-LM mean PPL from 248.6 to 111.4 while syntax-tree accuracy rises from 86.0 to 93.1.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use TTA instead of DExperts | diffusion text generation · RealToxicityPrompts | maximum toxicity: 32.0 -> 26.0 | Helped |
| Use TTA instead of LM-Steer | sentiment-controlled generation · Sentiment Control | accuracy: 85.4 -> 94.7 | Helped |
| Use TTA instead of Diffusion-LM | lexically-constrained generation · Lexically-Constrained Generation | Mean PPL: 248.6 to 111.4 | Helped |
| Add a token timestep schedule (ablation) | TTA (50), T=50 · Detoxification | toxicity: 14.0 to 12.5 | Helped |
| Add a token timestep schedule (ablation) | TTA (50), T=50 · Detoxification | PPL: 68.0 to 59.5 | Helped |
| Add a token timestep schedule (ablation) | TTA (50), T=50 · Sentiment Control | accuracy: 83.5% to 85.9% | Helped |
| Add a token timestep schedule (ablation) | TTA (50), T=50 · Sentiment Control | PPL: 44.0 to 40.2 | Helped |
| Add adaptive allocation (ablation) | Diffusion-LM · Sentiment Control | accuracy: 72.8% to 75.6% | Helped |
| Add adaptive allocation (ablation) | Diffusion-LM · Sentiment Control | PPL: 89.3 to 77.9 | Helped |
| Add adaptive allocation to D-CBG (ablation) | D-CBG, γ=1 · QM9 molecular property maximization | Valid (%): 989 -> 998 | Helped |
| Add adaptive allocation to D-CBG (ablation) | D-CBG, γ=10 · QM9 molecular property maximization | Mean Property: 0.585 -> 0.591 | Helped |
| Use adaptive rather than constant allocation (ablation) | sentiment control diffusion · Sentiment Control | fluctuation ratio | Helped |
Surprising findings
- Fixed-at-T and random token timestep allocations are suboptimal, with fixed-at-T showing large variance and occasional extreme perplexity values.
- Benefits of timestep allocation diminish as the number of denoising steps grows (e.g., at T=1000 gains are small).
Limitations
- Adaptive allocation relies on classifier gradients, so it requires a differentiable task classifier.
- Gains from allocation shrink at large step counts, and the theoretical analysis rests on assumed conditions (A1)-(A3).
Experimental setup
Datasets: C4 E2E IMDB Movie Reviews AGNews RealToxicityPrompts PPLM prompts QQP QM9
Benchmarks: RealToxicityPrompts (challenging subset) PPLM sentiment prompts PPLM topic prompts E2E lexical constraints QQP paraphrase QM9
Metrics: average toxicity maximum toxicity perplexity (GPT-2 Large) Dist-3 control accuracy BLEU BERTScore ROUGE-L
Baselines: PPLM GeDi DExperts Air-decoding LM-Steer Diffusion-LM SSD-LM LD4LG TESS
Builds on: SSD-LM TESS AR-Diffusion Diffusion-LM progressive distillation
Closest papers by meaning
- Denoising Time Matters: Diverse Generation in Diffusion Language Models (NeurIPS 2026 · similarity 0.90)
- Breaking AR’s Sampling Bottleneck: Provable Acceleration via Diffusion Language Models (NeurIPS 2025 · similarity 0.89)
- Scaling Diffusion Language Models via Adaptation from Autoregressive Models (ICLR 2025 · similarity 0.88)
- From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.88)
- Fast and Fluent Diffusion Language Models via Convolutional Decoding and Rejective Fine-tuning (NeurIPS 2025 · similarity 0.88)
- TEncDM: Understanding the Properties of the Diffusion Model in the Space of Language Model Encodings (AAAI 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.