Distilación de Agentes de Modelos de Lenguaje Grande en Modelos Pequeños con Herramientas de Recuperación y Código

Distilling LLM Agent into Small Models with Retrieval and Code Tools

May 23, 2025

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

cs.AI

Resumen

Los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés) sobresalen en tareas de razonamiento complejo, pero siguen siendo computacionalmente costosos, lo que limita su implementación práctica. Para abordar esto, trabajos recientes se han centrado en destilar capacidades de razonamiento en modelos de lenguaje más pequeños (sLMs) utilizando trazas de cadena de pensamiento (CoT) de LLMs maestros. Sin embargo, este enfoque tiene dificultades en escenarios que requieren conocimiento factual raro o cálculos precisos, donde los sLMs a menudo alucinan debido a su capacidad limitada. En este trabajo, proponemos Distilación de Agentes, un marco para transferir no solo la capacidad de razonamiento, sino también el comportamiento completo de resolución de tareas desde agentes basados en LLMs hacia sLMs equipados con herramientas de recuperación y código. Mejoramos la distilación de agentes en dos ejes complementarios: (1) introducimos un método de prompting llamado prefijo de primer pensamiento para mejorar la calidad de las trayectorias generadas por el maestro; y (2) proponemos una generación de acciones auto-consistente para mejorar la robustez en tiempo de prueba de los agentes pequeños. Evaluamos nuestro método en ocho tareas de razonamiento en dominios factuales y matemáticos, cubriendo tanto la generalización dentro del dominio como fuera de él. Nuestros resultados muestran que sLMs con tan solo 0.5B, 1.5B y 3B parámetros pueden alcanzar un rendimiento competitivo con modelos más grandes de 1.5B, 3B y 7B ajustados mediante distilación CoT, demostrando el potencial de la distilación de agentes para construir agentes pequeños prácticos que utilizan herramientas. Nuestro código está disponible en 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.