CaLR: Causal Latent Revision for Robust Diffusion Reasoning

Wei Cai; Jian Zhao; Yuchen Yuan; Xuelong Li

cs.AISep 17, 2026

CaLR: Causal Latent Revision for Robust Diffusion Reasoning

Authors: Wei Cai, Jian Zhao, Yuchen Yuan, Xuelong Li

Organizations: Peking University

Abstract

Autoregressive (AR) models suffer from local greediness, while diffusion language models (DLMs) often lack the strict causal structure required for reasoning. To combine the advantages and overcome the drawbacks of the dual, we propose Causal Latent Revision (CaLR), a framework that reformulates reasoning as constrained latent optimization. By adopting a causal topology matrix (CTM) from an expert model and implicit differentiation, CaLR performs gradient-guided ``thought revision" to enforce logical consistency, enabling dynamic self-correction of intermediate steps during parallel generation. Empirically, CaLR achieves SOTA DLM performance on complex benchmarks, surpassing strong AR baselines and demonstrating superior robustness in constrained tasks like Sudoku.

Figures & tables

Figure 1: Difference between AR, DLM, and CaLR: While AR models are limited by unidirectional constraints and DLMs disregard causal priors, CaLR explicitly injects causal meta-knowledge into the latent space, capturing intrinsic logical consistency for reasoning.

Figure 2: An overview of our proposed Causal Latent Revision (CaLR) framework. CaLR operates in two stages to enforce logical consistency in DLMs.

Model	GPQA		ARC-C		GSM8K	MMLU	HumanEval	Trip-Planning
Model	Acc	ThroughPut	Acc	ThroughPut	Acc	Acc	Acc	Acc	ThroughPut
AR model
Llama3.2-1B	0.22	0.13	38.50	0.35	28.50	40.30	28.60	0.12	0.15
SmolLM2-1.7B	0.24	0.12	48.30	0.36	54.10	60.10	40.80	0.13	0.19
Qwen2.5-0.5B	0.23	0.12	48.40	0.33	41.20	52.30	35.30	0.10	0.19
Qwen2.5-1.5B	0.26	0.15	53.50	0.38	76.40	66.50	41.50	0.14	0.20

Table 1: Comparison of different models on various benchmarks. Underline and bold indicate the best and second-best results of the current column, respectively.

Figure 3: Parallel decoding performance analysis across reasoning and coding benchmarks. We evaluate the trade-off between inference speed (Tokens/Step) and accuracy using the entropy bounded sampler.

Figure 4: Performance evaluation on the Sudoku task, where $n$ denotes the size of the training dataset.

Setting	GSM8K	MATH500	Sudoku	Avg
$\alpha=2$	81.43%	37.60%	77.73%	65.58%
$\alpha=3$	76.82%	34.00%	81.65%	64.16%
$\alpha=4$	81.96%	37.20%	78.64%	65.93%
w/o P.I	74.09%	33.20%	81.65%	62.98%
w/o P.II	68.64%	34.00%	81.35%	61.33%
$\alpha=5$	78.18%	36.80%	81.50%	65.49%

Table 2: Ablation study under different

\alpha

and scheduler settings. Gray underline (

\alpha=4

) stands for the optimal setting. w/o P.I denotes removing the Scheduler (Stage 1 dynamic control), and w/o P.II denotes removing Causal Alignment (Stage 2).

Reasoning Consistency in Language Models

Reasoning consistency in language models measures whether conclusions and reasoning remain reliable across equivalent prompts, samples, steps, and contexts.

See all 115 papers

Towards Efficient Reasoning: Learning Causal Shortcuts for Diffusion Language Models
Sep 23, 2026Dian Jin, Kairong Han, Baohong Li +5Diffusion Language Model Training Efficient Language Model Reasoning
Zhejiang University · Shanghai AI Laboratory
DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models
May 8, 2026Amin Karimi Monsefi, Dominic Culver, Nikhil Bhendawade +4Group Relative Policy Optimization Credit Assignment in RL
The Ohio State University · Apple
Introspective Diffusion Language Models
Apr 13, 2026Yifan Yu, Yuqing Jian, Junxiong Wang +12Diffusion Language Models Diffusion Language Model Decoding
Together AI · University of Illinois Urbana-Champaign · Princeton University +2

Sep 23, 2026J/K move · Enter open · S save

Towards Efficient Reasoning: Learning Causal Shortcuts for Diffusion Language Models

Dian Jin, Kairong Han, Baohong Li, Xinpeng Dong+4

Zhejiang University · Shanghai AI Laboratory

Diffusion Language Model Training Efficient Language Model Reasoning Diffusion Language Models Efficient Language Model Inference

Open paper AskPDF arXiv

Diffusion Language Models (DLMs) have attracted significant attention for their strong reasoning ability. However, under a bidirectional attention mechanism, DLMs operate over an exponentially large exploration space compared to autoregressive models (ARMs), making it challenging to focus on reasoning-guiding tokens under random masking. We define causal shortcuts as token chains that cover the full sequence and provide explicit guidance towards correct reasoning trajectories. We analyze the effects of causal shortcuts on the reasoning accuracy and convergence speed of DLMs, and find that they largely improve answer convergence efficiency and generation accuracy. Motivated by this, we propose a Causal Shortcut Learning (CSL) Framework for DLMs. Specifically, we introduce a step-by-step token extraction procedure to extract causal shortcuts from data, and apply parallel prioritized masking on these tokens during training to enable efficient and accurate convergence to correct answers via causal shortcuts. Extensive experiments across multiple reasoning benchmarks and two base models demonstrate that CSL consistently outperforms existing SFT-variant baselines, achieving an average improvement of

1.92\%

over SFT-only models, and up to

4.20\%

on MATH-500. The code is available at the \href{https://github.com/ZJUDianJin/Causal-Shortcuts-Learning}{https://github.com/ZJUDianJin/Causal-Shortcuts-Learning

Abstract

Reasoning Consistency in Language Models

Explore similar work

Towards Efficient Reasoning: Learning Causal Shortcuts for Diffusion Language Models

DACA-GRPO: Denoising-Aware Credit Assignment for Reinforcement Learning in Diffusion Language Models

Introspective Diffusion Language Models