Towards Efficient Reasoning: Learning Causal Shortcuts for Diffusion Language Models

2026-09-23

AI Summary

本文针对扩散语言模型(DLMs)在推理时难以聚焦关键信息的问题，提出了一种因果捷径学习(CSL)框架，通过提取和优先遮蔽关键令牌序列来提高推理效率与准确性。

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

Problem

• Diffusion Language Models
• bidirectional attention mechanism
• reasoning-guiding tokens
• causal shortcuts

Innovation

• Causal Shortcuts
• Diffusion Language Models
• Reasoning Accuracy
• Convergence Speed
• Prioritized Masking