Noise Your Prompt: Noising Conditioning Tokens in Continuous Diffusion Language Models

AI summary

Overview

• Research area: Generative modeling for language — specifically continuous diffusion and flow-matching language models, and the design of their conditional training objective.
• Technical level: Intermediate. The paper assumes familiarity with diffusion/flow models, denoising objectives, and classifier-free guidance, but the proposed change itself is conceptually simple.
• Scope: The paper introduces and empirically evaluates a single change to the diffusion language model training objective — also noising the conditioning prompt tokens — across two combinatorial reasoning benchmarks (Sudoku, N-Queens) and two natural language generation tasks (Gigaword summarization, open-ended dialogue).

What This Paper Is About

Continuous diffusion and flow language models generate text conditionally by being trained with the conditioning prompt tokens held clean — a practice the authors term "paste-in" conditioning — where a noised response is denoised given an un-noised prompt. The authors ask whether this default is actually optimal, and show that simply noising the conditioning tokens during training (with some probability, instead of always keeping them clean) improves generalization, solution diversity, and — in some but not all settings — natural language generation quality.

Key Contributions

1. A minimal training-objective modification. Instead of always training with clean conditioning tokens, the model keeps them clean with probability p_condclean and otherwise noises them to the same noise level as the response tokens. p = 1 recovers the standard paste-in objective; p = 0 recovers a fully unconditional joint-distribution objective; intermediate values interpolate.
2. Loss over the full sequence. The model is also asked to denoise the conditioning tokens, not just the response, which the authors argue makes the training objective more efficient by using every token.
3. Empirical evidence on combinatorial reasoning. Noised-prompt training substantially improves puzzle solve rates and solution diversity versus paste-in training, with the largest gains on harder puzzles.
4. Mixed evidence on natural language. Measurable improvements on Gigaword summarization, including against dropout-regularized baselines and at a larger (132M) model, but null results on open-ended dialogue — explicitly showing the gains do not transfer to all natural language tasks.
5. Guided sampling without default inference cost. Because the model is trained under both clean and noised prompts, its predictions can be affinely combined at sampling time in a classifier-free-guidance-inspired scheme. The method requires no additional inference cost by default. Code is publicly available.

Main Findings

