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 ( Ai ) given a question ( Q ). STAR-LDM utilizes both the answer text Ai and its latent embedding zi,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.20.4
42.20.2
44.00.1
50.50.1
69.30.1
27.10.3
37.60.5
26.30.3
41.50.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 ↑
Generative Perplexity ↓
Diversity ↑
GPT-2 Large
.77B
-
85.20.6
30.50.3
41.80.1
GPT-2 XL
1.5B
-
86.60.4
27.10.6
43.30.1
Pythia-1b
1B
-
85.50.1
36.50.7
45.60.1
Pythia 1.4b
1.4B
-
84.80.5
33.1.1
47.10.0
STAR-LDM
.96B
1.0
94.60.2
36.10.5
45.30.0
Table 2: Language generation evaluation on C4 validation examples. We report mean and standard error over three sets of generations.
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.
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.
Figure 13: Additional control-perplexity trade-offs beyond those shown in the main paper.
Figure 14: Relationship between n-gram diversity and sentiment control.
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.
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.
The success of large language models (LLMs) across diverse NLP tasks has elevated the importance of reasoning chain optimization as a critical step in aligning model behavior with task objectives. Existing reasoning chain tuning methods often rely on black-box heuristics or gradient-free search, which lack interpretability, generalization, and sample efficiency. In this work, we introduce Thoughts-as-Planning, a novel framework that formalizes reasoning chain optimization as a sequential decision-making process over a latent semantic space. We model the LLM as a partially observable environment and learn a latent world model that simulates the effect of reasoning chain edits on downstream outputs. A proximity-preserving embedding space is constructed to encode reasoning chain-response dynamics, enabling planning via gradient descent or reinforcement learning. Our method supports multi-scale abstraction, allowing reasoning chain edits at token, segment, and instruction levels to be integrated into a unified planner. Through extensive experiments on language understanding and generation tasks, we demonstrate that Thoughts-as-Planning outperforms state-of-the-art reasoning chain tuning baselines in efficiency, robustness, and generalization, while offering interpretability through its structured planning trajectory. Our code is available at https://github.com/FastLM/Thoughts-as-Planning.
Dong Liu, Yanxuan Yu, Ying Nian Wu
University of California, Los Angeles · Columbia University
Discrete diffusion language models offer a compelling alternative to autoregressive generation for tasks demanding bidirectional reasoning and global constraint satisfaction. Yet they share a structural bottleneck: when decoding in parallel, each token is sampled independently from its marginal, severing the statistical dependencies among the tokens decoded together. Continuous diffusion language models avoid this by denoising a shared continuous state, but their denoiser sees only that state, so nothing ties it to a valid token configuration until it is finally decoded. To address this, we propose Hierarchical Continuous Diffusion Language Models (HC-DLM), which couple discrete token generation with a continuous latent trajectory in a single, principled denoising process, whose training objective is derived from a variational bound on the token likelihood. In contrast to recent methods that attach continuous context to a self-contained discrete chain, HC-DLM makes the latent the only persistent generative state: tokens are read out from it at every step and feed back as a scaffold for the next latent update. On structured reasoning (Sudoku), mathematical planning (Countdown) and language modeling (LM1B), HC-DLM improves over discrete and continuous diffusion baselines at matched model size, in puzzle accuracy on Sudoku and Countdown and in generative perplexity on LM1B. Project page: https://hc-dlm.github.io/.
Hui Ren, Zihan Li, Chang Liu +2
University of Illinois Urbana-Champaign · Amazon.com, Inc.