STEPER: Step-wise Knowledge Distillation for Enhancing Reasoning Ability in Multi-Step Retrieval-Augmented Language Models

AI summary

Overview

• Research area: Natural Language Processing, specifically knowledge distillation and multi-step retrieval-augmented generation (RAG) for multi-hop question answering.
• Technical level: Advanced. The paper assumes familiarity with retrieval-augmented language models, teacher-student distillation, multi-hop QA benchmarks, and multi-task loss weighting.
• Scope: The paper introduces StepER, a training framework that distills step-specific reasoning abilities from a large teacher model into smaller student models for multi-step retrieval-augmented question answering, evaluated on three multi-hop QA benchmarks.

What This Paper Is About

Multi-step retrieval-augmented language models answer complex questions by retrieving evidence, reasoning over it, and retrieving again, but the strong reasoning needed for this is mostly found in very large models that are expensive to run. Existing knowledge distillation methods train a smaller student model to imitate the teacher's entire reasoning path at once, ignoring the fact that different steps require different reasoning abilities and different amounts of accessible information. StepER addresses this by supervising the student at each individual reasoning step and by adaptively reweighting training toward whichever step is currently most learnable.

Key Contributions

1. The paper categorizes the essential reasoning abilities required in multi-step retrieval settings into three stages—Reasoning Initialization, Reasoning Expansion, and Reasoning Aggregation—and argues that each needs separate support, illustrating this with a comparison of Vanilla-KD and StepER in Figure 1.
2. It proposes StepER, a framework that builds a step-wise dataset from a teacher multi-step retrieval-augmented LM and combines step-wise supervision with reasoning difficulty-aware training to learn these abilities.
3. It shows that StepER is flexible enough to integrate with different multi-step retrieval frameworks, demonstrated on both query-updating approaches (IRCOT-style) and question-decomposition approaches (Self-Ask), where it outperforms Vanilla-KD on HotpotQA and MuSiQue.
4. It reports extensive analyses—including GPT-4o-based evaluation of the three reasoning stages, ablation over step-data configurations, model scalability from 0.5B to 7B, rationale validity on SubQA, and out-of-domain transfer—showing improved overall performance and scalability.

Main Findings

• Best distillation result overall: With Llama3.1-Instruct as the base model, StepER achieves average scores of EM 46.06, F1 56.12, and Acc 53.69 across 2Wiki, HotpotQA, and MuSiQue, compared with Vanilla-KD at EM 42.46, F1 51.76, and Acc 49.03.
• Gap over Vanilla-KD: The paper reports an average accuracy improvement of approximately 9.5% over Vanilla-KD, and states StepER outperforms all baselines on 2Wiki and HotpotQA.
• Per-dataset numbers for StepER: 2Wiki EM 63.60, F1 69.45, Acc 66.00; HotpotQA EM 51.00, F1 62.80, Acc 61.00; MuSiQue EM 23.59, F1 36.13, Acc 34.07.
• Comparable to a much larger teacher: The abstract states that an 8B model achieves performance comparable to a 70B teacher model. In Table 1, StepER's 8B average Acc of 53.69 exceeds multi-step ICL baselines such as IRCOT 70B (51.45), Self-Ask 70B (51.40), and ReAct 70B (50.20).
• Step data matters for each ability: Using GPT-4o binary classification on HotpotQA, adding First-step data strengthens Reasoning Initialization, adding First-step plus First Mid-step data improves Reasoning Expansion, and StepER with all step data achieves the highest performance across all evaluation criteria.
• Difficulty-aware weighting beats fixed weights: On HotpotQA, the Difficulty-Aware strategy reaches EM 51.00, F1 62.80, Acc 61.00 versus Uniform weighting (50.40 / 61.57 / 58.40), Weight First (49.10 / 61.63 / 57.70), and Weight Last (48.80 / 60.78 / 58.00). On MuSiQue it reaches EM 23.59, F1 36.13, Acc 34.07 versus Uniform (21.67 / 33.28 / 33.58), Weight First (21.04 / 31.24 / 32.46), and Weight Last (21.91 / 33.85 / 33.37).
• Works with question decomposition too: Under Self-Ask, StepER achieves HotpotQA EM 49.80, F1 62.33, Acc 57.80 and MuSiQue EM 28.20, F1 40.52, Acc 38.80, beating both Self-Ask 8B and Vanilla-KD. The paper reports accuracy boosts of 9.6% on HotpotQA and 14.95% on MuSiQue over the Self-Ask 8B baseline.
• Scalability across Qwen2.5-Instruct sizes: StepER achieves the highest accuracy across 0.5B, 1.5B, 3B, and 7B on HotpotQA. The StepER 3B model nearly matches the Qwen2.5-Instruct 72B teacher while StepER 7B surpasses it; StepER 3B outperforms Vanilla-KD 7B, and StepER 1.5B surpasses Vanilla-KD 3B.
• More valid rationales: On SubQA, StepER scores Accuracy 55.70, SQA 72.00, and RI 87.90, compared with Vanilla-KD (50.30 / 62.95 / 81.00 vpc), IRCOT 8B (48.35 / 66.70 / 86.00), and IRCOT 70B (56.76 / 70.35 / 86.00).
• Better cross-domain transfer: In out-of-domain scenarios (HQ→2W, HQ→MQ, MQ→2W, MQ→HQ), StepER consistently outperforms Vanilla-KD with gains ranging from 1% to 4%.
• Retrieval step sensitivity of the teacher: IRCOT 70B teacher accuracy by maximum retrieval steps is 56.45 (3), 56.81 (4), 57.23 (5), 56.85 (6), and 56.43 (7), which the paper uses to note that answers generally appear within 2-3 steps.

