Find Tailored Step Example for Next Step: a Targeted Step-wise Retrieval Framework for Guiding LLM Reasoning

Abstract

Large language models (LLMs) have shown strong performance in mathematical reasoning, supported by approaches such as In-Context Learning (ICL) and Retrieval-Augmented Generation (RAG). However, existing methods often provide problem-level examples, which is too coarse-grained for multi-step reasoning to cause informational redundancy, and structural misalignment. To address this limitation, we propose Step-wise Training for In-context Reasoning (STIR) to provide step-synchronized and logically targeted guidance to enhance the model's mathematical reasoning capabilities. STIR enables a model to dynamically decide when to retrieve a single logically consistent next step, just using the current problem and its intermediate state as the query. First, We decompose expert solutions into Step-Level Reasoning Units inspired by human thinking patterns. Leveraging this data, a Step Retriever is trained for logical continuity to map current reasoning states to relevant subsequent steps. Then a Step Reasoner is trained to decide when to retrieve tailored step examples and incorporates this guidance into reasoning. We further extend STIR with a Process-aware Reinforcement Learning phase using Group Relative Policy Optimization to learn to self-formulate search queries and optimizes the decision-making policy. Experiments on seven benchmarks demonstrate that STIR achieves accuracy improvements ranging from 1.86% to 17.96%, maintaining lower token efficiency than baselines. Analysis via our proposed DSM, TCN and RCR metrics shows that STIR improves reasoning capability, achieving DSM scores ranging from 4.54 to 29.87 across backbones and significant improvements over the baseline in both TCN and RCR.