Stop-Think-AutoRegress: Language Modeling with Latent Diffusion Planning

Justin Lovelace; Christian Belardi; Sofian Zalouk; Adhitya Polavaram; Srivatsa Kundurthy; Kilian Q. Weinberger

cs.CLFeb 24, 2026

Stop-Think-AutoRegress: Language Modeling with Latent Diffusion Planning

Authors: Justin Lovelace, Christian Belardi, Sofian Zalouk, Adhitya Polavaram, Srivatsa Kundurthy, Kilian Q. Weinberger

Organizations: Department of Computer Science, Cornell University

Abstract

The Stop-Think-AutoRegress Language Diffusion Model (STAR-LDM) integrates latent diffusion planning with autoregressive generation. Unlike conventional autoregressive language models limited to token-by-token decisions, STAR-LDM incorporates a "thinking" phase that pauses generation to refine a semantic plan through diffusion before continuing. This enables global planning in continuous space prior to committing to discrete tokens. Evaluations show STAR-LDM significantly outperforms similar-sized models on language understanding benchmarks and achieves $>70\%$ win rates in LLM-as-judge comparisons for narrative coherence and commonsense reasoning. The architecture also allows straightforward control through lightweight classifiers, enabling fine-grained steering of attributes without model retraining while maintaining better fluency-control trade-offs than specialized approaches.

Figures & tables

Figure 1: Unified diffusion-guided language model architecture illustrating the training (left) and generation (right) processes. See text for details.

Figure 2: Visualization of the change in log likelihood due to conditioning on the clean continuation embedding versus pure noise. Tokens highlighted in red represent the information provided by the semantic plan.

Figure 3: Diagram comparing zero-shot NLU evaluation for standard LMs (left) and STAR-LDM (right). Both models score candidate answers ( $A_{i}$ ) given a question ( $Q$ ). STAR-LDM utilizes both the answer text $A_{i}$ and its latent embedding $\mathbf{z}_{i,0}$ from Sentence-T5 for scoring.

Model	Params	Scoring	CSQA	SIQA	HS	WG	PIQA	OBQA	ARC:E	ARC:C	Avg
Random Baseline	—		20.0	33.3	25.0	50.0	50.0	25.0	25.0	25.0	31.7
GPT2-Large	.77B	LM Loss	36.6	42.1	42.9	51.9	69.2	34.2	46.6	25.1	43.6
GPT2-XL	1.5B	LM Loss	36.3	42.3	47.9	53.1	70.5	34.4	51.1	28.5	45.5
Pytha-1b	1B	LM Loss	35.3	42.4	45.6	52.6	69.2	32.4	49.1	27.0	44.2
Pytha-1.4b	1.4B	LM Loss	35.7	43.2	50.9	54.2	71.1	36.0	54.0	28.3	46.7
STAR-LDM	.96B	LM Loss	$35.2_{0.4}$	$42.2_{0.2}$	$44.0_{0.1}$	$50.5_{0.1}$	$69.3_{0.1}$	$27.1_{0.3}$	$37.6_{0.5}$	$26.3_{0.3}$	$41.5_{0.1}$

Table 1: Zero Shot NLU Evaluation Results. We report mean and standard error over three monte-carlo evaluations for STAR-LDM.

Figure 4: LLM-Judge evaluation. Results presented with 95% confidence intervals.

			C4 Validation
	Params	CFG ( $w$ )	MAUVE $\uparrow$	Generative Perplexity $\downarrow$	Diversity $\uparrow$
GPT-2 Large	.77B	-	${85.2}_{0.6}$	${30.5}_{0.3}$	${41.8}_{0.1}$
GPT-2 XL	1.5B	-	${86.6}_{0.4}$	${27.1}_{0.6}$	${43.3}_{0.1}$
Pythia-1b	1B	-	$85.5_{0.1}$	$36.5_{0.7}$	$45.6_{0.1}$
Pythia 1.4b	1.4B	-	$84.8_{0.5}$	$33.1_{.1}$	$47.1_{0.0}$
STAR-LDM	.96B	1.0	$94.6_{0.2}$	$36.1_{0.5}$	$45.3_{0.0}$

