Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning
Authors: Hung Le
Dung Nguyen
Kien Do
Sunil Gupta
Svetha Venkatesh
Venue: ICLR 2025 Poster (2025)
Keywords: Reinforcement Learning
Memory
POMDP
Primary area: reinforcement learning
TL;DR: Stable Hadamard Memory introduces a novel, computationally efficient memory model for reinforcement learning, dynamically adjusting to erase irrelevant experiences and reinforce crucial ones.

Abstract: Effective decision-making in partially observable environments demands robust memory management. Despite their success in supervised learning, current deep-learning memory models struggle in reinforcement learning environments that are partially observable and long-term. They fail to efficiently capture relevant past information, adapt flexibly to changing observations, and maintain stable updates over long episodes. We theoretically analyze the limitations of existing memory models within a unified framework and introduce the Stable Hadamard Memory, a novel memory model for reinforcement learning agents. Our model dynamically adjusts memory by erasing no longer needed experiences and reinforcing crucial ones computationally efficiently. To this end, we leverage the Hadamard product for calibrating and updating memory, specifically designed to enhance memory capacity while mitigating numerical and learning challenges. Our approach significantly outperforms state-of-the-art memory-based methods on challenging partially observable benchmarks, such as meta-reinforcement learning, long-horizon credit assignment, and POPGym, demonstrating superior performance in handling long-term and evolving contexts.

Published as a conference paper at ICLR 2025
STABLE HADAMARD MEMORY: REVITALIZING MEMORY-AUGMENTED AGENTS FOR REINFORCEMENT LEARNING
Hung Le, Dung Nguyen, Kien Do, Sunil Gupta, and Svetha Venkatesh Applied AI Institute, Deakin University, Geelong, Australia (thai. .eii,ie,,,, Le,
k. do, sunil. gupta, svetha. weettaa..eaaateetheeeeeee... au
ABSTRACT
Effective decision-making partially observable environments demands robust memory management. Despite their success in supervised learning, current deep- learning memory models struggle in reinforcement learning environments that are partially observable and long-term. They fail to efficiently capture relevant past information, adapt flexibly to changing observations, and maintain stable updates over long episodes. We theoretically analyze the limitations of existing memory models within a unified framework and introduce the Stable Hadamard Memory, novel memory model for reinforcement learning agents. Our model dynamically adjusts memory by erasing no longer needed experiences and reinforcing crucial ones computationally efficiently. To this end, we leverage the Hadamard product for calibrating and updating memory, specifically designed to enhance memory capacity while mitigating numerical and learning challenges. Our approach significantly outperforms state-of-the-art memory-based methods on challenging partially observable benchmarks, such as meta-reinforcement learning, long- horizon credit assignment, and POPGym, demonstrating superior performance in handling long-term and evolving contexts. Our source code is available at www..///wwwwwwn/eeee.eoeccec
INTRODUCTION
Reinforcement learning agents necessitate memory. This is especially true in Partially Observable Markov Decision Processes (POMDPs (Kaelbling et al. 1998)), where past information is crucial for making informed decisions. However, designing a robust memory remains an enduring challenge, as agents must not only store long-term memories but also dynamically update them response to evolving environments. Memory-augmented neural networks -aatiua-partcuulaary those developed for supervised learning (Graves et al. 2016; Vaswani et al. 2017), while offering promising solutions, have consistently struggled in these dynamic settings. Recent empirical studies
(Morad et al. 2023 Ni al. 2024) have shown that MANNs exhibit instability and underperform simpler vector-based memory models such as GRU (Chung et al., 2014) or LSTM (Hochreiter) 1997). The issue exacerbated in complex and sparse reward scenarios where agents must selectively retain and erase memories based on relevance. Unfortunately, existing methods fail to provide a memory writing mechanism that is efficient, stable, and flexible meet these demands.
In this paper, we focus on designing a better writing mechanism to encode new information into the memory. To this end, we introduce the Hadamard Memory Framework (HMF), unified model that encompasses many existing writing methods as specific cases. This framework highlights the critical role of memory calibration, which involves linearly adjusting memory elements by multiplying the memory matrix with calibration matrix and then adding an update matrix. By leveraging Hadamard products that operate dlement-wise on memory matrices, we allow memory writing without mixing the memory cells in computationally efficient manner. More importantly, the calibration and update matrices are dynamically computed based on the input at each step. This enables the model to learn adaptive memory rules, which crucial for generalization. For instance, in meta-reint reinforcementt learning with varying maze layouts, fixed memory update rule may work for one layout but fail in another. By allowing the calibration matrix to adjust according the currett

---

Published as a conference paper at ICLR 2025
maze observation, the agent can learn to adapt to any layout configuration. A dynamic calibration matrix also enables the agent to flexibly forget and later recall information as needed. For example, an agent navigating room may need to remember the key's 's location, retain it during detour, and later recall it when reaching a door, while discarding irrelevant detour events.
Although most current memory models can be reformulated within the HMF, they tend to be either overly simplistic with limited calibration capabilities (Katharopoulos et al. 2020) Radford et al. 2019) or unable to manage memory writing reliably, suffering from gradient vanishing or exploding issues (Ba et al. 2016; Morad et al. 2024). To address these limitations, we propose a specific instance of HMF, called Stable Hadamard Memory, which introduces a novel calibration mechanism based on two key principles: (i) dynamically adjusting memory values in response to the current context input to selectively weaken outdated or enhance relevant information, and (ii) ensuring the expected value of the calibration matrix product remains bounded, thereby preventing gradient vanishing or exploding. Through extensive experimentation on POMDP benchmarks, including meta reinforcement learning, long-horizon credit assignment, and hard memorization games, demonstrate that our method consistently outperforms state-of-the-art memory-based models in terms of performance while also delivering competitive speed. We also provide comprehensive ablation studies that offer insights into the components and internal workings of our memory models.
In summary, our contributions are: (i) Unified Memory Framework: We propose unified memory framework, supported by theoretical analysis, to elucidate its properties as episode length increases. (ii) Novel Calibration Matrix: We design calibration matrix with random parameter selection, mitigating timestep dependencies and theoretically enhancing stability by regulating the expected gradient norm. (iii) Extensive Empirical Evaluation: Experiments on challenging benchmarks demonstrate the scalability and superior performance of our approach across various tasks.
2 BACKGROUND
2.1 REINFORCEMENT LEARNING PRELIMINARIES
A Partially Observable Markov Decision Process (POMDP) is formally defined tuple (S, A, O, RG 7) where S is the set of states, A is the set of actions, O is the observation space, R:S + R the reward function, P: A - (SS) defines the state transition probabilities, and y E [0, is the discount factor. Here, the agent does not directly observe the true environment state 8;. Instead, receives an observation o Os+)), which provides partial information about the state, often not enough for optimal decision making. Therefore, the agent must make decisions based on its current observation o and a history of previous observations, actions, and rewards (4o, To, O1...., ar-11 o.). The history may exclude past rewards or actions.
Let us denote the input context at timestep as It = (or, 01-1, rt-1) ) and assume that we can encode the sequence of contexts into a memory M, - (1,))/_1) The goal learn policy π (ay)MM that maximizes the expected cumulative discounted reward:
8 Σ
((x) E,
R(s at)aa T(ae)Mi),
8641 8t41
P(s ac), O(s;)
(1)
Thus, memory system capable of capturing past experiences is essential for agents to handle the partial observability the environment while maximizing long-term rewards.
2.2 MEMORY-AUGMENTER NEURAL NETWORKS
We focus on matrix memory M, and to simplify notation, we assume it is square matrix. Given memory M E # we usually read from the memory as:
h = Mig (xz)
(2)
where 9 is a query network RD R H map an input context de E R query vector. The read value h, later will be used as the input for policy/value functions. Even more important than the reading process is the memory writing: How can information be written into the memory
2

