Piensa antes de actuar: Transformadores de Decisión con Memoria de Trabajo 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

Resumen

Los agentes de toma de decisiones basados en modelos de lenguaje de gran escala (LLM) han demostrado la capacidad de generalizar en múltiples tareas. Sin embargo, su rendimiento depende de grandes cantidades de datos y recursos computacionales. Argumentamos que esta ineficiencia surge del fenómeno de olvido, en el cual un modelo memoriza sus comportamientos en los parámetros durante el entrenamiento. Como resultado, entrenar en una nueva tarea puede deteriorar el rendimiento del modelo en tareas anteriores. En contraste con el mecanismo de memoria implícita de los LLM, el cerebro humano utiliza un almacenamiento de memoria distribuido, que ayuda a gestionar y organizar múltiples habilidades de manera eficiente, mitigando el fenómeno de olvido. Inspirados por esto, proponemos un módulo de memoria de trabajo interno para almacenar, combinar y recuperar información para diferentes tareas posteriores. Los resultados de evaluación muestran que el método propuesto mejora la eficiencia del entrenamiento y la generalización tanto en juegos de Atari como en tareas de manipulación de objetos en meta-mundo. Además, demostramos que el ajuste fino de la memoria mejora aún más la adaptabilidad de la arquitectura propuesta.

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.