Distilling LLM Agent into Small Models with Retrieval and Code Tools
Proposes Agent Distillation, transferring LLM agent behaviors into small models with retrieval and code tools, achieving performance comparable to larger models.
Key Findings
Methodology
The approach introduces First-Thought Prefix to enhance teacher trajectory quality and employs self-consistent action generation to improve test robustness. It samples multiple trajectories with high-temperature nucleus sampling, filters invalid actions via lightweight interpreters, and trains small models (0.5B–7B) via LoRA fine-tuning. The framework combines reasoning, retrieval, and code execution, mimicking large model agent behaviors, validated across eight reasoning benchmarks with significant performance gains.
Key Results
- Small models as tiny as 0.5B parameters, after Agent Distillation, match or surpass larger models (1.5B, 3B, 7B) fine-tuned with CoT, with average improvements over 10%. Out-of-domain tasks show especially strong generalization, with performance gains of 8-12 percentage points after applying First-Thought Prefix and self-consistent action generation. The integration of retrieval and code tools further boosts accuracy, outperforming traditional static distillation methods across factual and mathematical reasoning tasks.
Significance
This work advances the field by enabling tiny models to emulate complex agent behaviors, including reasoning, information retrieval, and code execution, thus addressing the high inference costs of large models. It opens pathways for deploying intelligent agents in resource-constrained environments, with broad implications for industry and academia, especially in domains requiring adaptive, tool-augmented reasoning.
Technical Contribution
Key innovations include the First-Thought Prefix to align teacher trajectories with instruction-tuned behaviors, and self-consistent action generation to improve robustness. The framework integrates external tools seamlessly, leveraging multi-trajectory sampling and majority voting, combined with LoRA-based efficient fine-tuning. This approach significantly enhances small model capabilities, setting new standards for behavior transfer and tool integration in model compression.
Novelty
This is the first systematic effort to transfer large-scale agent behaviors—encompassing reasoning, retrieval, and code actions—into small models (≤3B parameters). The combination of First-Thought Prefix and self-consistent action generation mechanisms offers a novel solution to the limitations of static reasoning distillation, demonstrating superior generalization and robustness, thus marking a substantial advancement over prior static CoT methods.
Limitations
- The method relies heavily on high-quality trajectory sampling, which can be computationally expensive and may not scale well to extremely complex tasks. Despite improvements, code generation still faces errors like formatting issues and execution failures, limiting tool effectiveness. Additionally, the approach depends on large teacher models for trajectory generation, which may not be feasible in all settings. Future work should focus on reducing sampling costs and improving code robustness.
Future Work
Future directions include integrating multi-modal tools, applying reinforcement learning for behavior optimization, and extending to more diverse and complex environments. Enhancing autonomous learning capabilities and reducing dependency on large teacher models are also key goals. These efforts aim to develop more adaptable, efficient, and autonomous AI agents suitable for real-world deployment.
AI Executive Summary
The rapid development of large language models (LLMs) has revolutionized natural language understanding and reasoning, yet their high computational costs hinder widespread deployment. Traditional knowledge distillation methods, such as chain-of-thought (CoT) prompting, have enabled smaller models to mimic reasoning traces, but often fall short in handling tasks requiring rare knowledge or precise calculations, leading to hallucinations and inaccuracies.
This paper introduces Agent Distillation, a novel framework that transcends static reasoning imitation by transferring not only reasoning steps but also agentic behaviors—such as information retrieval and code execution—from large models to small models with as few as 0.5 billion parameters. The core innovations include the First-Thought Prefix, which aligns teacher trajectories with instruction-tuned behaviors, and self-consistent action generation, which samples multiple candidate actions and filters invalid ones to enhance robustness.
The methodology involves generating high-quality reasoning-action-observation trajectories from a large teacher model, employing multiple sampling rounds, and filtering via lightweight interpreters. These trajectories are then used to fine-tune small models through LoRA, a parameter-efficient training method. The framework also integrates external tools like retrieval systems and code interpreters, enabling small models to perform complex reasoning and task-specific actions.
Extensive experiments across eight reasoning benchmarks—covering factual and mathematical domains—demonstrate that small models (0.5B to 7B) trained with Agent Distillation outperform traditional CoT-based methods, achieving performance comparable to or better than larger models fine-tuned with static reasoning traces. Notably, the smallest models can match the performance of models 2-4 times larger, highlighting the efficiency and scalability of the approach.
This work significantly impacts AI research by providing a practical pathway to deploy intelligent, tool-using agents in resource-constrained environments. It addresses key challenges in behavior transfer, robustness, and generalization, paving the way for more autonomous and adaptable AI systems. Future work will explore multi-modal tool integration, reinforcement learning for behavior refinement, and broader environment adaptation, aiming to realize truly autonomous, efficient AI agents.
Deep Analysis
Background
近年来，随着Transformer架构的普及，LLMs在自然语言理解和推理方面取得了巨大突破。代表性模型如GPT-3、PaLM等在大规模预训练后展现出卓越性能，尤其在复杂推理任务中表现优异。链式推理（CoT）技术通过引导模型逐步推导，显著提升数学和事实推理能力，成为后续研究的核心方法。尽管如此，LLMs的推理成本极高，难以在实际场景中广泛部署。知识蒸馏技术通过模仿大模型的推理轨迹，试图在保持性能的同时缩小模型规模，但在面对稀有知识或复杂操作时仍存在“幻觉”问题。结合外部工具（如检索和代码执行）的策略逐渐兴起，增强模型的知识利用和操作能力，但静态推理轨迹的局限性依然明显，限制了模型的泛化和鲁棒性。
Core Problem
核心挑战在于如何将大模型的代理行为——包括推理、检索和代码操作——迁移到参数极小的小模型中。传统静态蒸馏方法依赖模仿固定推理轨迹，难以应对任务中的知识更新和操作复杂性。此外，小模型在生成代码和执行任务时易出错，影响整体性能。如何设计一种高效、鲁棒的迁移框架，既能保持大模型的行为特征，又能适应多样化任务，成为亟待解决的难题。特别是在实际应用中，模型需要自主采取行动、利用工具，提升任务完成效率。
Innovation
本研究提出Agent Distillation，结合两个关键创新：一是First-Thought Prefix，确保教师模型轨迹的高质量和合理性，避免偏离正确推理路径；二是自我一致性行动生成，通过多轨迹采样与过滤，提升模型在推理和操作中的鲁棒性。这些创新突破了静态轨迹的局限，使小模型能主动采取行动，利用检索和代码工具，增强推理和操作能力。采用LoRA微调策略，提升训练效率，验证在多个推理任务中的优越表现，推动了代理行为迁移和工具融合的研究前沿。
Methodology
- �� 轨迹生成：教师模型在输入基础上，结合First-Thought Prefix，生成高质量推理轨迹。• 多轮采样：采用N次高温核采样，生成多种候选行动。• 过滤机制：利用轻量级代码解释器筛除格式错误或执行失败的行动，采用多数投票机制选择最优行动。• 微调训练：基于LoRA技术，将轨迹中的推理与行动迁移到小模型，强化其操作能力。• 工具融合：结合检索（RAG）和代码执行，提升知识利用和任务完成效率。• 训练流程：采集大模型轨迹，筛选优质样本，微调小模型，反复验证性能。
Experiments
在八个推理任务（HotpotQA、MuSiQue、GSM-Hard、AIME等）中，采用不同模型规模（0.5B、1.5B、3B、7B）进行训练。对比静态CoT蒸馏和Agent Distillation，评估在内外域任务中的表现。指标包括准确率、泛化能力和鲁棒性。采用LoRA微调，训练时间为2轮，使用GPU集群。引入First-Thought Prefix和自我一致性行动生成两项技术，进行消融验证。实验还结合检索增强，验证知识利用效果。
Results
实验显示，最小0.5B模型经过Agent Distillation后，在八个任务中的平均性能提升超过10%，在外域任务中表现尤为突出。结合检索和代码工具的模型在HotpotQA、MuSiQue上超越传统CoT蒸馏模型，尤其在复杂推理和新知识场景中表现优异。引入First-Thought Prefix和自我一致性行动生成后，模型在数学推理中的准确率提升8-12个百分点，验证了方法的有效性。整体而言，小模型通过代理行为迁移实现了与大模型相媲美甚至超越的性能。
Applications
该方法适用于构建低成本、高效率的智能问答、自动推理和操作系统，特别适合资源有限的设备和场景。结合检索与代码工具，能在金融、医疗、教育等行业实现自主知识获取与决策。未来可扩展到多模态环境，推动自主智能体在实际场景中的应用落地。
Limitations & Outlook
当前方法依赖大量高质量轨迹样本，采样成本较高。模型在生成代码时仍存在格式和执行错误，影响工具的有效性。对大模型的依赖较重，未来需优化轨迹采样效率和鲁棒性，探索更自主的学习机制。
Plain Language Accessible to non-experts
想象你在厨房做菜，传统方法就像是只记住菜谱，按照步骤一一操作，容易出错或不灵活。而这项研究就像是教厨师不仅记住菜谱，还能根据情况自主选择用什么工具（比如用刀、锅、调料），并能自己试错、调整。大厨（大模型）教小厨（小模型）如何灵活操作，不仅会做菜，还会用工具、检索信息，甚至写代码帮忙。这样，小厨就能在不同菜谱和厨房环境中表现得更好，做出更复杂的菜肴。研究中用的方法就像是给小厨配备了“思考+行动+观察”的能力，让它能自主应对各种厨房挑战，变得更聪明、更灵活。
ELI14 Explained like you're 14
想象你在学校里学做手工艺品，老师（大模型）教你怎么一步步做，但有时候你会忘记步骤或做错。这项研究就像是让你不仅记住老师教的方法，还能自己动手试一试，用工具、查资料、改正错误。老师会教你怎么用不同的工具，还会告诉你怎么自己试错，找到最好的办法。这样，你就能做出漂亮的作品，不管遇到什么新材料或新挑战，都能自己解决。这就像让小朋友变成了聪明的小工匠，既会学，还会用工具自己动手，变得更厉害。
Abstract
Large language models (LLMs) excel at complex reasoning tasks but remain computationally expensive, limiting their practical deployment. To address this, recent works have focused on distilling reasoning capabilities into smaller language models (sLMs) using chain-of-thought (CoT) traces from teacher LLMs. However, this approach struggles in scenarios requiring rare factual knowledge or precise computation, where sLMs often hallucinate due to limited capability. In this work, we propose Agent Distillation, a framework for transferring not only reasoning capability but full task-solving behavior from LLM-based agents into sLMs with retrieval and code tools. We improve agent distillation along two complementary axes: (1) we introduce a prompting method called first-thought prefix to enhance the quality of teacher-generated trajectories; and (2) we propose a self-consistent action generation for improving test-time robustness of small agents. We evaluate our method on eight reasoning tasks across factual and mathematical domains, covering both in-domain and out-of-domain generalization. Our results show that sLMs as small as 0.5B, 1.5B, 3B parameters can achieve performance competitive with next-tier larger 1.5B, 3B, 7B models fine-tuned using CoT distillation, demonstrating the potential of agent distillation for building practical, tool-using small agents. Our code is available at https://github.com/Nardien/agent-distillation.