Interpreting and Controlling LLM Reasoning through Integrated Policy Gradient

TL;DR

This paper tackles the opacity of multi-step LLM reasoning by introducing Integrated Policy Gradient (IPG), a training-free, outcome-oriented attribution method that extends gradient-based analysis from parameter space to representation space and propagates long-horizon signals through inference trajectories via path-integrated gradients. IPG supports both neuron-level states and k-Sparse Autoencoder (k-SAE) feature representations, enabling robust identification of reasoning components and subsequent causal interventions by scaling selected activations with a factor $\gamma$. Empirically, IPG identifies sequentially influential components, shows strong cross-dataset transferability, and demonstrates precise control of reasoning capabilities and reasoning length across open-source models and reasoning benchmarks. The work also shows that IPG components transfer robustly to distilled models, supporting mechanisms that persist under model compression, and provides fine-grained interpretability by linking components to distinct reasoning aspects. Overall, IPG offers a scalable, interpretable, and transferable framework for understanding and steering LLM reasoning in practical settings, with potential extensions to broader behavioral domains.

Abstract

Large language models (LLMs) demonstrate strong reasoning abilities in solving complex real-world problems. Yet, the internal mechanisms driving these complex reasoning behaviors remain opaque. Existing interpretability approaches targeting reasoning either identify components (e.g., neurons) correlated with special textual patterns, or rely on human-annotated contrastive pairs to derive control vectors. Consequently, current methods struggle to precisely localize complex reasoning mechanisms or capture sequential influence from model internal workings to the reasoning outputs. In this paper, built on outcome-oriented and sequential-influence-aware principles, we focus on identifying components that have sequential contribution to reasoning behavior where outcomes are cumulated by long-range effects. We propose Integrated Policy Gradient (IPG), a novel framework that attributes reasoning behaviors to model's inner components by propagating compound outcome-based signals such as post reasoning accuracy backward through model inference trajectories. Empirical evaluations demonstrate that our approach achieves more precise localization and enables reliable modulation of reasoning behaviors (e.g., reasoning capability, reasoning strength) across diverse reasoning models.

Figures (12)

• Figure 1: Paradigms of interpretability method for reasoning in LLMs and comparison results on reasoning datasets. (a) Illustration of the three paradigms, regarding (i) text-pattern methods that associate high activations with special tokens (e.g., "Wait"); (ii) control-vector methods that rely on human-craft contrastive input pairs; and (iii) our proposed Integrated Policy Gradient method. $t$ is the reasoning time step. (b) Intervene reasoning component: selected components are scaled by a factor $\gamma$ (see Section \ref{['Sec:intervene']}). (c) Control reasoning behavior with both suppressing and enhancing.
• Figure 2: Effects of targeted interventions on reasoning performance. Left: impact of the intervention scaling factor on Qwen2.5-Math-1.5B-Instruct yang2024qwen25mathtechnicalreportmathematical. Right: impact of the number of selected features on reasoning accuracy.
• Figure 3: Reasoning strength control in GSM8K cobbe2021trainingverifierssolvemath for DeepSeek-R1-Distilled-Qwen-1.5B guo2025deepseek. IPG identifies neurons that is related with reasoning length, suggested by effective control while maintaining high accuracy.
• Figure 4: Responses generated by Qwen2.5-Math-1.5B-Instruct yang2024qwen25mathtechnicalreportmathematical on one example in GSM8K cobbe2021trainingverifierssolvemath, including both original and intervened outputs, with the model's reasoning ability elicited to arrive at the correct answer.
• Figure 5: Dual radar plots of neuron- and feature-level interventions: enhancing neuron #940 and feature #8053 on layer 20.