Title: Don’t Let It Fade: Preserving Edits in Diffusion Language Models via Token Timestep Allocation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion Language Models

Controllable Text Generation
Abstract: Classifier guidance is a widely adopted technique in diffusion language models, used to steer generation toward desired attributes. However, such guidance often introduces instability during the generation process, where token-level updates fluctuate across timesteps. We identify and formally characterize this phenomenon as update-forgetting. This instability disrupts the refinement process by overwriting semantic edits, ultimately degrading fluency and coherence, which is particularly problematic in tasks like controllable text generation. To address this, we propose TTA-Diffusion, a novel inference-time approach that dynamically allocates timesteps per token based on refinement needs. Unlike conventional diffusion models that apply uniform updates, TTA-Diffusion employs structured timestep allocation, preserving stable tokens while allowing uncertain tokens to undergo progressive adjustment. Experimental results across diverse tasks demonstrate that TTA-Diffusion significantly outperforms both diffusion-based and auto-regressive baselines in fluency and control accuracy while improving computational efficiency by reducing the number of required timesteps. On the sentiment control task, TTA-Diffusion achieves over 20\% higher accuracy and nearly half the perplexity of prior diffusion models, using less than one-fifth the denoising steps. This work highlights the importance of mitigating fluctuations in token updates and promoting a balanced refinement process, thereby enhancing stability and controllability in controllable language modeling.


Peer reviews:
comments: SUMMARY:
The paper investigates why classifier-guided diffusion language models become unstable, tracing the problem to two factors—excessive fluctuations and update-forgetting—and proposes adaptive per-token timestep allocation as an inference-time fix. Additionally, they bridge the discrete–continuous gap through logit-simplex mapping and uses timestep distillation to reduce sampling steps. Extensive experiments and ablations confirm that the method yields more stable and controllable generation.

STRENGTHS AND WEAKNESSES:
Strengths
1. The paper’s clear empirical diagnosis (fluctuation-driven update-forgetting) naturally motivates the adaptive per-token timestep strategy, making the solution intuitive and clear.
2. The method shows consistent improvements in overall metrics compared to other baselines even with reduced sampling steps.
3. The author provides extensive experiments in the main paper and the appendix.

Weaknesses
1. While adaptive token-level timesteps for classifier-guided inference are a thoughtful extension, I feel this idea is incremental, as it builds on earlier token scheduling concepts like those in AR-Diffusion. Moreover, other components such as the logit-simplex mapping and timestep reduction have already been proposed in prior work.
2. From my understanding, the time distillation is used for compensating the time overhead from using the classifier guidance. I think there should be a comparison regarding the time consumption for generating the text.
3. Several informative studies in the appendix are not mentioned in the main paper making it easy for readers to miss evidence that actually strengthens the claims.

QUESTIONS:
Please refer to the weakness section.

