Think Before You Act: Decision Transformers with Internal Memory

Q5

Are $\omega$  in Line 202, 215, and 232 the same?

A: Yes, the address $\omega$ remains the same in these lines. Line 202 defines how we compute the address. Line 215 defines how we write information to the memory, and line 232 defines how we read from the memory.

The reasons are as follows:

1. Content-based Addressing: DT-Mem uses a content-based addressing mechanism, where the content to be read or written influences the memory address. Since we are calculated the address using current input embeddings $E$, this mechanism aligns with using the same address for both operations, as the content being processed is directly tied to where it is stored or retrieved from.
2. Memory Consistency and Coherence: This approach allows the model to immediately read back what it has just written, ensuring that the memory reflects the latest changes made by the network.
3. Simplified Memory Management: Using the same address for reading and writing simplifies the memory management process. It reduces the complexity of the controller's operations, as it doesn't need to maintain separate mechanisms or algorithms for addressing reads and writes.

Q6

What is the motivation to compute the strength $\beta$ using attention? Why do we need to use $\beta$ in both erasing and adding vectors?

A: The primary motivation behind employing attention mechanisms is to achieve Context-Awareness. By using cross-attention between the input embeddings, $E$, and the memory, $M$, the model can selectively concentrate on pertinent aspects of the input. This attention-based computation ensures that memory updates are context-dependent, enhancing both the dynamism and efficiency of the process.

The computation of $\beta$ plays a crucial role in regulating the flow of information to the memory. As $\beta$ represents a balance between the information in the addition vector $\epsilon^a$ and the erasure vector $\epsilon^e$, it effectively governs the extent of new information added to, and old information removed from, the current memory. This parameter allows for dynamic control over memory updates. For instance, when the input information is highly relevant to specific memory segments of interest (as determined by the content-address $\omega$), $\beta$ assumes a larger value. Consequently, only a small portion of memory is erased (calculated as 1-$\beta$, as outlined in line 216) and predominantly mixed with the addition vector (as described in line 222). Conversely, if the input information significantly differs, the model erases more from the current memory, adding only a minimal amount of new information.

Q7

Is $t$  in Line 222 the step in the trajectory? Can you provide an algorithm to explain clearly how memory read and write are executed within a trajectory?

A: Thank you for highlighting this point of confusion. In this context, the symbol $t$ represents the $N$th timestep of memory operations that occur within the interactions between the transformer module and memory. It specifically refers to the internal loop of memory operations and is not associated with the trajectory. For a detailed explanation, we have included the relevant algorithm in Appendix Algorithm 3.

Q8

Is Step 1 Line 187 important? Do you have an ablation study on Step 1?

A: Step 1 serves two main purposes:

1. To map various input dimensions into a consistent dimension for storage. Given that our memory is a matrix with fixed dimensions (i.e., number_of_slots * dimensions), it's imperative to harmonize input dimensions for storage.
2. As highlighted in our paper, the input sequence comprises multiple steps of a tuple $< \hat{r}_t, s_t, a_t >$. We use the timestep $t$ to arrange these inputs, thereby preserving the temporal relationships among them.

We've conducted results without including input sequence organization:

Though minor differences exist between the two sets of inputs, the variance in outcomes between the two methodologies is not pronounced.

Thank you for your insightful feedback!

Q1

The novelty is limited. The main idea is to integrate an external memory for Transformers to improve memorization and reasoning. There have been many works along this line of research. The proposed memory read and write mechanism is also straightforward and heavily based on the mechanism (content-based attention, add, erase, ...) of NTM and DNC.

A: Internal Memory vs. NTM: While both the internal memory module and NTM aim to introduce memory mechanisms to neural architectures, our internal memory is inspired by human cognitive processes and is specifically tailored to decision-making agents. The design and integration of this module into DT, and its subsequent application, is a novel contribution in the context of our work. Besides, there are other papers inspired by NTM: $1$ Memory Augmented Neural Networks with Wormhole Connections. ArXiv:1701.08718 $2$ Hybrid computing using a neural network with dynamic external memory. Nature 2016. Same as the above mentioned papers, despite that our work inspired by NTM, we design a new memory reading and writing techniques for DT.