Methodology in Plain English

The researchers start with a teacher model (Llama3.1-Instruct 70B) and a set of multi-hop questions. For each question, the teacher retrieves passages and reasons in stages. The framework splits this into three stages:

1. Reasoning Initialization – the teacher retrieves the first passages using the question itself and generates an initial reasoning step from limited information.
2. Reasoning Expansion – at each later step, a step-search query derived from the partial reasoning chain (either a previous reasoning step or a decomposed question) retrieves additional passages, and the teacher produces the next reasoning step given the question, all passages retrieved so far, and the reasoning so far. If a reasoning step already contains an answer flag such as "So the answer is:", the chain is recorded and expansion stops early.
3. Reasoning Aggregation – the teacher combines all prior reasoning and passages and appends the correct answer.

Any sample where the teacher's final statement does not match the ground-truth answer is filtered out, so the step-wise dataset contains only correct reasoning processes.

Training then uses a multi-task objective with three loss terms, one for each stage: an initialization loss, a summed expansion loss over intermediate steps, and an aggregation loss where the previous reasoning is concatenated with the gold answer. Instead of weighting these equally, the method treats the difficulty of each stage as a trainable parameter (sigma) and scales each loss by 1/(2·sigma²) plus a log-sigma regularization term, so harder tasks receive higher sigma and easier ones lower sigma. This lets the model shift focus across stages as training progresses.

Implementation details: BM25 retriever, maximum S = 5 retrieval steps, top-K = 4 passages per step, learning rate 5×10⁻⁶, 2 total epochs, cosine scheduler with linear warmup, trained on 4×A100 GPUs with DeepSpeed ZeRO Stage 3 and gradient checkpointing. Training data was built by sampling 50,000 instances from each dataset's training split (the full MuSiQue training set, which has fewer than 50,000), with final filtered training samples of 33,584 for 2WikiMultiHopQA, 30,572 for HotpotQA, and 5,515 for MuSiQue; validation and test sets were 500 randomly sampled instances each.

Why This Matters

Impact on research. The paper argues that prior distillation work overlooks the fact that reasoning demands change from step to step, and shows that decomposing supervision by stage produces measurable gains in reasoning quality, rationale validity, and cross-domain transfer. It also frames step-wise distillation as a way to close the gap between small student models and much larger teachers without adding the extra models that some prior approaches require.

Real-world applications:

• Multi-hop question answering systems over large document collections, where a question requires chaining facts across several sources.
• Cost-constrained deployment of reasoning assistants, since StepER reports that an 8B model can perform comparably to a 70B teacher and that smaller students can bridge model-scale gaps.
• Domain-specific expert support tools that need staged evidence gathering and final synthesis, mirroring the paper's motivating analogy of a doctor diagnosing ankle pain through initial symptoms, additional tests, and final diagnosis.
• Enterprise or research search assistants that must generate justifiable intermediate rationales, given the SubQA results on reasoning integrity.

Industry relevance. The method reduces inference cost by enabling smaller models to handle complex queries, and the paper reports it is broadly compatible with multi-step retrieval frameworks (both query-updating and question-decomposition styles) and across model families and sizes, which matters for teams that must choose between a large expensive model and a smaller deployable one.

Future Directions

• More fine-grained, step-wise filtering of the distilled dataset based on the validity of the reasoning path, rather than filtering solely on final-answer correctness, which the paper notes currently allows wrong reasoning that coincidentally reaches a correct answer.
• Applying parameter-efficient fine-tuning together with StepER to improve training efficiency and make the framework more practical.
• Extending the approach and evaluating it across a wider range of multi-step retrieval frameworks beyond the query-updating and question-decomposition approaches tested here, including combinations with methods that learn when to retrieve.
• Further investigation of how the learned difficulty parameters evolve during training and how that evolution relates to final performance, which the paper points to in an appendix figure rather than analyzing in the main text.

Target Audience

Researchers and practitioners working on retrieval-augmented generation, knowledge distillation, and multi-hop question answering who want smaller models to inherit complex reasoning behavior; engineers deploying reasoning systems under compute or latency budgets; and readers interested in how training objectives can be structured around the distinct stages of a reasoning process.

Authors’ abstract

Answering complex real-world questions requires step-by-step retrieval and integration of relevant information to generate well-grounded responses. However, existing knowledge distillation methods overlook the need for different reasoning abilities at different steps, hindering transfer in multi-step retrieval-augmented frameworks. To address this, we propose Stepwise Knowledge Distillation for Enhancing Reasoning Ability in Multi-Step Retrieval-Augmented Language Models (StepER). StepER employs step-wise supervision to align with evolving information and reasoning demands across stages. Additionally, it incorporates difficulty-aware training to progressively optimize learning by prioritizing suitable steps. Our method is adaptable to various multi-step retrieval-augmented language models, including those that use retrieval queries for reasoning paths or decomposed questions. Extensive experiments show that StepER outperforms prior methods on multi-hop QA benchmarks, with an 8B model achieving performance comparable to a 70B teacher model.