Research map / Large language models / Reasoning and chain of thought
Distilling LLM Agent into Small Models with Retrieval and Code Tools
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Distills LLM agent trajectories (thought-code-observation with retrieval and code tools) into small Qwen2.5 models using a first-thought prefix for teacher trajectories and self-consistent action generation at test time, matching next-tier larger CoT-distilled models.
Problem
Small language models hallucinate on rare factual knowledge and precise computation; CoT distillation does not teach them to use tools, so they fail on out-of-distribution knowledge and calculation tasks.
Approach
Distill full agent trajectories (reason-act-observe with Wikipedia retrieval and Python code execution, CodeAct style) from a 32B teacher into 0.5B-7B students, improve teacher trajectories by prefixing the agent's first thought with a CoT first step, and at test time sample multiple actions per step, filter failed code, and pick the majority-consistent observation.
Key results
- The 3B agent with FTP+SAG averages 36.60 across eight tasks, versus 29.27 for the 3B CoT prompting baseline and 27.72 for CoT distillation.
- The 0.5B agent with FTP+SAG reaches 21.90 average, matching the 1.5B CoT-distilled model (21.28) and exceeding it; the 7B agent with FTP+SAG reaches 42.68 average, above the 32B teacher's CoT prompting (39.54).
- On Llama-3.2-1B-Instruct, FTP+SAG raises average from 13.23 (CoT distill) to 23.97; on Phi-4-mini, from 27.41 (CoT distill) to 35.79.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| distill reason-act-observe agent trajectories instead of static CoT traces | Qwen2.5-1.5B-Instruct · Average across 8 tasks | average score: 21.28 -> 30.55 | Helped |
| add first-thought prefix to teacher agent trajectories (ablation) | Qwen2.5-1.5B-Instruct agent distillation · Average across 8 tasks | average score: 28.06 -> 29.11 | Helped |
| add self-consistent action generation (ablation) | Qwen2.5-1.5B-Instruct agent distillation · Average across 8 tasks | average score: 28.06 -> 30.29 | Helped |
| add both first-thought prefix and self-consistent action generation (ablation) | Qwen2.5-1.5B-Instruct agent distillation · Average across 8 tasks | average score: 28.06 -> 30.55 | Helped |
| add first-thought prefix to teacher agent trajectories (ablation) | Qwen2.5-3B agent distillation · MATH500 | exact match: 54.0 -> 54.4 | Helped |
| add self-consistent action generation (ablation) | Qwen2.5-3B agent distillation · MATH500 | exact match: 54.0 -> 57.4 | Helped |
| add first-thought prefix to teacher agent trajectories (ablation) | Qwen2.5-0.5B agent distillation · Bamboogle | LLM-as-a-judge score: 17.6 -> 24.0 | Helped |
| add first-thought prefix to teacher agent trajectories (ablation) | Qwen2.5-1.5B agent distillation · MuSiQue | LLM-as-a-judge score: 9.0 -> 8.0 | Hurt |
| use Qwen2.5-Coder-32B-Instruct as the teacher instead of the general 32B model (ablation) | Qwen2.5-1.5B agent distillation · Average across 8 tasks | average score: 30.55 -> 30.91 | Helped |
| use a code-specific student instead of a general student (ablation) | Qwen2.5-1.5B agent distillation · Average across 8 tasks | average score: 30.55 -> 29.52 | Hurt |
| set SAG temperature to 1.0 instead of 0.4 (ablation) | Qwen2.5-1.5B-Instruct agent · Average across math benchmarks | average score: 30.59 -> 31.24 | Helped |
| replace LoRA with full fine-tuning (ablation) | Qwen2.5-1.5B-Instruct agent distillation · Average across 8 tasks | average score: 29.11 -> 26.24 | Hurt |
Surprising findings
- Prompting Qwen2.5 instruction-tuned models as agents underperforms CoT prompting on MATH500 for the 32B teacher (69.2 vs 79.2), and naive agent distillation hurts 3B MATH500 relative to CoT distillation.
- The first-thought prefix reduces retrieval calls, and can increase hallucination by encouraging the model to generate facts internally instead of retrieving them.
Limitations
- Experiments limited to Qwen2.5 series for main results and a single teacher model (Qwen2.5-32B), with no proprietary teacher tested.
- First-thought prefix can degrade performance when the agent generates facts internally instead of using retrieval.
- Agents underperform CoT on analytic math problems such as precalculus; only retrieval and code tools are studied, and safety of code execution is not addressed.
Experimental setup
Datasets: HotPotQA MATH Bamboogle MuSiQue 2WikiMultiHopQA GSM-Hard AIME OlymMATH Wikipedia 2018
Benchmarks: HotPotQA MATH500 MuSiQue Bamboogle 2WikiMultiHopQA GSM-Hard AIME OlymMATH
Metrics: accuracy exact match average accuracy code parse error rate code execution error rate retrieval call count token count
Baselines: CoT prompting CoT distillation CoT distillation + RAG agent prompting self-consistency CoT FireAct
Builds on: CodeAct ReAct reasoning distillation (CoT distillation) self-consistency smolagents
Closest papers by meaning
- Progress-Aware Distillation for Mitigating Stagnation in Small Language Model Agents (NeurIPS 2026 · similarity 0.91)
- MAGDi: Structured Distillation of Multi-Agent Interaction Graphs Improves Reasoning in Smaller Language Models (ICML 2024 · similarity 0.91)
- AgentBench: Evaluating LLMs as Agents (ICLR 2024 · similarity 0.90)
- StrategyLLM: Large Language Models as Strategy Generators, Executors, Optimizers, and Evaluators for Problem Solving (NeurIPS 2024 · similarity 0.90)
- LLM+AL: Bridging Large Language Models and Action Languages for Complex Reasoning About Actions (AAAI 2025 · similarity 0.90)
- Eliciting Reasoning in Language Models with Cognitive Tools (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.