Reinforcement Learning with Fast and Forgetful Memory
Authors: Steven Morad
Ryan Kortvelesy
Stephan Liwicki
Amanda Prorok
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
partially observable
POMDP
memory
rnn
transformer
TL;DR: We propose a drop-in alternative to RNNs for model-free recurrent RL

Abstract: Nearly all real world tasks are inherently partially observable, necessitating the use of memory in Reinforcement Learning (RL). Most model-free approaches summarize the trajectory into a latent Markov state using memory models borrowed from Supervised Learning (SL), even though RL tends to exhibit different training and efficiency characteristics. Addressing this discrepancy, we introduce Fast and Forgetful Memory, an algorithm-agnostic memory model designed specifically for RL. Our approach constrains the model search space via strong structural priors inspired by computational psychology. It is a drop-in replacement for recurrent neural networks (RNNs) in recurrent RL algorithms, achieving greater reward than RNNs across various recurrent benchmarks and algorithms _without changing any hyperparameters_. Moreover, Fast and Forgetful Memory exhibits training speeds two orders of magnitude faster than RNNs, attributed to its logarithmic time and linear space complexity. Our implementation is available at https://github.com/proroklab/ffm.

Reinforcement Learning with Fast and Forgetful Memory
Steven Morad University of Cambridge Cambridge, UK sm2558@cam. ac uk
Ryan Kortvelesy University of Cambridge Cambridge, UK rk627@cam ac.uk uk
Stephan Liwicki Toshiba Europe Ltd. Cambridge, UK stephan.1iwicki www.liiiitiittsbbbe.....
Amanda Prorok University of Cambridge Cambridge, UK asp450cam ac uk
Abstract
Nearly all real world tasks are inherently partially observable, necessitating the use of memory in Reinforcement Learning (RL). Most model-free approaches marize the trajectory into a latent Markov state using memory models borrowed from Supervised Learning (SL), even though RL tends exhibit different train- ing and efficiency characteristics. Addressing this discrepancy, we introduce Fast and Forgetful Memory, an
memory model designed specifi- cally for RL. Our approach constrains the model search space via strong structural priors inspired by computational psychology.
drop-in replacement for re- current neural networks (RNNs) in recurrent RL algorithms. achieving greater reward than RNNs across various recurrent benchmarks and algorithms without changing any hyperparameters. Moreover, Fast and Forgetful Memory exhibits training speeds two orders of magnitude faster than RNNs, attributed its log- arithmic time and linear space complexity. Our implementation is available at https: /ttp//ggttttb/tt wwt.coo//ppoooo hoom/poroolabbffm..
Introduction
Reinforcement Learning (RL) was originally designed solve Markov Decision Processes (MDPs) [Suttoon and Barto, 2018), but many real world tasks violate the Markov property, confining su- perhuman RL agents to simulators. When MDPs produce noisy or incomplete observations, they become Partially Observable MDPs (POMDPs). Using memory. we can summarize the trajectory into a Markov state estimate, extending convergence guarantees from traditional RL approaches to POMDPs [Kelbbing et al., 1998].
In model-free RL, there two main approaches to modeling memory: (1) RL-specific architec- tures that explicitly model a probabilistic belief over Markov states [Kaelbling et al., 1998] and (2) general-purpose models such RNNs or transformers that distill the trajectory into a fixed-size latent Markov state. Ni et al. [2022] and Yang and Nguyen [2021] reveal that with suitable hyperpa- rameters, general-purpose memory often outperforms
specialized belief-based memory.
Most applications of memory to RL tend to follow Hausknecht and Stone [2015], using RNNs like Long Short-Term Memory (LSTM) or the Gated Recurrent Unit (GRU) to summarize the trajectory, with other works evaluating transformers, and finding them tricky and data mata-hungry to train [Parisotto al., 2020]. Morad al. [2023] evaluate large collection of recent memory models across many
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

partially observable tasks, finding that the GRU outperformed all other models, including newer models like linear transformers. Interestingly, they show that there is little correlation between how well memory models perform in SL. and RI. Even recent, high-profile RI. work like Hafner et al. [2023], Kapturowski et al. [2023] use RNNs over more modern alternatives, raising the question question: why do older memory models continue to overshadow their contemporary counterparts in modern RL?
Training model-free RL policies today is sample inefficient, alchemical, and prone to collapse [Schulman al., 2015, Zhang et al., 2018], and reasoning over the entire trajectory magnifies these issues. SL can better utilize scale, compute, and dataset advancements than RL, largely remov- ing the need for strong inductive biases. For example, transformers execute pairwise comparisons across all inputs, in a somewhat "brute force" approach to sequence learning. Similarly, State Space Models [Gu et al., 2021] or the Legendre Memory Unit [Voelker et al., 2019] are designed to retain information over tens or hundreds of thousands of timesteps, growing the model search space with each additional observation.
Through strong inductive biases, older recurrent memory models provide sample efficiency and sta- bility in exchange for flexibility. For example, the GRU integrates inputs into the recurrent states in sequential order, has explicit forgetting mechanics, and keeps recurrent states bounded using satu- rating activation functions. These strong inductive biases curtail the model search space, improve training stability and sample efficiency, and provide more "user-friendly" training experience. If inductive biases are indeed responsible for improved performance, can we better leverage them to improve memory in RI.?
Contributions We introduce a memory model that summarizes a trajectory into latent Markov state for a downstream policy. To enhance training stability and efficiency, we employ strong in- ductive priors inspired by computational psychology. Our model can replace RNNs in recurrent RL algorithms with single line of code while training nearly two orders of magnitude faster than RNNs. Our experiments demonstrate that our model attains greater reward in both on-policy and off-policy settings and across various POMDP task suites, with similar sample efficiency to RNNs.
2 Related Work
RL and SL models differ in their computational requirements. While in SL. the model training duration is primarily influenced by the forward and backward passes, RI. produces training data through numerous inference workers interacting with the environment step by step. Consequently, it's imperative for RL memory to be both fast and efficient during both training and inference. However, for memory models, there is often an efficiency trade-off between between the training and inference stages, well as trade-off between time and space complexity.
Recurrent Models Recurrent models like LSTM [Hochreiter and Schmidhuber, 1997], GRUs [Chung et al., 2014], Legendre Memory Units (LMUs) [Voelker et al., 2019], and Independent RNNs [Li et al., 2018] are slow to train but fast at inference. Each recurrent state S over a sequence must be computed sequentially
y₁ S, = ST-1), jE [I.....n]
where F My are the inputs and outputs respectively at time j. and f updates the state S incremen- tally. The best-case computational complexity of recurrent models scales linearly with the length of the sequence and such models cannot be parallelized over the time dimension, making them prohibitively slow to train over long sequences. On the other hand, such models are quick at infer- ence, exhibiting constant- time complexity per inference timestep and constant memory usage over possibly infinite sequences.
Parallel Models Parallel or "batched" models like temporal CNNs [Bai et al., 2018) or transformers [Vaswani et al., 2017] do not rely on a recurrent state and can process entire sequence in parallel i.e.
y f [ 2,), E [1....,]]
Given the nature of GPUs, these models exhibit faster training than recurrent models. Unfortunately, such models require storing all or portion of the trajectory, preventing their use on long or infinite

---