Use of LoRA: While the concept of LoRA might not be new, its application in conjunction with our internal memory module and within the realm of LLM-based decision-making agents brings a novel perspective. The combination of these components, and the synergies they produce, offers a unique approach to addressing the generalization and adaptation in MDT.

The novelty of our work isn't merely in the individual components but in the holistic approach and the results achieved. By blending inspiration from human cognition with state-of-the-art machine learning techniques, our paper is the first paper that introduces internal memory in DT improved training efficiency and generalization in a variety of tasks, as evidenced in our experiments.

Contextual Novelty: It's crucial to recognize that novelty doesn't always stem from introducing entirely new concepts. Often, the innovative combination, adaptation, and application of existing ideas to new contexts or challenges can bring significant advancements to a field.

Thank you for your insightful feedback!

Q1

Some explanations in the paper are unclear. For instance, the reference to "Large Language Model based decision making agents" needs clarification. It's unclear if this refers to Transformer-based agents or models within the Decision Transformer family.

A: This phrase generally refers to all the large language model-based agents, which includes both decision transformer, trajectory transformer [1] and the following works. In this paper, we mainly focus on decision transformer family. We’ve revised this sentence in the paper revision line 1.

[1] Offline Reinforcement Learning as One Big Sequence Modeling Problem. https://arxiv.org/abs/2106.02039v4

Q2

The motivation for this work is not entirely clear. While the paper argues that memorization through neural network parameters can lead to weaker performance in forgetting during training, it remains unclear how this relates to generalization performance for unseen tasks.

A: Memorization and Forgetting in Neural Networks: When a neural network is trained on a series of tasks, it tends to "memorize" the specifics of those tasks through its parameters. This process can lead to a phenomenon known as "catastrophic forgetting," where learning new tasks impairs the network's performance on previously learned tasks [1]. This is a significant issue in the context of decision-making models, where adaptability and the ability to handle a variety of scenarios are crucial.

Generalization to Unseen Tasks: Generalization refers to a model's ability to perform well on new, unseen tasks, not just the ones it was trained on. The problem with the traditional approach of memorizing through parameters is that it often leads to overfitting on the training data, which can diminish the model's ability to generalize. When a model is overfitted, it becomes too tailored to the specifics of its training data and loses the flexibility needed to adapt to new situations or tasks [2].

Relation Between Forgetting and Generalization: The link between the forgetting phenomenon and generalization is that both are impacted by how a model learns and stores information. A model that is prone to forgetting may struggle with generalization because it cannot effectively retain and utilize the broad range of knowledge required to handle new tasks [3]. Conversely, a model designed to minimize forgetting (such as through an internal memory mechanism like DT-Mem) can potentially maintain a more diverse and generalizable knowledge base.

Motivation of DT-Mem: The motivation for introducing DT-Mem, as such, is to address these interrelated issues. By incorporating an internal memory module, DT-Mem aims to reduce the reliance on parameter-based memorization, thereby mitigating the effects of catastrophic forgetting. This is hypothesized to improve the model's ability to generalize to new tasks, as it can more effectively retain and apply knowledge from a broader range of experiences. Section 5.3 Table 1 the results of model generalization support this findings.

[1] Overcoming catastrophic forgetting in neural networks. National academy of sciences 2017.

[2] Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530.

[3] Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 114(13), 3521-3526. 2017

Q3

Is the memory not initialized throughout the entire training process? Clarifying this point could help readers better understand the novelty of this approach, as memory initialization is typically performed per episode (e.g., NTM).

A: The memory is initialized only once during the entire training process. This approach ensures that the knowledge acquired from one task can be leveraged in subsequent tasks, leading to improved pre-training performance, as evidenced in Fig. 6.