Table 2: Language generation evaluation on C4 validation examples. We report mean and standard error over three sets of generations.

Panel of Figure 5 — Figure 5: Relationship between perplexity and content attributes across guidance scales ( $s$ ).

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Visualization of the change in log likelihood due to conditioning on a semantic plan. Shown for two different generations. Positive values (red) indicate tokens that become more likely when conditioned on the plan, while negative values (blue) indicate tokens that become less likely.

Figure 7: Perplexity at various noise levels on a subset of 100 examples from the C4 evaluation set (left), or 100 sampled generations (right).

Figure 8: Perplexity on a subset of 100 examples from the C4 evaluation set (left), or 100 sampled generations (right).

Figure 9: Visualization of the token wise log likelihood by log(SNR) (top). Visualization of the token wise change in log likelihood by log(SNR) (bottom).

Figure 10: Quality-Speed Tradeoff in Language Generation. We plot generative perplexity (lower is better) against the wall-clock time to generate a 96-token continuation. The STAR-LDM curve shows performance as the number of diffusion sampling steps is varied. Baselines are standard autoregressive models.

Panel of Figure 11 — Figure 11: Prompt template used for LLM-as-Judge evaluations with Claude 3.7 Sonnet on StoryCloze generations.

Context: Rick grew up in a troubled household. He never found good support in family, and turned to gangs. It wasn’t long before Rick got shot in a robbery. The incident caused him to turn a new leaf.
STAR-LDM (A): Ricks attempt to change his life was long and hard. However, the process has been worth it. The way he lived changed when he started a new chapter…
GPT2-XL (B): Rick now looks after his younger brother, who is disabled.

The two are the heart of what is perhaps the most important project this summer: a massive…
Evaluation of Story Continuations After analyzing the four-sentence context about Rick and the two competing continuations, here is my assessment based on the specified criteria: Narrative Coherence Continuation A follows the context’s redemption arc naturally, describing Rick’s difficult but rewarding journey to change his life after the shooting incident. Continuation B introduces Rick caring for a disabled younger brother (new information) and then abruptly shifts to discussing a ”project” without clear connection to the established narrative. A maintains consistent focus on Rick’s transformation while B introduces disconnected elements. [[Coherence A]] Commonsense Reasoning Continuation A demonstrates understanding of how personal transformation is typically difficult but rewarding, showing appropriate causal relationships. Continuation B presents a plausible caregiving scenario but fails to connect this to the ”important project” mentioned afterward, breaking logical flow. The second sentence in B feels disconnected from both the context and the first sentence of the continuation. [[Reasoning A]] Language Quality Continuation A has minor issues (e.g., ”Ricks” missing apostrophe) but maintains the narrative’s tone and flows naturally from the context. Continuation B starts cohesively but abruptly transitions to an incomplete thought about a ”project,” creating a jarring stylistic break that doesn’t match the established narrative voice. A feels more like a complete thought while B seems fragmented. [[Quality A]] Emotional/Psychological Plausibility Continuation A acknowledges the psychological reality that changing one’s life after trauma is challenging but worthwhile, reflecting plausible emotional growth. Continuation B presents a plausible caring relationship but fails to connect this to Rick’s emotional journey established in the context, and the sudden mention of a ”project” lacks psychological grounding in the character’s established motivations. [[Plausibility A]]

Appendix

Figure 12: Qualitative example of story completions from the StoryCloze dataset. The LLM judge evaluated both continuations, where Continuation A is STAR-LDM and Continuation B is GPT2-XL. The judge preferred STAR-LDM’s continuation across all evaluation criteria.

Panel of Figure 13 — Figure 13: Additional control-perplexity trade-offs beyond those shown in the main paper.

Panel of Figure 14 — Figure 14: Relationship between n-gram diversity and sentiment control.

Panel of Figure 15 — Figure 15: Relationship between n-gram diversity and toxicity control.

