Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning

Thank you for your thoughtful review. We address your concerns point by point below.

Weakness

In the original manuscript, we have presented results and discussions regarding the change in performance across task complexity levels. The answers to your questions below will provide more details.

Questions

1. We tested SHM on PopGym tasks with different memory sizes $H$, and the results are shown in Fig. 3b. Overall, performance improves as $H$ increases; however, in one task, performance peaks at $H = 128$. Thus, we selected $H = 128$ for other PopGym tasks to balance both performance and computational efficiency.

2. In these POMDP tasks, the episode length often determines task complexity, with longer episodes typically being more challenging. We evaluated performance (measured by success rate) as episode length increased and reported the results in Fig. 2. Additionally, in the meta-RL setting, we studied task complexity by varying the train-test task ratio (lower ratios indicate higher complexity). The performance under different train-test ratios is presented in Fig. 1. Overall, as task complexity increases, the performance of all methods declines. However, our SHM maintains reasonable performance and clearly outperforms other approaches. Regarding computational requirements, we have added Table 5 in the appendix of the revised manuscript to show the running time of our method vs GPT-2 as task complexity changes. As shown, the running time of our method remains comparable to that of the other baseline (see the table in our answer to your question 3). Please note our current implementation of SHM is naive without any parallel support while GPT-2 implementation has Pytorch-optimized attention.

3. Yes, in Fig. 2, we reported the performance of SHM over extended episodes, including lengths of 250 and 500 steps. (The benchmark in [1] also only tested up to 500 steps for Key-to-Door tasks). Beyond this limit, testing requires significant computational and memory resources for all methods, which are beyond our current capacity. That said, we agree that it is valuable to present results with more diverse episode lengths. To address this, we have added results with an episode length of 100 in the appendix. These results further demonstrate the dominance of SHM over GPT-2 (best performer reported in [1]). Please see example comparison below for both time and performance as episode length increases for Key To Door (see more in Appendix Table 5). | Episode Length | SHM Total Training Time (hours) | GPT-2 Total Training Time (hours) | SHM Success Rate (%) | GPT Success Rate (%) | |----------------|---------------------------------|-----------------------------------|----------------------|-----------------------| | 100 | 20 | 18 | 100 | 59 | | 250 | 27 | 23 | 78 | 21 | | 500 | 41 | 36 | 55 | 25 |

We hope these revisions adequately address your concerns. Please feel free to let us know if there is anything else we can clarify or enhance to improve your score on our paper.

Reference

[1] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in rl? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36, 2024

Thank you for the response. The authors have addressed all of my concerns, and I maintain that this paper should be accepted. I agree with the other reviewer, in that perhaps there is a more descriptive name than "Stable Hadamard Memory". But ultimately, this is up to the authors.

Thank you for your positive review. We address your concerns point by point below.

Weakness

"The authors do not provide code": Due to time constraints, we did not submit the code at the submission time. However, as mentioned in the original manuscript, we committed to release the code. In this revision, we have attached the code for SHM and the baselines in Popgym tasks. We plan to complete and release the full code upon publication.

"The initialization of ...": Thank you for your suggestions. We have revised this part accordingly.

"Proposition 3 is poorly worded ...": Thank you for pointing it out. We have changed it to:

"Proposition 3. If calibration is enabled such that $C_{\theta}(x_{t})\neq\boldsymbol{1}$, and the calibration matrix remains fixed and independent of the input $x_{t}$ (i.e., $\forall t:C_{\theta}(x_{t})=\theta\in\mathbb{R}^{H\times H}$ ), this will lead to either numerical instability or learning difficulties"

"Section 3.3 is misspelled ...": Thank you, we have fixed it.

Questions

"Why is 1+tanh⁡(x) ....": Because the proof in Appendix A.3 requires an odd function like tanh (see Line 753 in the original manuscript)

"How come the recurrent state ...": Thank you for your insightful comments. Our method does not utilize complex-valued states; instead, it employs real-valued recurrent states like standard models such as GRU. Incorporating complex-valued states to track phase information could be an orthogonal contribution to our approach. We believe that integrating this aspect could potentially enhance performance further. However, we will leave this exploration for future work.

