Structured Agent Distillation for Large Language Model

Cet article propose la « Structured Agent Distillation », une méthode qui compresse les agents basés sur les grands modèles de langage en modèles plus petits en segmentant les trajectoires en phases de raisonnement et d'action pour améliorer la fidélité décisionnelle tout en réduisant les coûts d'inférence.

L'essentiel

Imaginez que vous avez un chef étoilé (le grand modèle d'IA) qui est capable de cuisiner des plats complexes en suivant une recette précise. Il ne se contente pas de donner le plat final ; il explique chaque étape : « Je coupe l'oignon, je fais chauffer la poêle, je verse l'huile... ».

Le problème, c'est que ce chef est énorme, très lent et coûte très cher à embaucher pour chaque petit repas. Vous voulez donc embaucher un apprenti (le petit modèle d'IA) pour faire le travail, mais vous voulez qu'il soit aussi intelligent que le chef, sans avoir besoin d'un cerveau de géant.

Jusqu'à présent, la méthode habituelle pour apprendre à l'apprenti était de lui montrer le plat final et de lui dire : « Copie-moi mot à mot ». C'est comme si vous lui disiez : « Répète après moi : 'Je coupe l'oignon' ». Le problème ? L'apprenti apprend à répéter les mots, mais il ne comprend pas pourquoi le chef coupe l'oignon avant de chauffer la poêle. Il devient un perroquet : il peut dire les bonnes choses, mais s'il doit improviser, il perd le fil.

La Solution : La "Distillation Structurée"

Les auteurs de cette recherche proposent une nouvelle méthode appelée Distillation Structurée d'Agents. Voici comment cela fonctionne, avec une analogie simple :

1. Séparer la "Réflexion" de l'"Action"

Au lieu de traiter la conversation du chef comme un long bloc de texte, cette méthode la découpe en deux parties distinctes, comme si on séparait le carnet de notes du chef de ses mouvements de mains.

• La partie [RAISONNEMENT] (Le Carnet de Notes) : C'est là où le chef explique sa logique. « Je dois couper l'oignon car il va brûler si je le mets trop tard. »
• La partie [ACTION] (Les Mains) : C'est là où le chef agit. « Coupe l'oignon. »

2. Entraîner l'Apprenti avec deux professeurs différents

Dans la méthode classique, l'apprenti écoute tout d'un coup. Ici, on utilise une approche en deux temps :

• Le Professeur de Logique : Il ne regarde que le carnet de notes. Il dit à l'apprenti : « Tu dois apprendre à penser comme le chef. Si le chef explique pourquoi il coupe l'oignon, tu dois aussi expliquer pourquoi, pas juste dire 'coupe'. »
• Le Professeur d'Action : Lui, il ne regarde que les mouvements de mains. Il dit : « Peu importe comment tu l'expliques, l'important est que tu saches quand et comment couper l'oignon pour obtenir le bon résultat. »

3. Pourquoi c'est magique ?

Imaginez que vous apprenez à conduire.

• L'ancienne méthode (Distillation par token) : Vous apprenez à répéter les ordres du moniteur : « Tourne à gauche, tourne à droite ». Si la route change, vous paniquez car vous ne savez pas pourquoi on tourne.
• La nouvelle méthode (Distillation Structurée) : Le moniteur vous dit : « Regarde, je tourne à gauche PARCE QUE il y a un virage (Raisonnement), et ENSUITE je tourne le volant (Action). »

Grâce à cette séparation, l'apprenti comprend la structure de la décision. Il apprend à distinguer le moment où il doit réfléchir du moment où il doit agir.

Les Résultats Concrets

Grâce à cette méthode, les chercheurs ont montré que :

1. L'apprenti est plus intelligent : Même avec un cerveau plus petit (un modèle plus petit), il réussit mieux ses tâches car il a compris la logique, pas juste les mots.
2. Il est plus rapide : Il ne perd pas de temps à réfléchir à des choses inutiles ou à faire des détours, car il a bien intégré le schéma de décision.
3. Il est plus fiable : S'il doit résoudre un problème nouveau, il sait comment s'y prendre car il a appris à séparer la pensée de l'action, tout comme le chef.

En résumé

Cette recherche nous dit : Ne faites pas juste copier-coller à vos IA. Pour créer de petits agents intelligents et efficaces, il faut leur apprendre à structurer leur pensée en séparant clairement le moment où ils réfléchissent du moment où ils agissent. C'est comme passer d'un apprenti qui répète des phrases à un apprenti qui comprend la recette et sait cuisiner par lui-même.

Résumé technique

1. Problématique

Les modèles de langage (LLM) utilisés comme agents autonomes (via des cadres comme ReAct) excellent dans la prise de décision en intercalant raisonnement et actions. Cependant, leur déploiement pratique est entravé par deux facteurs majeurs :

• Coût d'inférence élevé et taille des modèles.
• Limites de la distillation classique : Les méthodes actuelles de distillation de connaissances (Knowledge Distillation - KD) traitent les trajectoires d'agents comme de simples séquences de tokens plates. Elles appliquent une supervision au niveau du token (ex: minimisation de la divergence KL token par token), ignorant la structure compositionnelle de l'agent.

Les limites de l'approche "Token-Level" :

1. Elle néglige les dépendances à long terme entre le raisonnement et l'action.
2. Elle efface la distinction sémantique entre la phase de planification (raisonnement) et la phase d'exécution (action).
3. Elle provoque un dérive sémantique lors des simulations (rollouts), dégradant la cohérence et le taux de réussite des tâches.

