Skip to content

Distilling LLM Agent into Small Models with Retrieval and Code Tools

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

Agent Distillation transfers full agentic behavior—reason-act-observe trajectories with retrieval and code tools—from large LLMs (32B) into small language models (0.5B–7B), going beyond static chain-of-thought distillation. Two techniques improve trajectory quality (first-thought prefix) and test-time robustness (self-consistent action generation). Distilled 0.5B–3B agents match CoT-distilled models that are 2–4× larger.

Problem

Chain-of-thought (CoT) distillation from large teacher models fails for sLMs on tasks requiring rare factual knowledge or precise arithmetic, because sLMs must memorize knowledge and computations rather than learning reusable problem-solving strategies. When facts or numbers fall outside the training distribution, distilled CoT models hallucinate. Existing tool-augmented distillation methods rely on static demonstrations and do not teach interactive, environment-grounded behavior.

Method

Agent Distillation fine-tunes sLMs on teacher-generated reason-act-observe trajectories (CodeAct format: Thought → Python code → interpreter Observation), training on thoughts and actions only (observations excluded from loss). The student learns to call retrieval and code execution tools adaptively rather than memorizing answers.

Two complementary improvements:

  1. First-thought prefix (ftp): Instruction-tuned teachers suffer distributional drift under agent prompts, producing lower-quality first reasoning steps. ftp first prompts the teacher with a CoT prompt to elicit a high-quality initial reasoning step y₁, then uses y₁ as a prefix when sampling the full agent trajectory. This anchors the teacher's reasoning before it enters the agentic format. The prefix is only used at trajectory generation time, not at student inference.

  2. Self-consistent action generation (sag): At inference, the student samples N=8 candidate thought-action sequences via nucleus sampling (temperature 0.4). Sequences with parse or execution errors are filtered by a code interpreter. Majority voting over the observations of surviving candidates selects the final action; if all fail, one is kept and its error is fed back for self-correction.

Training uses LoRA (rank 64) on Qwen2.5-Instruct series (0.5B–7B) with a 32B teacher, ~2,000 filtered trajectories per domain, 2 epochs on 4× A100-80GB.

Key Contributions

  • Agent Distillation framework: distilling interactive agentic trajectories (not static CoT) from a 32B LLM into sLMs as small as 0.5B, teaching tool-augmented problem-solving rather than knowledge memorization.
  • First-thought prefix: a training-time trajectory improvement that resolves distributional drift in instruction-tuned teacher agents, improving trajectory quality without teacher fine-tuning.
  • Self-consistent action generation: a test-time compute method that filters invalid code actions and applies majority voting over observations to boost sLM robustness.
  • Empirical demonstration across 8 benchmarks (4 factual multi-hop QA, 4 math) covering in-domain and out-of-domain generalization, with ablations across 4 Qwen2.5 sizes, Llama-3.2-1B, and Phi-4-mini.

Results

  • Agent vs. CoT distillation (same model size): Agent distillation consistently outperforms CoT distillation on out-of-domain tasks; gains are most pronounced at smaller sizes and harder tasks.
  • Cross-tier comparison: 0.5B agent (avg 19.24) ≈ 1.5B CoT (avg 21.28); 1.5B agent (avg 28.06) ≈ 3B CoT (avg 27.72); 3B agent (avg 33.60) > 7B CoT (avg 33.54); 7B agent (avg 39.85) > 32B CoT teacher (avg 39.54).
  • ftp+sag (3B): average 36.60 vs. bare agent distillation 33.60, vs. CoT 27.72.
  • AIME (3B, ftp+sag): 15.6 vs. CoT 5.6; GSM-Hard (3B, ftp+sag): 65.4 vs. CoT 56.8.
  • RAG-enhanced CoT is not competitive with agent distillation despite also having retrieval access.
  • Cross-family generalization: Llama-3.2-1B agent (ftp+sag): avg 23.97 vs. CoT 13.23; Phi-4-mini (3.8B) agent (ftp+sag): avg 35.79 vs. CoT 27.41.
  • Code-specific models (Qwen2.5-Coder) yield only marginal improvement over general models; code knowledge is not the primary bottleneck.
  • Token counts are comparable between agent and CoT; agents generate slightly more tokens on factual tasks (multiple retrieval calls) but fewer on math tasks (code replaces verbose calculation).

Limitations

  • ftp can cause hallucination: by seeding the first reasoning step with CoT-style factual statements, it reduces retrieval calls and can lead smaller models to rely on internal (potentially wrong) knowledge.
  • sag adds inference cost: N=8 parallel samples per step increases compute at test time; there is no analysis of latency overhead.
  • MATH500 regression for larger students: 3B and 7B agents underperform CoT-distilled counterparts on MATH500, likely because Qwen2.5 is heavily instruction-tuned on college math aligned with CoT format; analytic problem types (e.g., precalculus) that require symbolic reasoning rather than numeric computation resist code-based solutions.
  • Does not improve core reasoning: agent distillation teaches tool use but not deeper inference; RL in tool-augmented environments is identified as a needed next step.
  • Smaller models underuse retrieval: 0.5B/1.5B agents over-rely on initially retrieved documents rather than re-querying, limiting factual accuracy on multi-hop tasks.
  • Static Wikipedia knowledge base: experiments use Wikipedia 2018 corpus rather than live search, limiting ecological validity for time-sensitive factual queries.

Relevance to Agentic AI / LLM Agents

This paper directly addresses the scalability bottleneck for deploying LLM agents: large models (≥30B) are expensive at inference, while small models cannot reliably act as agents when trained only on CoT traces. Agent distillation shows that the ReAct/CodeAct interactive loop—the core architecture of most LLM agent frameworks—can be compressed into sub-3B models with competitive capability, opening a path to on-device or low-latency agents. The first-thought prefix addresses a general failure mode of instruction-tuned models when operating under agent prompts (distributional drift), which is relevant to any agent pipeline that repurposes chat-tuned models. Self-consistent action generation is a practical test-time compute strategy that complements process reward models and is immediately applicable to existing agentic frameworks that expose a code interpreter. Together, these results shift the practical frontier for tool-using agents toward much smaller and cheaper models.