Q4

Could you investigate whether MDT can be fine-tuned using the LoRA technique? As LoRA is applicable to the general Transformer architecture, it would be insightful to assess the potential for fine-tuning MDT using this approach.

A: The results shown in Section 5.5 Figure 5 are all fine-tuned using the LoRA technique. As we concluded in lines 342-351. The results show that “the consistent superior performance of DT-Mem across most games suggests that this method might have a more adaptable approach.”

We evaluated both DT-Mem and MDT using the Memory Maze offline probing benchmark to determine if DT-Mem enhances performance in tasks requiring long-term dependency. Performance was measured using wall prediction accuracy (%) for the Walls benchmarks (higher the better) and mean-squared error for the Objects benchmarks (lower the better). We report the average score across three training runs. The results are as follows:

These results indicate that DT-Mem achieves superior performance in both environments across all four metrics. This suggests that the incorporation of a memory module in DT-Mem effectively stores past information for future decision-making. It highlights the method's efficacy and its advantages in handling long-term dependency tasks.

Q2.c

About the memory. There are some related methods, neural episodic control (https://arxiv.org/abs/1703.01988) for the writing and lookup. The two methods share many similarities, so maybe add a detailed comparison of the proposed methods and all related methods, so we can fully understand the contributions.

A: Thank you for the comments. Due to the page limit, we have added these comparisons in the paper revision Appendix Section E

Q2.d

Still about the memory. What is exactly the difference between the external memory and the internal memory? Could you provide an example about the two kinds of memories, as well as the advantages of the proposed memory.

External Memory vs. Internal Memory

1. External Memory:
   • Definition: External memory in AI and machine learning models refers to a separate storage unit that is not part of the core neural network architecture. It's like an additional database or repository that the model can access.
   • Example: A neural Turing machine (NTM) or Differentiable Neural Computer (DNC), which uses an external memory matrix that it can read from and write to, separate from its neural network parameters.
   • Advantages:
     • Scalability: Can store large amounts of data beyond the capacity of the model's parameters.
     • Flexibility: Allows the model to store and retrieve information in a more structured way, similar to how a computer uses RAM and disk storage.
2. Internal Memory:
   • Definition: Internal memory refers to the capacity of a neural network to store information within its own architecture, such as in its weights and activations. It does not use a separate storage unit but relies on the network's inherent structure.
   • Example: LSTM (Long Short-Term Memory) networks that use their recurrent connections to store information over time within the network.
   • Advantages:
     • Efficiency: More efficient in terms of retrieval speed, as the information is stored within the network.
     • Integration: Better integrated with the model's learning process, enabling more seamless information flow and processing.

Advantages of the Proposed Memory

1. Contextual Retrieval: Utilizing mechanisms like attention to retrieve and store information more contextually and relevantly.
2. Enhanced Generalization: By having a more flexible memory system, the model generalize better to new, unseen tasks.

Q3.i

Some issues about the experiments. I generally think the experiments are sufficient, but with some suggestions: i) can we use the prompting for the DT-Mem? as we know prompting is much easier than fine-tuning.

A: Thank you for the suggestion. The primary objective of the fine-tuning in Section 5.5 is to investigate whether modifications to the memory module can boost the performance of DT-Mem without altering other transformer parameters. This is intended to assess the role of the memory module's stored information in aiding decision-making for new tasks. Since the memory module engages in read and write operations through cross-attention, the use of prompt tuning, which does not alter DT-Mem's stored information, is not aligned with our approach.

Nevertheless, we have conducted a comparative analysis of prompt-tuning DT (PDT) and DT-Mem in Table 2. The results demonstrate the effectiveness of our proposed method and highlight the advantages of fine-tuning the memory module.

Q3.ii

And ii) what is the limit of the internal memory, given the fixed size, i.e., parameters, of the memory?