"Line 233: ...": This is a typo. We have revised it to m,k. Here, m and k are the row and column indices of a matrix.

Thank you for your constructive comments. We address your concerns point by point below.

Weaknesses

1. Thank you for your insightful comments on the novelty and naming of our approach. “Hadamard Memory” is a unified framework for matrix-based memory, distinct from vector-based memories in RNNs and GRUs. Although "Matrix Hadamard Memory" might be clearer, we chose a more concise name. While the Hadamard product is used in other architectures, we leverage it for memory calibration, enabling matrix-level memory updates. Please note that a key contribution of our work is identifying this unified framework, and analyzing its theoretical properties, and limitations (Sections 3.1-3.2). Based on the analysis, we propose a solution to improve the Hadamard Memory Framework. We named it “Stable Hadamard Memory” to emphasize our focus on stability, particularly through the calibration matrix’s random parameter selection. We used Section 3.3 to write about this stability component. To make it clearer, we have revised the introduction and add more descriptions to Sec.3 ending to clarify the contributions regarding designing the calibration matrix with random parameter selection for stable updates.

2. Although our approach might appear similar to linear attention in Linear Transformer [1], there is a clear distinction between our Hadamard Memory Framework (Eq. 7-8) and linear attention (Eq. 5):

• Linear attention lacks the calibration matrix C, which is central to our framework, as explained in Section 3.1.
• Linear attention does not include the update gate η​, which enhances memory control in our method.

We have clarified these points in Sec. 3.1 of the revision. We did not initially report the performance of Linear Transformer because, according to the Popgym benchmark [2], it underperforms against baselines like GRU and FFM. As a result, we prioritized comparisons with stronger or more recent baselines (for your reference, we have added the reported results of Linear Transformer to Appendix Table 6 in this revision). Moreover, we have also added mLSTM [3], a recent and more related matrix-based LSTM variant, to our experiments. New results in Fig. 1 and 2 demonstrate that SHM still significantly outperforms mLSTM. We also summarize the results between our SHM and mLSTM in the table below.

3. The random selection mechanism does not substantially increase memory requirements for storing the model's parameter. The additional parameters introduced by our calibration matrix $\theta \in \mathbb{R}^{L \times H}$, with $L = 128$ and $H$ typically set to 24, resulting in around 3,000 extra parameters—less than 3% of the total memory used by the standard policy and value networks in our experiments (see Appendix Table 2). Thus, the effect on memory usage is minimal, and comparable to other baselines in term of both number of parameters and model's memory capacity (see results in Sec. 4.1)

4. We followed the benchmark from Morad et al., (2023) [2] and trained each model 3 times per task. Please note that the RL tasks we used are long-term by design, with each run taking several days to complete. Given the large number of tasks (24 in total) and resource limitations, we believe that 3 runs are reasonable, and the results still show clear statistical differences between the methods. For details on the tuning process, we provided an explanation in Appendix C1 and C2 in the original manuscript. Basically, we focus on tuning $H$, the memory dimension.

5. Thank you for pointing it out. Although this paper does not focus on the matrix memory update aspect, it is indeed related. We have included this paper in the related work of the revision.

Thank you for your encouraging reviews. We address your concern below.

"The paper is overly complex ...": We have simplified the notation on $\theta$ and added a simpler and more detailed description at the end of Sec. 3 to highlight and explain the benefit of the memory update rule as follows:

"The proposed update rule improves memory management and stability through several key benefits. The dynamic calibration matrix $C_t$ adjusts memory elements based on the input context, allowing the model to prioritize relevant data while forgetting outdated information. This ensures that transient details, like detour events, are discarded while critical information, like landmarks, is retained. Moreover, the random parameter selection in the calibration matrix breaks timestep dependencies and controls the expected gradient norm, preventing numerical instabilities like gradient explosion or vanishing. This stable update mechanism ensures consistent learning and robust performance over long episodes."

We hope these changes have effectively addressed your concerns. Please let us know if there is anything further we can clarify or improve to support a higher evaluation.