Distilling LLM Agent into Small Models with Retrieval and Code Tools

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Kang, Minki, Jeong, Jongwon, Lee, Seanie, Cho, Jaewoong, Hwang, Sung Ju
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866914138255523840
author Kang, Minki
Jeong, Jongwon
Lee, Seanie
Cho, Jaewoong
Hwang, Sung Ju
author_facet Kang, Minki
Jeong, Jongwon
Lee, Seanie
Cho, Jaewoong
Hwang, Sung Ju
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2505_17612
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Distilling LLM Agent into Small Models with Retrieval and Code Tools
Kang, Minki
Jeong, Jongwon
Lee, Seanie
Cho, Jaewoong
Hwang, Sung Ju
Computation and Language
Artificial Intelligence
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.
title Distilling LLM Agent into Small Models with Retrieval and Code Tools
topic Computation and Language
Artificial Intelligence
url https://arxiv.org/abs/2505.17612