A: The major limitation of internal memory is Capacity Limitation: The capacity of internal memory is inherently limited by the size of the neural network, i.e., the number of its parameters and the architecture's complexity. Once these parameters are set, the capacity to store and recall information is fixed. The ability to store information depends on how densely information can be encoded within the network's parameters, which has practical limits.

Thank you for your insightful feedback!

Q1.a

For Section 3.1. It seems that you focus on offline RL, rather than RL. The two fields have differences. You may need a brief introduction of offline RL, as well as existing methods as baselines, such as DT, MGDT, RMDT, HDT.

A: We appreciate the insightful suggestion. Accordingly, a concise overview of offline RL has been incorporated into Section 3.1 (lines 141-146). Additionally, due to page constraints, we have briefly introduced the baselines in Section 5.2.

Q1.b

Figure 2 is somehow misleading. Figure 2a, does the memory module is a layer of the transformer? I think they are separated, however, the plot seems that the memory is stacked with the transformer. Figure 2b, the retrieved memory is another memory? I think should be the retrieved experiences, or other terms. Please make the terms unique and clear.

A: We appreciate the suggestion. To clarify, in Figure 2a, the memory module is depicted as an independent module, distinct from the transformer module. Conversely, in Figure 2b, the memory that is retrieved serves as an intermediate input to the transformer module, integrating with it more directly. We have updated Figure 2 in the revised version of the paper to better illustrate these configurations and their respective operational dynamics.

Q2.a

The introduction uses large language model as the motivation, however, even using GPT-2 architecture, the embedding and the tokens I believe is not about words, it should be game-specific tokens. so please using transformer or decoder-only transformer to avoid any confusion.

A: Thank you for the comments. We have revised this term in line 16 and line 20.

Q2.b

It seems that the largest model used in the paper is 50M. Compared with LLM, it still very small. Does LoRA is necessary? LoRA can be used to any models, which is not a technical contribution of this paper...

A: Model Size: The choice of a 50M model size was dictated by the computational resources available to us. Nonetheless, in Figure 3, we explore the scaling law of DT-Mem, which reveals an interesting trend: as the number of parameters increases, so does performance. Importantly, when compared to the 200M MDT model, our DT-Mem, with only a quarter of the parameter size, demonstrates superior performance. This not only underscores the efficiency of DT-Mem in generalizing across games but also empirically supports its effectiveness even with fewer parameters.

Motivation of LoRA: The motivations of using LoRA to fine-tune the model can be concluded in following two reasons:

1. According to a paper $1$ , compared to other adapter methods, LoRA method not only don't introduce inferences latency nor reduce input sequence length while retaining high model quality. It allows for quick task-switching when deployed as a service by sharing the vast majority of the model parameters
2. Parameter-efficient fine-tuning (PEFT) approaches (such as LoRA) only fine-tune a small number of (extra) model parameters while freezing most parameters of the pretrained LLMs, thereby greatly decreasing the computational and storage costs $1$ . This also overcomes the issues of catastrophic forgetting $2$ . PEFT approaches have also shown to be better than fine-tuning in the low-data regimes and generalize better to out-of-domain scenarios $3$ .

Full Fine-tuning (FFT) vs. LoRA: To assess whether the use of LoRA adversely affects performance, we conducted experiments contrasting Full Fine-Tuning (FFT) of memory parameters with LoRA. In this context, FFT-single refers to fine-tuning all parameters exclusively on a single game, whereas FFT-All represents fine-tuning on the entire set of games simultaneously. Results are DQN-normalized score.

Based on above results, we conclude the following observations:

• LoRA appears to be the most consistently effective strategy across the games provided.
• While FFT-Single occasionally outperforms PEFT (like in Alien), FFT-All consistently trails behind the other two.

The reason full fine-tuning is not comparable to PEFT comes from the following parts: 1. Fine-tuning dataset size. Note that we only use 50k data in LoRA and full fine-tuning compares on 500k used in MDT paper 2. The benefits of LoRA listed above: "This approach also addresses catastrophic forgetting [2] and has outperformed standard fine-tuning in low-data and out-of-domain situations [3]”