2. Méthodologie : Structured Agent Distillation (SAD)

Les auteurs proposent SAD, le premier cadre de distillation conçu spécifiquement pour les agents ReAct, qui préserve à la fois la fidélité du raisonnement et la cohérence des actions.

A. Segmentation Structurée des Trajectoires

Contrairement aux approches plates, SAD segmente explicitement chaque trajectoire d'agent en deux spans (segments) distincts :

• Span [REASON] : Contient les traces de raisonnement (Chain-of-Thought, étapes intermédiaires).
• Span [ACT] : Contient les actions exécutées (appels d'outils, commandes, réponses finales).
• Note : Les observations de l'environnement sont également identifiées mais exclues de la fonction de perte principale pour éviter le surapprentissage sur des feedbacks déterministes.

B. Objectifs d'Apprentissage Spécifiques aux Segments

Le modèle étudiant (élève) est entraîné avec deux pertes distinctes, appliquées via des masques binaires ($m_r$ et $m_a$) :

1. Alignement de la Politique CoT ($L_{CoT}$) : Minimise la divergence KL entre les distributions du professeur et de l'élève uniquement sur les tokens de raisonnement. Cela force l'élève à imiter la logique de pensée du professeur.
2. Cohérence des Actions ($L_{Act}$) : Minimise la divergence KL sur l'espace d'actions discret uniquement sur les tokens d'action. Cela garantit que l'élève prend les bonnes décisions d'exécution.

La perte totale est une somme pondérée :
$$L_{total} = \lambda_r \cdot L_{CoT} + \lambda_a \cdot L_{Act}$$
(Où $\lambda_r = \lambda_a = 1.0$ par défaut).

C. Perspective d'Optimisation (Projection de Gradient)

L'article offre une analyse théorique montrant que la distillation token-level couple des gradients hétérogènes (raisonnement fréquent vs actions rares mais critiques), créant un angle de conflit dans l'espace des paramètres.
SAD résout ce problème par une projection orthogonale des gradients :

• Les gradients du raisonnement sont projetés sur le sous-espace $V_r$.
• Les gradients d'action sont projetés sur le sous-espace $V_a$.
• Cela élimine les interférences inter-segments et permet une mise à jour géométrique plus précise, séparant l'apprentissage de la logique haute niveau de l'exécution basse niveau.

D. Apprentissage par Curriculum

Pour améliorer la stabilité, les trajectoires sont triées par complexité (basée sur la longueur des segments et l'entropie du professeur) et présentées progressivement à l'élève, des exemples simples aux plus complexes.

3. Contributions Clés

1. Première distillation structurée : SAD est la première méthode à segmenter les trajectoires ReAct en spans de raisonnement et d'action pour une supervision fine, surpassant la distillation token-level naïve.
2. Découplage sémantique : La méthode permet aux agents compacts d'apprendre indépendamment le raisonnement symbolique et l'exécution ancrée, évitant le compromis habituel entre les deux.
3. Validation empirique robuste : Résultats démontrés sur trois benchmarks majeurs (ALFWorld, WebShop, HotPotQA-ReAct) avec différents modèles professeurs (GPT-2, OPT, LLaMA, Orca2) et tailles d'élèves (de 120M à 760M paramètres).

4. Résultats Expérimentaux

Les expériences montrent que SAD surpasse systématiquement les baselines (Token-level KD, SeqKD, MiniLLM) :

• Taux de réussite des tâches (Task Success) : Améliorations significatives, notamment sur les petits modèles. Par exemple, sur ALFWorld avec un modèle de 120M, SAD atteint 43,7 % contre 39,4 % pour la meilleure baseline token-level.
• Efficacité du raisonnement : Les agents entraînés avec SAD génèrent des chaînes de pensée plus courtes et plus concises (réduction de la longueur moyenne des tokens de raisonnement).
• Cohérence CoT (Chain-of-Thought) : Taux de correspondance plus élevé avec le raisonnement du professeur, indiquant une meilleure fidélité structurelle.
• Latence : Réduction du nombre d'étapes (actions + raisonnement) nécessaires pour compléter une tâche, accélérant ainsi l'exécution.
• Évolutivité (Scaling) : Les gains sont particulièrement prononcés pour les modèles de petite taille (120M, 340M), où SAD atténue la dégradation de performance souvent observée avec la distillation classique.

5. Signification et Impact

Ce travail est significatif car il démontre que la structure de la trajectoire est aussi importante que le contenu des tokens pour l'apprentissage des agents.

• Déploiement pratique : SAD permet de compresser des agents complexes en modèles légers et déployables sans sacrifier leur capacité de planification.
• Nouvelle direction pour la distillation : Il suggère que pour les tâches interactives et décisionnelles, les méthodes de distillation doivent évoluer au-delà de la prédiction de tokens pour adopter une supervision structurelle et fonctionnelle.
• Généralisation : La méthode est agnostique à l'architecture et s'applique aussi bien aux modèles decoder-only qu'aux modèles instruction-tuned, ouvrant la voie à la compression d'agents basés sur des LLMs de pointe (Frontier LLMs).

En résumé, Structured Agent Distillation établit un nouveau standard pour la compression d'agents LLM en traitant le raisonnement et l'action comme des composantes distinctes mais complémentaires, nécessitant des objectifs d'apprentissage spécialisés pour une imitation fidèle.