Guiding the Reasoning Ability of Diffusion Language Models via Sparse Autoencoder Feature Intervention

This article presents an innovative study that uses Sparse Autoencoder (SAE) feature intervention technology to guide the chain-of-thought reasoning behavior of Diffusion Language Models (DLMs) during the inference phase, significantly improving mathematical problem-solving ability without additional training. The core idea is to activate the reasoning-related features already encoded inside DLMs, and its effectiveness is validated using the GSM8K elementary school math word problem dataset.

After the success of diffusion models in the image domain, they were extended to NLP to form DLMs. Their iterative denoising generation paradigm has advantages such as good controllability and parallel decoding, but their performance on complex reasoning tasks is weaker than autoregressive models, and they lack an explicit chain-of-thought mechanism. Traditional improvements require extensive fine-tuning, which is costly and inflexible. This study proposes an SAE feature intervention scheme to activate internal reasoning features.

1. SAE Application: A Top-K sparse autoencoder is used with a dictionary size of 4× the model dimension, retaining Top-K activated features. It is deployed across multiple layers of the DLM to capture features at different abstraction levels.
2. Contrastive Feature Discovery: By comparing feature activations between Chain-of-Thought (CoT) and Direct prompts, reasoning-related features are identified using difference calculation, Welch's t-test, and Cohen's d.
3. Diffusion Time Step Intervention: An intervention X_{l,k}[s_k] += α × m_f × v_f is injected at each denoising step, without additional training.

Baseline: GSM8K dataset (approximately 8000 elementary school math problems). Metrics: GSM8K accuracy, reasoning score (reasoning markers/operations/structure), concept improvement C(f), steering score S(f). Comparative Experiments: Baseline, positive steering (α=2.0), negative steering (α=-2.0), random feature control group.

Architecture: A modular structure including models, data, training, analysis, steering, and other modules. Usage: Supports local/Colab execution, provides complete pipeline scripts, and allows specifying stages to run. Code repository: https://github.com/Pranaynk07/dlm-reasoning-steering.

Academic Significance: Proves that DLMs have interpretable reasoning features inside, enhancing controllability. Applications: Educational tutoring, code generation, scientific research assistance, controllable text generation. Limitations: Validated only on DiffuGPT-Medium; generalization needs further investigation. Future: Extend to large-scale models, automated feature discovery, multi-task transfer, integration with reinforcement learning.

This study achieves precise guidance of DLM reasoning ability during the inference phase, improving mathematical problem-solving ability without training. It has both academic value (revealing the reasoning mechanism of DLMs) and practical prospects; such technologies are expected to become standard tools for AI controllability.