• Sudoku hard gains are the headline result. On Sudoku hard (30/81 clues revealed, N = 10,000 training puzzles), the paste-in baseline solves 3.73% of 2,000 held-out puzzles versus 24.65% at p_condclean = 0.1.
• Sudoku easy improves too. At p = 0.3 with N = 10,000, solve rate is 91.18% versus 68.49% for the standard paste-in setup.
• Data-limited regime behaves differently. With only N = 1,000 Sudoku puzzles, p = 0.5 performs best (6.92% on easy), while paste-in (1.47%) and p = 0.2 (1.39%) perform similarly with no statistically significant difference.
• Fully unconditional training fails. p = 0.0 yields 1.41% on Sudoku easy and 0.00% on Sudoku hard — the model must see some clean prompts so that the clean inference-time prompt stays within the training distribution.
• N-Queens accuracy improves at the larger scale. At n = 10 (length 100), p = 0.2 reaches 98.32% accuracy versus 86.23% for paste-in; at n = 8 (length 64) accuracy is high throughout (99.37% at p = 0.2 versus 97.48% for paste-in).
• Solution diversity is where N-Queens shows the clearest effect. Coverage — the fraction of distinct valid solutions recovered across k = 20 samples per puzzle — is 73.79% at p = 0.2 versus 50.60% for paste-in at n = 10, and 92.39% versus 89.02% at n = 8, over 750 held-out puzzles.
• Clean-prompt training mode-collapses. In Figure 1, on a single held-out 8x8 N-Queens board admitting six distinct valid solutions, the noised-prompt model (p = 0.2) spreads five generations across 4/6 distinct solutions, while the paste-in model (p = 1.0) collapses onto a single completion.
• Noised prompts act as a strong regularizer on modest language data. On Gigaword, the 57M model's best validation NLL is 0.3707 +/- 0.0093 at p = 0.5, versus 0.4581 +/- 0.0041 for the p = 1.0, dropout 0.1 baseline — and better than the heaviest dropout baseline tried (0.4455 +/- 0.0080 at dropout 0.4). The 132M model shows the same pattern (0.3518 +/- 0.0112 at p = 0.5 versus 0.4573 +/- 0.0022 for paste-in).
• Gigaword quality-diversity frontier improves as measured by MAUVE against average unigram entropy, when sweeping softmax temperature on the denoised x_1 distribution over T in [0.2, 1.5], with the reference entropies of 2.935 bits (Gigaword) and 3.878 (dialogue) shown as targets.
• Dialogue shows null results. No significant effect is observed on open-ended dialogue under noising, at either the larger 3.38M-sample regime or the smaller 191k-sample regime, even at larger noising probabilities like p = 0.3.
• Guided sampling improves Sudoku solve rate further. On Sudoku hard (N = 10,000, p = 0.2), guidance weight w = 5 raises solve rate from 23.89% to 30.82% (+6.93); on Sudoku easy (N = 10,000) from 89.44% to 92.08% (+2.64); in the data-limited N = 1,000 easy setting from 1.66% to 3.94% (+2.28). As reported in prior classifier-free guidance work, this trades off against diversity: increasing guidance generally improves accuracy but decreases N-Queens solution coverage.
• Ablation: separate conditioning noise time. Drawing an independent time t' ~ U(0,1) for conditioning tokens reaches 93.23% on Sudoku easy (above the best single-time setting of 91.18% at p = 0.3) and 25.09% on Sudoku hard, but performs worse than frequent same-time noising on N-Queens (97.20% accuracy / 88.84% coverage at n = 8; 88.56% / 51.69% at n = 10).
• Ablation: response-only loss. Restricting the loss to response tokens gives 87.70% at p = 0.2 versus 63.34% at p = 1.0 on Sudoku easy — the gains mainly come from noising, not from the full-sequence loss, though the full-sequence objective performs similar or better.
• Ablation: per-token independent noising. Instead of jointly keeping or noising the whole prompt, per-token Bernoulli noising gains relatively more at higher clean probability (89.75% at p = 0.5 versus 75.97% at p = 0.2 on Sudoku easy) but is overall worse than joint noising on Sudoku hard (16.98% at p = 0.5, 8.99% at p = 0.2).
• Ablation: masked token noise. Noising the prompt by replacing a matched fraction of tokens with [MASK] — chosen so the mutual-information ratio under masking matches that under Gaussian noise at level t — performs statistically similar to Gaussian noising at p = 0.2 and p = 0.5. (The numeric results for this ablation are cut off in the available content.)
• Comparison to data scaling. On Sudoku easy, the default paste-in objective requires roughly 10x more training data to match the performance of noised-prompt training at p = 0.2 (single seed, 2,000 held-out samples drawn from the same generator).

Methodology in Plain English

The authors work with an in-place infilling layout: the entire puzzle board or text sequence is treated as one sequence, and every token carries its own noise level tau_i in [0, 1], where 1 means clean and 0 means full noise. The model is a flow / x_1-prediction denoiser that receives the corrupted sequence plus per-token times and predicts the clean tokens, trained with cross-entropy over the full sequence.

The only change is what happens to the conditioning tokens. For each example, they draw a shared time t, then flip a Bernoulli coin with probability p_condclean. If it comes up "clean", all conditioning tokens are held at noise level 1 while the response is noised to t — exactly the standard paste-in setup. Otherwise, all conditioning tokens are noised to the same level t as the response, so the model must denoise with less prompt signal. At generation time they follow standard practice: conditioning tokens are kept clean, clamped, and re-pasted at every step of a 128-step first-order Euler sampler. Keeping some fraction of training examples fully clean is necessary to avoid a train/inference mismatch.

Because the model has been trained on both clean and noised prompts, they can additionally sample with a guidance scheme inspired by classifier-free guidance: they take the log-softmax logits of the model's prediction under a noised prompt and under a clean prompt and affinely interpolate them with weight w, where w = 1 recovers default paste-in sampling and w > 1 makes the model lean harder on the prompt.

Setup details: the backbone is the Flow Language Model (FLM) from Lee et al. (2026), a DiT-style transformer with hidden size 512, 8 blocks, 8 attention heads, conditioning dimension 128, and dropout 0.1. Training uses AdamW (learning rate 3e-4, (beta_1, beta_2) = (0.9, 0.999), no weight decay), a constant schedule with 2,500 warmup steps, gradient clipping at 1.0, bf16-mixed precision, and EMA of 0.9999, on a single NVIDIA L40S GPU. Sudoku sequences have length L = 81; N-Queens has L = n^2. Language tasks use the bert-base-uncased tokenizer with vocabulary V = 30,522.