[1] LoRA: Low-Rank Adaptation of Large Language Models, ICLR 2022 [2] Adapter-Fusion: Non-destructive task composition for transfer learning.EACL, 2021. [3] Prefix-tuning: Optimizing continuous prompts for generation. ACL, 2021.

Thank you for your insightful feedback!

Q1

Fig. 2(b) is too simple, making it difficult to correspond one-to-one with the steps in the method part. The authors should make the figure more comprehensive and understandable.

A: Thank you for the comments. We have revised this figure in the paper revision Fig. 2

Q2

There are many papers demonstrating the ideas about internal memory, and the authors should explain the differences with similar methods in more detail in the method section.

A: Thank you for pointing this out. To the best of my knowledge, the work most closely related to ours (DT with internal memory) is RMDT, which we have discussed extensively in the paper. I would greatly appreciate it if the reviewer could share specific papers they believe should be compared to our work. We are more than willing to discuss these differences in the revised version of the paper.

Q3

Add more analysis about different situations, such as the input misleading by the content stored in the memory (i.e., noise or dissimilar pattern), how does the method eliminates this type of impact.

A: We thank the reviewer for their valuable suggestion. Following this, we conducted an experiment to assess the robustness of the proposed method against input distortion. This involved adding Gaussian noise to the input frames of Atari games. Specifically, we set the mean to 0 and experimented with various standard deviation values. The results are detailed in the table below:

From the results above, we conclude that the proposed DT-Mem demonstrates greater robustness to noisy inputs compared to the MDT method. This is evident as the DT-Mem consistently outperforms MDT under various levels of Gaussian noise. Notably, the performance with a standard deviation of 0.5 shows minimal difference compared to the no-noise scenario, illustrating DT-Mem's effectiveness in mitigating the impact of varying input distortions.

Dear Reviewers,

Thank you for your time. We know it’s a busy moment, but we would like to know if there’s anything else we can do to clarify your concerns about our paper.

Please see the summary of the paper revision above.

Reviewers gRQQ, hmxu: Thank you for your positive feedback, we hope the rebuttal helped in answering your questions. Please let us know if you have other comments and concerns.

Reviewer qzjQ: Your main concern was about the clarity of our approach, its underlying motivation, and its implementation details, particularly regarding the memory initialization process and the application of the LoRA technique.

In the response, we provided detailed explanations to clarify these points. We elaborated on the conceptual framework connecting memorization and forgetting in neural networks with generalization performance, and how DT-Mem is designed to improve this. We also clarified the single initialization process of the memory module during training, which is a distinctive aspect of our approach.

We also implemented your proposed Memory Maze benchmark to evaluate the performance of DT-Mem in long-term dependency tasks. The results from this benchmark were included to demonstrate the superior performance of DT-Mem in handling complex tasks, thereby addressing your concerns about its effectiveness and generalizability.

Is there anything else that we can do to address your concerns?

Reviewer NgFj: Your main concern was about the novelty of our approach, specifically regarding the integration of an external memory into Transformers, the comparison with other memory-augmented Transformer models, and the need for more detailed background content and related work in our paper.

In the response, we highlighted the unique aspects of our working memory module inspired by human cognition, specifically designed for decision-making agents. We emphasized the novel application of LoRA in conjunction with our memory module in LLM-based decision-making agents. We also provided a detailed comparison of our memory read and write mechanism with other models like ∞-former, LONGMEM, and KNN-transformer, and addressed the lack of sufficient background content by expanding on Decision Transformers and offline RL settings in our revised paper.

We also expanded related work section to cover more memory-based Transformer papers. We conducted additional experiments and ablation studies as requested, such as the comparison of Full Fine-Tuning (FFT) with LoRA using Atari games and provided visualizations to demonstrate how our model effectively uses memory in decision-making.

Is there anything else that we can do to address your concerns?

Thanks!