Time-Annealed Perturbation Sampling: Diverse Generation for Diffusion Language Models

1. Problem Statement

Diffusion Language Models (Diffusion-LMs) have emerged as a powerful alternative to autoregressive models, offering explicit temporal control over text generation. However, they suffer from a critical limitation: output homogenization. Despite instructions to be diverse, Diffusion-LMs often converge to repetitive or conservative semantic paths, similar to autoregressive models.

Existing diversity control methods face significant trade-offs:

• Training-time methods (e.g., modifying loss functions) are costly, require curated datasets, and lack flexibility.
• Decoding-time methods (e.g., temperature scaling, top-p, min-p) often degrade generation quality (fluency, coherence) when aggressively tuned to increase diversity.
• Prompt-based methods rely heavily on the model's inherent capabilities and can cause semantic drift.

The core challenge is to develop a training-free inference strategy that leverages the unique temporal dynamics of diffusion models to encourage semantic branching (diversity) without sacrificing generation quality or instruction adherence.

2. Key Insight: Temporal Division of Labor

The authors identify a fundamental characteristic of Diffusion-LMs analogous to image diffusion models:

• Early Denoising Steps: Primarily determine the global semantic structure and high-level narrative flow.
• Later Denoising Steps: Focus on local lexical refinement, fluency, and stylistic details.

This insight suggests that injecting randomness early in the process can alter the semantic trajectory (branching), while preserving the later steps ensures the output remains coherent and follows the original prompt.

3. Methodology: Time-Annealed Perturbation Sampling (TAPS)

TAPS is a training-free inference strategy that perturbs the conditioning signal (prompt embeddings) during the denoising process. The perturbation magnitude is time-dependent, decaying as the generation progresses.

Core Mechanism

1. Perturbation Injection:
   Instead of using a static prompt embedding $E$, TAPS introduces stochastic noise $\epsilon$ at inference step $t$:
   $$\tilde{E}(t) = E + \sigma(t)\epsilon$$
   where $\sigma(t)$ is a monotonically decreasing annealing schedule. Strong perturbations are applied in early steps ($t_{start}$) to encourage semantic exploration, while perturbations vanish in later steps ($t_{end}$) to ensure alignment with the original prompt.

2. Quality Preservation Mechanisms:
   To prevent the injected noise from causing distributional drift or degrading quality, TAPS employs two stabilization techniques:

   • Rescaling and Mixing: The perturbed embedding is rescaled to match the mean and standard deviation of the original embeddings. It is then interpolated with the original embedding using a mixing coefficient $\psi$:
     $$\hat{E} = \psi E' + (1-\psi)E$$
     This ensures the perturbed signal stays within a valid statistical range.
   • Norm-Preserving Projection: Additive noise can alter the magnitude (norm) of token embeddings, disrupting the model's learned geometry. TAPS projects the perturbed embeddings back onto the hypersphere defined by the original token norms:
     $$E^*_i = \frac{\hat{E}_i}{\|\hat{E}_i\|_2} \cdot \|E_i\|_2$$
     This preserves the directional variation (diversity) while maintaining the original embedding magnitudes (stability).

4. Key Contributions

1. Empirical Discovery: Identified the temporal semantic structure in Diffusion-LMs, proving that early steps govern semantics while later steps handle lexical refinement.
2. Novel Algorithm (TAPS): Proposed a simple, scalable, and training-free inference strategy that leverages this temporal structure to induce semantic branching without retraining.
3. Comprehensive Validation: Demonstrated consistent improvements across diverse backbones (LLaDA, TraDo) and tasks (creative writing, instruction following, mathematical reasoning) with negligible computational overhead.

5. Experimental Results

The authors evaluated TAPS on LLaDA-8B-Instruct and TraDo-8B-Instruct across four benchmarks: WritingPrompts, NoveltyBench, Arena-Hard-Auto, and GSM8K.

• Diversity Gains: TAPS consistently outperformed baselines (Top-k, Top-p, Min-p, Diverse Prompt, EDT) across lexical (IntraDistinct, Div-Blue), semantic (Sent-BERT), and embedding-level (EAD) metrics.
  • Example: On NoveltyBench, TAPS achieved a Sent-BERT score of 36.04 vs. 25.80 for the base model on LLaDA, indicating significantly deeper semantic variation.
• Quality Preservation: Unlike aggressive sampling methods that degrade coherence, TAPS maintained or slightly improved quality scores (creativity, coherence, writing quality) as judged by GPT-4o and human preference models.
• Reasoning Robustness (GSM8K):
  • Pass@1: TAPS showed a mild drop in single-sample accuracy compared to the base model due to exploration.
  • Majority Voting: TAPS significantly outperformed all baselines in majority-vote accuracy. The diverse reasoning paths generated by TAPS contained less correlated errors, allowing aggregation to recover the correct answer more effectively.
• Ablation Studies:
  • Noise Scale ($\sigma$): Higher noise increases diversity but can hurt quality if too high; optimal balance found at $\sigma=0.2$.
  • Injection Window: Injecting noise only in early steps ($t \in [0.9, 0.3]$) was crucial for semantic diversity; late-stage noise only affected lexical variation.
  • Mixing Coefficient ($\psi$): A value of $\psi=0.9$ provided the best trade-off, stabilizing quality without sacrificing diversity.

6. Significance

• Paradigm Shift: TAPS demonstrates that Diffusion-LMs possess a unique "temporal division of labor" that can be exploited for controlled generation, a property not present in autoregressive models.
• Practical Utility: As a training-free method, TAPS can be immediately applied to existing Diffusion-LMs to enhance creativity and reasoning exploration without the cost of fine-tuning.
• Balanced Trade-off: It successfully resolves the traditional diversity-quality trade-off, showing that structured exploration (via time-annealed noise) yields diverse outputs that remain coherent and instruction-following.
• Future Direction: The paper highlights Diffusion-LMs as a promising direction for controlled text generation, suggesting that future research should focus on leveraging temporal dynamics for more sophisticated control mechanisms.