Think Before You Act: Decision Transformers with Working Memory

• URL: http://arxiv.org/abs/2305.16338v3
• Date: Tue, 28 May 2024 19:03:33 GMT
• Title: Think Before You Act: Decision Transformers with Working Memory
• Authors: Jikun Kang, Romain Laroche, Xingdi Yuan, Adam Trischler, Xue Liu, Jie Fu,
• Abstract summary: Decision Transformer-based decision-making agents have shown the ability to generalize across multiple tasks. We argue that this inefficiency stems from the forgetting phenomenon, in which a model memorizes its behaviors in parameters throughout training. We propose a working memory module to store, blend, and retrieve information for different downstream tasks.
• Score: 44.18926449252084
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Decision Transformer-based decision-making agents have shown the ability to generalize across multiple tasks. However, their performance relies on massive data and computation. 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. Inspired by this, we propose a 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 Atari games and Meta-World object manipulation tasks. Moreover, we demonstrate that memory fine-tuning further enhances the adaptability of the proposed architecture.

Related papers

• MemSifter: Offloading LLM Memory Retrieval via Outcome-Driven Proxy Reasoning [78.46301394559903]
  Large Language Models (LLMs) are increasingly used for long-duration tasks.<n>Current methods face a trade-off between cost and accuracy.<n>MemSifter is a novel framework that offloads the memory retrieval process to a small-scale proxy model.
  arXiv  Detail & Related papers  (2026-03-03T02:57:38Z)
• Memory Caching: RNNs with Growing Memory [56.25483647131372]
  We introduce Memory Caching (MC), a technique that enhances recurrent models by caching checkpoints of memory states (a.k.a. hidden states)<n>We propose four variants of MC, including gated aggregation and sparse selective mechanisms, and discuss their implications on both linear and deep memory modules.<n>The results indicate that while Transformers achieve the best accuracy, our MC variants show competitive performance, close the gap with Transformers, and performs better than state-of-the-art recurrent models.
  arXiv  Detail & Related papers  (2026-02-27T18:53:41Z)
• AtomMem : Learnable Dynamic Agentic Memory with Atomic Memory Operation [40.1709026042412]
  We propose AtomMem, which reframes memory management as a dynamic decision-making problem.<n>By combining supervised fine-tuning with reinforcement learning, AtomMem learns an autonomous, task-aligned policy to orchestrate memory behaviors.<n> Experimental results across 3 long-context benchmarks demonstrate that the trained AtomMem-8B consistently outperforms prior static-workflow memory methods.
  arXiv  Detail & Related papers  (2026-01-13T08:22:28Z)
• Dynamic Affective Memory Management for Personalized LLM Agents [1.7600011132381626]
  We propose a new memory management system for affective scenarios.<n>Our system achieves superior performance in personalization, logical coherence, and accuracy.<n>Our work offers new insights into the design of long-term memory systems.
  arXiv  Detail & Related papers  (2025-10-31T12:12:51Z)
• Memo: Training Memory-Efficient Embodied Agents with Reinforcement Learning [53.72709564555407]
  Memo is a transformer-based architecture and training recipe for reinforcement learning.<n>It incorporates the creation and retrieval of memory by interleaving periodic summarization tokens with the inputs of a model during training.<n>We demonstrate Memo's effectiveness on a gridworld meta-RL benchmark and a multi-object navigation task in photo-realistic indoor settings.
  arXiv  Detail & Related papers  (2025-10-22T16:24:47Z)
• Memp: Exploring Agent Procedural Memory [72.41472703974935]
  Large Language Models (LLMs) based agents excel at diverse tasks, yet they suffer from brittle procedural memory that is manually engineered or entangled in static parameters.<n>We propose Memp that distills past agent trajectories into both fine-grained, step-by-step instructions and higher-level, script-like abstractions.<n>We show that as the memory repository is refined, agents achieve steadily higher success rates and greater efficiency on analogous tasks.
  arXiv  Detail & Related papers  (2025-08-08T16:20:56Z)
• How Memory Management Impacts LLM Agents: An Empirical Study of Experience-Following Behavior [49.62361184944454]
  Memory is a critical component in large language model (LLM)-based agents.<n>We study how memory management choices impact the LLM agents' behavior, especially their long-term performance.
  arXiv  Detail & Related papers  (2025-05-21T22:35:01Z)
