Distillation d'un agent LLM en petits modèles avec des outils de recherche et de code

Distilling LLM Agent into Small Models with Retrieval and Code Tools

May 23, 2025

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

cs.AI

Résumé

Les grands modèles de langage (LLMs) excellent dans les tâches de raisonnement complexe, mais restent coûteux en termes de calcul, limitant leur déploiement pratique. Pour remédier à cela, des travaux récents se sont concentrés sur la distillation des capacités de raisonnement dans des modèles de langage plus petits (sLMs) en utilisant des traces de chaîne de pensée (CoT) provenant de LLMs enseignants. Cependant, cette approche rencontre des difficultés dans les scénarios nécessitant des connaissances factuelles rares ou des calculs précis, où les sLMs ont souvent tendance à halluciner en raison de leurs capacités limitées. Dans ce travail, nous proposons la Distillation d'Agent, un cadre pour transférer non seulement la capacité de raisonnement, mais aussi le comportement complet de résolution de tâches, des agents basés sur LLM vers des sLMs équipés d'outils de recherche et de code. Nous améliorons la distillation d'agent selon deux axes complémentaires : (1) nous introduisons une méthode d'invite appelée préfixe de première pensée pour améliorer la qualité des trajectoires générées par l'enseignant ; et (2) nous proposons une génération d'actions auto-cohérente pour améliorer la robustesse des petits agents lors des tests. Nous évaluons notre méthode sur huit tâches de raisonnement dans les domaines factuels et mathématiques, couvrant à la fois la généralisation en domaine et hors domaine. Nos résultats montrent que des sLMs aussi petits que 0,5B, 1,5B, 3B paramètres peuvent atteindre des performances compétitives avec des modèles plus grands de 1,5B, 3B, 7B paramètres affinés par distillation CoT, démontrant le potentiel de la distillation d'agent pour construire des petits agents pratiques utilisant des outils. Notre code est disponible à l'adresse 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.