Résumé technique : Attribution de crédit basée sur un arbre pour un système de mémoire multi-agent (TreeMem)

Énoncé du problème

Les systèmes de mémoire sont de plus en plus utilisés pour étendre les modèles de langage de grande taille (LLM) au-delà de leurs fenêtres de contexte pour des tâches à long horizon, généralement organisés sous forme de pipelines multi-agents impliquant des agents Constructeur, Résumé et Récupération. Bien que l'apprentissage par renforcement (RL) ait été exploré pour optimiser ces agents, les paradigmes existants rencontrent des limitations significatives :

1. Récompenses finales uniformes : L'application uniforme de la récompense finale de la tâche en aval (par exemple, la précision des questions-réponses) à tous les agents fournit une supervision grossière et ambiguë. Cela échoue à distinguer les contributions spécifiques d'agents hétérogènes, limitant leur capacité à se spécialiser.
2. Récompenses spécifiques à la tâche : Concevoir des récompenses intermédiaires pour des sous-tâches spécifiques (par exemple, la couverture des preuves pour les Constructeurs, la qualité de l'abstraction pour les Résumés) nécessite des annotations manuelles coûteuses et est sujet à l'instabilité. De plus, des récompenses conçues à la main peuvent induire un « piratage de la récompense », où les agents optimisent la métrique sans améliorer les performances en aval.

Il existe un manque de méthodes capables de dériver des signaux de crédit fins et spécifiques à l'agent directement à partir de la récompense finale en aval, sans dépendre d'annotations spécifiques à la tâche coûteuses.

Méthodologie : TreeMem

Les auteurs proposent TreeMem, un cadre d'entraînement par RL basé sur un arbre qui dérive un crédit spécifique à l'agent à partir de la récompense finale grâce à un déroulement en structure d'arbre et à une moyenne de Monte Carlo.

1. Déroulement en structure d'arbre

Au lieu d'un pipeline linéaire, TreeMem étend le flux de travail multi-agent en une structure d'arbre pendant l'entraînement :

• Phase Constructeur : Étant donné un historique $H$, l'agent Constructeur échantillonne $G$ actions candidates de mémoire ($a_1^i$).
• Phase Résumé : Pour chaque action du constructeur, l'agent Résumé échantillonne $J$ branches de mémoire résumée ($a_2^{i,j}$).
• Phase Récupération : Pour chaque paire constructeur-résumé, l'agent Récupération échantillonne $K$ réponses finales ($a_3^{i,j,k}$).
  Cela résulte en $G \times J \times K$ trajectoires complètes, formant un arbre de chemins candidats allant de l'historique brut aux réponses finales.

2. Attribution de crédit par moyenne de Monte Carlo

TreeMem attribue un crédit aux actions intermédiaires en propageant la récompense finale vers le haut à travers l'arbre :

• Agent Récupération : Reçoit la récompense finale directe $R$ de sa réponse générée.
• Agent Résumé : Son crédit est estimé en moyennant les récompenses finales sur les $K$ branches de réponse dépendant de son résumé spécifique.
• Agent Constructeur : Son crédit est estimé en moyennant les récompenses finales sur toutes les $J \times K$ branches dépendant de sa construction de mémoire spécifique, avec une pénalité de longueur ajoutée pour décourager la copie mot à mot d'historiques longs.

Ce processus convertit un signal sparse, uniquement final, en signaux d'optimisation granulaires et spécifiques à l'agent.

3. Optimisation de la politique

Le cadre construit un groupe d'entraînement en sélectionnant aléatoirement une trajectoire complète parmi chacun des $G$ sous-arbres enracinés aux actions du constructeur. Ces trajectoires sélectionnées, appariées à leurs crédits spécifiques à l'agent calculés, sont utilisées pour mettre à jour simultanément les politiques des agents Constructeur, Résumé et Récupération. L'optimisation suit un principe relatif au groupe (inspiré de GRPO), normalisant les crédits au sein du groupe pour calculer les avantages. Cela permet aux agents hétérogènes de se spécialiser efficacement tout en restant alignés sur l'objectif global de la tâche. Notamment, l'expansion de l'arbre est utilisée uniquement pendant l'entraînement ; à l'inférence, le système fonctionne directement avec les politiques optimisées, sans exécuter de branches d'arbre.

Contributions clés

• Nouveau cadre : Introduction de TreeMem, le premier cadre dérivant un crédit spécifique à l'agent pour les systèmes de mémoire multi-agent directement à partir de récompenses finales sans annotations spécifiques à la tâche.
• Estimation en structure d'arbre : Un mécanisme pour modéliser le pipeline multi-agent comme un arbre, utilisant la moyenne de Monte Carlo pour isoler l'impact des actions d'agents intermédiaires sur le résultat final.
• Spécialisation sans annotation : Démontre qu'une attribution de crédit fine peut être réalisée sans l'étiquetage manuel coûteux requis par les méthodes précédentes de récompenses spécifiques à la tâche.
• Optimisation simultanée : Une stratégie pour mettre à jour conjointement les politiques d'agents hétérogènes en utilisant des crédits estimés sur des sous-arbres, favorisant une spécialisation des rôles efficace.

Résultats expérimentaux

TreeMem a été évalué sur des benchmarks à long horizon incluant PersonaMem (contextes de 32K, 128K, 1M tokens), LongMemEval et LOCOMO, en utilisant les modèles de base Qwen et Llama.

• Performance : TreeMem a constamment surpassé des bases solides, y compris des méthodes basées sur le contexte, des systèmes de mémoire pilotés par des invites et des méthodes existantes optimisées par RL (par exemple, Mem1, Memory-R1, CoMAM). Sur PersonaMem, il a réalisé une amélioration moyenne de +7,5 % par rapport à la meilleure base RL.
• Études d'ablation :
  • Signaux de récompense : TreeMem a surpassé les variantes utilisant uniquement des récompenses finales ou des récompenses spécifiques à la tâche conçues manuellement, validant que l'attribution de crédit basée sur un arbre fournit des signaux d'entraînement supérieurs.
  • Impact des agents : La suppression du signal basé sur un arbre pour n'importe quel agent unique (Constructeur ou Résumé) a entraîné une dégradation des performances, confirmant que tous les agents bénéficient d'un crédit spécifique à l'agent.
• Efficacité : Bien que le déroulement en arbre introduise une surcharge de calcul pendant l'entraînement, la méthode accélère la convergence (nécessitant moins d'étapes) et améliore l'efficacité des tokens (produisant des mémoires plus compactes). À l'inférence, elle n'engendre aucune surcharge supplémentaire par rapport aux bases standard.
• Sensibilité aux hyperparamètres : Le système a démontré des performances robustes sur une large gamme de tailles de groupes ($G$) et de nombres de branches ($J, K$), ne nécessitant aucune recherche exhaustive sur grille.

Importance et affirmations

L'article affirme que TreeMem répond à un goulot d'étranglement critique dans l'optimisation de la mémoire multi-agent : l'incapacité à fournir des signaux d'entraînement spécifiques aux agents hétérogènes sans annotations coûteuses. En transformant des récompenses finales grossières en signaux granulaires et spécifiques à l'agent grâce à des déroulements en structure d'arbre, TreeMem permet aux agents de mémoire de se spécialiser efficacement dans leurs rôles respectifs (construction, résumé, récupération).

Les auteurs affirment que cette approche valide l'importance de l'attribution de crédit fine dans les systèmes complexes multi-agents de LLM. Les résultats suggèrent que TreeMem peut améliorer considérablement les performances des systèmes de mémoire à long horizon, les rendant plus fiables et efficaces pour des tâches nécessitant un stockage et une récupération dynamiques d'informations sur des interactions étendues. Ce travail se positionne comme une solution pratique qui équilibre l'alignement au niveau du système avec le besoin d'un comportement d'agent spécialisé, évitant les écueils tant de la supervision uniforme que de l'ingénierie de récompense manuelle peu fiable.