• Breaking Memory Limits: Gradient Wavelet Transform Enhances LLMs Training [45.225732322141994]
  Large language models (LLMs) have impressive performance across a range of natural language processing tasks. Their vast number of parameters introduces significant memory challenges during training. Existing memory-efficient algorithms often rely on techniques such as singular value decomposition projection or weight freezing. We propose a novel solution called Gradient Wavelet Transform (GWT), which applies wavelet transforms to gradients in order to significantly reduce the memory requirements.
  arXiv  Detail & Related papers  (2025-01-13T11:35:09Z)
• Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning [64.93848182403116]
  Current deep-learning memory models struggle in reinforcement learning environments that are partially observable and long-term. We introduce the Stable Hadamard Memory, a novel memory model for reinforcement learning agents. Our approach significantly outperforms state-of-the-art memory-based methods on challenging partially observable benchmarks.
  arXiv  Detail & Related papers  (2024-10-14T03:50:17Z)
• Parameter-Efficient and Memory-Efficient Tuning for Vision Transformer: A Disentangled Approach [87.8330887605381]
  We show how to adapt a pre-trained Vision Transformer to downstream recognition tasks with only a few learnable parameters. We synthesize a task-specific query with a learnable and lightweight module, which is independent of the pre-trained model. Our method achieves state-of-the-art performance under memory constraints, showcasing its applicability in real-world situations.
  arXiv  Detail & Related papers  (2024-07-09T15:45:04Z)
• PARMESAN: Parameter-Free Memory Search and Transduction for Dense Prediction Tasks [5.5127111704068374]
  This work addresses flexibility in deep learning by means of transductive reasoning. We propose PARMESAN, a scalable method which leverages a memory module for solving dense prediction tasks. Our method is compatible with commonly used architectures and canonically transfers to 1D, 2D, and 3D grid-based data.
  arXiv  Detail & Related papers  (2024-03-18T12:55:40Z)
• Spatially-Aware Transformer for Embodied Agents [20.498778205143477]
  This paper explores the use of Spatially-Aware Transformer models that incorporate spatial information. We demonstrate that memory utilization efficiency can be improved, leading to enhanced accuracy in various place-centric downstream tasks. We also propose the Adaptive Memory Allocator, a memory management method based on reinforcement learning.
  arXiv  Detail & Related papers  (2024-02-23T07:46:30Z)
• Recurrent Action Transformer with Memory [39.58317527488534]
  This paper proposes a novel model architecture that incorporates a recurrent memory mechanism designed to regulate information retention. We conduct experiments on memory-intensive environments (ViZDoom-Two-Colors, T-Maze, Memory Maze, Minigrid-Memory), classic Atari games, and MuJoCo control environments. The results show that using memory can significantly improve performance in memory-intensive environments, while maintaining or improving results in classic environments.
  arXiv  Detail & Related papers  (2023-06-15T19:29:08Z)
• A Memory Model for Question Answering from Streaming Data Supported by Rehearsal and Anticipation of Coreference Information [19.559853775982386]
  We propose a memory model that performs rehearsal and anticipation while processing inputs to important information for solving question answering tasks from streaming data. We validate our model on a short-sequence (bAbI) dataset as well as large-sequence textual (NarrativeQA) and video (ActivityNet-QA) question answering datasets.
  arXiv  Detail & Related papers  (2023-05-12T15:46:36Z)
• A Memory Transformer Network for Incremental Learning [64.0410375349852]
  We study class-incremental learning, a training setup in which new classes of data are observed over time for the model to learn from. Despite the straightforward problem formulation, the naive application of classification models to class-incremental learning results in the "catastrophic forgetting" of previously seen classes. One of the most successful existing methods has been the use of a memory of exemplars, which overcomes the issue of catastrophic forgetting by saving a subset of past data into a memory bank and utilizing it to prevent forgetting when training future tasks.
  arXiv  Detail & Related papers  (2022-10-10T08:27:28Z)
• A Model or 603 Exemplars: Towards Memory-Efficient Class-Incremental Learning [56.450090618578]
  Class-Incremental Learning (CIL) aims to train a model with limited memory size to meet this requirement. We show that when counting the model size into the total budget and comparing methods with aligned memory size, saving models do not consistently work. We propose a simple yet effective baseline, denoted as MEMO for Memory-efficient Expandable MOdel.
  arXiv  Detail & Related papers  (2022-05-26T08:24:01Z)
• Learning to Learn Variational Semantic Memory [132.39737669936125]
  We introduce variational semantic memory into meta-learning to acquire long-term knowledge for few-shot learning. The semantic memory is grown from scratch and gradually consolidated by absorbing information from tasks it experiences. We formulate memory recall as the variational inference of a latent memory variable from addressed contents.
  arXiv  Detail & Related papers  (2020-10-20T15:05:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.