sequences encountered during inference. Furthermore, certain parallel models like transformers require quadratic space and n comparisons each timestep during inference. This limits both the number of inference workers and their speed, resulting in inefficiencies, especially for on-policy algorithms.
Hybrid Models Linear transformers IKatharopoulos et al., 2020, Schlag et al., 2021, Su al., 2021] or state space models [Gu et al., 2021. 2020] provide the best of both worlds by providing equivalent recurrent and closed-form (parallel) formulations
S = f(z,, S5-1) = g g(.11...., zz), jE [E.....,,
(3)
Training employs the batched formula to leverage GPU paralletism, while inference exploits recur- rent models' low latency and small memory footprint. Thus, hybrid models well-suited for RL because they provide both fast training and fast inference, which is critical given the poor sample ef- ficiency of RL. However, recent findings show that common hybrid models typically underperform in comparison to RNNs on POMDPs [Morad et al., 2023].
3 Problem Statement
We given sequence of actions and observations (01, 0), (02, a1). (i))....(..., (on; an-1) up to time n. Let the trajectory X be some corresponding encoding æ = E(0, a) of each action-observation pair
X E I....., .n = [=(01. , an-1)
(4)
Our goal is to summarize X, into latent Markov state yn using some function f. For the sake of both space and time efficiency, we restrict our model to the space of hybrid memory models. Then, task is to find an f. such that
Ya S. s(X kin Sk-1),
where Xim shorthand for #t #k+,, Fy. Note that by setting k: = n, we achieve a one-step re- current formulation that one would see in an RNN, i.e., Yn, Sn f(X n.n Sn-1) = f( Sn
4 Background
In the search for inductive priors to constrain memory model search space, we turn the field of computational psychology. In computational psychology, we use the term frace to refer to the phys- ical embodiment of a memory that is, the discernable alteration in the brain's 's structure before and 's formation. Computational psychologists model long-term memory as a collection
after a memory's of traces, often in matrix form
S, , Fn.
(6)
where x, are individual traces represented as column vectors and S, is the memory at timestep n [Kahana, 2020].
Composite Memory Composite memory [Galton, 1883] approximates memory as lossy blend- ing of individual traces æ via summation, providing an explanation for how a lifetime of experiences can fit within fixed-volume brain. Murdock [1982] expresses this blending via the recurrent for- mula
S, = + B,
where 7 E (0, 1) is the forgetting parameter and B, is diagonal matrix sampled from Bernoulli distribution, determining a subset of x, to add to memory. In Murdock's ' model, a vector, not a matrix. We can expand "unroll" Murdock's recurrence relation, rewriting it in closed form
S = 7 1 +... B,. 0 0<7
making it clear that memory decays exponentially with time. Armed with equivalent recurrent and parallel formulations of composite memory, we can begin developing hybrid memory model.
3

---

It is vital to note that incoming traces capture raw sensory data and lack con- text. Context, be it spatial, situational, or temporal, differentiates between seemingly identical raw sensory inputs, and is critical to memory and decision making. The prevailing theory among com- putational psychologists is that contextual and sensory information are mixed via outer product [Howard and Kahana, 2002]
Contextual Drift
Zn = In n
(9)
Wn = pwin-1 + Tm: 0 0<p<l
(10)
where 77, is the contextual state at time n,, and P ensures smooth or gradual changes between texts. Although context spans many domains, we focus solely on temporal context.
5 Fast and Forgetful Memory
Fast and Forgetful Memory (FFM) is hybrid memory model based on theories of composite mem- ory and contextual drift. It is composed of two main components: cell and aggregator. The cell receives an input and recurrent state x, S and produces corresponding output and updated re- current state S S (Figure 1). The aggregator resides within the cell and is responsible for updating S (Figure 2). We provide the complete equations for both the aggregator and cell and end with the reasoning behind their design.
Aggregator The aggregator computes a summary S, of Xim. given recurrent state S and inputs Xkin
S. =gn-ht1 o Sk-1 +
Sn E Cmc
(11)
j=k
Tepp (-l(a) + iwi))
exp(-l(an + +iw)))
7¹ = (exp (-a) (-xxeepp exp (-iw) 61
E Cmxe
exp (-taat + twi)) exp (-t(am + iwe))
(12)
where is the Hadamard product (or power), is the trace size, is the context size, and R w E Ie are trainable parameters representing decay and context respectively. Multiplying column vector by __ "broadcasts" repeats the column vector e times.
In Appendix C. derive memory-efficient and numerically stable closed form solution to com- pute all states Skn in parallel over the time dimension. Since our model operates over relative time, map absolute time k,...,m n to relative time p O,...,t, where t k. The closed form to
compute a state Sktp given S | is then
Sktp pp @
+ +p
Y J (xx+) Skyp E CTxe.
(13)
1=0
We can rewrite the aggregator's closed form (Equation 13) in matrix notation to highlight its time- parallel nature, computing all states Skm
Sk, Sk+ Sk4n at once
Σ En_ x i (m)+)12))
x'
S
Sam
+
Skm E ((++1)xmxe
S
(Σj= x 10 (xx+)
(14)
The cumulative sum term in the rightmost matrix can distributed across processors, each requir- ing Olog t) time using a prefix sum or scan [Harris et al., 2007].
Cell The aggregator alone is insufficient to produce a Markov state y. S is complex-valued and contains large amount of information that need not be present in y at the current timestep, making

---

21
e ta
(e (T
R(11)
F1
R(-211)
3(x1)
n
3(-21)
n 1 n
Figure detailed example of aggregator dynamics for m 1, 2, showing how one- dimensional input Vj contributes to the recurrent state S over time t = [1. At the first timestep, Si is simply æ1. Over time, the contribution of æ S undergoes exponential decay via the (forgetting) and oscillations via the w term (temporal context). In this example, the contribution from 21 to S, approaches zero time n.
o(a(+)) 1 (1 (1- ((4-)( @ (s(-)
Z
(1))) O (f2())
- LN((s())
S e7 +
S
Figure 2: A visualization of an FFM cell running recurrently for a single timestep (i.e., inference mode). Inputs Fg. Fg.S S go through various linear layers (f), hadamard products (O), addition (+), normalization (LN),, and sigmoids (a) to produce outputs yk, Sy. The boxed region denotes the aggregator, which decays and shifts Sk-1 1 in time via 7. During training, the aggregator computes Se,..., S, in parallel so that we can compute all y ........ Yn in a single forward pass.
cumbersome interpret for a downstream policy. Furthermore, X could benefit from additional preprocessing. The cell applies input gating to X and extracts a real-valued Markov state from S.
Tk.n & (z @ (E2((()))))
Tken R (++1)xm
(15)
S Skm Agg(E Sk-1),
km C(t+1)xmxe
(16)
Zke Ey (Flatten R Sk:n || ((Skm)))),
Zken R (t+1))x
(17)
yken LN(zkmm) @ (fa(we))) ls (a @ (1 (fs
yken R (t+1)xd
(18)
Agg represents the aggregator (Equation 14) and l represents linear layers with mappings f Rd - m lg Rmm22 -> Rdd l, : R -> Rd. R, S extract the real and imaginary components of complex number as reals, Flatten reshapes matrix (m X c c-> me) and || is the concatenation operator. LN is nonparametric layer norm, and a is sigmoid activation. Equation 15 applies input gating, Equation 16 computes the recurrent states, Equation 17 projects the state into the real domain, and Equation 18 applies output gating.
Modeling Inductive Biases To justify FFM's architecture, detail its connections to computa- tional psychology. Drawing from the theory of composite memory, we integrate sigmoidal gating to in Equation 15, approximating the sparsity-driven Bernoulli term, B, presented in (Equation 7). This inductive bias represents that only a small subset of each sensory experience is stored in mem- ory.
Following the blending constraint imposed by composite memory (Equation 8), we sum the traces together (Equation 11). We suspect this blending prior yields model that is more robust to the bad policy updates that plague RI,, compared to RNNs that apply nonlinear and destructive transforms (e.g. deletion) to the recurrent state.
5

---

In Equation 12, the exponential component exp (-ta) ensures that traces are forgotten according to composite memory decay from Equation 8. We choose an exponential decay e e- approximating Murre and Dros [2015]. This inductive bias asymptotically decays traces, reducing the size of the model optimization space. This is in direct contrast to recent SL. memory models which aim to store as much information as possible for as long as possible [Gu et al., 2021, Voelker et al., 2019, Schlag et al., 2021]. In our ablations, we show that forgetting is critical in RL.
Traditionally, contextual coding and composite memory are separate mechanisms, but we incorpo- rate both through 7 The exp (-iwl) component of in Equation 12 applies temporal context, enabling relative-time reasoning. Following contextual coding theory, we take the outer product of a trace vector, in our case decayed by a. with a context vector produced by w. We can rewrite y = -/= C y, mirroring the gradual recurrent context changes from Equation 10.
The recurrent state S is a complex- valued matrix, so we project it back to the real domain as z (Equation 17). can be large for dimensions where is near zero (minimal decay), we find it crucial to apply layer normalization in Equation 18. Finally, we apply a gated residual connection to improve convergence speed, letting gradient descent find the optimal mixture of input memory z for the Markov state y-
The Mathematics Underpinning FFM The goal of FFM is produce a memory mechanism which can be parallelized across the time dimension. As discussed in section 4, we can achieve this through summation of sensory information mixed by some arbitrary function & with temporal context Wy vi),, generating an aggregator of the form
S Σ (b()),
(19)
Any implementation under this formulation would provide memory with temporal context for abso- lute time. However, in order to maximize generalization, we wish to operate over relative time. This introduces a problem, as the relative temporal context ('n- (in--] j) associated with each æj must be up- dated at each timestep. Therefore, we must find function fa that updates each temporal context with some offset: (()(i), (*+)) vi + k). Solving for h, we get: h(j, k)
(z) + s (k)). If we also set ø = h, then applying temporal shift h(,, k) to each term in the sum yields the same result as initially inserting each x, with vi + k):
h()(#()) z,), k) o(b + k),
(20)
Consequently, it is possible update the terms each timestep reflect the relative time n - j. However, it is intractable to recompute relative temporal contexts at each timestep, as that would result in n² time and space complexity for sequence of length n. To apply temporal context in batched manner in linear space, we can leverage the distributive property. If we select ((a. b) = a a- then the update distributes over the sum, updating all of the terms simultaneously. In other words, we can update the context of n terms in constant time and space (see Appendix C for the full derivation). As sb defined as a function of ø, we can solve to find (1) e (although any exponential constitutes solution, we select the natural base). This yields an aggregator of the form
S. Σ E(nn 1))
(21)
1=0
g E R,, + 0, then the recurrent state will explode over long sequences: 1im ((nn))xx) x, = Do, so the real component of & should be negative. Using different decay rates Ej. & E R_. we can deduce the temporal ordering between terms æy and Ir- Unfortunately, this requires both 2j, JA eventually decay to zero what if there are important memories we do not want to forget, while simultaneously retaining their ordering? If is imaginary, then we can determine the relative time between æj By as the terms oscillate indefinitely without decaying. Thus, we use a complex & with negative real component, combining both forgetting and long-term temporal context. In the fol- lowing paragraphs, we show that with a complex ç, the FFM cell becomes a universal approximator of convolution.
Universal Approximation of Convolution Here, we show that FFM can approximate any tempo- ral convolutional. Let us look at the term from Equation 17, given the input ă with a slight change

---

Model
Training
Inference
Parallel Time
Space
Time
Space
RNN
O(n)
(n)
O(1)
O(1)
Transformer
(log n)
Onn")
O(n)
(
FFM (ours)
(log n)
O(n)
0(1)
0(1)
Table 1: The time and space complexity of memory models for a sequence of length n (training), or
computing a single output recurrently during a rollout (inference).
notation to avoid the overloaded subseript notation: z(r) = -r. z(n) = Zn. For a sequence 1 to n, =nn), the precursor to the Markov state Yn. can be written
(n) b Σ a(T) e exp (-ra) exp --TiwT)
(22)
where A, is the weight and bias from linear layer lz. indexed by subscript. Looking at a single input dimension & of E; we have
(k+1)em
72
(n2) b +
Σ
A
(x(r) (-oxpp(-(((e))) (-T(iw ax))
(23)
keem
(1++1)om
(T)
Σ
A exp (-T(iw + +x))
(24)
j=kem
Equation 24 is a temporal convolution of it (t) using a Fourier Series filter with e terms (w E C?), with an additional learnable "filter extent" term The Fourier Series is universal function ap- proximator, so FFM can approximate any convolutional filter over the signal #(7). Appendix A further shows how this is related Laplace transform. Unlike discrete convolution, we do not need to explicitly store prior inputs or engage in zero padding, resulting better space efficiency. Fur- thermore, the filter extent ag is learned and dynamic it can expand for sequences with long-term temporal dependencies and shrink for tasks with short-term dependencies. Discrete temporal con- volution from methods like Bai et al. [2018] use fixed-size user-defined filter extent.
Interpretability Unlike other memory models, FFM is interpretable. Fach dimension in S has known decay rate and contextual period. We can determine trace durability (how long a trace lasts) ta and the maximum contextual period tu of each dimension in via
log(3) a1
2x
ta
tu
(25)
1,
where B determines the strength at which trace is considered forgotten. For example, , = 0.01 would correspond to the time when the trace Fy contributes 1 % of its original value Sn,- FFM can measure the relative time between inputs up modulo of tw..
6 Experiments and Discussion
We evaluate FFM on the two largest POMDP benchmarks currently available: POPGym [Morad et al., 2023] and the POMDP tasks from [Ni et al., 2022], which we henceforth refer to as POMDP- Baselines. For POPGym, we train a shared-memorg actor-critic model using recurrent using Proxi- mal Policy Optimization [Schulman et al., 2017]. For POMDP-Baselines models for the actor and critic using recurrent Soft Actor Critic (SAC) [Haarnoja et al., 2018] and recurrent Twin Delayed DDPG (TD3) [Fujimoto et al., 2018]. We replicate the experiments from Baselines papers as-is without changing any hyperparameters. We use a
train separate memory
the POPGym and single FFM configuration across all experiments, except for varying the hidden and recurrent sizes to match the RNNs, and initialize a; w based on the max sequence length. See Appendix for further details.
Table lists memory baselines, Table 3 ablates each component of FFM, Figure 3 contains sum- mary of training statistics for all models (wall-clock train time, mean reward, etc), and Figure
7

---

Name
Type
Paper/Description
Name
Type
Paper/Descriptton
MLP
Basic
2 layer perceptron
LMU
RNN
Voelker et al. [2019]
PosMLP
Basic
MLP w/ positional encoding
IndRNN
RNN
Li al. [2018]
Stack
Basic
Mnih et al. [2015]
Elman
RNN
Elman [1990]
TCN
Conv
Bai al. [2018]
GRU
RNN
Chung al. [2014]
FWP
XFormer
Schlag et al. [2021]
LSTM
RNN
Hochreiter and Schmid-
FART
XFormer Katharopoulos et [2020]
huber [1997]
S4D
Hybrid
Gu et al. [2021]
DNC
MANN
Graves et al. [2016]
Table 2: The memory models used in the POPGym experiments, Morad et al. [2023] for further
details. Conv stands for convolution, XFormer stands for transformer, and MANN stands for mem- ory augmented neural network.
543
198
DM2
Rvdd
71M
(87
FART
516
0R%.
77M
Dron
7C4
PART
IYRRN
15
1.5
100
TNN (79)
1INERTER Latenny (19))
Paxt Memory (010
Figure 3: Training statistics computed over ten trials and one epoch of PPO, with FFM in orange. FFM trains nearly two orders of magnitude faster on the GPU than the fastest RNNs (note the log- scale x axis). Even on the CPU, FFM trains roughly one order of magnitude faster than RNNs on the GPU. FFM has sub-millisecond CPU inference latency, making it useful for on on-policy algorithms. FFM memory usage is on-par with RNNs, providing scalability to long sequences. Despite effi- ciency improvements, FFM still attains greater episodic reward than other models on POPGym.
contains detailed comparison of FFM against the best performing RNN and transformer POP- Gym. Figure 5 evaluates FFM on the POMDP-BB BaiellB-s tasks and Figure 6 investigates FFM inter- pretability. See Appendix E and Appendix F for more granular plots. Error bars in all cases denote the 95% bootstrapped confidence interval.
Practicality and Robustness: FFM demonstrates exceptional performance on the POPGym and benchmarks, achieving the highest mean reward as depicted in Figure 3 and Fig- ure 5 without any changes to the default hyperparameters. FFM also executes a PPO training epoch nearly two orders of magnitude faster than GRU (Figure 3, Table 1). Without forgetting, FFM underperforms the GRU, showing that forgetting is the most important inductive prior (Table 3). We use single FFM configuration for each benchmark, demonstrating it is reasonably insensitive to hyperparameters. All in all, FFM outperforms all other models on average, across three RL algo- rithms and 52 tasks. Surprisingly, there are few occurrences where FFM is noticeably worse than others, suggesting it is robust and a good general-purpose model.
Explainability and Prior Knowledge: Figure 6 demonstrates that FFM learns suitable and in- terpretable decay rates and context lengths. We observe separate memory modes for the actor and
FFM
FFM-NI
FFM-NO
FFM-NC
FFM-FC
FFM-ND
FFM-FD
FFM-HP
0.399
0.392
0.395
0.382
0.383
0.333
0.390
0.395
a 0.005
0.001
0.006
0.005
0.005
0.005
0.006
0.010
Table 3: Ablating FFM components over all POPGym tasks. FFM-NI removes input gating, FFM- NO removes the output gating, FFM-NC does not use temporal context, FFM-FC uses fixed (non- learned) context, FFM-ND does not use decay, and FFM-FD uses fixed (non-learned) decay. FFM- HP uses the Hadamard product instead of the outer product to compute 7 in Equation 12. FFM-HP requires additional parameters but ensures independence between rows and columns of the state. FFM-ND (no decay) validates our assumption that forgetting is an important inductive bias.
8

---

Reward Reletvve to GRU
10%
1%
0%%
-1%
100%
Env
Figure 4: FFM compared to the GRU, the best performing model on POPGym. See Appendix F for comparisons across all other POPGym models, such linear transformers. Error bars denote the bootstrapped 95% confidence interval over five trials. FFM noticeably outperforms state of the art seventeen tasks, while only doing noticeably worse four tasks.
5000
RNN-BAC
Rowere
2000
FRM+SAC
RNN-TD3
1000
Waleer
Ant-P
HaltCheetan-/
Walken
Hopgen-V
Figure Comparing FFM with tuned RNNs on the continuous control tasks from POMDP- Baselines. Ni et al. [2022] selects the best performing RNN (either GRU or LSTM) and tunes hyperparameters on per-task basis. Error bars denote the bootstrapped 95% confidence inter- val over five random seeds. The task suffix V or P denotes whether the observation is velocity-only (masked position) or position-on (masked velocity) respectively. Untuned FFM outperforms tuned RNNs on all but one task for SAC. For TD3, FFM meets exceeds RNN performance all but one task.
Model 100)
Modelf
Rowerd
200
FTM-P
FFM-11104
30
1101
155
2000
501
200
150
20C
200
150
200
Fpozh
Fpom
Fpors
Figure 6: case study on FFM explainability. We investigate the task RepcauPreviousMedduum where the episode length is 104 timesteps and the agent must output the observation from 32 timesteps ago. (Left two) We visualize 7 using trace durability ta, = 0.1 and context period tw. Both t.. tw demonstrate modes for the actor and critic, which we highlight in the figures. The critic requires the episode's 's current timestep to accurately predict the discounted return, while the actor just requires 32 timesteps of memory. (Left) The ty critic mode converges to the maximum episode length of 104. The actor mode converges to the necessary period (perhaps because cos( 22 the real component of exp -xp, is monotonic on the half-period 0.5t...). (Middle) The ta terms con- verge to 32 for the actor and large value for the critic. (Right) We initialize ax to encapsulate ideal actor and critic modes, such that t, E [32, 104), tw E [32, 104). We find FFM with informed initialization outperforms FFM with the original initialization of t, E [1, 1024], tw E [1, 1024].

---

critic, and we can initialize a, w using prior knowledge to improve returns (Figure 6). This sort of prior knowledge injection not possible in RNNs, and could be useful for determining the number of "burn-in" steps for methods that break episodes into fixed-length sequences, like R2D2 [Kaptur- owski et al., 2019] or MEME [Kapturowski al., 2023].
Limitations and Future Work
We found that large context sizes e resulted in decreased performance. We hypoth- esize that a large number of periodic functions results in a difficult to optimize loss landscape with many local extrema. Interestingly, transformers also employ sinusoidal encodings, which might ex- plain the difficulty of training them in model-frre RL. We tried multiple configurations of FFM cells in series, which helped in some tasks but often learned much more slowly. In theory, serial FFM cells could learn a temporal feature hierarchy similar the spatial hierarchy in image CNNs. In many cases, serial FFM models did not appear fully converged, so it is possible training for longer could solve this issue.
Scaling Up:
Hyperparameter tuning would almost guarantee better performance, but would also result in biased comparison to other models. We evaluated FFM with on-policy and off-policy algorithms but did not experiment with offline or model-based RI algorithms. In theory, FFM can run continuous time or irceuularly-ppaced intervals simply by letting be continuously or irregular, but in RL we often work with discrete timesteps at regular intervals so we due not pursue this further.
Additional Experiments:
Numerical Precision: FFM experiences loss of numerical precision caused by the repeated mul- tiplication of exponentials, resulting in very large or small numbers. Care must be taken to prevent overflow, such as upper-bounding a. Breaking sequence into multiple forward passes while prop- agating recurrent state (setting n - less than the sequence length) fixes this issue, but also reduces the training time efficiency benefits of FFM. We found FFM performed poorly using single preci- sion floats, and recommend using double precision floats during multiplication by 7. We tested a maximum sequence length of 1024 per forward pass, although could go higher by decreasing the maximum decay rate. With quadruple precision floats, we could process sequences of roughly 350,000 timesteps in single forward pass with the decay rate used. Unfortunately, Pytorch does not currently support quads.
7 Conclusion
The inductive priors underpinning FFM are key its success, constraining the optimization space and providing parallelism-enablig structure. Unlike many memory models, FFM is interpretable and even provides tuning opportunities based on prior knowledge about individual tasks. FFM pro- vides low-engineering cost "plug and play" upgrade to existing recurrent RI algorithms, improving model efficiency and reward in partially observable model-free RL with a one-line code change.
Acknowledgements
Steven Morad and Stephan Liwicki gratefully acknowledge the support of Toshiba Europe Ltd. Ryan Kortvelesy and Amanda Prorok were supported in part by ARL DCIST CRA W911NF-17-2-0181
10

---

References
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Leslie Pack Kaelbling. Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2)99-134, 1998. Publisher: Elsevier.
Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent Model-Free RL Can Be Strong Baseline for Many POMDPs. In Proceedings of the 39th International Conference on Machine Learning, pages 16691 16691-16723, PMLR, June 2022. URL http:: mml.
press/v162/m122a.. html.
ISSN: 2640-3498.
Zhihan Yang and Hai Nguyen. Recurrent Off- Off-policy Baselines for Memory-baser Continuous Con- trol. Deep RL Workshop, NeurlPS 2021, 2021.
Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable mdps. In 2015 aaai fall symposium series, 2015.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphaël Lopez Kaufman, Aidan Clark, Seb Noury, Matthew Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing Transformers for Reinforcement Learning, In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 7487-7498. PMLR, July 2020. URL https: http///prrceeenneeeeeeeeeeeesset mlr. press/ wttppp//l1l/eariiiveeeeoeooeeeee ww1p/p/raittoo2aaa...m htmm..
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym: Benchmarking Partially Observable Reinforcement Learning. In The Eleventh International Conference on Learning Representations, 2023. URL
chDrutUTsOK.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering Diverse Domains through World Models. arXiv preprint arXiv:2301.04104 2023.
Steven Kapturowski, Víctor Campos, Ray Jiang, Nemanja Rakicevic, Hado van Hasselt, Charles Blundell, and Adria Puigdomenech Badia. Human-level atari 200x faster. In The Eleventh In- lernational Conference on Learning Representations, 2023. URL https://openreviee .eteeneeeeeeeee/eeeeeee.e. forum?id=Jt Korumnid=-JCGGOHHRJJ.
John Schulman, Sergey Levine, Picter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference on machine learning, pages 1889-1897. PMLR, 2015.
Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A Study Overfitting in Deep Reinforcement Learning. arXiv, 2018. URL https: http://twwwooopooo//oo//////
org/ /////dff11004 0w8.06888.... pdf.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining Recurrent, Convolutional, and Continuous-time Models with Linear State Space Lay- ers. Advances in Neural Information Processing Systems, volume 34, pages 572-585, Cur- ran Associates, Inc., 2021. URL https://proceedi sagrocceeinggs.nuuii.s.... www.nnww......ppppeeeeeee (ineer/202//hhhh 1al46fo0e8aa917fccdd0eeeeeeeeee Atbaaececbbbbbb--bbbbttt....... html.
Aaron Voelker, Ivana Kajić, and Chris Eliasmith. Legendre Memory Units: Continuous- Time Representation in Recurrent Neural Networks. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: https://pappee: /tt///aaprrr.r/// www..nip....... cc/paper/////cpccc
Sat.://A.........oo--aAattttttt html.
Sepp Hochreiter and Jürgen Schmidhuber. Long Short-Term Memory. Neural Computation, (8):1735-1780, November 1997. ISSN 0899-76677 doi: 211162/neco..997..938 http::::/doo...r//// wrg//o..16.//ecc.. 1997 ...9.998.. 8. 1735.
URL
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical Evaluation of Gated Recurrent Neural Networks on Sequence Modeling. Technical Report arXiv: 14T.:1111.3555 arxiv, December 2014. URL UtL.ttp: http::t/t//xii....//// wrw.:/s//////aaeeae..e..ossssses /412e 3555. arXiv:1412.355. [es] type:
article.

---

Shuai Li, Wanqing Li. Chris Cook, Ce Zhu, and Yanbo Gao. Independently Recurrent Neu- ral Network (IndRNN): Building Longer and Deeper RNN. 2018 IEEE/CVF Confer- ence on Computer Vision and Pattern Recognition, pages 5457-5466, Salt Lake City, UT, June 2018. IEEE. ISBN 978-1-5386-6420-9, doi: 20.1009/CVPR.2018.0000 //iseexplore ieeeeeerreeeeeeeee.eee.oeeeeeeee org/ /ttp///cmmme///7//7////////
URL https:
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv: /80.00001,.1.12000
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention All you Need. In Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL https: https://proceeddin htwp:/diwi............eepeeeeeee cc/paper/221//hash -baaaasaaaaaaaaaattttt.. .tmaaa.ttt................m.....
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: fast autoregressive transformers with linear attention. In Proceedings of the 37th Interna- fional Conference on Machine Learning, (CML,'20, pages 5156-5165. JoL.Rorg. July 2020.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers Are Secretly Fast Weight Programmers. In Proceedings of the 38th International Conference on Machine Learning, pages PMLR, July 2021. URL httph:ttt/// hatttss:/proceeeiingg html. ISSN: 2640-3498.
wwwpppreeiiiaa schlag21a.
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09980 2021.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. НiРPО: Recurrent Memory with Optimal Polynomial Projections. Advances Neural Information Processing Systems, volume 33, pages 1474-1487. Curran Associates, Inc., 2020. URL https: httpp::/ppooeeeeeeeeeeeeee http///errceeee..../o neurips. cc/ppp/e//////0/0//// hwwhh/aaaa/aaa//aa http:////0000aaaae//000bebeeeee 12220006eeeea11222777000aaaa0002
Abstract html.
Michael J Kahana. Computational models of memory search. Annual Review of Psychology, 71:
107-138, 2020. Publisher: Annual Reviews.
Francis Galton. Inquiries into human faculty and its development. Maemillan, 1883.
Bennet B Murdock. A Theory for the Storage and Retrieval of Item and Associative Information. Psychological Review, 89(6):609, 1982. Publisher: American Psychological Association.
Mare W. Howard and Michael Kahana. A Distributed Representation of Temporal Con- text. Journal of Mathematical Psychology, 46(3):269-20) 299, June 2002. ISSN 0022-2496. doi: 10.10066/mmps220011333. URL https://www, sci wwweeenneeeeettccooo wwt.eacccteet.eccceeeeeeeeeeeeee c pi1/$8002249661913888.
Mark Harris, Shubhabrata Sengupta, and John Owens. Parallel prefix sum (scan) with CUDA. GPU gems, 3(39):851-876, 2007.
Jaap M. J. Murre and Joeri Dros. Replication and Analysis of Ebbinghaus' Forgetting Curve. PLOS ONE, (007)xe0120644 July 2015. ISSN 1932-6203, doi: 10.137l/jourmll .onne.012064.. URL
https://dx plos. org/10. 1011000. 1371 harr//ooraalll pone. 0120644.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 1707:07777 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pages 1861-1870. PMLR, 2018
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-
critic methods. In International conference on machine learning, pages 1587-1596. PMLR, 2018.
12

---

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, and others. Human-level control through deep reinforcement learning. nature, 518(7540):52-533, 2015. Publisher: Nature Publishing Group.
Jeffrey L Elman. Finding structure in time. Cognitive science, 14(2)119-211, 1990. Publisher:
Wiley Online Library.
Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwińska, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, and others. Hybrid computing using neural network with dynamic external memory. Nature, (58(7626)447147760 2016. Publisher: Nature Publishing Group.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney, Recurrent ex- perience replay in distributed reinforcement learning. In International conference on learning representations, 2019.
13

---

A Relationship to the Laplace Transform
We often think of RL terms of discrete timesteps, but the sum Equation 11 could be replaced with integral for continuous-tin time domains, where ø is function of time:
S - L a(T) e exp (-ra) exp (-Tiw) T dr.
(26)
If we look just single dimension ã, we have
S(n) (n)=J/ / ((r))x iw)) dr,
(27)
which is equivalent to the Laplace Transform of a function ((T). One way to interpret the connec- tion between our aggregator and the Laplace transform is that the aggregator transforms the memory task into a pole-finding task in the S-plane, a well-known problem in control theory. When we train FFM for RL task, are attempting find the poles and zeros that minimize the actor/critic loss for a batch of sequences.
B Weight Initialization
We find that our model is sensitive to initialization of the or, w values. We compute the upper limit for a such that memory will retain = 0.01 = 1% of its value after some elapsed number of timesteps call tg. This can be computed mathematically via
log 3 t
(28)
The lower limit is set based on the maximum value double precision float can represent, minus small
log 1.79 x 10308
(29)
te
Note that if needed, we can choose a smaller lower bound, but the input must be chunked into sequences of length le and run minibatches. Since we can compute minibatches recurrently, the gradient spans across all minibatches rather than being truncated like a quadratic transformer. Nonetheless, we did not need to do this in any of our experiments. We set a to be linearly spaced between the computed limits
a = [an.... ......j..... anl
(30)
j log 3
3 m
(log 1.79 x 10308
= m te
+11
(31)
te
We initialize the w terms using a similar approach, with linearly spaced denominators between one and t
2T/ [Wi, W2, .,,
(32)
wj j + (1 1 2)te
(33)
All other parameters are initialized using the default Pytorch initialization.
C Computing States in Parallel with Linear Space Complexity
Here, we show how the recurrent formulation can be computed in parallel. naive implementation would (n)2) space, but using trick, accomplish this in O(n) space.
Assume we have already computed hidden state S 1 some sequence. Our job now to compute the next n - recurrent states Sk, Skt1..... S, in parallel. We can rewrite Equation 11 in matrix
14

---

form:
n0 (x tgl @ Sk-1
S
70 (#x+1 tg о (x + + o R-
Sk+1
Sk+2
7
(#k+2le + o (#s+ile) t o (ax1e + r Sk-1
Skm
(34)
Sn
>" @ (#*+)1.)) + gpt ① Sk-1
We can write the closed form for the pth row of the matrix, where 2 p 2
Sktp = Σ Enn @ (Bn+s17) )
+ g o Sk-1
(35)
(n- k+)
Unfortunately, it appears we will need to materialize
terms:
2
y o (xx1.) y o (xr1)),7 @ (xx+1l))
(36)
(37)
(38)
y 3 (xx1.) pm ....,, (x,1.) )
(39)
However, we can factor out y P from Equation 35 via combination of the distributive property and the product of exponentials
Skop - 7PO Σr (r'o(w))) + o Sk-1
(40)
Now, each 7 is associated with a single æ, requiring just n - k + terms:
1 ( @ (xx+11.)
(41)
(42)
(43)
(m.1 r' o 1......)...... 7" (x,11)
(44)
We can represent each of these rows as a slice of single n + length tensor, for a space complexity of (n k) or for 1, O(n).
Finally, we want swap the exponent signs because provides better precision when working with floating point numbers. Computing small values, then big values is more numerically stable than the other way around. We want to compute the inner sum using small numbers (y* results in small numbers while y produces big numbers; n k is positive and k- is negative). We can factor g out of the sum and rewrite Equation 40 as
S Skip = A +
+ +xptt o Sx-1
(45)
If we let k, this is equivalent to Equation 13:
Sktp p 2 Sk-1 + - Σ (er+)ll))
Skip E Cmxxe
(46)
We also need to compute 2(n 1: + 1) gamma terms: g =n ,.. g..... resulting in linear space complexity O(n) for sequence of length
15

---

D Benchmark Details
We utilize the episodic reward and normalized episodic reward metrics. We record episodic rewards as single mean for each epoch, and report the maximum reward over trial.
We do modify the hyperparameters for PPO, SAC, TD3 from the original benchmark papers. We direct the reader to [Morad et al., 2023] for PPO hyperparameters and [Ni et al., 2022] for SAC and TD3 hyperparameters.
D.1 Hardware and Efficiency Information
We trained on a server with Xeon CPU running Torch version 1.13 with CUDA version 11.7, with consistent access to two 2080Ti GPUs. Wandb reports that we used roughly 161 GPU days of compute to produce the Basell-aselines results, 10 GPU days for the POPGym results, and 45 GPU days for the FFM ablation using POPGym. This results in total of 216 GPU days. Since we ran four jobs per GPU, this corresponds to 54 days of wall-clock time for all experiments.
D.2 PPO and POPGym Baselines
Morad et al. [2023] compares models along the recurrent state size, with recurrent state size of 256. For fairness, We let can be represented as a 256 dimensional real vector. We initialize a, w following Appendix for te 1024, 8 0.01. We run version 0.0.2 of POPGym, and compare FFM numerically with the MMER score from the paper Table 4.
32 and 4, which results in a complex recurrent state of 128, which
D.3 SAC, TD3, and POMDP-Baseline
[Ni et al., 2022] does not compare along the recurrent state size, but rather the hidden size, while utilizing various hidden sizes h. In other words, the LSTM recurrent size is twice that of the GRU. We let c = h/32 and m h/c. so that mc = h. This produces an equivalent FFM configuration to the POPGym baseline when h 128, with equivalent recurrent size in bytes the LSTM (or equivalent dimensions to the GRU). The paper truncates episodes into segments of length 32 or 64 depending the task, we let fc 128 ensure that information can persist between segments. Thanks to the determinism provided in the paper. readers should be able to reproduce our exact results using the random seeds 1, 2, 4. We utilize separate memory modules for the actor and critic, as done in the paper. We base our experiments off of the most recent commit at the time of writing, 4d9cbf1.
Model
MMER
MLP
0.067
PosMLP
0.064
FWP
0.112
FART
0.138
S4D
-0.180
TCN
0.233
Fr.Stack
0.190
LMU
0.229
IndRNN
0.259
Elman
0.249
GRU
0.349
LSTM
0.255
DNC
0.065
FFM
0.400
Table 4: MMER score comparison from the POPGym paper
16

---

E Benchmark Lineplots
We provide cumulative max reward lineplots in Figure 9, Figure 7, and Figure 8 for all the experi- ments we ran.
Exriogel
Figure 7: SAC lineplots for POMDP-Baselines, where each trial is plotted separately.
15
30
Figure 8: TD3 lineplots for MaaDl-Basellnes, where each trial is plotted separately.
17

---

Figure 9: Lineplots for POPGym, where the shaded region represents the 95% bootstrapped confi- dence interval.
18

---

POPGym Comparisons by Model
We provide Figure 10, Figure 11, Figure 12 showing the relative FFM return compared to the other
12 POPGym models, including temporal convolution, linear transformers, and more.
Reward Rolative DNC
1000%
100%
10%
0%
-1%
-100%
Roward Relatve Elman
100%
Roward Retative to FART
100%
10%
0%
-1%
Converentcn caanenn
Rewerd Relative FVP
100%
10%
0%
-1%
Env
Figure 10: Relative POPGym returns compared to FFM.
19

---

Rewerd Relative to IndRNN
100%
10%
1%
0%
-1%
-10%
Env
Reward Relaivve ta LMU
100%
10%
1%
0%
-1%
-10%
Eny
Reward Reletivs to LSTM
100%
10%
1%
0%
-1%
-10%
Env
Reward Roletiie MUP
Env
Figure 11: Relative POPGym returns compared to FFM.
20

---

Roward Relative PoaMLLF
100%
11%
0%
-10%
Reward Relalivve to S4D
1000%
100%
0%
-100%
Reward Relative Slack
10%
1%
-10%
Doge
Reward Relative to TON
10%
1%
0%
-1%
-10%
Env
Figure 12: Relative POPGym returns compared to FFM.
21

---

G Efficiency Statistic Details
This section explains how we computed the bar plots showing efficiency statistics. We construct the POPGym models and perform one PPO training epoch. The train time metric corresponds to the time spent doing forward and backward passes over the data. One epoch corresponds the epoch defined in POPGym: 30 minibatches of 65,336 transitions each, with an episode length of 1024, hidden dimension of 128, and a recurrent dimension of 256. We do this times and compute the mean and confidence interval. Torch GRU, LSTM, and Elman networks have specialized CUDA kernels, making them artificially faster than LMU and IndRNN which are written Torch and require the use of for loops. We utilize the pure python implementation of these models, wrapping them in a for loop instead of utilizing their hand-designed CUDA kernels. We consider this a fair comparison since this makes the GRU, Elman, and LSTM networks still run slightly faster than IndRNN and LMU (Torch-native RNNs). FFM is also written Torch and does not have access to specialized CUDA kernels. CUDA programmers more skilled than us could implement custom FFM kernel that would see speedup similar to the
kernels.
To compute inference latency, we turn off Autograd and run inference for 1024 timesteps, computing the time for each forward pass. do this 10 times and compute the mean and 95% confidence interval.
To compute memory usage, we utilize Torch's GPU memory tools and record the maximum memory usage at any point during the training statistic. Memory usage is constant between trials.
For reward, we take the mean reward over all environments, split by model and trial. We then report the mean and 95% confidence interval over trials and models.
22


Reviews:
[{"review_id": "9bVLStmEp1", "reviewer": "Reviewer_ARGZ", "summary": "The paper aims to propose an algorithm-agnostic memory model specifically designed for reinforcement learning (RL). The proposed neural network (NN) model is claimed to offer advantages over existing sequence-based memory models like Recurrent Neural Networks (RNNs) and can be seamlessly substituted for RNNs. The key advantage over RNNs is that the model can process a sequence in parallel, making it more efficient compared to traditional RNNs. The main design components of the model include:\nA context drift formula that computes the embedding of sensory information with context.\nA decaying mechanism incorporated into the recurrent formula to compute a summary of the embedding.\nGating mechanisms used to extract the final output from the learned summary of the embedding.", "questions": "see weaknesses.", "limitations": "na", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Appropriate architecture may be suitable for a particular type of tasks. I agree that there might be some interesting special NN architecture that is particularly beneficial for RL tasks. Studying RL-specific NN structure is an important and practically useful research direction. \n\n2. The proposed NN architecture seems to be new. ", "weaknesses": "The main weaknesses are: clarity of presentation, lack of justifications of some design choices, experiments are not persuasive enough. \n\nClarification questions: \n\n1. What is the difference between subscript t and n? I found it to be confusing in eq(9). What should be the relation between n and t there? \n\n2. eq(4) already defined x_n as the embedding of observation-action pair, why eq(9) redefined this notation? \n\nAnd what is hat{x}? It is said that it is the raw sensory information, isn’t it part of observation variables? \n\nWhat is the context variable in RL? Do you simply define omega as a constant one vector? \n\nI strongly suggest the authors to use a particular RL environment to explain all these notations. I found it to be confusing when reading notations in the paper. \n\nJustify some designs: \n\n1. As a new NN architecture,  it is critical to study which component of the proposed methods matters. \nAt least, the paper should study the necessity of exponential decay mechanism, outer product design (eq 9), etc. One possible baseline in the experiments could be simply adding a decaying mechanism in the hidden state of a regular RNN so that it can somehow show the importance of the decay mechanism.  \n\nI have doubts on why decaying is good. Why it makes sense to set up an expiration for each feature unit? In existing literature, it is critical to always maintain a so-called coverage in the state space to learn an optimal policy (please refer to references [1][2][3]). If setting expiration for some feature units, doesn’t it mean manually discard some of the states? \n\n2. The paper should also include discussions with successor features, as the method of computing the hidden state (i.e. eq(7)) is highly similar to how a successor representation is learned. An in-depth study of successor representation may result in better understanding of the proposed architecture. \n\nExperiments. \n\nI found the experiments a bit disappointing. As a general NN architecture, why not also test it on those commonly used RL domains like mujoco, Atari, etc. as mentioned in the abstract, “it is a drop-in replacement for RNN.” There are already existing work using RNN in a variety of RL benchmark tasks. It is better to test the proposed architecture on those well-tested RL domains.  \n\nI would like to emphasize that i don’t think it is necessary to restrict the test on POMDP. In function approximation setting, a MDP can be always viewed as a POMDP. \n\n[1] An equivalence between loss functions and non-uniform sampling in experience replay by Scott Fujimoto et al. \n[2] Understanding and mitigating the limitations of prioritized experience replay by Pan et al. \n[3] regret minimization ER in off-policy RL et al. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. We believe a reject rating for the issues you raised is unduly harsh, and we believe there might also be some misunderstanding regarding the scope of our paper. The goal of this paper is to design an efficient replacement for RNNs in model-free RL, which summarizes observations from a POMDP into a latent Markov state.\n\n### Possible Misconceptions\n- There might be some confusion about the word \"memory\". Your citations [1,2,3] discuss replay buffers for MDPs and do not mention POMDPs or RNNs. More than half of our experiments (POPGym) are on-policy and do not use a replay buffer. We are concerned with \"memory\" as it applies to an online summary of observations, rather than \"memory\" in the sense of storing and replaying state-action-reward tuples. We will add a sentence to the introduction to better explain our definition of \"memory\".\n- Solving MDPs provides no signal on the capacity of memory models to solve POMDPs. This is because MDPs can be solved without memory [A], while POMDPs cannot [B]. The goal of this paper is to develop a memory model to solve POMDPs.\n- We only use pre-existing benchmarks in our experiments. Our experiments from the POMDP study by Ni et al. (Fig. 5) are the MuJoCo tasks you are thinking of (Walker, Ant, HalfCheetah, Hopper), but replicated in the PyBullet engine and modified to be partially observable. We also test on POPGym, which is the largest POMDP benchmark available. Atari is expensive to train recurrent policies on, and we have already used over 50 wall-clock days training. Furthermore, many Atari games are MDPs or can be made MDPs by concatenating the last four observations, which is insufficient to evaluate long-term memory.\n- We perform an extensive ablation of seven different FFM architectures in Table 2, including four decay variants. See the following rebuttal section for more information.\n\n### The Case for Decay and Other Architectural Decisions \nDecay is already built into most RNNs. LSTMs and GRUs explicitly decay the recurrent state using the forget gate [C]. The forget gate activation $f_t = \\sigma(\\dots)$ in an LSTM would produce elementwise decay of the recurrent cell state: $c_t = f_t \\odot c_{t-1} + \\dots$ just like FFM's elementwise decay \n.\n- We ablate decay in Table 2, showing that without decay our model performs worse than the GRU (Fig. 3). Thus, decay is absolutely critical in our model.\n- If forgetting is detrimental to a specific task, our model will learn $\\alpha = 0$ through gradient descent, which corresponds to no decay. In Fig. 6, we record the decay rate and show it decreases over time for a long-term memory task.\n- The outer product architecture follows from computational psychology (Eq. 9-10). We tried variants without the outer product (e.g., where each dimension had its own $\\alpha, \\omega$), but we did not report them.\n- We are not familiar with successor features. Perhaps we could use them to investigate recurrent states produced by memory models, but we already investigate how FFM builds recurrent states in Fig. 6.\n\n### Notation\nOur notation is different than a normal RL paper because we need to represent both computational psychology and RL using a single notation. That said, we still try to follow standard ML conventions (e.g. $x$ is NN input, $y$ is output, $s$ is the recurrent state, $t,n$ are time/sequence indices, $o$ is observation, etc.). We will do a second pass through all the notation to see if anything can be improved. If you have specific suggestions, we would be happy to make changes.\n\n- In Eq. 9-10, the $t$ is just a placeholder for the index/timestep. We can rewrite Eq. 9-10 replacing $t$ with $n$ if this resolves the issue. \n- We introduce the hat operator in Eq. 9 to make it clear that Eq. 7-8 would store the trace after applying the context ($x$) rather than the observation without context ($\\hat{x}$). We will rewrite equations 9 and 10 to be more explicit.\n- The context variable is $\\exp(i\\omega)$, which is combined with the decay term and defined in Eq. 12. The $\\omega$ term is a learned weight of size $c$.\n\n[A] Sutton, Barto. _Reinforcement Learning: An Introduction._ Section 4.4 (http://incompleteideas.net/book/ebook/node44.html)\n\n[B] https://cs.brown.edu/research/ai/pomdp/tutorial/pomdp-background.html\n\n[C] https://en.wikipedia.org/wiki/Long_short-term_memory"}, {"review_id": "xToDWC0k4c", "reviewer": "Reviewer_Zkgc", "summary": "The paper introduces Fast and Forgetful Memory (FFM), a memory model designed for reinforcement learning. FFM is based on theories of composite memory and contextual drift from computational psychology. It consists of two main components: an aggregator and a cell. The aggregator computes complex-valued summaries of inputs, while the cell performs input gating and extracts real-valued outputs from the summaries. Conceptually, the aggregator implements composite memory by blending individual traces and incorporating an exponential forgetting mechanism. It further enables relative-time reasoning by applying a temporal context. On a suite of POMDP tasks from POPGym and POMDP-Baselines, FFM-enabled recurrent RL algorithms outperform those running on standard RNN. Moreover, since FFM belongs to the class of hybrid memory models, which provides equivalent recurrent and closed-form formulae, it achieves linear space complexity and logarithm time complexity with parallelism. This makes FFM an efficient and performant memory model that can be used as a drop-in replacement for RNN in recurrent RL algorithms.", "questions": "Besides drawing inspiration from computational psychology, are there intuitive or formal reasons that the proposed inductive biases lead to better performance on reinforcement learning tasks compared to the various instantiations of RNN? In particular:\n- Why is storing memory as individual traces more desirable than having a single context vector?\n- Is there a benefit to using imaginary temporal context over standard position encoding?\n- Is there evidence that temporal context enables relative-time reasoning?", "limitations": "The authors address the following limitations in their paper:\n- Performance decreases with larger context sizes, presumably because of a larger number of periodic functions resulting in more local extrema in optimization.\n- Limitation in the number of experiments.\n- Repeated multiplication of exponentials can lead to a loss of numerical precision. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper proposes a novel memory mechanism for reinforcement learning by drawing inspiration from the field of computational psychology. It represents its cell state as a collection of traces and features numerous inductive biases for composition, forgetting, and relative-time reasoning. These inductive biases collectively lead to improved performance on RL tasks.\n- The experiments are substantial and thorough, providing ample evidence that FFM can be used as a drop-in replacement for RNN in recurrent RL algorithms. The ablation experiments and analysis provide further insight into the significance of each design choice.\n- The contributions of this paper are significant, especially considering the importance of memory models in partially-observable RL.", "weaknesses": "- The intuition behind some inductive biases is elusive. Please refer to the Questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper, and thank you for your interest. We hope we are understanding your first question correctly: we only ever store the __sum__ of time-contextualized traces. We assume psychologists like using distinct traces like Eq. 6 because they are conceptually easier to reason over, but from a computational standpoint, they are very inefficient and scale poorly to large sequences. Note that in our paper, every time we say \"context\" we are explicitly referring to the temporal context (i.e., $e^{iw}$). We can think about the distinct traces being present in the sum, which allows us to do tricks like multiplying by exponentials to \"shift\" the recurrent state. Please see the section titled __An Explanation of Temporal Context__ in the __Author Rebuttal__ section for why an imaginary context is beneficial. \n\nWe do not report it in the paper, but the RepeatPreviousEasy task where the agent must output the observation from $k$ timesteps ago was unsolvable without temporal context. We believe our context-free model had trouble determining which observation it saw precisely $k$ timesteps ago. This extended to other tasks (e.g. Autoencode and others that we do not recall) that depended on the order of observations. Adding context fixed this issue. Fig. 6 provides some evidence for this conclusion. \n\nIn Fig.6, FFM empirically learned a context period of $2k = 64$. With this context period, the recurrent state will \"spike\" at half the period $k=32$. In other words, at $t=32$, the imaginary component of $x_{t-k}$ in the recurrent state will be $i \\sin(\\omega t) x_{t-k} = x_{t-k}$ and the real component will be $\\cos(\\omega t) x_{t-k} = 0$. The downstream layers should be able to pick up on this \"spike\" and output the observation $x_{t-k}$. Put another way -- it is mighty suspicious that the periods our learned context parameters dropped from various periods of $>400$ to precisely $2k$. One might conclude that gradient descent is pushing the context period to the environment period because it is beneficial!"}, {"review_id": "Tt5DWKBpHh", "reviewer": "Reviewer_qKYK", "summary": "This work introduces a new recurrent neural network layer called Fast and Forgetful Memory (FFM), which, differently from traditional architectures such as GRUs and LSTMs, FFM is explicitly designed for deep reinforcement learning", "questions": "- Do you foresee any possible issues of instability?\n- How long were the temporal dependencies tested in the experiments?", "limitations": "I think authors did fine work addressing the limitations of their contribution at the end of the paper", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The authors present a novel and promising memory layer for reinforcement learning, the core ideas of the layer are intuitive. The theoretical foundation is well explained and easy to follow, I didn't spot any errors. Also the empirical analysis is extensive. Authors cover time and space complexity of the layer which is very beneficial to get a general idea of its advantages.\n\nIn general I think is a good work and a strong accept", "weaknesses": "No major weaknesses that I can tell, just a few minor points:\n- In the abstract I would not say so broadly that all model-free approaches are using architectures from SL I think that applies to RNNs but mainstream transformers used in RL were specifically redesigned for this framework\n- Line 104, I would explain what traces are\n- Line 117, remove brackets from citations that are part of the sentence\n- Page 4, I would do an effor to inegrate the footnote in the main text, it would make the section much easier to follow", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking time to review, we hope you enjoyed the paper.\n\nThank you for the corrections, we will fix these at once. The footnote was originally in the main text, but it took up a lot of space and we were worried it might put some readers to sleep. We will see where we can fit it back in. We briefly cite the Stabilized Transformers for RL paper in the introduction, but we will also add it to the abstract.\n\nAs described in the limitations, floating-point instabilities can result from learning very large values of $\\alpha$. However, once we clamped $\\alpha$, we did not observe any other instabilities unique to FFM. It occasionally overestimated Q values on the off-policy tasks, but not any more often than the LSTM or GRU. As we note in the experiments section, we did not change any hyperparameters for any benchmarks, so FFM should be stable enough for an average user to apply it to standard RL tasks.\n\nThe temporal dependencies varied by experiment. On the longer side, the POPGym navigation experiments have a maximum episode length of 1024 timesteps, but it is not clear if the dependencies span all 1024 timesteps. More strictly speaking, the POPGym RepeatFirstHard environment has a temporal dependency of 832 timesteps. We would guess that the POPGym benchmark has an average temporal dependency measured in low hundreds of timesteps, but we do not have a hard number. The POMDP-Baselines benchmark tends to stick to BPTT lengths of 32 or 64, so we suspect it has a lower temporal dependency than POPGym, but we cannot be sure."}, {"review_id": "BCPB5mdXXQ", "reviewer": "Reviewer_AuBN", "summary": "The paper proposes a new memory model for reinforcement learning agents, targeting POMDP setting. The memory (FFM) aggerates past data via an exponentially weighted summation operation where the weights are learnable parameters. The memory is used to compute the Markov state through sophisticated gating mechanisms. The paper evaluates FFM on 2 POMDP benchmarks and shows improvement over baselines such as GRU. ", "questions": "- It is unclear why the inductive bias (composite memory and contextual drift) is designed specifically for RL setting. The bias is very general and can be used for both SL and RL. This raises the questions such as whether FFM can show good performance under SL? Why only test it on RL?\n- It is hard to link Eq. (11) to the contextual drift motivation where $\\omega$ should be recurrently computed. \n- It is also unclear why the decay should be complex-valued. The authors should present the explanation for the changes compared to the original motivation Eq. 8-10. \n- The choice of gating for the Cell seems arbitrary. Could you explain the reasons behind these designs, or do you follow any prior papers/ideas? ", "limitations": "More investigation is needed on the initialization of the decay. Now, it seems to be heuristic and unclear if it still works for other tasks or settings.  ", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The memory design is motivated by computational psychology \n- The proposed memory is light-weight and fast", "weaknesses": "- The novelty is questionable. The main contribution is Eq. 11, and it is very similar to Eq. 1 in [1]. The difference is that the decay term here is learnable and consists of real and imaginary components. Table 2 reveals that non-learned decay also shows decent performance, which questions the need of learning the decay. \n- The literature review lacks a line of research using MANN on RL [2,3,4]. They are memory models specifically designed for RL. It is fairer to compare FFM with them rather than comparing with general memory models such as DNC or GRU. Please note that general memory models may require heavy hyperparameter tuning and adjustment to make them work on RL, thus comparing with well-established memory-based RL methods is better. For example, Transformer needs some tricks to work on RL [5]. Please consider using this baseline instead of general Transformer models.   \n- The experimental results are not convincing. In POPGym, the gap between FFM and GRU is not clear. The ablation study in Table 2 indicates that the full model does not outperform other configurations significantly, especially when the std is high and only 3 runs are executed. It would suggest that the complicated design might not add much to the performance gain. In Fig. 5, there are many tasks FFM shows similar or worse results than RNN. It questions the benefit of using FFM. \n\n[1] Ba, Jimmy, Geoffrey E. Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin Ionescu. \"Using fast weights to attend to the recent past.\" Advances in neural information processing systems 29 (2016).\n\n[2] Wayne, Greg, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska-Barwinska, Jack Rae et al. \"Unsupervised predictive memory in a goal-directed agent.\" arXiv preprint arXiv:1803.10760 (2018).\n\n[3]  Fortunato, Meire, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charles Deck, Joel Z. Leibo, and Charles Blundell. \"Generalization of reinforcement learners with working and episodic memory.\" Advances in neural information processing systems 32 (2019).\n\n[4] Lampinen, Andrew, Stephanie Chan, Andrea Banino, and Felix Hill. \"Towards mental time travel: a hierarchical memory for reinforcement learning agents.\" Advances in Neural Information Processing Systems 34 (2021): 28182-28195.\n\n[5] Parisotto, Emilio, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg et al. \"Stabilizing transformers for reinforcement learning.\" In International conference on machine learning, pp. 7487-7498. PMLR, 2020.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review. Our primary concern in this paper was to design a __more efficient RNN for model-free RL__, with the idea that anyone with an existing recurrent RL implementation (e.g. CleanRL's recurrent PPO) could easily replace their LSTM/GRU with FFM, and train a comparable policy using fewer resources. We believe we have succeeded in that regard.\n### Novelty and the Relationship to [1]\nThe $\\alpha$ term in FFM's aggregator produces similar decay to $\\lambda$ from [1], but we disagree with the premise that our approach is a simple modification of [1]. \n- We benchmark FFM against modern variants of [1], such as the Fast Autoregressive Transformer (FART) and the Fast Weight Programmer (FWP) (Fig. 3 and Appendix Fig. 10).\n- [1] only studies a single RL task which we would argue is insufficient to draw conclusions on RL performance.\n- [1] does not appear parallelizable over time like FFM, meaning it will be less efficient and on-par with RNN efficiency (Eq. 2 in [1] shows $h_{s+1}$ has a nonlinear dependence on $h_s$)\n- We compute an outer product, but it is not using the recurrent state $h(\\tau) h(\\tau)^\\top$ as in [1], but rather with a function of input the and the context vector $f(x_\\tau) \\omega(\\tau)^\\top$ (using their notation).\n\nThe temporal context fixes some of the theoretical shortcomings of [1], and is a non-trivial addition that requires numerous additional considerations (see __An Explanation of Temporal Context__ in the __Author Rebuttal__). Let us summarize and inspect the ablations from Table 2:\n\n|Temporal Context ($\\omega$)|Score|Std.|\n|-|-|-|\n|None|0.382|0.005|\n|Fixed|0.383|0.005|\n|Learned|0.399|0.005|\n\nWe can see that a learned context is important and statistically significant ($p=0.0141$ according to T-test computed via [A]). Furthermore, if the context $\\omega$ is not learned, then equations 20-22 fall apart and FFM can no longer approximate any temporal convolution.\n\nThe case behind learning decay $\\alpha$ is a bit weaker, but still provides a marginal benefit at the cost of 32 additional parameters. Readers are free to use our model with a fixed $\\alpha$ if they find the results unconvincing -- it does not diminish our contributions.\n### MANNs and Transformers\nMany of your cited papers are indeed missing from our related work. Given our paper's scope and focus on efficiency, we will explain why we did not include or compare to them:\n- [2] is not peer-reviewed and exceedingly complex -- there is not one working public implementation after 5 years. Furthermore, the backbone is an RNN which we show is much less efficient than FFM.\n- [3] is not parallelizable over time due to writes and reads to external memory, making it slow like the RNNs we tested. Furthermore, it is not model-free RL as it requires predicting the next observation (i.e., it is a dynamics model).\n- [4] and [5] combine an LSTM/GRU with a transformer, resulting in quadratic space complexity while also being slow to run because it cannot be parallelized over time. We cite [5] and tried to run it, but our GPU did not have enough memory to fit it given the batch sizes used by POPGym.\n\nIt is possible that some of these methods could outperform FFM on a per-sample basis, but not on a time or hardware efficiency basis.\n\nFurthermore, the LSTM and GRUs from the POMDP-Baselines experiments used hand-tuned initialization and task-specific hyperparameters from Ni et al. Our FFM model outperformed those models while using the exact same hyperparameters.\n### Experimental Gaps Between FFM and the GRU\nFig. 3 shows that on average, we exceed the returns of the best POPGym model (GRU) by ~10%, without overlapping 95th percentile error bounds. Fig. 4 breaks this down further, showing that we perform noticeably worse on 4 tasks and noticeably better on 18. For POMDP-Baselines, FFM outperforms the LSTM/GRU on 12 out of 14 tasks. We argue that this paints a clear picture of FFM outperforming RNNs, even before considering efficiency.\n### Inductive Biases\nYou are correct in that FFM's inductive bias could be useful on SL tasks, and in fact, FFM could be run without modification on SL tasks. Frankly, we are less familiar with the wider SL literature, and as you point out in [1], there are already many good papers on the subject (which we will add to our related work). The POPGym paper shows that there is an unpredictable performance gap between memory in SL and RL, so we argue that RL-focused experiments are especially useful to the RL community. \n### Recurrent Temporal Context\nThe purpose of Eq. 10 is to show that in psychology, context changes gradually. We can reformulate our context implementation recurrently like in Eq. 10 as $\\omega_{t+1} = e^{iw} \\odot \\omega_t$. Please see the section entitled __An Explanation of Temporal Context__ in the __Author Rebuttal__ response for why it is crucial that the temporal context be complex-valued.\n### Cell Gating\nFFM output gating is very similar to the GRU output gating: $h_t = (1 - z_t) \\odot h_{t-1} + z_t \\odot g_t; \\quad z_t = \\sigma(W x_t + b + \\dots)$. Accessing the previous state $h_{t-1}$ in FFM would break the parallel nature of FFM, so we replace the $h_{t-1}$ with the input $x_t$, resulting in a skip-connection that improves gradient propagation upstream of the FFM module. It also provides a shortcut for downstream layers to access the current observation.\n\nThe input gating approximates the product of the Bernoulli random variable $B$ with input $x$ from Eq. 7-8. We tried sampling from an actual Bernoulli distribution using the reparameterization trick, as well as approximating sparsity via $B_t x_t = ReLU(x_t)$. Our current approximation worked best.\n\n[A] https://www.statskingdom.com/140MeanT2eq.html"}]

Author rebuttal:
Thank you to all the reviewers and ACs who spent time reviewing this paper. We are happy to hear that our proposed method is "novel and promising", as well as "inuitive...well explained, and easy to follow", that "the empirical analysis is extensive" and "The contributions of this paper are significant, especially considering the importance of memory models in partially-observable RL."

We believe some readers might benefit from a better explanation of temporal context, so we will move the subsection "Efficient Convolution with Infinite-Extent Filters" to the appendix to make room for an expanded explanation of temporal context. In the meantime, let us explain the rationale behind our implementation of temporal context.

### An Explanation of Temporal Context
The context is traditionally real-valued in computational psychology (Eq. 9-10), but our implementation leverages the complex domain to break symmetries and improve relative-time reasoning. Since we are effectively summing inputs to compute the recurrent state, we can run into ordering symmetries in the real domain because addition is commutative. 

Let us begin with a simple example very similar to FFM aggregation, but without context and of a single dimension. Our goal is to compute a recurrent state $s_n$

$$ s_n = x_1 e^{(n-1)\alpha} + x_2 e^{(n-2)\alpha} + \dots $$

where $\alpha$ is the decay term. Consider the case where we learn a real decay term of $\alpha = 0$ (no decay), perhaps because two inputs $x_j, x_k$ are very important and cannot be forgotten. For clarity and without loss of generality, assume all $x \notin \{x_j, x_k\}$ are worthless and we learn to map them to zero. Then, our recurrent state is

$$ s_n = e^{(n - j) \alpha} x_j + e^{(n - k) \alpha} x_k = x_j + x_k. $$

How are we to know whether we observed $x_j$ first or $x_k$ first, or how long ago they occurred? In the current case, we cannot know either. We could add a transformer positional encoding, but this would not solve the ordering ambiguity, as we would have

$$ s_n = x_j + \textrm{pos}(j) + x_k + \textrm{pos}(k) = x_j + \textrm{pos}(k) + x_k + \textrm{pos}(j). $$

Using an imaginary context $e^{i\omega}$ breaks these symmetries. With imaginary context, our recurrent state becomes

$$ s_n = x_j e^{(n - j) i\omega} + x_k e^{(n - k) i\omega}$$

Now, we can tell whether $x_j$ or $x_k$ came first, and how old each observation is. Thanks to Euler's Formula ($e^{i\omega} = i \sin \omega + cos \omega$), we can shift these inputs infinitely far into the past or future without losing the relative distance between $j$ and $k$. In other words, as $n \to \infty$, the phase shift between $e^{(n - j) i\omega}$ and $e^{(n - k) i\omega}$ remains constant.

Let us demonstrate how we "shift" the recurrent state with a concrete example. At timestep $n+1$, we shift the recurrent state one step into the past via multiplication: 

$$ s_{n+1} = e^{iw} (x_j e^{(n - j) i\omega} + x_k e^{(n - k) i\omega}) = x_j e^{(n + 1 - j) i\omega} + x_k e^{(n + 1 - k) i\omega}. $$ 

In practice, this shift requires computing very large exponents and quadratic memory complexity. Appendix C derives tricks to avoid these issues. The use of complex exponentials also makes FFM able to approximate any convolutional filter (subsection Efficient Convolution with Infinite-Extent Filters), providing a strong theoretical reason for learning the context parameter $\omega$.

Meta-review: I am thrilled to convey my recommendation for the acceptance of your paper to be presented at NeurIPS.

This is effectively a new neural net architecture drop-in replacement for RNN's. It's well written and appears novel enough. It could have gone either way. Nothing really exciting, nor dramatically disappointing.

NeurIPS is a prestigious platform that attracts the brightest minds in the field, and your paper's acceptance adds to the conference's reputation for excellence. I am genuinely excited to see your work presented and discussed among peers who share your passion for pushing the boundaries of knowledge.
Decision: Accept (poster)