🤖 AI Summary
This work addresses the instability and poor interpretability of large reasoning models in multi-step reasoning, often caused by disrupted information flow. The study identifies and formally defines two previously unrecognized failure modes: Shallow Lock-in and Deep Decay. To mitigate these issues without retraining, the authors propose StepFlow, a test-time intervention that constructs Step-Saliency maps by integrating attention and gradient signals. StepFlow corrects shallow-layer saliency via an Odds-Equal Bridge and restores deep-layer information flow through Step Momentum Injection, which reintroduces residual signals into deeper layers. Evaluated across multiple large reasoning models, StepFlow consistently improves accuracy on mathematical, scientific, and programming tasks, demonstrating that targeted repair of information flow effectively recovers model reasoning capabilities.
📝 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 measured by Step-Saliency 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.