---

Published as a conference paper at ICLR 2025
ensure efficient and accurate memory reading? A general formulation for memory writing is Mr - (M with f as the update function that characterizes the memory models.
The simplest form of memory writing traces back Hebbian learning rules: Mr M,-1 + % It where X is outer product (Kohonen & Ruohonen 1973) Hebb 2005). Later, researchers have
proposed "fast weight" memory (Marr & Thach 1991 Schmidhuber 1992 Ba et al. 2016
M₁ = M,- A + ny (M,- 1, we) @ g (,)
where g a non-linear function that take the previous memory and the current input data as the input; and are constant hyperparameters. On the other hand, computer-inspire memory architectures such as Neural Turing Machine (NTM, (Graves al. 2014)) and Differentiable Neural Computer (DNC, (Graves et al. 2016)) introduce more sophisticated memory writing:
= M @ (1 w w (1) & (M)) + w (M-11, 21) @ (M-11,
(4)
where w, e and u are non-linear functions that take the previous memory and the current input data as the imput to produce the writing weight, crase and value vectors, respectively. 3 is the Hadamard (element-wise) product. The problem with non-linear f w.r.t M is that the computation must be done in recursive way, and thus being slow. Therefore, recent memory- such as Linear Transformer models adopt simplified linear update (Katharopoulos al. 2020):
@ (k (2x))
M (x,) 2
(5)
E, o (k (xx)))
where d is an activation function; k and v are functions that transform the input to key and value. In another perspective inspired by neuroscience, Fast Forgetful Memory (FFM, (Morad et al. 2024) employs parallelable memory writing, which processes a single step update as follows:
M, M, +7+ (xx) @ 1 0-
where y is a trainable matrix, is an input transformation function, and n is the last timestep (see more memory models in Appendix D).
METHODS
In this section, begin by introducing unified memory writing framework that incorporates several of the memory writing approaches discussed earlier. Next, we examine the limitations of current memory writing approaches through analysis of this framework. Following this analysis, we propose specialized techniques to address these limitations. For clarity and consistency, all matrix and vector indices will be referenced starting from rather than 0. Constant matrices are denoted by bold numbers.
3.1 HADAMARD MEMORY FRAMEWORK (HMF)
We propose a general memory framework that uses the Hadamard product as its core operation. The memory writing at time step is defined as:
M, = M-1 2 (oz)) + +Up(i
Un
where Co IRD h H U. R H are parameterized functions that map the current input to two matrices Ct (calibration matrix) and U. (update matrix). Here, 0 and are referred to as calibration and update parameters. Intuitively, the calibration matrix C, determines which parts of the previous memory M,-1 should be weakened and which should be strengthened while the update matrix U, specifies the content to be encoded into the memory. We specifically choose the Hadamard product (0) the matrix operator because it operates individually. We avoid using the matrix product to prevent mixing the content of different memory cells during calibration and update. Additionally, the matrix product is computationally slower.
each memory element
3

---

Published as a conference paper at ICLR 2025
There are many ways to design C₁ and Uy- Given proper choices of C2 and Ug. Eqs. 3 can be reformulated into Eq. 7 Inspired by prior "fast weight" works, we propose simple update matrix:
Up (xx) Tu (Ie) [o (xx) @ k (2x)]
where k and are trainable neural networks that transform the input to key and value representations. 7hp : IRD + is a parameterized function that maps the current input dt to an update gate that controls the amount of update at step t. For example, if Tp (ax) = 0, the memory will not be updated with any new content, whereas nh (xt) = 1 the memory will be fully updated with the content from the 1-th input like Linear Transformer. We implement me (24) network with sigmoid activation function.
neural
We direct focus towards the design of Ct. which is the contribution of work. The calibration matrix selectively updates the memory by erasing no longer important memories and reinforcing ongoing critical ones. In a degenerate case, if C. not forget or strengthen any past information, and will only memorize new information over time, similar the Hebbian Rule and Linear Transformer. To analyze the role of the calibration matrix. is useful to unroll the recurrence, leading to the closed- equation (see proof in Appendix A.I A.I
for all the memory will
M, Mo II
C.
Σ
U, O II C
(9)
j-1-
where II represents element-wise products. Then, h4 = Mig = Mo II Cig (ax) +
Ei= U. @ (i=+t1Cyq
Calibrating the memory is important because without calibration (G, = VI), the read value becomes: h; Mog (Ix) + Ei= Usg (aza). In this case, making he to reflect a past context at any step j requires that q (xx) * and Eizy Uiq (zx) = Eigy Thp (xx) (zx) [A: (I:) q (Ia)] Ro O, which can be achieved we can find (It) such that k (zz) (xz) N 0 Vi A j. Yet, this becomes hard when T >> H and Th (r;) o as it leads to an overdetermined system with more equations than variables. We note that avoiding memorizing any i-th step with 7p (x₂) = 0 is suboptimal since may be required for another reading step t' o t.
Therefore, at certain timestep it is critical to eliminate longer relevant timesteps from M, by calibration, i.e., U, O IT C2 0 for unimportant i (forgetting). For example, an agent may first encounter an important event, like seeing a color code, before doing less important tasks, such as picking apples. When reaching the goal requiring to identify a door matching the color code, it would be beneficial for the agent erase memories related the apple-picking task, ensuring a clean retrieval of relevant information the color code. Conversely. if timestep i becomes relevant
again at a later timestep 1', we need to recover its information, ensuring U, ④ IT' C₁ # (strengthening), just like the agent, after identifying the door, may return to collecting apple task.
Remark 1. In the Hadamard Memory Framework, calibration should be enabled (C, * 1) and conditioned on the input context.
Regarding computing efficiency, if C, and U, not functions of Mas, we compute the memory using Eq. parallel, ensuring fast execution. In particular, the set of products {II)_ C2
can be calculated in parallel in (log t) time. The summation can also be done parallel in O (log t) time. Additionally, since all operations are element-wise, they can be executed in parallel with respect to the memory dimensions. Consequently, the total time complexity is O (log t). Appendix Algo. Tillustrates implementation supporting parallelization.
Remark 2. In the Hadamard Memory Framework, with optimal parallel implementation, the time complexity for processing sequence of steps is O (logt). ) By contrast, without parallelization, the time complexity is O 21.12)..

---

Published as a conference paper at ICLR 2025
3.2 CHALLENGES ON MEMORY CALIBRATION
The calibration matrix enables agents to either forget or enhance past memories. However, complicates learning due the well-known issues of gradient vanishing or exploding. This can be observed when examining the policy gradient over 7 steps, which reads:
VeJ (Te) = E Es
V e log ne (4.) Adu (st. ac,
(10)
(e
where è is the set of parameters, containing {0, v), Ade represents the advantage function, which
integrates reward information R (st, and G. (e) B (e M,) 0M: captures information
GM:
de
related to the memory. Considering main gradient step required to learn 0 and yp:
0M,
0 a Co(wi)
i II Co(w))
Mo
Σ 1=1
Up
(11)
00
00
00
G1(4,10)
0M,
2
a (x:)
II
Co(og)
(12)
0e
Op
j=C+1
G2(LA.0)
We collectively refer G1 (i, t, 0) and G2 (i,t, 0) as G1,2 (i, t, 0) "I These terms are critical as they capture the influence state information timestep training challenges arise from these two terms as the number of timesteps & increases: (i) Numerical Instability (Gradient Exploding): if 3m, k [1, H] s.t. G1,2 2 (i, t, 0) [m, k) - oc, this leads to overflow, causing the gradient to become "nan"; (ii) Learning Difficulty (Gradient Vanishing): if < io, meaning no signal from timesteps < in contributes to
the learning parameters and yp. The
t >> fg, ||G1,2 (i,t, learning the parameters. This is suboptimal, especially when important observations occur early in the episode, and rewards are sparse and given at the episode end, .ie., R (st, a,) 4 T.
How to design the calibration matrix Co (I) to overcome the training challenges? A common approach is to either fix as hyperparameters or make it learnable parameters independent on the demonstrate that this leads to either numerical
input xı (e.g., Eqs. and Unfortunately, instability or learning difficulties as formalized Proposition 3 In the next section, we will provide better design for the calibration matrix.
Proposition 3. If calibration is enabled such that Ca(xe) + 1, and the calibration matrix remains fixed and independent of the imput Is (i.e., Vt Co(xs) = E R H HH,, this will lead to either numerical instability learning difficulties.
Proof. See Appendix A.2
3.3 STABLE HADAMARD MEMORY (SHM)
To avoid numerical and learning problems, is important to ensure each element C2 is not always greater than or smaller than 1, which ends up in their product will be bounded such that
E I[= C, # {0, )o) as T + cc. At the same time, we want C, to be function of x; to enable calibration conditioned on the current context. this end, we propose the calibration matrix:
Co (xx) = tanh (0, @ ve (xt))
(13)
where v, R + R is a mapping function, and 6, E "f represents the main calibration parameters. Here, implement v. linear transformation to map the input memory space.
5

---

Published as a conference paper at ICLR 2025
Notably, the choice of 0g determines the stability of the calibration. We propose to design 0, as trainable parameters that is randomly selected from a set of parameters In particular, given 0 = [0t)- where 0: E Ie and L is the set size, we sample uniformly a random parameter from 0: u (0). We name the design as Stable Hadamard Memory (SHM). Given the formulation, the range of an element 24 = Co (xx) [m, ki is [0, 2] where m, E [1. We can show one under certain conditions.
important property of this choice that E
II C2
Proposition 4. Assume: (i) N (0, 2.), (i) vy is linear transformation, i.e. (x,) Wa (iii) {z+ Og @ de (a=))== are independent across Given Co (xt) defined in Eq. 13 then:
E
m
E
Co (24) [m,
1 1V 0,1 km,
Proof. see Appendix A.3
Assumption (i) is reasonable as Fr can treated as being drawn from a Gaussian distribution, and LayerNorm be applied to normalize 2; to have mean of zero. Assumption (ii) can be realized implement U as linear transformation. Assumption (iii) is more restrictive because {a:)]- are often dependent in RI. setting, which means { = 8g @ v. (rz)) are not independent and thus, [T=
2t m.k # 1. However, by reducing the linear dependence between 2; through random selection of O., we can make [II m.,k closer to 1, and thus being bounded. Specifically, we will prove that by using 02 ~ L (0), the Pearson Correlation Coefficient between timesteps is minimized, as stated in the following proposition:
Proposition 5. Let 24 2im. uf where _m, (0, @ ve (21)) [m, km, ut" 0s [ma] and of ve (3rx) [k]. Given the Pearson correlation coefficient of two random variables X and Y is defined then Vot, vb:
as p (X,Y)= VVar(X) War(Y) (Wrrl)
Co
10 (u]" unl ul) S ( uf)))))
Proof. See Appendix A.4
As result, our choice of O, outperforms other straightforward designs for minimizing dependencies between timesteps. For instance, fixed Og (0, = 0 E R²) results in higher dependencies because p(0 [m] of A [m]] o = p(uř, utt 2 p (uf" of up o In practice, even when [ITT_ 2t m., bounded, the cumulative product occasionally become very large for certain episodes and timesteps, leading to overflow and disrupting the learning process. This can be avoided by clipping the gradients. In experiments, memory is then integrated into policy-gradient RL algorithms to optimize Eq. with the read value he (Eq. 2) used as imput for value/policy networks.
implement SHM using nonparallel recursive form (Eq. 7 The
The proposed update rule improves memory management and stability through several key benefits. The dynamic calibration matrix C, 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 vanishing, and ensureing robust performance over long episodes.
4 EXPERIMENTAL RESULTS
We evaluate our method alongside notable memory-augment agents POMDP environments. Unless stated otherwise, the context consists of the observation and previous action. All training

---

Published as a conference paper at ICLR 2025
FFM
FWP
GPT-2
GRU
SHM (Ours)
S6
mLSTM
wind-50-150
point rebot-50-150
wind-10-190
point robot-10-190
0.75
Sucoess Rate
0.4
0.4
0.50
0.5
0.25
0.2
0.2
0.00
0.0
0.0
0.0
400K
200K
400K
400K
200K
Figure 1: Meta-RL: Wind and Point Robot learning curves. Mean + std. over 5 runs.
FFM
FWP
GPT-2
CRU
SHM (Ours)
S6
mLSTM
visual match-250
visual, match-500
keytodoor-250
keytodoor-500
Success Rate
1.0
1.0
0.75
0.50
0.5
0.5
0.25
0.2
0.0
0.0
0.00
0.0
Figure 2: Credit Assignment: Visual Match, Key-to- toor learning curves. Mean + std. over runs.
uses the same hardware (single NVDIA H100 GPU), RL architecture, algorithm, training protocol. and hyperparameters. The baselines differ only their memory components: GRU (Chung al. 2014). FWP (Schlag et al. 2021). GPT-2 (Radford et al. 2019, S6 (Gu & Dao 2023), mLSTM (Beck et al. 2024b), FFM (Morad et al., 2024) and SHM (Ours). For tasks with a clear goal, we measure performance using the Success Rate, defined as the ratio of episodes that reach the goal to the total number of evaluation episodes. For tasks in POP-Gym, we use Episode Return as the evaluation metric. We fix SHM's number of possible 0,, L 128, across experiments.
4.1 SAMPLE- EFFECCEET META REINFORCEMENT LEARNING
Meta-RL targets POMDPs where rewards and environmental dynamics differ across episodes, representing various tasks (Schmidhuber 1987; Thrun & Pratt 1998). To excel in all tasks, memory agents must learn general memory update rules that can adapt to any environments. We enhanced the Wind and Point Robot environments from Ni et al. (2022) to increase difficulty. In these environments, the observation consists of the agent's 2D position Pt, while the goal state Pg is hidden. The agent takes continuous actions 0z by moving with 2D velocity vector. The sparse reward is defined as R(p++1, Py) = 1 (llp++1 Pyll2 2 S r) where r = 0.01 is the radius. In Wind, the goal fix P, [0, 1], yet there are noises in the dynamics: Pt+1 Pt + a; + w, with the "wind" w sampled from 1, 0.1] at the start and fixed thereafter. In Point Robot, the goal varies across episodes, sampled from U[-10, 10). To simulate real-world conditions where the training tasks are limited, we create 2 modes using different number of training and testing tasks: [50, 150) and [10, 190), respectively. Following the modifications, these simple environments become significantly more challenging to navigate toward the goal, so we set the episode length to 100.
We incorporate the memory methods to the Soft Actor Critic (SAC, (Haarnoja et al. 2018)), using the same code base introduced in Ni et al. (2022). We keep the SHM model sizes and memory capacities small, at around 2MB for the checkpoint and 512 memory elements, which is roughly equivalent to GRU (see Appendix C.I). We train all models for 500,000 environment steps, and report the learning curves in Fig. achieves the best performance, with a near-optimal success rate, while other methods underperform by 20-50% on average in Wind and Point Robot, respectively. In the hard mode (10-190), SHM continues to outperform other methods by approximately 20%, showing earlier signs of learning.
In the easy mode (50-150), our SHM method consistently
7

---

Published as a conference paper at ICLR 2025
Model Average Return
GRU 1.3.4.1.3
FFM
SHM (Ours)
-22
-5.1166.3
Table I: PopGym Hardest Tasks: Mean return + std. (x 100) at the end of training over 3 runs. The range of return (x (x100) is [-100, 100].
4.2 LONG-TERM CREDIT ASSIGNMENT
In this task, we select the Visual Match and Key-to-Door environments, the most challenge tasks mentioned in Ni et al. (2024). Both have observation as the local view of the agent, discrete actions and sparse rewards dependent on the full trajectory, requiring long-term memorization. In particular, the pixel-based Visual Match task features an intricate reward system: in Phase 1, observing color codes yields no rewards, while in Phase 2. picking an apple provides immediate reward of one, relying on short-term memory. The final reward-a bonus for reaching the door with the matching color code is set to 10. Key-to-Door also involves multiple phases: finding the key, picking apples, and reaching the door. The terminal reward is given if the key was acquired in the first phase and used to open the door. Both tasks can be seen as decomposed episodic problems with noisy immediate rewards, requiring that in the final phase, the agent remembers the event in the first phase. We create 2 configurations using different episode steps in the second phases: 250 and 500, respectively.
We use SAC- AiC-Discrete (Christodoulou 2019) the RL algorithm. We train the same set of baselines as in Sec. 4.1 for 2 million environment steps. The results in Fig. 2 show that SHM significantly outperforms all other methods in success rate. Notably, SHM is the only method that can perfectly solve both 250 and 500-step Visual Match while the second-best performer, FFM, achieves only a 77% and 25% success rate, respectively. In Key-To-Door, our method continues showing superior results with high success rate. By contrast, no meaningful learning is observed from the others, which perform similarly to GPT-2, as also noted by Ni al. 2024).
4.3 POPGYM HARDEST GAMES
We evaluate SHM on the POPGym benchmark (Morad et al. 2023), the largest POMDP benchmark to date. Following previous studies (Morad al. 2023] Samsami et al. (2024), we focus the most memory-intensive tasks: Autoencode, Battleship. Concentration and RepeatPrevious. These tasks require ultra long-term memorization, with complexity increasing across Easy, Medium, and Hard levels. All tasks use categorical action and observation spaces, allowing up to 1024 steps.
For comparison, evaluate SHM against state-of-the-art model-free methods, including GRU and FFM. Other memory models, such as DNC, Transformers, FWP, and SSMs, have been reported to perform worse. The examined memory models integrated into PPO (Schulman al.) (017). trained for 15 million steps using the same codebase as Morad et al. (023) to ensure fairness. The models differ only in their memory, controlled by the memory dimension hyperparameter H. We tune it for each baseline, adjusting it to optimize performance, as larger H values typically improve results. The pers-peeformin configurations are reported Table and Appendix Table 6 where SHM demonstrates a relative improvement of 100-12% over FFM and GRU on average. Notably, the learning curves in Appendix Fig. 4 show that only SHM demonstrates signs learning in several tasks, including RepeatPrevious- Propauurevviuu-MMiiu/iHHtttee Medium//ardd, and Autoencode- Autoenoode-EayyMMddin.. Detailed hyperparameter setting and additional results are provided in Appendix
In terms of running time, the average batch inference time in milliseconds for GRU, FFM, and SHM is 1.6, 1.8, and 1.9, respectively, leading to a total of 7, 8, and hours of training per task. While SHM is slightly slower than GRU and FFM, the difference is reasonable trade-off for improved memory management in partially observable environments. Last but not least, our model's 's runtime was measured using con-parallel implementation, while GRU benefits from hardware optimization in the PyTorch library. SHM's running time could be further improved with proper parallelization.
4.4 MODEL ANALYSIS AND ABLATION STUDY
Choice of Calibration Prop. 5 suggests that selecting random 0, E R in Eq. 13 will reduce the dependencies between Ct, bringing l Ct closer avoid gradient issues. In this section,
8

---

Published as a conference paper at ICLR 2025
Rondom 6/Durs)
(b)
R 128
H 128
Episodic Return
H 156
H 156
Episodic Retura
0.5
Qumglative Preduct
1.0
0.50
0.8
A
H
8.6
0.00
-0.40
0.5
-0.50
15M
Step
SM
100
Autoencode--aay, Fasy
Battleship-Fasy
Sep
Important memory
t=20
(c)
1=
100
M
-100
1-0
1-15
(-50
1-50
[-100
1-115
1-120
C
Phase
Forget Remember
Phase 2
Phase
Figure 3: (a) Left: Return of calibration designs over 3 runs; Right: Calibration matrix cumulative
product over 100 episodes. (b) Return of memory sizes on Autoencode- (left) and Battleship-
Easy (right). (c) Memory (M, top) and calibration (C, bottom) matrices over timesteps in Visual
Match: SHM erases memory that are no longer required and strengthens the important ones.
empirically verify that by comparing our proposed Random Og with the following designs: C 1. no calibration is used; Random C, where random calibration matrix sampled from normal distribution each timestep, having E [II" m.,k under mild assumptions, but preventing learning meaningful calibration rules; Fixed C, simpler method for learning calibration, but prone gradient problems (Prop. 33; Fixed 0., where we learn fixed parameter fr, which is proven to be less effective than Random O, in reducing linear dependencies between timesteps (Prop. [): Neural O,, where 0, = FFW (zx), generated by a feedforward neural network like mLSTM, but with no guarantee of reducing timestep dependencies. We trained RL agents using the above designs of the calibration matrix with H = 72 on the Autoencode-Ease Easy task and present the learning curves in Fig. (a, left). The results show that our proposed Random 0 outperforms the other baselines by substantial margin, with at least a 30% improvement in performance. This confirms the effectiveness of our calibration design in enhancing the final results.
Vanishing Behavior practice, exploding gradients can be mitigated by clipping. Thus, we focus on the vamishing gradient, which depends on the cumulative product C, I C.. Our theory suggests that Random should be less susceptible to the vanishing phenomenon compared to other methods such as Fixed C, Fixed 0g and Neural 0, To verify that, for each episode, we compute the average value of elements in the matrix Cj that are smaller than (C, [< 1). as those larger than are prone exploding and are not appropriate for averaging with the vanishing values. We plot C for = 1,2, 100 over 100 episodes Fig (a, right).
The results are consistent with theoretical predictions: Fixed C leads to rapid vanishing of the cumulative product in just 10 steps. Neural 8, also perform badly, likely due to more complex dependencies between timesteps because z now becomes non-linear function of 2t, causing I m.k to deviate further from 1. While Fixed 0, is better than Neural 0,, it still exhibits quicker vanishing compared to our approach Random 0,. As expected, Random C' shows little vanishing, but like setting C' 1, it fails to leverage memory calibration, resulting in underperformance (Fig. (a, left).. Random 0, although its C, also deviates from 1, shows the smallest deviation among the calibration learning approaches. Additionally, the vanishing remain manageable after 100 timesteps, allowing gradients to still propagate effectively and enabling the calibration parameters to be learned.
Memory Size The primary hyperparameter of our method is H, determining the memory capacity. We test SHM with H (24, 72, 128, 156) on the Autoencode-Eand Easy and Battleship Mattleship-Easy tasks. Fig. 3 (b) shows that larger memory generally results in better performance. In terms of speed,

---

Published as a conference paper at ICLR 2025
the average batch inference times (in milliseconds) for different H values are 1.7, 1.8, 1.9, and 2.1, respectively. We choose H = 128 for other POPGym tasks to balance performance and speed.
Forgetting and Remembering We study the learned calibration strategy of SHM on Visual Match with 15,100, and 10 steps in Phase I, 2 and, 3, respectively. We sample several representative M, and C₁ from 3 phases and visualize them in Fig. [ (c). In Phase 1, the agent identifies the color code and stores it in M, possibly in two columns of M, marked as "important memory". In Phase unimportant memory elements are erased where C. 0 (e.g., those within the blue rectangle). However, important experiences related to the Phase I's code are preserved across timesteps until Phase (e.g., those within the red rectangle where C2 21), which is desirable.
RELATED WORKS
Numerous efforts have been made to equip RL agents with memory mechanisms for handling POMDPs. Memory models can be broadly categorized into vector-based and matrix-base approaches. Vector-based memory, like RNNs (Elman 1990), processes inputs sequentially and stores past inputs in their hidden states. While RNNs are slower to train, they are efficient during inference. Advanced variants, such GRU and LSTM, have shown strong performance in POMDPs, often outperforming more complex RL methods (Ni et al. 2022 Morad et al., 2023). Others focus on the use of permutation-innaaiinn sequence models meta-reinforceemet learning. demonstrating their advantage even without task inference objectives (Beck et al. (2024a). Recently, faster alternatives like convolutional and structured state space models (SSM) have gained attention (Bai et al., 2018 Gu et al., 2020), though their effectiveness in RI. is still under exploration. Initial attempts with models like $4 underperformed in POMDP tasks (Morad et al. 2023), but improved SSM versions using S5, S6 world models have shown promise (Lu et al., [024] Gu & Dao [0233
Samsami et al. 2024). Despite these advancements, vector-based memory is limited, as compressing history into single vector makes it challenging scale for high-dimensional memory space.
Matrix-based memory, on the other hand, offers higher capacity by storing history in a matrix but at the cost of increased complexity. Attention-bast based models, such Transformers (Vaswani ct al. al..||2.000 (017), have largely replaced RNNs SL, also delivering good results in standard POMDPs (Parisotto et al., 2020). However, their quadratic memory requirements limit their use in environments with long episodes. Empirical studies have also shown that Transformers struggle with long-term memorization and credit assignment tasks (Ni et al. (024). While classic memory- augmented neural networks (MANNs) demonstrated good performance in well-crafted long-term settings (Graves et al., 2016 Hung et al. 2019 Le et al. 2020), they are slow and do not scale well in larger benchmarks like POPGym (Morad et al. (023). New variants of LSTM (Beck et al.) 2024b), including those based on matrices, have not been tested in reinforcement learning settings and lack theoretical grounding to ensure stability.
Simplified matrix memory models (Katharopoulos et al., 20200 Schlag et al. 2021 offer scalable solutions but have underperformed compared to simple RNNs in the POPGym benchmark, highlighting the challenges of designing efficient matrix memory for POMDPs. Recently, Fast and Forgetful Memory (FFM, (Morad et al. 2024)) ), incorporating inductive biases from neuroscience, has demonstrated better average results than RNNs in the benchmark. However, in the most weemory-intensio environments, the improvement remains limited. Compared to approach, these matrix-based memory methods lack a flexible memory calibration mechanism and do not have robust safeguards to prevent numerical and learning issues in extremely long episodes.
DISCUSSION
In this paper, we introduced the Stable Hadamard Framework (SHF) and its effective instance, the Stable Hadamard Memory (SHM), a novel memory model designed to tackle the challenges of dynamic memory management in partially observable environments. By utilizing the Hadamard product for memory calibration and update, SHM provides an efficient and theoretically grounded mechanism for selectively erasing and reinforcing memories based on relevance. Our experiments on the POPGym and POMDP benchmarks demonstrate that SHM significantly outperforms state-of- the-art memory- based models, particularly in long-term memory tasks, while being fast to execute.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin lonescu. Using fast weights attend to the recent past. In Advances in Neural Information Processing Systems, pp. 4331-4339, 2016.
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint 1uXiv:1803.0127 2018.
Jacob Beck, Matthew Jackson, Risto Vuorio, Zheng Xiong, and Shimon Whiteson. Splagger: Split aggregation for meta-reinforcemmet learning. In Reinforcement Learning Conference, 2024a.
URL htttpss://ppeeeeeeee wwwp:/eeeeooommiioeooooooooo Sodd0l10uuuuaaae 4RVW
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint arXiv:2405.04517 2024b.
Petros Christodoulou. Soft actor-critic for discrete action settings. arXiv preprint arXiv: 2019.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv: 2014.
Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2):199-11,, 1990.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv: 2014.
arXiv preprint
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiñska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538 (7626):4714776 2016.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. 2023.
Albert Gu, Tri Dao. Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Donald Olding Hebb. The organization of behavior: A neuropsychological theory. Psychology press, 2005.
Hochreiter. Long short-term memory. Neural Computation MIT-Press, 1997.
Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Abuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):5223, 2019.
Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Plamning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134, 1998.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Teuvo Kohonen and Matti Ruohonen. Representation of associated data by matrix operators. IEEE Transactions on Computers, 100(7):701-702, 1973.
Hung Le, Truyen Tran, and Svetha Venkatesh. Self-attentive associative memory. In International conference on machine learning, pp. 5682-5691. PMLR, 2020.

---

Published as a conference paper at ICLR 2025
Chris Lu, Yannick Schroecker, Albert Gu. Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Behbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
David Marr and W Thomas Thach. A theory of cerebellar cortex. In From the Retina to the Neocorler, pp. 11-50. Springer, 1991.
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. Popgym: Benchmarking partially observable reinforcement learning. arXiv preprint arXiv:2303011559 2023.
Steven Morad, Ryan Kortvelesy, Stephan Liwicki, and Amanda Prorok. Reinforcement learning with fast and forgetful memory. Advances in Neural Information Processing Systems, 36, 2024.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be strong baseline for many pomdps. In International Conference on Machine Learning, pp. 16691 16723. PMLR, 2022.
Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Lue Bacon. When do transformers shine in rl? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36, 2024.
Emilio Parisotto, Francis Song, Jack Rac, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference machine learning, pp. 7487-7498. PMLR, 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mastering memory tasks with world models. arXiv preprint 1.0Xv::2003042553, 2024.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers secretly fast weight programmers. In International Conference on Machine Learning, pp. 9355-9366. PMLR, 2021.
Jürgen Schmidhuber. Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta.... hook. PhD thesis, Technische Universität München, 1987.
Jürgen Schmidhuber. Learning to control fastht memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139 1992.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.066 2017.
Sebastian Thrun and Lorien Pratt. Learning to learn: Introduction and overview. In Learning to learn pp. 3-17. Springer, 1998.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Neural Information Processing Systems, 2017. URL https: : /ppi. semanticscholas htt.a//coooo.ca.ooooooooooooaaa.
12

---

Published as a conference paper at ICLR 2025
APPENDIX
A THEORETICAL RESULTS
A.1 CLOSED-FORM MEMORY UPDATE
Given
M, = M M
(14)
then
M, = Mo II C. + Σ
U II c j=-1
(15)
t=
Proof. We prove by induction.
Base case: for = 1, the equation becomes for both Eqs. 14 and 15
M1 My C1 U1
Thus, the equation holds for / 1.
Inductive hypothesis: Assume the equation holds for = n.
n
12 Σ U I1 C
M, My II C,
j=it1
Inductive step: We now prove the equation holds for t = n 1. From the update rule in Eq. 14
M = Mn @ Cn+1 + Un+1
Substitute.M, using the inductive hypothesis:
a II
n U II C,
Mn+1
Mo
C
Σ
Cn+1 + Un+1
1=1
j=++1
n+1
n+1
Mo
II
C
Σ
U
II
C,
Un+1
j=it1
Mo
II
C
Σ
U
II j=++1
C
1=1
This matches the form of the closed-form Eq. for completing the proof by induction.
A.2 PROPOSITION 3
Definition 6. Critical Memory Gradients: In the Hadamard Memory Framework, we define the critical memory gradients of memory rules as follows:
O II1 -1+1 Co(o,)
aCo (xz)
G G1 (1,t,))
Σ
II 1,
Co(ok)
00
08
j=it1
G₂ (1.1,0) =
II
Co(ar)
Now can proceed the proof for Proposition
3
13

---

Published as a conference paper at ICLR 2025
Proof. In the case Co (xt) E Rx*" the critical gradients read:
G1 (1,1,0) Σ II 0
j=i+1 k=
= =(t i) g
G2 (i,t,) Ap
30 m, < s.t. 10 [m, ki] 1, 21 (i, t, 0, [m, kj ~ o (0 [m,, k)" )))++~)) and thus + 0o, i.e.
numerical problem arises. Let ||-||| denote the infinity norm, we also have:
||G (i, t, 0)l [| 0 0-1 S (t |101
GG (i, 8)1 ||0'011 - ||0||*
Note that if Vm, & 10 [m, k3] < 1, |||| < 1, both terms become 0 as f -i increases, thus learning problem always arises. conclusion, in this case, to avoid both numerical and learning problems, V', le 0 (m, k]] = 1, which is not optimal in general.
A.3 PROPOSITION 4
Proof. Let 2: m..k = ut"h uff'ff where 2; 2/ (0, @ ve (z+)) [m, kl. wrth 0, [m] and ot = ue (xx) [k]]. Using assumption (i) and (ii), ey is Gaussian variable, i.e., of N (0, p.)) By definition, uf" is categorical random variable that can take values {0 [m, [e]] with equal probability 1/L. For now, we drop the subscripts m, k and t for notation ease. The PDF as a mixture distribution since u is categorical and can take discrete values {wi)err The PDF of 2, denoted as ff (z), is given by:
wv can be expressed
f(=) = Σ
P(u ut)um
1=1
I Σ (2)
where [ujo(2) is the PDF of uyv, and is constant for each 1. Thus, [ugu (2) is the scaled PDF of w:
2
(2)
fo
Since v ~ (N, p). the PDF of v. denoted as f. (z), is symmetric about 0, we have:
(x)(2)) fo = = fo
W = (aje(-2).
This shows that (2)(2) is symmetric around for each I. Therefore, the PDF f(2) is also symmetric:
1 L =())= = - f2) = (m) f(-2)
L
Since tanh an odd function and the PDF of mm.k is symmetric about [ann (m
0 and
thus [1 + tanh (z = 1. Finally, using assumption (iii),
[T= m.,k = -I F m.k
T
14

---

Published as a conference paper at ICLR 2025
Task
Input type
Policy/ networks
RL algorithm
Batch size
3-layer FFW Memory
Meta-RL
Vector
SAC
32
layer (128, 128, 128, H)
Credit Assignment
2- 2-layer CNN 2-layer FFW
Image
SAC-D
32
+ Memory layer (128, 128, H)
3-layer FFW + Memory layer (128, 64, H, 64)
POPGym
Vector
PPO
65,536
Table 2: Network architecture shared across memory baselines.
A.4 PROPOSITION
Proof. Without loss of generalization, we can drop the indice m and is for notation ease. Since (0), it is reasonable assume that each of {Me, Je {u,, Ve}. {us ve {ur' och. {ue, up are independent. In this case, let us denote X = U2 and X = te we have
Cov(X,X)) Cov(u ue ve
Elusve ver ] - E (1 Elurvvy)
B( (we IE ) (uu (B(e))
E( (4 (E))) (E(ur) E(u [[())
E[ B(ue Cov(us, Br))
The variances read:
Var(X) Var(ugu; 1E( Eu Bluil (B[e²) B[(e)²)
Bluu]] Var(v:)
Similarly:
Var(X') Var(urr
E Var(ue))
Given P as the Pearson Correlation Coefficient, consider the ratio:
B Ms Bl Coo( / M.r
(((XX')) 1p(ve, vrr
B [uu Var(er) )-B Var(v.))
Cov(vs.o,)))
E u, E(uy
E E
By the independence of ue and uy and the Cauchy-Schward inequality:
E B[ue]] E [u
V E E(u])
which implies
Ip(ueve, uy lo(o,, we)
< > lp (urre, 10 (vr, vy)
The equality holds when ur Bug.
DETAILS ON METHODOLOGY
In this section, describe RL frameworks used across experiments, which adopted exactly from the provided benchmark. Table summarizes the main configurations. Further details can be
found in the benchmark papers (Ni al. 2022 Morad et al. 2023)
15

---

Published as a conference paper at ICLR 2025
Algorithm 1: Theoretical Parallel Hadamard Memory Framework.
Input: Mo E BR C e RBBT*H>H U E B
Operator: @ parallel prefix sum, parallel prefix product, @ Hadamard product
Ouput: M = [M] E BxTxxXX
/* Parallel prefix product along T. Complexity: Ollog(t)).
Cp & di)
/* Concatenation along т. Complexity: 0(1).
+/
D concat([MMo CJ dim
/* Parallel prefix product along т. Complexity: O(log(())
Dp @DD, dim 1)
/* Parallel Hadamard product (Cp # 0) Complexity: 0(1).
E = U @ C
Parallel prefix sum along T. Complexity: O(log(t)).
Ep ((E, dim 1)
Parallel sum. Complexity: 0(1).
M = Dpl: 1:
Task Meta-RL Credit Assignment
URL
License
https://ggthu...cmmmmmm.m/////// baseeinessesseeeeeeee
ines
MIT
MIT
POP-Gym
MIT
Table 3: Benchmark repositories used in our paper.
C DETAILS OF EXPERIMENTS
We adopt public benchmark repositories to conduct our experiments. The detail is given in Table
C.1 HAMPLE-EFFICIEE ASSIGNMENT DETAILS
META REINFORCEMENT LEARNING AND LONG-TERM CREDIT
We report the choice memory hyperparameter in Table. 4 Due to the lack of efficient parallel processing in the codebase from Ni et al. (2022), rumning experiments with larger memory sizes is prohibitively slow. As a result, we were only able to test models with 512 memory elements, limiting the potential performance of our SHM. This constraint contrasts with the POPGym benchmark with better parallel processing, where our method scales comfortably with larger memory sizes, as demonstrated later in
We also report the running time and final performance for Visual Match and Key-to-Dook experiments for different episode lengths. The results in Table [ demonstrate that SHM can scale reasonably in both computation efficiency and performance as the task complexity increases.
C.2 POPGYM HARDEST GAMES DETAILS
In this experiment, we tuned the memory model size to ensure optimal performance. Specifically, for GRU, tested hidden sizes of 256, 512, 1024, and 2048 on the Autoencode- task. Increasing GRU's hidden size did not lead to performance gains but resulted a significant rise in computation
Model
GRU
FWP
GPT-2*
S6
mLSTM
FFM
SHM (Ours)
512
24
512
512
24
128
24
Memory elements
512
512
512 x T
512
512
512
512
Table 4: Memory dimension for Meta-RL. and Credit Assignment tasks. GRU and S6 use vector memory of size H. GPT-2 does not have a fixed size memory and attend all previous T timesteps in the episode. H = 512 is the dimension of Transformer's hidden layer. FFM's memory shape is 2 x m X where = H = m 128. FWP and SHM's memory shape is x H.
16

---

Published as a conference paper at ICLR 2025
Total Training
Success Rate (%)
Success Rate (%)
Episode Length
Time (hours)
Visual Match
Key To Door
SHM
GPT-2
SHM
GPT-2
SHM
GPT-2
100
20
18
100
83
100
59
250
27
23
100
72
78
500
41
36
100
47
55
22
Table 5: Average running time and success rate across episode lengths.
FFM
GRU
SHM (Ours)
Avg. Episods Return
euta
Battleship-easy
Concentration-easy
RepestPrevtous-eayy
0.5
-0.2
0.0
-0.4
-0.5
-022
10M
10M
10M
10M
Step
Step
Step
Step
Avg. Episode Retum
Autoencode-medlum
Battleship-medlum
Concentration-mmedum
RepeatPrevious
0.5
0.2
0.20
-0.3
0.3
-0.4
-0.4
-0.5
-0.5
10M
10M
10M
10M
Step
Step
Step
Step
Avg. Episode Return
Autoencode-hard
Battleship-hard
Concentration-hart
RepeatPrevious-aard
002
-0.45
-0.2
-0.850
-0.3
-0.775
-0.4
-0.50
-04
10M
10M
10M
10M
Step
Step
Step
Step
Figure POPGym learning curves: Mean + std. over runs.
cost. For instance, at Н = 2048, the average batch inference time was 30 milliseconds, compared 1.6 milliseconds at H 256. Thus, we set GRU's hidden size to 256, as recommended by the POPGym documentation. For FFM, we set the context size to 4, as the authors suggest that a larger context size degrades performance. We tuned the trace size m E {32. 128, 256, 512} on the Autoencode-Eaie value for all experiments. For our method, we tested H E {24, 72, 128, 156) on the same task and observed that larger values of H led to better performance, as shown in the ablation study in Sec. 4.4 However, to balance runtime, memory cost, and performance, we set H = 128 for all POPGym experiments. The learning curves are given in Fig. 4 The final results are reported fully in Table
task and found that m 128 provided slightly better results, so used this
D MORE RELATED MEMORY MODELS
Recently, Beck et al. (2024b) have proposed matrix- LSTM (mLSTM):
M, = f (ax)) M + (ax) v (ax) @ (xx) where f and i are the forget and input gates, respectively.
(16)
17

---

Published as a conference paper at ICLR 2025
Task
Level
Linear Transformer"
GRU
FFM
SHIM (Ours)
Easy
44.7411.4
-3.997.77
-2.770066
49.5::23.3
Autoencode
Medium
47.8+0..2
43.623.5
332.710.6
-28.8=14.4
Hard
48.110.11
48.1120.7
47.740.55
43.9-0.9
Easy
.1.330..5
41.121.
-3.0047.1
-12.322.4
Battleship
Medium
-3,.2+0.3
03.4
-37.1+3.1
-1..8=0.6
Hard
-38.410.22
-3.55...5
-38.8.0.3
221.2.1..
Easy
-18.5+0.2
110.9
10.7+112
-1.9=2.4
Concentration
Medium
-18.610.2
21.40..5
22..7.0.1
21.0
Hard
-3..010.1
84.0.1..3
-8.55
-3.
Easy
6.0#4.00
99.910.0
98,4
88.9
RepcatPrevious
Medium
1..
33.721..7
24.340..4
48.2=7.2
Hard
-8..510.33
41.7718.8
-33.911.0
-19.4
Average
-3..140.4
-28.4+1.3
-24.211.2
-5.1=6.3
Table PopGym: Mean return std. (x 100) the end of training over 3 runs. The range of return (x 100) is [- 100, 100). is reported from Morad et al. (2023).
18


Reviews:
[{"review_id": "gJdpUahzVZ", "reviewer": "Reviewer_1jQd", "summary": "This work presents a memory-based architecture for partially observable RL, named Stable Hadamard Memory (SHM). SHM is evaluated on a range of POMDP environments and POPGym against recurrent and attention-based baselines, as well as some ablations and analysis.", "questions": "* Typo in the title of Section 3.3 - Stable Hadamrad Memory.\n\n* How were the baselines tuned for each of the experiments? I could only find partial detail in the appendix. What efforts were made to normalize computational and memory cost between baselines and SHM?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well presented and the method is explained in detail. It is evaluated against a range of baselines and on multiple benchmarks.", "weaknesses": "1. The novelty of the method is not directly apparent from reading the paper. Stable Hadamard Memory is a confusing description for the method, as many existing RNNs, such as GRUs, make use of a Hadamard product as a core component. Furthermore, the Hadamard product is not the defining feature of the method, any more than addition. Instead, it would seem that the design of the calibration matrix, with random selection/dropout of update parameters, is the primary contribution of the approach. Based on this, more focus on this component in the paper and method naming would benefit the reader.\n\n2. The method's update rule is highly similar to linear attention, which is cited but not discussed in the work. A discussion of how this method relates to linear attention and empirical evaluation against it would be compelling.\n\n3. Due to the random selection of update parameters in SHM, the memory usage of SHM should be higher than methods without parameter dropout. Discussing this and adding experiments comparing SHM to baselines with the same memory usage would be useful.\n\n4. The experiments have very few repeats (3 seeds) and it is unclear how some baselines were tuned.\n\n5. A paper from RLC 2024 exploring memory aggregation in Meta-RL is not cited [1].\n\n[1] Jacob Beck, Matthew Jackson, Risto Vuorio, Zheng Xiong, and Shimon Whiteson. Splagger: Split aggregation for meta-reinforcement learning. Reinforcement Learning Conference, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2TUu2rLyKE", "reviewer": "Reviewer_Av1T", "summary": "This work addresses memory in partially observable RL environments. The authors state that current methods fail to capture relevant past information, making it difficult to retain important memories and know when to dynamically update them in response to changes in the environment. This issue is particularly significant for continual learning.\n\nThe authors introduce a Hadamard memory framework, suggesting that it performs memory calibration and linearly adjusts memory elements by maintaining a calibration matrix used to update these elements via an update matrix. The Hadamard product allows these elements to be learned using neural networks, and this method is well-supported in the literature. The memory matrix is then formulated using the standard policy gradient method, and the authors also demonstrate stability guarantees.\n\nThis approach is based on the Hadamard product applied to the memory matrix, enabling adaptive memory rules to be learned.\n\nThe method is evaluated on various Meta-RL tasks and challenging credit assignment tasks, with results showing dramatic improvements. The authors propose that this dynamic memory writing method facilitates more efficient credit assignment. They also evaluate the method on a set of partially observable gym environments, visually demonstrating the memory and calibration matrices.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* As said in my summary, the results show improvements on a broad set of tasks and appear to be extremely significant.\n\n* The method is mathematically well motivated, justified well and theoretically guarantees are given. I especially appreciate the stability guarantees.  It appears to be sound to use this for updating agent memories.\n\n* The credit assignment benefits are justified well and empirically shown.\n\n* The work is very well motivated and many of the theoretical guarantees are strong and supported well.\n\nOverall this work is solid and presents and interesting way of updating memory.  I currently support acceptance of this work.  I do not have high confidence in my review as this work is a bit orthogonal to the previous works I have done on partially observably environments.", "weaknesses": "* The paper is overly complex at times going into too many mathematical details.  I would like a very clear description of the benefits of using this update rule and why it empirically works well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2OIafjI0uF", "reviewer": "Reviewer_SMQZ", "summary": "The authors propose a relatively simple recurrent model for reinforcement learning under partial observability. Their model is parallelizable in time due to a linear recurrent state update. They demonstrate that their method outperforms prior work across various benchmarks, both with PPO SAC algorithms.\n\nAll in all, this is a good paper tackling an important problem with a novel approach and thorough experimentation. In my opinion, the biggest weakness is that the authors have not provided their code.", "questions": "- Why is $1 + \\tanh(x)$ used instead of $\\textrm{sigmoid}(x)$?\n- How come the recurrent state is real-valued compared to methods like FFM? Did the authors try using complex-valued states to track phase information across sequences?\n- Line 233: What is $m.k$?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper is nicely structured\n- The authors provide a thorough literature review\n- The authors' method is very straightforward to implement\n- The derivation of the linear matrix C is much shorter and easier to read than the derivation of the A and B matrices from the State Space Models paper (and yet, the authors' method performs better)\n- Rather than designing linear operations to keep inputs around in the recurrent state as long as possible (State Space Models), the authors take relatively novel approach to design a linear operation that has well-behaved gradients\n- Randomly sampling $\\theta$ to break the dependence between inputs is a novel and interesting idea\n- The authors run numerous experiments, across multiple benchmarks and RL algorithms", "weaknesses": "- The authors do not provide code\n- The writing is a bit awkward in places, hampering my understanding -- I think the paper would benefit from having another proofread\n    - The initialization of $C$ is quite important (lines 259-264), but it took me a few reads to understand the process -- perhaps the authors can consider writing something like $\\theta = [\\theta_1, \\theta_2, \\dots, \\theta_{L}]; \\quad \\theta_t \\sim \\mathcal{U}(\\theta)$\n    - Proposition 3 is poorly worded, and should be rewritten\n    - Section 3.3 is misspelled (STABLE HADAMRAD MEMORY (SHM))", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xlFZLhbZDN", "reviewer": "Reviewer_3seu", "summary": "This paper introduce\"SHM\"model, spcefically solving memory management problem in RL. SHM use a Hadamard product-based memory update mechanism to dynamically retain important information and discard irrelevant data, aiming to enhance memory stability and efficiency. The authors propose SHM as a solution to the limitations of existing memory models, such as numerical instability and inefficient memory retention, particularly in complex and long-horizon tasks. Experimental results demonstrate SHM’s superior performance compared to state-of-the-art memory models across multiple benchmarks, including meta-reinforcement learning and the POPGym suite, indicating its potential as a robust memory mechanism for RL applications.", "questions": "Hopefully, author interpret HaM in aspect of  \"scaling law\".\n\n① Performance vs. Memory Size: How does the performance of the Hadamard Memory model scale with increases in memory size? Is there an optimal memory size at which the model's performance peaks?\n\n② Impact of Task Complexity: As task complexity increases, how do the computational requirements and performance of the model change? Can the authors provide metrics to measure the impact of task complexity on model performance?\n\n③ Scalability with Episode Length: Does the Hadamard Memory model maintain stable and efficient updates in longer episodes? Have the authors considered the model's stability and efficiency over extended periods?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "This paper's got some fresh ideas that really matter in the real world. First off, they've come up with this nifty thing called the Stable Hadamard Memory (SHM) model. It's like a brain upgrade for AI in situations where it can't see everything. The unique twist is this memory update method that uses something called the Hadamard product. It's all about keeping the good memories and dumping the useless ones, which makes the AI's memory super stable and efficient.\n\nBut that's not all. The paper really shows off the SHM's chops by putting it through its paces in a bunch of tests.  It outperforms the current bigwigs in memory models by a mile. This model's got the chops to handle those tough, long-term memory tasks that other models can't touch, especially when it comes to memory-intensive tasks in reinforcement learning.\n\nSo, in a nutshell, the research in this paper could really move the needle in the field of reinforcement learning.", "weaknesses": "In reinforcement learning, especially when dealing with \"partially observable\" environments, effective memory management is really crucial. As task complexity increases, memory demand also skyrockets, and challenges grow alongside it. So, discussing how models \"hold up\" as tasks become more complex and expansive is pretty important.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)