DuoMem: Rumo a Agentes de Memória em Dispositivo Capazes via Destilação de Espaço Duplo

Resumo

Agentes baseados em modelos de linguagem de grande porte (LLM) podem resolver tarefas processuais complexas interagindo com ambientes ao longo de múltiplas rodadas, mas essa capacidade geralmente depende de modelos grandes, contextos longos e chamadas repetidas de inferência. Isso torna agentes avançados com aumento de memória difíceis de implantar em dispositivos com recursos limitados. Apresentamos o DuoMem, uma estrutura de destilação em dois espaços que transfere a capacidade de resolução de problemas processuais de um modelo professor grande para modelos alunos compactos. O DuoMem destila em dois espaços complementares: (1) destilação no espaço de contexto, que substitui as memórias geradas pelo aluno por memórias processuais de maior qualidade geradas pelo professor, anexadas à entrada do aluno; e (2) destilação no espaço de parâmetros, que ajusta finamente adaptadores LoRA leves em trajetórias bem-sucedidas do professor. Avaliado no ALFWorld, um desafiador benchmark de tomada de decisão corporificada, o DuoMem eleva um modelo de 4 bilhões de parâmetros de uma taxa de sucesso de 4,3% para 77,9%, fechando a maior parte da lacuna em relação ao modelo professor de 72 bilhões de parâmetros (87,1%), enquanto adiciona menos de 10 milhões de parâmetros treináveis e apenas alguns megabytes de memórias do professor pré-computadas. Além disso, o modelo de 4B aprimorado pelo DuoMem conclui tarefas mais de 3 vezes mais rápido que o professor de 72B em tempo real de relógio, tornando-o viável para implantação em tempo real na borda, o que seria desafiador para o professor. Ablações extensivas em oito modelos que abrangem de 2B a 72B parâmetros revelam que ambos os eixos de destilação contribuem de forma complementar.

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