LIMITATIONS:
Please refer to the weakness section.
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n* The authors propose TTA-Diffusion, an inference-time technique to reduce token fluctuations when applying classifier guidance.\n* The authors observe that semantics of generated sequences are perturbed and modified during diffusion and formalize this problem as diffusion fluctuation and update-forgetting.\n* The authors propose a dynamic timestep allocation strategy that leverages gradient signals from the guiding classifier to ensure that recently guided tokens are not immediately revered or perturbed by the denoising model.\n* The authors show that TTA-Diffusion outperforms autoregressive and diffusion baselines on controllability tasks, such as detoxification, sentiment control, topic control, etc.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n* Diffusion generation, and in particular timestep allocation strategies and fine-grained steering, is a relevant topic of interest to the community.\n* Presented claims and hypotheses are empirically appropriately supported by experiments and ablations.\n* The paper is clean and easy to follow.\n\nWeaknesses:\n* The paper focuses on a rather specific problem (steering simplex diffusion models via classifier guidance), limiting its general applicability.\n* The definition of token disruption and update-forgetting seems rather simple and is unintuitive. It is based on strict token matching, which at a glance seems like a poor measure of semantic coherence.\n* Some claims are weakly supported by theoretical justifications drawn from the discrete diffusion literature. While the rough intuition may translate to simplex diffusion models, some details and rigor are likely lacking.\n\nQUESTIONS:\nComments and questions are not necessarily in order of importance (largely chronological with the flow of the paper).\n\n1. Are the examples in Figure 1 actual observed samples / examples of update-forgetting? The transition seems drastic. Diffusion disruption and update-forgetting are defined by strict token-equality checks, which seem brittle and susceptible to over-estimation due to minor word substitution to synonyms or modifications to sentence structure. Actual observed samples from the experiments in Appendix B.1 would be helpful in contextualizing the problem (maybe alongside TTA-Diffusion outputs in Table 12 for clearer comparison).\n\n2. The authors write that TTA-Diffusion \"build[s] upon the Simplex-based diffusion language model [11], and extend it into a fully non-autoregressive framework [28] that naturally supports external classifier guidance.\" However, the non-autoregressive extension simply follows from using TESS, so it seems slightly misleading to suggest this as a contribution. A rephrasing seems necessary.\n\n3. It would be helpful if Definition 2 includes details on how key tokens from classifier guidance are identified (highest gradient norm, as written in Hypothesis 2. Moreover, Table 12 seems to suggest that context / prompt tokens are included in computing the top $K$ most important tokens -- is this correct? I would expect context tokens to be excluded from this measurement as well as updates from classifier guidance or diffusion.\n\n4. Sec. 2.2. indicates that TESS without self-conditioning seems to be used. Does enabling self-conditioning make a difference? The simplex averaging self-conditioning method is effectively an exponential moving average of predicted probabilities, so I would expect the predictions to remain more consistent without excessive disruptions.\n\n5.  At a high level, it seems like update-forgetting occurs because the backward diffusion step erases the effect of classifier guidance. Is the guidance factor at all relevant to this? Update-forgetting seems to suggest that classifier guidance was perhaps too aggressive that it causes a \"tug-of-war\" between guidance updates and backward diffusion. Can a well-tuned guidance factor alleviate this? It seems like $\\alpha_\\mathrm{smooth}$ explicitly modulates this, but really a fine-tuned $\\lambda$ should suffice (and perhaps the adaptive timestep allocation strategy in Sec. 3.2 is effectively controlling $\\lambda$ per timestep per token).\n\n6. Hypothesis 1 seems to suggest token fluctuations is the cause of reduced fluency and coherence. Is it possible that the disruptions are due to model's low confidence, i.e., the reduced fluency is more the byproduct of the model struggling to converge on a stable output sequence due to the difficulty of generating a sample? From this point of view, the reduced fluency is not due to the fluctuations per se, but due to low confidence / more uncertainty in the model's predictions?\n\n7. Is the adaptive timestep in Sec. 3.2. have a monotonicity property, i.e., $f(i, t) < f(i, t + 1)$? It seems unnatural and contradictory for a token that was assigned a small timestep (\"this is a stable, finished token\") to later have a larger timestep (\"now we need to aggressively refine\"). The norm-based convex combination with time doesn't seem to guarantee this. \n\n8. How does progressive step reduction relate to the earlier hypotheses and problem definitions? It's not clear to me what the motivation for this is, especially given the narrative that TTA-Diffusion is an inference-time technique.\n\n9. Is update-forgetting unique to the continuous-to-discrete nature of simplex modeling, or do you think there are continuous analogues?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
The authors propose to improve the stability of classifier-guided semantic edits in diffusion language models by proposing context-dependent token updates. They show that their approach significantly improves both computational efficiency and generation quality.

STRENGTHS AND WEAKNESSES:
**Strengths**
- The authors clearly formulate the "update-forgetting" failure mode with supporting empirical evidence and thorough ablations
- The proposed context-dependent token updates leads to significant improvements in both computational efficiency and generation quality
- The range of tasks and ablations are thorough and convincing 

