DuoMem : Vers des agents mémoire sur appareil performants via une distillation à double espace

Résumé

Les agents basés sur de grands modèles de langage (LLM) peuvent résoudre des tâches procédurales complexes en interagissant avec des environnements sur plusieurs tours, mais cette capacité repose généralement sur des modèles de grande taille, des contextes longs et des appels d'inférence répétés. Cela rend difficile le déploiement d'agents à mémoire augmentée sur des appareils aux ressources limitées. Nous présentons DuoMem, un cadre de distillation bi-espace qui transfère la capacité de résolution de problèmes procéduraux d'un grand modèle enseignant vers des modèles étudiants compacts. DuoMem distille dans deux espaces complémentaires : (1) la distillation dans l'espace contextuel, qui remplace les mémoires générées par l'étudiant par des mémoires procédurales de meilleure qualité générées par l'enseignant, préfixées à l'entrée de l'étudiant, et (2) la distillation dans l'espace paramétrique, qui affine de légers adaptateurs LoRA sur des trajectoires réussies de l'enseignant. Évalué sur ALFWorld, un benchmark exigeant de prise de décision incarnée, DuoMem améliore le taux de réussite d'un modèle de 4 milliards de paramètres, le faisant passer de 4,3 % à 77,9 %, comblant ainsi la majeure partie de l'écart avec le modèle enseignant de 72 milliards (87,1 %), tout en ajoutant moins de 10 millions de paramètres entraînables et seulement quelques mégaoctets de mémoires pré-calculées de l'enseignant. De plus, le modèle amélioré par DuoMem (4 milliards) termine les tâches plus de 3 fois plus rapidement que le modèle enseignant de 72 milliards en temps réel, ce qui le rend viable pour un déploiement en périphérie en temps réel, là où l'enseignant serait difficile à utiliser. Des ablations approfondies sur huit modèles allant de 2 à 72 milliards de paramètres révèlent que les deux axes de distillation apportent des contributions complémentaires.

English

Large Language Model (LLM)-based agents can solve complex procedural tasks by interacting with environments over multiple turns, but this ability typically depends on large models, long contexts, and repeated inference calls. This makes advanced memory-augmented agents difficult to deploy on resource-constrained devices. We introduce DuoMem, a dual-space distillation framework that transfers procedural problem-solving ability from a large teacher model to compact student models. DuoMem distils in two complementary spaces: (1)context-space distillation, which replaces student-generated memories with higher-quality teacher-generated procedural memories prepended to the student's input, and (2)parameter-space distillation, which fine-tunes lightweight LoRA adapters on successful teacher trajectories. Evaluated on ALFWorld, a challenging embodied decision-making benchmark, DuoMem boosts a 4B-parameter model from 4.3% to 77.9% task success rate, closing most of the gap to a 72B teacher model (87.1%), while adding fewer than 10M trainable parameters and only a few megabytes of pre-computed teacher memories. Moreover, the DuoMem-enhanced 4B model completes tasks over 3x faster than the 72B teacher in wall-clock time, making it viable for real-time edge deployment, which would be challenging for the teacher.Extensive ablations across eight models spanning 2B-72B parameters reveal that both distillation axes contribute complementary