Distillare un Agente LLM in Modelli Piccoli con Strumenti di Recupero e Codice

Distilling LLM Agent into Small Models with Retrieval and Code Tools

May 23, 2025

Autori: Minki Kang, Jongwon Jeong, Seanie Lee, Jaewoong Cho, Sung Ju Hwang

cs.AI

Abstract

I grandi modelli linguistici (LLM) eccellono in compiti di ragionamento complesso, ma rimangono computazionalmente costosi, limitandone l'implementazione pratica. Per affrontare questo problema, lavori recenti si sono concentrati sulla distillazione delle capacità di ragionamento in modelli linguistici più piccoli (sLM) utilizzando tracce di catena di pensiero (CoT) da LLM insegnanti. Tuttavia, questo approccio incontra difficoltà in scenari che richiedono conoscenze fattuali rare o calcoli precisi, dove gli sLM spesso producono allucinazioni a causa delle capacità limitate. In questo lavoro, proponiamo la Distillazione di Agenti, un framework per trasferire non solo la capacità di ragionamento, ma l'intero comportamento di risoluzione dei compiti da agenti basati su LLM a sLM dotati di strumenti di recupero e codice. Miglioriamo la distillazione di agenti lungo due assi complementari: (1) introduciamo un metodo di prompting chiamato prefisso di primo pensiero per migliorare la qualità delle traiettorie generate dagli insegnanti; e (2) proponiamo una generazione di azioni auto-consistente per migliorare la robustezza al test-time dei piccoli agenti. Valutiamo il nostro metodo su otto compiti di ragionamento in domini fattuali e matematici, coprendo sia la generalizzazione in-dominio che out-of-domain. I nostri risultati mostrano che sLM con soli 0,5B, 1,5B, 3B parametri possono raggiungere prestazioni competitive con modelli più grandi di livello successivo (1,5B, 3B, 7B) fine-tuned utilizzando la distillazione CoT, dimostrando il potenziale della distillazione di agenti per costruire piccoli agenti pratici e dotati di strumenti. Il nostro codice è disponibile all'indirizzo https://github.com/Nardien/agent-distillation.

English

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.