Causal shortcuts improve reasoning accuracy and speed in diffusion language models
Towards Efficient Reasoning: Learning Causal Shortcuts for Diffusion Language Models
Computation and Language
Summary
Diffusion language models are good at reasoning but face a huge challenge because they look at many possible token arrangements, which slows them down. The authors introduce the idea of 'causal shortcuts,' which are key token sequences that help guide the model toward the right answers faster and more accurately. They develop a method to identify these shortcuts and train models to pay special attention to them. Their experiments show this approach improves performance on several reasoning tasks.
What this means in practice
- •For natural language processing engineers: Improve the efficiency and accuracy of reasoning tasks in language models by integrating causal shortcut learning during training.
- •For educational technology developers: Enhance automated math problem-solving systems, particularly for complex benchmarks like MATH-500, through more focused reasoning pathways.
Authors
Dian Jin, Kairong Han, Baohong Li, Xinpeng Dong, Zijing Hu, Nuanqiao Shan, Fei Wu, Kun Kuang
Abstract
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