Pense Antes de Agir: Transformadores de Decisão com Memória de Trabalho Interna

Think Before You Act: Decision Transformers with Internal Working Memory

May 24, 2023

Autores: Jikun Kang, Romain Laroche, Xindi Yuan, Adam Trischler, Xue Liu, Jie Fu

cs.AI

Resumo

Agentes de tomada de decisão baseados em modelos de linguagem de grande escala (LLMs) demonstraram a capacidade de generalizar em múltiplas tarefas. No entanto, seu desempenho depende de grandes volumes de dados e poder computacional. Argumentamos que essa ineficiência decorre do fenômeno de esquecimento, no qual um modelo memoriza seus comportamentos nos parâmetros ao longo do treinamento. Como resultado, o treinamento em uma nova tarefa pode deteriorar o desempenho do modelo em tarefas anteriores. Em contraste com o mecanismo de memória implícita dos LLMs, o cérebro humano utiliza um armazenamento de memória distribuído, que ajuda a gerenciar e organizar múltiplas habilidades de forma eficiente, mitigando o fenômeno de esquecimento. Assim inspirados, propomos um módulo de memória de trabalho interno para armazenar, combinar e recuperar informações para diferentes tarefas subsequentes. Os resultados da avaliação mostram que o método proposto melhora a eficiência do treinamento e a generalização tanto em jogos Atari quanto em tarefas de manipulação de objetos no meta-mundo. Além disso, demonstramos que o ajuste fino da memória aprimora ainda mais a adaptabilidade da arquitetura proposta.

English

Large language model (LLM)-based decision-making agents have shown the ability to generalize across multiple tasks. However, their performance relies on massive data and compute. We argue that this inefficiency stems from the forgetting phenomenon, in which a model memorizes its behaviors in parameters throughout training. As a result, training on a new task may deteriorate the model's performance on previous tasks. In contrast to LLMs' implicit memory mechanism, the human brain utilizes distributed memory storage, which helps manage and organize multiple skills efficiently, mitigating the forgetting phenomenon. Thus inspired, we propose an internal working memory module to store, blend, and retrieve information for different downstream tasks. Evaluation results show that the proposed method improves training efficiency and generalization in both Atari games and meta-world object manipulation tasks. Moreover, we demonstrate that memory fine-tuning further enhances the adaptability of the proposed architecture.