Architecture
Base Language Model	GPT-2 Large (774M parameters)
Sentence Embedding Model	Sentence-T5 XL (768 dimensions)
Total Trainable Parameters	956M
DiT Architecture — Prompt Encoder and Diffusion Prediction
DiT Layers per Module	6
DiT Hidden Dimension	1024

Appendix

Table 3: Implementation details for STAR-LDM training.

MLP Architecture
Input Dimension	768
Hidden Dimension	1536
Number of MLP Blocks	4
Number of Hidden Layers	1 hidden layer
Normalization	Pre-norm with adaptive RMSNorm
Residual Connections	Applied after each MLP

Appendix

Table 4: Implementation details for the Noise-Conditioned MLP classifier.

Figure 16: Classifier performance across noise levels for sentiment (left), and toxicity (right) classification.

Autoregressive Language Modeling

Autoregressive language modeling predicts sequences one step at a time, using prior tokens to generate text, speech, code, and other data.

See all 128 papers

Think First, Diffuse Fast: Improving Diffusion Language Model Reasoning via Autoregressive Plan Conditioning
Feb 20, 2026Earl J St SauverDiffusion Language Models LLM Planning
Independent Researcher
Thoughts-as-Planning: Latent World Models for Chain-of-Thoughts Optimization via Reinforcement Planning
Apr 27, 2026Dong Liu, Yanxuan Yu, Ying Nian WuCoT Reasoning Large Language Model-Guided Optimization
University of California, Los Angeles · Columbia University
Hierarchical Continuous Diffusion Language Models
Oct 1, 2026Hui Ren, Zihan Li, Chang Liu +2Structured Reasoning in Language Models Diffusion Language Model Decoding
University of Illinois Urbana-Champaign · Amazon.com, Inc.

Feb 20, 2026J/K move · Enter open · S save

Think First, Diffuse Fast: Improving Diffusion Language Model Reasoning via Autoregressive Plan Conditioning

Earl J St Sauver

Independent Researcher

Diffusion Language Models LLM Planning

Open paper AskPDF arXiv

Diffusion large language models (dLLMs) generate text via iterative denoising but consistently underperform on multi-step reasoning. We hypothesize this gap stems from a coordination problem: AR models build coherence token-by-token, while diffusion models must coordinate all positions simultaneously. We propose plan conditioning, a training-free method that prepends a short (~100-token) natural-language plan from an AR model to the diffusion model's prompt. The plan serves as a frozen scaffold -- globally visible context that every token position can attend to from the first denoising step. On GSM8K, plan conditioning improves LLaDA-8B-Instruct from 75.6% to 87.2% (+11.6 percentage points), matching a same-size AR model (LLaMA 3.1 8B, 87.7%) despite a 6.4pp weaker baseline. On HumanEval, the gain is +12.8pp (37.2% to 50.0%), showing plans generalize to code. The same plans improve LLaMA by only +5.7pp on GSM8K and +1.3pp on HumanEval -- diffusion models benefit 2-10x more, supporting the coordination-problem hypothesis. Across 5 random seeds, plan-conditioned GSM8K accuracy has zero standard deviation, making diffusion inference highly stable. Ablations reveal the model follows plan strategy (wrong-strategy plans cause -16.3pp) but is robust to plan values (perturbed numbers: -1.1pp), and that planner quality has a sharp threshold: smaller Llama-class plans hurt (-1.6 to -6.8pp) while frontier plans provide the full lift. Attention analysis confirms the mechanism: plan tokens receive 1.8x excess attention during early denoising, declining to uniform as completion tokens solidify. Plan conditioning costs ~$0.002 per problem and adds ~2s of latency.

Abstract

Autoregressive Language Modeling

Explore similar work

Think First, Diffuse Fast: Improving Diffusion Language Model Reasoning via Autoregressive Plan Conditioning

Thoughts-as-Planning: Latent World Models for Chain-of-Thoughts Optimization via Reinforcement Planning

Hierarchical Continuous Diffusion Language Models