**Weaknesses**
- Empirical speedups of proposed approach under fewer sampling steps are not presented
- Proposed approach is grounded in Gaussian diffusion. However, SOTA diffusion LMs operate in a discrete state space, with recent work proposing guidance mechanisms under this setting [1]
- The authors do not compare to discrete diffusion LMs that support classifier guidance [1]


[1] Schiff, Y. "Simple Guidance Mechanisms for Discrete Diffusion Models." ICLR 2025.

QUESTIONS:
- Can the authors report the empirical speedups of proposed approach under fewer sampling steps?
- Can the context-dependent token updates be applied to existing discrete diffusion frameworks? [1]

[1] Schiff, Y. "Simple Guidance Mechanisms for Discrete Diffusion Models." ICLR 2025.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes TTA-Diffusion, which is an inference time method to adaptively assign a differing number of diffusion time steps to every token in the sequence. The authors argue that this mechanism helps perform better classifier-guided controllable generation. The number of steps for a token is set by using signal from the classifier. The authors also show that their method is effective when diffusion models are progressively distilled to reduce the number of diffusion steps.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- TTA-Diffusion works with just 50-100 time steps. Compared to baselines, this should make TTA-Diffusion much faster.\n- The analysis done for hypothesis 1 is a convincing motivating story for TTA-Diffusion.\n- TTA-Diffusion is an inference-time method that can potentially be easily applied to many different diffusion methods.\n\nWeaknesses: \n- There are a few missing baselines. DGLM ([1]) is a very related method and there are no comparisons to DGLM. AR-Diffusion [2] proposes using something similar to linear schedule for assigning time steps, but there are also no comparisons to AR-Diffusion. \n- The main contribution of the paper is to propose an adaptive strategy to assign time steps to each token. However, for their detoxification experiments they use a linear allocation schedule as mentioned in caption of table 1. Does this mean the adaptive strategy is only useful sometimes?\n- The writing clarity can be improved. Some details are unclear; for e.g. is the time step allocation done once at the start of denoising or is the allocation changed at every step. Are the time steps for all tokens set according to the equation in 199 or only a subset of tokens?\n\n[1] Lovelace, J., Kishore, V., Chen, Y., & Weinberger, K. (2024, August). Diffusion Guided Language Modeling. In Findings of the Association for Computational Linguistics ACL 2024 (pp. 14936-14952).\n[2] Wu, T., Fan, Z., Liu, X., Zheng, H. T., Gong, Y., Jiao, J., ... & Chen, W. (2023). Ar-diffusion: Auto-regressive diffusion model for text generation. Advances in Neural Information Processing Systems, 36, 39957-39974.\n\nQUESTIONS:\nHere are some questions I have about the paper:\n- Does forgetting, as illustrated in figure 1, empirically occur? Doesn't the classifier guidance applied at every denoising step prevent \"love\" to change to \"hate\" when stearing for positive sentiment?\n- The experimental setting for hypothesis 2 is unclear. Is classifier guidance being applied? If so, why does the probability drop?\n- What is the input to the classifiers? Is it tokens, is it vectors? Appendix B2 mentions off-the-shelf classifiers are used for the different tasks, but these classifier operate on clean text. Are you using these same classifiers on noisy intermediate sequences during diffusion?\n- Section 4.4 mentions for T=1000 the sentiment accuracy improves from 92.2 to 92.6, but these numbers are different from the ones reported in the table in the main paper and the appendix? Which table are these numbers from?\n- Why are the results from TTA (5000) bad? Table 3a shows a very high perplexity for  TTA (5000). Table 9 also shows that TTA (5000) is worse than TTA (1000).\n- In figure 2, how many samples correspond to each mean fluctuation ratio? Is the distribution close to uniform?\n\nLIMITATIONS:\nYes, the authors address limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}