🔥🔥Reasoning Fails Where Step Flow Breaks

🎉 Accepted to ACL 2026 Main Conference (CCF-A) 🎉

Paper  |  PDF  |  Code

Xiaoyu Xu^1,2, Yulan Pan², Xiaosong Yuan^3‡, Zhihong Shen², Minghao Su², Yuanhao Su², Xiaofeng Zhang^1†

¹Shanghai Jiao Tong University    ²Fuzhou University    ³Jilin University

^†Corresponding author    ^‡Project leader

---

Abstract

Large reasoning models (LRMs) that generate long chains of thought now perform well on multi-step math, science, and coding tasks. However, their behavior is still unstable and hard to interpret, and existing analysis tools struggle with such long, structured reasoning traces. We introduce Step-Saliency, which pools attention–gradient scores into step-to-step maps along the question–thinking–summary trajectory. Across several models, Step-Saliency reveals two recurring information-flow failures: Shallow Lock-in, where shallow layers over-focus on the current step and barely use earlier context, and Deep Decay, where deep layers gradually lose saliency on the thinking segment and the summary increasingly attends to itself and the last few steps. Motivated by these patterns, we propose StepFlow, a saliency-inspired test-time intervention that adjusts shallow saliency patterns via Odds-Equal Bridge and adds a small step-level residual in deep layers via Step Momentum Injection. StepFlow improves accuracy on math, science, and coding tasks across multiple LRMs without retraining, indicating that repairing information flow can recover part of their missing reasoning performance.

Key Contributions

1. Step-Saliency: A step-level diagnostic tool that aggregates token-level attention–gradient saliency into step→step maps, making long reasoning traces interpretable at the step level.

2. Two failure patterns identified: We discover two depth-wise information-flow failures in large reasoning models:

   • Shallow Lock-in — shallow layers over-attend to the current step and under-use earlier reasoning context.
   • Deep Decay — deep layers progressively lose saliency on the thinking segment; the summary increasingly attends only to itself and the last few steps.

3. StepFlow: A test-time intervention (no retraining required) that repairs information flow via two mechanisms:

   • Odds-Equal Bridge (OEB) — redistributes attention mass in shallow layers to restore cross-step flow.
   • Step Momentum Injection (SMI) — injects a small step-level residual in deep layers to counteract saliency decay.

   StepFlow consistently improves accuracy across DeepSeek-R1-Distill (7B/14B/32B), GPT-OSS-20B, and QwQ-32B on AIME24, AIME25, AMC23, MATH-500, GPQA-Diamond, and LiveCodeBench.

Figures

Figure 1: Token-level saliency maps are dense and noisy; Step-Saliency pools them into question/thinking/summary blocks. Correct traces show smooth information flow, while errors exhibit shallow lock-in and weak thinking→summary links.

Figure 2: Step-Saliency patterns for shallow vs. deep layers and correct vs. error traces. Shallow Lock-in (narrow local flow in error traces) and Deep Decay (faster saliency loss in deep layers).

Project Structure

saliency/
  src/
    generate_saliency_maps.py   # Generate step-level saliency maps
    saliency_extractor.py       # Universal attention saliency extractor
    model_config.py             # Model configurations (GPT-OSS, DeepSeek, QwQ)
    interventions/
      attention_manager.py      # Attention hook manager for interventions
      bridge_guard_oeb.py       # Odds-Equal Bridge (OEB) implementation
      smi.py                    # Step Momentum Injection (SMI) implementation
      state_controller.py       # State tracking for channel segments
  scripts/
    analyze_step_saliency.py    # Step saliency analysis

Installation

pip install -r requirements.txt

GPT-OSS model support requires a custom transformers build that includes transformers.models.gpt_oss. DeepSeek-R1-Distill and QwQ models work with the standard transformers package.

Usage

Generate Saliency Maps

python src/generate_saliency_maps.py \
    --model-path /path/to/model \
    --dataset math \
    --output-dir outputs/saliency

Hyperparameters

Default StepFlow configuration (Table 8 in paper):

Model	OEB layers	SMI layers	tau_max	alpha
R1-Distill-7B/14B/32B	bottom 1/4	top 1/4	0.15	0.06
GPT-OSS-20B	bottom 1/4	top 1/4	0.15	0.06
QwQ-32B	bottom 1/4	top 1/4	0.15	0.06

Override via CLI:

--smi-strength 0.06       # SMI residual scale alpha
--oeb-layers 1,3,5,7     # Specific OEB layers
--tau-max 0.15            # OEB bridge mass upper bound

License

This project is released under the Apache License 2.0.