Evaluation uses puzzle accuracy and (for N-Queens) coverage across k = 20 generations per puzzle; for language, MAUVE plus average unigram entropy. Unless otherwise stated, results report 95% two-sided t-distribution confidence intervals across n = 5 training seeds (dialogue uses n = 3), with bootstrap intervals from 10,000 resamples for Sudoku. Sudoku and N-Queens are evaluated with the final checkpoint, while language models — which overfit early — are evaluated at the best validation NLL checkpoint, validated every 2,000 steps.

Why This Matters

The paper challenges a default that the continuous diffusion language modeling literature largely treats as given. If simply noising the prompt during training yields gains this large on hard constraint-satisfaction puzzles, then the standard "keep conditioning clean" recipe may be leaving substantial performance on the table — for a change that costs one line of code and, by default, nothing at inference time. The paper is also careful to report where the method does not help (open-ended dialogue), which is useful for anyone deciding whether to adopt it.

Real-world applications:

• Constraint-satisfaction and planning assistants — scheduling, logistics, and resource allocation problems share the structure of Sudoku and N-Queens: hard combinatorial constraints with many valid solutions, where both correctness and diversity of proposals matter.
• Diverse candidate generation in design and engineering — situations where a single valid answer is not enough and a spread of distinct valid solutions is the deliverable.
• Summarization pipelines, particularly in data-scarce or domain-specific settings, where the paper reports measured gains against standard dropout regularization, including at a 132M-parameter model.
• Compute-constrained training — if noised prompts genuinely substitute for roughly 10x more training data, smaller teams with limited data and hardware budgets gain the most.

Industry relevance: The modification is trivial to implement, adds no inference cost by default, and slots into existing flow/diffusion language model codebases. The guided-sampling option provides a tunable accuracy–diversity knob at inference time for practitioners who want one.

Future Directions

• Understanding why the gains do not transfer to dialogue. The paper establishes the null result on open-ended dialogue at both 3.38M and 191k training samples but does not explain the mechanism, leaving open what property of a task determines whether noised-prompt training helps.
• Tuning p_condclean per task and per data regime. The optimal value varied across settings (0.1 best on Sudoku hard at N = 10,000, 0.5 best at N = 1,000, 0.2 best on N-Queens), and no rule for selecting it is offered.
• Extending beyond Gaussian one-hot noise. The masked-noise ablation performed statistically similarly, suggesting the method may generalize to other corruption processes and potentially to discrete-state diffusion families like D3PM, SEDD, MDLM, and LLaDA.
• Combining with related conditioning-corruption work. Prompt infilling for masked diffusion language models (Fujinuma and Sakaguchi, 2026) is the closest prior work; how the two approaches interact is unexplored, as is the interplay with per-token noise-level schemes like diffusion forcing and AR-Diffusion.
• Scaling and model size. The paper tests a 132M model only on Gigaword validation NLL, so behavior at larger scales on reasoning tasks is not reported.

Target Audience

Researchers and engineers working on diffusion or flow-based language models who are implementing conditional training objectives — the paper's value is highest for those who already have such a model and want a low-cost change to try. It is also relevant to readers interested in the theory of conditional generation and classifier-free guidance, since the paper reframes conditioning strength as a continuous quantity controlled by noise level rather than a binary drop/don't-drop decision. Practitioners applying diffusion language models to structured reasoning or summarization will find the concrete benchmark numbers directly actionable. Readers without a diffusion or flow-matching background will need to acquire that first.

Authors’ abstract

We revisit a standard accepted practice in the continuous diffusion language model literature of fixing conditioning prompt tokens clean during training. We make a very simple modification: also noise the conditioning prompt tokens during training. We demonstrate that under this modified training objective, we achieve better generalization in combinatorial reasoning tasks such as Sudoku and N-Queens, with the largest gains on harder variants ($3.73\% \to 24.65\%$ solve rate on Sudoku Hard), and increased diversity of generated solutions ($50.60\% \to 73.79\%$ coverage on 10x10 N-Queens). We also show measurable improvements to natural language generation quality in modest dataset regimes with Gigaword summarization, but notably demonstrate that gains do not transfer to all natural language tasks (e.g open ended dialogue generation). Our method is a single line change to the training objective, requires no additional inference costs by default, and provides the flexibility of classifier-free guidance inspired guided sampling. Our \href{https://github.com/LateralIntelligence/noise-your-prompt} {code} is publicly available.