AGaLiTe: Approximate Gated Linear Transformers for Online Reinforcement Learning
Authors: Subhojeet Pramanik
Esraa Elelimy
Marlos C. Machado
Adam White
Venue: Accepted by TMLR ()

Abstract: In this paper we investigate transformer architectures designed for partially observable online reinforcement learning. The self-attention mechanism in the transformer architecture is capable of capturing long-range dependencies and it is the main reason behind its effectiveness in processing sequential data. Nevertheless, despite their success, transformers have two significant drawbacks that still limit their applicability in online reinforcement learning: (1) in order to remember all past information, the self-attention mechanism requires access to the whole history to be provided as context. (2) The inference cost in transformers is expensive. In this paper, we introduce recurrent alternatives to the transformer self-attention mechanism that offer context-independent inference cost, leverage long-range dependencies effectively, and performs well in online reinforcement learning task. We quantify the impact of the different components of our architecture in a diagnostic environment and assess performance gains in 2D and 3D pixel-based partially-observable environments (e.g. T-Maze, Mystery Path, Craftax, and Memory Maze). Compared with a state-of-the-art architecture, GTrXL, inference in our approach is at least 40% cheaper while reducing memory use more than 50%. Our approach either performs similarly or better than GTrXL, improving more than 37% upon GTrXL performance in harder tasks.

Published in Transactions on Machine Learning Research (10/2024)
AGaLiTe: Approximate Gated Linear Transformers for Online Reinforcement Learning
Subhojeet Pramanik" Esraa Elelimy": Marlos C. Machado Adam bh
"Department of Computing Science, University of Alberta, Canada
"Alberta Machine Intelligence Institute (Amii), Canada "Canada CIFAR AI Chair (spramanik, elelimy, machado, amve)oualberta.ca
Reviewed on OpenReview:
Abstract
In this paper we investigate transformer architectures designed for partially observable on- line reinforcement learning. The self-attention mechanism in the transformer architecture is capable of capturing long-range dependencies and it is the main reason behind its effective- ness in processing sequential data. Nevertheless, despite their success, transformers have two significant drawbacks that still limit their applicability online reinforcement learning: (1) in order to remember all past information, the self-attention mechanism requires access to the whole history to be provided as context. (2) The inference cost in transformers is ex- pensive. In this paper, we introduce recurrent alternatives to the transformer self-attention mochanism that offer context-independeen inference cost, leverage long-range dependencies effectively, and perforiss well in online reinforcement learning task. We quantify the impact of the different components of our architecture in diagnostic environment and assess per- formance gains in 2D and 3D pixel-based partially-observable environments (e.g. T-Maze, Mystery Path, Craftax, and Memory Maze). Compared with a state-of-the-art architecture, GTrXL, inference our approach al least 40% cheaper while reducing memory use more than 50%. Our approach either performs similarly or better than GTrXL, improving more than 37% upon GTrXL performance in harder tasks.
1 Introduction
In many real-world settings agents often have limited observability of the environment making decision making extra challenging. For example, an agent designed to drive cars must remember the road signs it saw a few minutes ago to adjust its velocity if there are any major changes to the road. A naive approach would to store the entire history of camera observations. However, such an approach is not scalable as the history of observations can be longer than the memory available to the agent (MeCallum) 1996). Alternatively, the agent can learn compressed representation of the history of obscrvations, and use make decisions. This approach, however, not feasible in continuing problems where agent's 's face an unending stream of experience and information critical for decision making occurred in the distant past. Therefore, can incrementally update their internal representation of the environment state with computation that does not growth as a function of total experience. In this paper, we investigate incremental state construction in the context of partially observable online reinforcement learning (RL), where agents learn while interacting with the world and thus computation and memory require special consideration.
need agents that
In RL, incremental state construction longestudied problem with many possible solution methods. Re- current neural network (RNN) architectures provide framework for learning such representations due to

---

Published in Transactions on Machine Learning Research (10/2024)
their ability to automatically learn relationships about the past. RNNs, such as LSTMs (Hochreiter & Schmidhuber 1997) and GRUs (Gao & Glowacka 2016), handle sequential data by maintaining a vector of hidden states that capture dependencies between consecutive elements in the sequence. RNNs have been applicd to a wide range of partially observable RL environments such as Atari 2600 games (Hausknecht & Stone 2015) and Starcraft. (Vinyals al.] 2017). The inference cost RNNs the cost of processing a single element a sequence of data-is independent of the length of the sequence making them an attractive choice for online RL. Unfortunately, RNNs such as LSTMs are notoriously difficult to train (Bakker 2001; Khandelwal et al., 2018), and their computations over the imput cannot be parallelized.
Transformers Vaswani et al. 2017 have achieved state-of-the-art performance in many sequential data processing problems, but have seen limited application in online RL (Parisotto et al. 2020). Transformer architectures have been widely used in natural language processing (e.g., Brown et al. 2020, Devlin et al. 2018) and computer vision (e.g., Petit et al.] 2021: Zhong et al., 2020). These successes are often attributed to the transformers' sclf-attention mechanism which can capture long-range dependencics, but they cannot learn relationships that exceed this fixed-length memory. In addition, the inference costs are higher than an RNN: linear in the length of the context window. The linear transformer architecture reduces computational complexity of the self-attention mechanism (Katharopoulos et al. 2020).
In this paper, we introduce two new approaches designed for partially observable RL problems based on the linear transformer's self-attention mechanism. Our new approach to self-attention, based on linear transform- ers, was designed to achieve the following: (1) a self-attention mechanism that can add and delete previous information, (2) a learned feature map, and (3) a self-attention mechanism that requires computation linear in the size of the embedding dimension. Our first contribution, Gated Linear Transformer (GaLiTe), uses gated structure that allows it to uncover relationships far in the past. It also uses a different self-attention mechanism that can learn a highly parallelizable feature map that is amenable to sequential computation with a context-independenn inference cost. Our second contribution, Approximate Gated Linear Transformer (AGaLiTe), introduccs an approximate version of GaLiTe's self-attention mechanism, climinating the neod to mainlain a matrix as a recurrent state.
We demonstrate the utility of our proposed approaches in several partially observable RL problems. Our experiments show that both GaLiTe and AGaLiTe can match the performance of more computationally expensive transformer architectures in small diagnostic T-Maze enviromment. In pixel-based navigation task, we find that our approach outperforms the ofatte-f-thheart transformer architecture, GTrXL (Parisotto et al. 2020), by more than 37%. Our AGaLiTe-based agent achieves higher rewards than a GTr-X- -base agent and higher performance across various in-game skills in Craftax Symbolic (Matthews et al. 2024
a symbolic adaptation of the 2D survival game Crafter (Hafner) (001). Tn 3D pixel-based navigation tasks, AGaLiTe's performance is close to GTrXL while reducing the computation and memory by 40% and 50% respectively. Our results in Craftax and 3D navigation provide promising initial evidence of the scalability of our new approach; step towards better transformer-basee architectures for online, partially observable RL. Code and implementation for this work is publicly available
2 Preliminaries
In this section, we provide a brief background on Transformers. We first discuss the canonical transformer architecture and then we discuss the lincar transformer approach, which is the basis of our approach.
The transformer architecture was introduced for supervised next token prediction tasks (Vaswani et al. 2017). Our main contribution is a new self-attention mechanism: this section provides the background required to understand the self-attention mechanism in transformers.
Self-attention is mechanically simple. For given query token i (embedded in x; : X(i, we output an embedded context vector that weights each input token's importance (attention weighted) to the query token. The input to the self-attention layer is a matrix X E Nxd an embedding of each input token (1 to N) into a vector, Rdd The output is a matrix A E N where d, is the head dimension. Algorithm shows a single self-attention layer with learnable parameters Wo, W K, Wr E Rdxd,
http:///uuhhub.cm//ahhoo////////
2

---

Published in Transactions on Machine Learning Research (10/2024)
2.1 Canonical Transformer Architecture
We can think of the process in two steps. In step one we cal- culate the attention weights. We compare each token in the context to all other tokens in the context (QK"). The weights are then scaled the size of the embedding dimension and nor- malized with an element-wise softmar. In step two, we compute and return the eatentiion-wiihhee context vectors, one for cach input in X.
Algorithm Canonical Self-Attention
Input: X R
Parameters: Wo, WK, Wv
1: Q < XWQ
XWK K
V <- XWv
The self-attention mechanism in Algorithm D is computation- ally expensive. We define the inference cost of self-attention as the cost. for processing a single element in a sequence. To generate representations for single element, a query vector is calculated instead of query matrix using a single input (Algorithm 1 step 1). In canonical self-attention mechanism, processing single element requires having the full sequence as input for generating the value and key matrix (Algorithm [ step 2 and 3). Thus, the inference cost depends on the input sequence length N. For a naive implementation, the inference cost has O(Nd²) time and ((Nd) space complexity; increasing the sequence length linearly increases the computat ional complexity. A simple miligation is to limit the size of the imput sequence by maintaining a window the history of input activations in memory (Dai et al. 2019), but doing so limits the past information the self-attention mechanism can recall.
4: < softman OK
Output: A R
2.2 Recurrent Attention with Linear Transformers
The linear transformer architecture (Katharopou- los et al. 2020) introduces a general way of for- mulating sell-attention as a recurrent neural net- work by replacing the softmar with a kernel func- tion, leveraging its equivalence to applying kernel smoothing over inputs (see work by Tsai et al.) 2019).
Algorithm 2 Linear Transformer's Self- Attentinn
Input: x E Rdd C.-1 Rddxde 8t-1 E Rdd Parameters: We, WK. Wy Rd xd So t O. + 0.
1: q: + (Wox))
2: k, t- ((xx))
A single time-step of inference of the linear trans-
3: V, Wrx;
former self-attention is described in Algorithm 2 Note that we present the algorithm for processing a single input vector in the case the lincar trans- former. This in contrast. to Algorithm [ which presents the algorithm for processing a sequence. Let k(a, b) ø(a) <(b), where & R + Rd. is non-linear feature map, dy is the output dimension of the feature map and k Rd Rd R Additionally, let & be defined as the outer product vector operation. At a given timestep the linear transformer self-attention maintains matrix, C,-1 E Rddx and a vector, St E Rdd, as a recurrent state, which is updated iteratively using the current input vector, X Different from Algorithm [ Algorithm 2 applies the feature map, Φ to generate the query and key for a given time-step (lines and 2). The linear transformer self-attention stores the outer product of value and key vectors as a recurrent state matrix, C, (line 4). Additionally, the sum of the key vectors is stored as recurrent normalization vector S, (line 5). The attention output vector, a,, is calculated by multiplying the recurrent state with the query vector, and normalizing it using the product of the normalization vector, St, and the query vector, q (line 6).
4: C. + C--1 + vi k;
5: 8t + 8 - k;
6: a + (C, q:)
Output: a Rd ,,, Rdd, St E Rdd
The linear transformer's self-attention has a context-independent inference cost, unlike the canonical self-attention mechanism. In Algorithm| 2 processing a single input vector (xx) has a space and time complexity of (dd)), assuming d, the embedding dimension (of the input), greater than d,, which the size of the attention- attention-weighted context vector, a,. Unlike vanilla self-attention, the computational complexity does not depend the context length, making it more efficient. for longer sequences.
3

---

Published in Transactions on Machine Learning Research (10/2024)
3 Gated Linear Transformers (GaLiTe)
In this section, we introduce GaLiTe to address two of the limitations of linear transformers. Specifically, (1) the recurrent equations in Algorithm 2 (lines 5 and 6) add positive values to the recurrent state, without any mechanism to delete past information. (2) Performance critically depends on the choice the kernel feature map & (lines and 2); clement-wise functions such as the Exponential Linear Unit (ELU) typically perform worse than softmax [Katharopoulos et al. 2020).
GaLiTe mitigates these two issues by introducing gating mechanism and parameterized feature map. The gating mechanism controls the flow of information at each index of C (the location of the recurrent states of the self-attention mechanism), allowing arbitrary context memory (inducing a trade-off with precision). The parameterized feature map is used to calculate the key and query vectors in the self-attention mechanism, eliminating the choice of the kernel feature map Ø
3.1 Gating Mechanism to Control the Flow of Information
In the linear transformer self-attention, at a given time-step t, Algorithm 2 increments the recurrent state, Ct-1, and the normalization vector, St-1, (lines and 5). Assuming Cg and Sy are initialized to zero, recall the update equations for C, and 8, are recursively defined as follows:
= C V ki,
(1)
St-1 kt-
(2)
As k, is a function of & and under the assumption of a positive feature map Φ, Equation 2 adds arbitrary positive values to 8--1- Similarly, Equation [ adds arbitrary positive values to C,-1 if the elements in value vector V, are positive. Both Equation [ and 2 have no way to control the flow of past information and the values in the recurrent state could grow. Instead, we use a normalized exponential average with element-wise learned decay parameters which smoothly reduces the impact of past information.
Gating mechanisms pose a learned outer-produet-bassed gating mechanism that decays every element of C.-1 and s allowing the network to learn the decay for each element (also known as the memory location). We introduce learnable parameters W3 E Rddadd W, E Rdexd, and gating vectors Be, and Ye. Let 0g be a sigmoid function defined as az(x) i 1/1te-s we define 3r and Vt as follows:
be used to control the flow of information in recurrent updates. We pro-
3, - (W
(3)
- og ( Wx)
(4)
Let O be the element-wise product, we use the outer product of 3, and % control the flow of past information in recurrent states C, and s,, modifying Equations
I andd as follows:
C = ((1- Bx) @ x)) 3 C.-1 + (B. @ (ne @k),
(5)
= (1 - m1 3 9: + no 2xOk kr.
(6)
We use outer products to learn the decay rate for each index of Ct, without requiring individual parameters for each index. The outer product assumes the decay rate at each index is independent from each other.
3.2 Learnable Feature Map for Self-Attention
Recall that the self-attention mechanism of the linear transfomer uses a kernel feature map calculate the key and query vectors:
k (Wxx))
(7)
인 ( (Qx:).
(8)
We consider a deterministic approach to learn the key and value vectors in the linear transfomer sell-atiention mechanism. We introduce modifications to k;, qr, and gating vectors calculation described in Equations 7 and respectively. We start by introducing hyperparameter n that controls the dimension of the feature maps used construct k, and qı. Let Wp,, W.. Wp.. Wps E Raxd be learnable parameters. We modify the dimensions of W,, W. Rdxxd getting rid of dR, the kernel feature map dimension. Let f() be function
S
4

---

Published in Transactions on Machine Learning Research (10/2024)
that flattens matrix into a vector. redefine k, and qi (previously defined in Equations 7 8
follows:
k, - Krelul W,, p,, x.) @ relul W xx))))
q: - Arelul W Xt) relul W (xx)).
(10)
We also modify the gating vectors calculation in Equation 4 as follows:
2 = fo (W (W X.) & (W,x)))
(11)
Using the modified key, query, and gating vectors, the recurrent states C2 E Rdax and S, R are calculated according to Equations 5 and D Tt is important to note that the feature map dimension, de = ndd,, is now controlled by the hyperparameter 7). Equations 9 and 10 use outer products to learn multiplicative interactions in the key and query vectors. Learning multiplicative interactions in the feature vectors allows learning complex non-linear relationships through training instead of relying on an explicit non-linear element-wise function or on random feature maps. Using outer products to generate an expansive feature map allows us to have a large feature map output dimension. Having a large feature map output dimension is essential it correlates with the memory capacity (see the work by [Schlag et al. 2021).
Finally, use the relu actívation function to ensure the output of the feature map is positive. A positive feature map is a common assumption in the linear transformer literature as it is simple way to ensure that similarity scores produced by the underlying kernel function are positive.
The Gated Linear Transformer (GaLiTe) self-attention incorporates the changes discussed above into the linear transfomer self-attention. The pseudo-code for GaLiTe is available in Appendix A GaLiTe has similar space and time complexity as the linear transfomer. For processing a single element a sequence, GaLiTe has space and time complexity of O (nd²) and O (nd²), respectively. In comparison, the linear transfomer requires O (ded) and O (dxd). Notice d& is defined to be the output dimension of the kernel feature map, which is qdd, in GaLiTe. Similar to the linear translomer, the space and time complexity of GaLiTe is independent of N and only depend on static hyperparameters d and 17.
4 Approximate Gated Linear Transformer (AGaLiTe)
Operating large matrices is expensive. Recall that GaLiTe stores matrix of dimension di,1 recurrent hidden state. This becomes more problematic with the use of multiple heads and layers; which are typically required to improve stability during the training (see the work by Michel et al., 2019). For example, previous applications of transformers in RL by Parisotto et al. (2020) use 8 heads and 12 layers; 96 heads in total. Second, the update to C₁ makes use of expensive and memory heavy operations: an outer product, clement- wise matrix sum, and multiplication.
Our goal is to approximate the recurrent state update in Equation 5 with an approximation that uses less space than O(1d²). Recall that Equation 5 replaces C, with C.-1 plus new outer product. To derive an approximation, we want to replace C,-1 with a matrix that has a lower rank. Also, we want to derive an update rule that is an approximation of Equation 5 but instead of updating the full-rank matrix, C,-1, it updates the low-rank approximation.
Our second approach, called Approximate Gated Linear Transformer (AGaLiTe), uses a low-rank approx- imation to roduce the space complexity of GaLiTe. We replace the previous recurrent state matrix, C,-1, with a set of vectors, reducing the space complexity of GaLiTe by d. We introduce an approximation of the Kronecker delta function using a sum of cosine functions and we use this to approximate Ct-1-
We introduce an approximation that uses sum of cosine functions to approximate sum of outer products. This approximation is deterministic, it does not introduce variance in the approximation, and il keeps incre- mental updates to the state end- to-end differentiable. Our approach is inspired by the rank-1 approximation introduced by Ollivier et al. (2015), but instcad of using random numbers to approximate a Kronecker delta function, we use a trigonometric identity that relates a Kronecker delta function to an integral over cosines. Recall that the Kronecker delta function is defined for integers m and such that 8pm = 1 if n, and
5

---

Published in Transactions on Machine Learning Research (10/2024)
6mn 0 if m # n. We present an approximation smm of 8mn such that Swn is defined follows:
2 Σ 1=0
2Ti
2 (- m ))
8mn
cos
(12)
It can further be shown that lim 8mn 8ma. The derivation for this result is presented in Appendix B.I. We use the approximation of the Kronecker delta function in Equation 12 to approximate the recurrent state update in Equation 5 Briefly, the approximation introduces the approximate Kronecker delta function to approximate C₁ as a sum of r outer-procd pitter-produts,, where each of the vectors in the outer-produce is defined recursively and updated using the value and key at the current timestep. For a given r, we maintain recurrent states v and k for k = 0,1,...,A For wi = 2nk and assuriin v and kj are initialized as zeros, we di- rectly calculate the attention output, a,, in replacement of Cr, considering the recurrent updates to vi and kj:
V = cos(wat)3: O v., +11-3)) ViL
(13)
K = cos(uxt)n @ k, + (1 -1)) @ k
(14)
Ek K= V (k)) 91)
(15)
2r(s, q:)
Due to space constraints, the rationale behind these approximations is presented in Appendix B.1.1
The pseudocode for AGaLiTe can be found in Ap- pendix [ Unlike Equation 5 Equations 13 and [4 define recurrence over vectors instead of matri- Des. If r < d, then the recurrence is more effi- cient in space than the recurrence Equation 5 In Appendix E. we provide an empirical evaluation of the impact of different values of r in the qual- ity of the approximation, showing that, in practice, it seems small values of do not compromise the quality of the approximation or the overall perfor- mance. The computational complexity of AGaLiTe is (rnd) and O (d2 + rnd) in space and time. With AGaLiTe, we have significantly improved the com- plexity of the self-attention mechanism and these show next.
Space and time complexity of AGaLiTe, GaLiTe, linear transformer, and GTrXL for processing a single element in a streaming sequence (M: memory size in GTrXL, d: representation dimension, d2 fea- ture map dimension in the linear transformer, q): fea- ture map hyperparameter in GaLiTe and AGaLiTe, approximation parameter in AGaLiTT).
Table 1:
Space
Time
GTrXL
O(Md)
O(M d))
Linear Transformer
O (d+d)
O (dad)
GaLiTe
(nd²)
1 (nd²)
AGaLiTe
O (rnd)
o (d + rrjd)
differences manifest in experiments
We compare the computational complexities of our proposed approaches and GTrXL (Parisotto et al. 2020 in Table 1 ments of forward pass using the AGaLiTe architecture and compare it to GTrXL in Appendix K We also discuss the parallelization GaLiTe and AGaLiTe over a sequence of data in Appendix F
We provide empirical latency measure-
5 Empirical Evaluation
This section investigates our proposed approaches in several partially observable reinforcement learning (RL) control problems. The memory requirements vary across the environments we consider. In T-Maze (Bakker 2001), the agent must remember single cue signal. In CartPole, the agent must. estimate the hidden state by integrating information over time. In Mystery Path (Pleines et al. 2023), the agent must remember multiple locations in grid environment. In Craftax (Matthews et al. 2024), 2D survival game, the agent faces with partial observability as it can only observe a limited portion of a large 2D map. In Memory Maze environment [Pašukonis el al. 2023 the agent must retain the layout of a 3D maze in addition to several locations across the maze.
Additionally, we also provide results in Long Range Arena Tay 2021 in Appendix D classical benchmark used to evaluate the ability to learn long-range dependencies in supervised learning scenaroo. We

---

Published in Transactions on Machine Learning Research (10/2024)
evaluate in two of the tasks from the benchmark: ListOps and IMDB. We found that AGaLiTe outperforms transformers and linear transformers across both of these tasks, despite using a smaller number of learnable parameters.
Diagnostic MDP. The T-Maze environment is used evaluate an agent's ability to learn long context dependencies in reinforcement learning scenario (Bakker) 2001). In this environment, the agent must remember cue shown only at the beginning of an episode in order to decide which way to turn al the end of a hallway (inset plot in Figure D. The cue is only included in the observation on the first timestep. The difficulty of this environment can be increased by increasing the corridor length. The agent's actions are NSEW, and the observation a binary encoding of the current cell (gray code), the cue (on the first step), and several random distractor bits. The full details are provided in Appendix G.1
We trained seven agents for five million steps in the T-Maze environment, for corridor lengths 120- 200. The network architecture for each agent has shared representation learning layer, either an RNN or a transformer, which is then followed by sepa- rate actor and critic heads. Two of these agents were trained using an RNN as the shared represen- lation layer, namely LSTM (Hochreiter & Schmid- huber, 1997) and GRU (Cho et al. 2014]. The other two agents used transformer, particularly the GTrXL architecture (Parisotto et al. 2020), and the linear transformer architecture (Katharopoulos it al.] 2020 In GTrXL, the memory size hyper- parameter, defined as the amount of stored history, controls the context length. We train two GTrXL agents, G1TXL-112 and GTYXLL-266, corresponding to memory sizes 128 and 256. Note that for the corridor lengths considered, GTrXL-256 has the en- tire episode provided as input. We also evaluate GaLiTe (17 4) and AGaLiTe (n = 4, = 1); we do so by replacing the XT~attention (Dai et al. 2019))
Success Rate
T-Maze
1.0
GTrXL-256
GaLiTe
AGaLiTe
0.8
Carrider Length
0.6
GTrXL-1228
GRU
LSTM
120
140
160
180
200
Corridor Length
Figure 1: Success rate in the last 100K timesteps av- eraged over 50 runs in T-Maze (shown inset). The shaded regions represents the standard error.
of GTrXL with one of the two approaches, while preserving the order of the layers and the gating of GTrXL. This allows us to evaluate exactly the impact of the newly introduced self-attention mechanisms without other confounders. The base RL algorithm for all agents use Advantage Aclor-Critic (A2C) (Wu et al.
2017 Architecture-specifee hyperparameters and tuning strategies are described in Appendix G.1
Figure D summarizes the main results. We report the success rate, the percentage of correct decisions, averaged over the last 100K timesteps of the experiment. An agent that chooses randomly at the intersection would achieve a success rate of 0.5. In this experiment, GTrXL sensitive the amount of history provided as input: GTyXL-128 (brown) fails for corridor lengths greater than 120, whereas GTrXL-256 (orange) works well across all corridor lengths. GaLiTe (purple) and AGaLiTe (red) match the performance of GTrXL-256 despite not having access to the entire episode as input. Note that AGaLiTe performs close to GaLiTe even with r r= (the approximation parameter). GRU (green) and lincar transformer (grey) outperform LSTM (blue), but their performance drop in the longest. corridor lengths. AGaLiTe is more computationally efficient T-Maze. For a single attention head, AGaLiTe uses roughly 125.1 times fewer operations
than GTFXLL-256 than GT-XL-256, and 36.57 times less space. Also, AGaLiTe uses roughly 62.67 times fewer operations and 18.28 times less space than GTYXL-128.
Partially Observable Classic Control. Inspired by previous work (Morad et al., 2022 Duan el al. 2016), we explored two variants partially observable CartPole (Barto et al.) 1983). In the first variant, we mask out the velocity information from the observation voctor and only allow positional information. This modification makes the problem difficult as the agent now needs to estimate these velocities itself. The second modification induede additional challenge by adding noise to the posit ional information communicated
7

---

Published in Transactions on Machine Learning Research (10/2024)
to the agent. We sampled the noise from normal distribution with zero mean and 0.1 standard deviation. This problem is qualitatively different from the T-Maze because of the different, requirements imposed by the environment. In CartPole, the agents must integrate information over time to construct a reasonable estimate of the underlying state of the MDP, whereas in T-Maze the agent must learn the cue was important and remember for a long period of time.
We used both GRU and GTrXL baselines for this prob- lem. GRU-based approaches perform best these partially observable classical control tasks, even compared to transform- ers (Morad et al., 2022). We used PPO (Schulman et al. 2017) and trained all agents for 5M steps on the two variants of Cart- pole. We also performed an extensive sweep of the hyperparam- eters of PPO and GRU, which is described Appendix G.2
Total Noisy Stateless CartPole Reward
AGaLiTe
GRU GTrXL
Figure sunmarizes the results of our experiment in the Noisy stateless CartPole, the second variant from above. The AGaLiTe agent learns faster and finds better-balancing pol- icy than the GRU and GTrXL agents. The results for the other variant. of partially observable Polt (without noise) is qualitatively similar and can be found in Appendix G.2
Env Steps
The total reward is binned over 10 timesteps and averaged over 30 different seeds + standard error.
Figure 2:
Mystery Path. In Mystery Path (Pleines et al. 2023 the
agent required to remember multiple cue signals for long periods time in a 2D pixel-based environment. In this environment, the agent's goal is to reach target position by traversing through a random invisible path. Episodes have fixed length and the agent reset back to the start location (along with a feedback observation) upon deviating from the path. We consider two configurations of this environment: MPGrid and MP; MP is more difficult. In MP, there six actions and smoother motion dynamics compared to MPGrid, with grid-like movements and four actions. MPGrid has maximum episode length of 128, while MP's is 512. Appendix G.3 describes the enviromment and the configurations considered.
Success Rate
Success Rate
MPGrid
MP
1.0
AGaLiTe-8 0.8
0.8
GTrXL-128
0.6
GTrXL-64
0.6
GTrXL-32
0.4
0.4
AGaLITe-4
Path
0.2
0.2
Agent
Start
0.0
0.0
0.5
1.0
1.5
2
4 Steps
Steps
x 10
x 10
Figure 3: Left: Learning curves MPGrid (averaged over 15 seeds + 95% bootstrapped CI) along with an inset figure showing a possible ground truth maze layout. Right: Learning curves in MP (averaged over 5 seeds +95% bootstrapped CI) along with inset figure depicting the agent's observation. The agent does not observe the path to the goal (left); red cross is shown feedback if the agent deviates off from the path, with the agent being reset to the start tile (right).
8

---

Published in Transactions on Machine Learning Research (10/2024)
We trained three GTrXL agents with memory sizes {32, 64, 128). and two AGaLiTe agents with feature map dimension 7 E {4,8), and = 1. The architecture sizes for GTrXT, and AGaLiTe were chosen to be similar to the ones used in the T-Maze experiments. PPO was the base RL agent used. We used a standard agent network architecture (c.g., Mnih et al. 2016 Schulman et al. 2017 for all agents. Details on hyperparameterss sweeps can be found in Appendix
G.3
Figure 3 summarizes the main results. Again we report success rate, the percentage of episodes the agent reaches the goal before an episode timeout, calculated over a window of one million steps. Across both configurations, MPGrid and MP, we observe that AGaLiTe matches the performance of GTrXL-128 when n 4 and surpasses GTT-XL-1128 in mcan performance when 1 = 8. Also, similar to T-Maze, we observe that reducing the memory size of GTTXL drastically impacts its performance.
We observe again that AGaLiTe is more computationally efficient than GTrXL. For a single attention head, AGaLiTe-8 uses roughly 55.75 times fewer operations than GTrXI--128, and uses 9.84 times less space. In other words, we observe performance al least as good GTrXL, in both variants of pixel-based control, fraction of the cost.
Craftax. Crafter (Hafner 2021 is a 2D open-world survival game where agent needs to forage for food and water, find shelter to sleep, defend against monsters, collect materials, and build tools. This environment is designed to evaluate an agent's ' ability to perform a wide range tasks purely from a scalar reward signal. The environment is partially observable as the agent can only observe a portion of the large randomiy gener- ated map that it navigates. Our hypothesis that an agent with better memory capabilities and longer context should achieve higher performance through navigating and utilizing the cues in the environment effectively. We consider the sym- bolic variant of the environment, Craftax symbolic, detailed in Matthews et al. 2024)). The symbolic variant simplifies the observations of original pixel-based crafter by encoding each pixel as one-hot vector and an intensity value.
Total Reward
AGaL
40
PPO+RNN
30
GT/LL-128
20
10
8.00
0.25
0.50
0.75
1.00
Steps
X 109
Figure 4: Learning curves of GT:XTL-128 and AGaLiTe the Craftax symbolic en- vironment (averaged over 15 seeds 95% bootstrapped CT). The inset plot shows the result of rendering sample observation (left) and the full map (right).
We trained GTrXL agent with memory size 128 and AGaLiTe agent with feature map dimension 1 8, and ==1 for 1B steps. The architecture sizes for GTrXL and AGaLiTe were chosen similar to the ones used in the T-Maze and Mystery Path experiments. PPO was the base RL agent used. Details on hyperparameters sweeps can be found in Appendix G.5
Figure 4 shows the total episodic reward achieved by the AGaLiTe-based agent compared with GTrXL- based agent. We observe that the AGaLiTe achieves a higher reward than both GTrXL and previously reported PPO+RNN bascline (Matthews et al. 2024). In Appendix J we compare the performance across the varions gamaerelated achievements present in Crafter. We find that AGaLiTe achieves higher scores in several of these achievements.
We also considered a 3D navigation environment called Memory Maze (Pašukonis et al., 2023) that has a fixed horizon and that requires the agent to remember multiple cue signals for long periods of time. At the beginning of each episode, new maze is generated randomly and several objects of different colors are distributed across the maze. The agent perceives 64 x 64 RGB image with colored border indicating the color of the current object of interest. Once the agent touches the object, it gets a +1 reward and the borders' colors change. The agent's 's goal is to maximize rewards within the fixed time budget. Thus, the agent must remember the objects' locations to travel through the maze as quickly as possible. Figure 5 (inset) provides an illustration of the Memory Maze environment. In the main paper, we report results on the largest maze size, 15 15, with an episode duration of 4,000 steps. Results for other maze sizes can be found in Appendix H
9

---

Published in Transactions on Machine Learning Research (10/2024)
Memory Maze. We trained GTrXL agent and Total Reward
AGaLiTe agent, each with 22M learnable parameters, for 100M steps using the Asyne-PPO algorithm (Petrenko et al.] 2020). The GTrXL agent had a memory size of 256, and the AaaiiTe agent had a Teature map 7 = and an approximation hyper- parameter = 7. We based our architectures for both the policy and the critic on the work by Petrenko et al. (2020). In this work, ResNet (He et al. 2016) is used to extract the features from the input image, then a sequence of features are fed into an RNN or a transformer. We detail the hyperparam- eters used, the archilecture sizes, and the luning strategy in Appendix G.4
MemoryMaze 15x15
25
GTxXL-256
AGaLITe
15
10
0.0
0.2
0.4
0.6
0.8
1.0
Figure shows the total opisodic reward achieved by our (GaLiTe-basee agent compared with GTrXL-based agent. The asymptotic performance of all the three agents is similar, but the GTrXL-based agent cxhibits faster learning carly on. Additionally, we explore the impact of smaller GTTXL memory sizes in Appendix D We find that reducing the memory size of GTrXL did not affect the performance in this environment, suggesting that the agents might not be utilizing their memory capabilities effectively.
Steps
1e8
Figure Learning curves in MemoryMaze The bold lines report total episodic reward averaged over three seeds; the other lines are individual seeds. The inset plot shows a sample observation.
Finally, we looked at the agents' utilization of the computational resources. For a single attention head, AGaLiTe uses roughly 125 times fewer operations than GTrXL-256 and it uses 46 times less space. Addi- tionally, we measured the frames per second (FPS) and the memory usage from 12 AGaLiTe and GTTXL agents. Overall, AGaLiTe achioves 535.63 = 0.52 FPS while GTrXL achieves 373.63 + 0.49 FPS, correspond- ing to a 43.36% improvement. Further, AGaLiTe uses 52.37% less memory than the GTrXL agent. The number of operations and space used are asymptotic information that highlight the benefits one can expect when using even bigger neural network architectures, such as those now common in industry. The FPS rate demonstrates the performance gain when AGaLiTe is instantiated in particular network architecture.
6 Ablation Study
In this section, we present ablations for each of the three modifications proposed to the linear transformer architecture in AGaLiTe, highlighting the importance cach modification. We present the results for these ablations in Figure
Our first ablation evaluates the impact of proposed gating mechanisms in AGaLiTe. We conduct this ablation in the MPGrid environment. This environment requires agent to selectively filter and remember multiple pieces of information throughout an episode. Our proposed gating mechanism significantly outperforms an AGaLiTe without the gating mechanism in Figure 6a
The other two ablations compare the proposed feature map and approximation approach to other approaches in the literature. We conducted these ablations in the T-Maze environment with the corridor length set to 200. We use the same hyperparameter tuning strategy as described Scction G.I but focusing on the best hyperparameter on corridor length 200. We report the success rate while training an agent for 5M steps over 50 seeds.
To evaluate the impact of different feature maps o in AGaLiTe, we consider two alternatives, proposed in the existing literature. The first uses an element-wise feature map ELU + +1 (Clevert et al. 2016). which was used originally in the linear transformer architecture. The second is the deterministic parameter free projection (DPFP) introducod by Schlag et al. (2021). which was shown to outperform exisiting feature map approaches in language modelling tasks. We present these results in Figure 6b We observed that our proposed feature map outperform both of these methods.
10

---

Published in Transactions on Machine Learning Research (10/2024)
Success Rate
Success
Success Rale
Rate
1.0
0.6
0.8-
0.6
0.4
0.4
0.4
AdaU le
AGaLTe out galing
0.2
AGaLITe
AGal with DPFP
AdwliTe with Rank-1 asprex
0.0
0.0
0.0
0.5
1,0
L55
4
0
Steps
x10"
Steps
x10
Steps
x10
(a) Gating mechanism
(b) Feature map
(c) Approximation approach
Figure 6: Ablation of various components of the AGaLiTe architecture
Finally, we compare AGaLiTe's approximation to an alternative incremental low-rank approximation method. We consider the rank-1 trick introduced by Ollivier et al. (2015). The rank-1 trick approximates a Kronecker delta function using random signs drawn from a uniform distribution. Similar to our proposed approximation approach, the rank-1 trick could be applied to derive incremental updates to a low-rank decomposition of matrix. We derived an approximation using the rank-1 trick and compared it to our proposed approximation (with r = 1) in Figure Gi trick.
We observe that our proposed approximation approach outperforms the rank-1
7 Related Work
Similar to our approach, several previous works have explored extensions to the linear transformer architec- ture, addressing its limitations. The DeltaNet architecture (Schlag et al. 2021 implements a scalar gating mechanism to accumulate the value vectors over time, and then uses an error-correcting delta rule to update the recurrent states. In contrast, our proposed approach utilizes an element wise gating mechanism applied directly over the recurrent states. DeltaNet also introduces deterministic feature map called DPFP have found to perform worse compared to our proposed feature map in Figure G RecurrentDeltaNce ar- chitecture [rie et al. 2021 proposed improvements to the DeltaNet architecture by introducing additional recurrence and non- non-linearit, to the updates of the key, value, and query vectors. However, the non-linear update rule in RecurrentDeltaNNtt limits its parallelizability over an input sequence. In contrast, our pro- posed approach uses linear update rules, and could be parallelized over an input sequence (Appendix FT] In concurrent and independent work, Aksenov al. (2024) explored a learnable kernel function inspired from Taylor expansion of exponential functions, demonstrating impressive in-context learning performance. Notably, none of these work improve upon the computational complexity or introduce approximations to recurrence mechanism of the linear transformer.
Gating mechanisms such as the one we used in GaLiTe and AGaLiTTe are commonly used in RNNs to control the flow of information and mitigate the impact vanishing gradients (Hochreiter Schmidhuber 1997)
Often, scalar gating mechanisms have been applied, such as in the linear transformer (Peng et al. 2021
However, using a single learned cocfficient could be sub-optimal as does not allow for a more fine-grained control of the flow of past. information from each index location in a recurrent stale.
The choice of the feature map ó can have significant impact on the overall performance (Schlag et al, 2021). For example, a non-expansive map based on ELU+1 can be used (Katharopoulos et al. (020). however, element-wise activation functions are limited in their ability to learn complex non-linear rela- tionships and using them as a feature map limits the memory capacity of the architecture (Schlag et al. 2021 Alternatively, random feature maps can be used approximate a softmax function (Peng et al.
2021 Choromanski al. (2021). Although randomized feature maps are equivalent to softmax function expectation, they introduce additional variance. Our model is deterministic.
11

---

Published in Transactions on Machine Learning Research (10/2024)
In the context of AGaLiTe, there are other incremental approaches approximating large matrices. In- cremental singular value decomposition (SVD) (Brand] 2002 2006) provides way to perform additive modifications to a low-rank singular value decomposition of a matrix. Previous applications of incremen- tal SVD in RL, however, suggest that sensitivity to the rank parameter is a significant issuc (Pan et al., 2017 The rank-1 approximation introduced by Ollivier et al. (2015) uses random numbers to approximate a Kronecker delta function producing an unbiased approximation of a matrix represented as sum of outer products. The use of random numbers, however, introduces variance in the approximation (Cooijmans & Martens 2019); our results in the T-Maze suggest the proposed approximation leads to better results than the rank-1 approxination.
Various approaches have been proposed that increase the context length of the transformers or reduce the computational complexity of self-attention. The LongFormer architecture (Beltagy et al. 2020 selectively calculates sparse attention matrix over a large context and the Transformer-XI, architecture (Dai et al. 2019) caches previous activations to attend over larger context. Alternatively, the RMT architecture
Bulatov al. 2022) introduced segment-level recurrence to pass global information over larger context.
Other approaches reduce the computational complexity of self-attention by proposing approximations (Kiteev
ct al.
2020
Choromanski et al. 2021 Wang et al. 2020).
2023 ), and state-space models such as LRU (Orvieto et al. 2023 (2023) offer context-independent inference cost while leveraging parallelization over a sequence. The $4 architecture was recently applied to offline in-context RL (Lu et al. 2024) and model-based RL (Samsami et al. 2024), demonstrating superiority over RNNs. However, none of these approaches have yet been explored in the model-free RL setting, which has been the main focus of this
Other methods such as RWKV (Peng et al.
S4 (Gu et al. 2021), and $5 Smith et al.
paper.
Several works have explored using transformers in RL. Parisotto & Salakhutdinov (2021) used transform- ers to learn policies in an asynchronous setting relying on policy distillation to make interaction with the environment feasible. Others have explored transformers in model-based, fully-observable RL, such as the TransDreamer architecture which replaces the GRU used inside Dreamer V2 (Hafner et al. 2020 with a transformer (Chen et al. 2022..
While focus of this paper has been towards online RL settings, several previous works have applied transform- ers to offline RL and in-context RT. Chen et al.
(0021) demonstrated that transformers could learn single-task
policies from offline RL data through imitation learning. Extensions to Chen et al. (022)) demonstated that
transformers could also derive multi-task policies from offline RL data in both same-domain (Lee et al., 2022)
and cross-domain environments (Reed et al.) 2022). Laskin et al. (2023) and Lin et al. (2024) demonst that transformers could be trained on offline datasets to learn RL policies in-context on novel tasks. Offline training of transformers can also be combined with online RL and has demonstrated application in discov- ering high value molecules (Ghugare et al. 2024). We believe that proposed approach could be also be useful across offline and in-context RI. settings by offering larger context with reduced computational requirements.
8 Conclusion and Future Work
Transformers have revolutionized many branches of AI research, but their computational requirements make extension to other domains such as online RI. difficult. In this paper, we have introduced two recurrent alternatives of the sclf-attention mechanism in transformers, called Gated Lincar Transformer (GaLiTe) and Approximate Gated Linear Transformer (AGaLiTe). We demonstrate the efficacy of both approaches in several partially observable reinforcement learning tasks (e.g., T-Maze, Mystery Path, Craftax, Memory Maze). When compared to a stato-of-the-art architecture GTrXL, the inference cost of our approach is more than 40% cheaper while reducing memory use more than 50%.
Future work could explore algorithmic improvements to AGaLiTe such as using updates based on efficient real-time recurrent learning (Zucchet et al.] 22023] Williams & Zipser 1989). Furthermore, the application of AGaLiTe to model-basee RL algorithms such as the Dreamer V3 (Hafner al. (023)) could be exciting, Finally, Morad et al. 2024)) leverages the properties of linear transformer approaches to propose a batching
12

---

Published in Transactions on Machine Learning Research (10/2024)
method that improves sample efficiency, increases the return, and simplifies the implementation of recurrent loss functions in RL. Application of some of these ideas to AGaLiTe would be exciting, In addition, previous work has found RNN-based approaches are best in some tasks and transformers better in others. There is much to be understood empirically partially observable RL.
Acknowledgements
We would like to thank Martha White, Dale Schnurmans, and Michael Bowling for providing valuable feedback and for their helpful discussions. We would like to thank Martha White for also providing access additional computational resources. We would like to thank Vincent Liu for providing feodback on the derivations presented in this paper. The research is supported in part by the Natural Sciences and Engineering Research Council of Canada (NSERC), the Canada CIFAR AI Chair Program, the University of Alberta, Google Cloud Incubator, TPU Research Cloud Program, and the Digital Research Alliance of Canada.
References
Yaroslav Aksenov, Nikita Balagansky, Sofia Lo Cicero Vaina, Boris Shaposhnikov, Alexey Gorbatovski, and Daniil Gavrilov. Linear transformers with learnable kernel functions are better in-context models. Association for Computational Linguistics, 2024.
Bakker. Reinforcement learning with long short-term memory. In Newral Information Processing Systems, 2001.
Andrew G Barto, Sutton, and Charles W Anderson. Neuronlike adaptive elements that solve difficult learning control problems. IEEE Transactions Systems, Man, and Cybernetics, 1983.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.055: 051100 2020.
Blelloch. Prefix sums and their applications. Technical report, 1990.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adamn Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: compos- able transformations of Python-NumPy programs, 2018.
Brand. Fast low-rank modifications of the thin singular value decomposition. Linear Algebra and its Appli- cations, 2006.
Matthew Brand. Incremental singular value decomposition of uncertain data with missing values. In European Conference Computer Vision, 2002.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Nee- lakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krucger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are fow-shot learners. In Neural Information Processing Systems, 2020.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. In Neural Information Processing Systems, 2022.
Changgu Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. orXiv:2202.09482 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Leo, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforeement learning via sequence modeling. In Neural Information Processing Systems, 2021.
13

---

Published in Transactions on Machine Learning Research (10/2024)
Kyunghyun Cho, Bart Van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statis- tical machine translation. 4r6.v:1406.1078, 2014.
Krzysztof Marcin Choromanski, Valeríi Tikhosherstov, David Dohan, Xingyou Song, Andreca Gane, Tamas Sarlos, Peter Hawkins, Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Represenlations, 2021.
Ajork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (ELUs). In International Conference on Learning Representations, 2016.
Tim Cooijmans and James Martens. On the variance of unbiased online recurrent optimization. arXiv:1902.09// 02400, 2019.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer- XL: Attentive language models beyond a fixed-length context. arXiv:1901.02966, 2019.
Devlin, Chang, Lee, and Toutanova, BERT: Pre-training of deep bidirectional transformers for language understanding, Xrxi::1810.00 04805, 2018.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement learning for continuous control. In International Conference Machine Learning, 2016.
DI engine Contributors. DI-engine: OpenDILab decision intelligence engine. https://github.. com/ opendilab/DI-enginn 2021.
Yuan Gao and Dorota Glowacka. Deep gate recurrent neural network. In Asian Conference on Machine Learning, 2016.
Raj Ghugare, Santiago Miret, Adriana Hugessen, Mariano Phielipp, and Glen Berseth. Searching for high- value molecules using reinforcement. learning and transformers. In International Conference on Learning Representations, 2024.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021.
Danijar Hafner. Benchmarking the spectrum of agent capabilities. arXiv preprint arXin:2109006780 2021.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world modcls. In International Conference on Learning Representations, 2020.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world modcls. arXiv:2301.00 04104, 2023.
Hausknecht and Stone. Deep recurrent Q-learning for partially observable MDPs. In Association for the Advancement of Artificial Intelligence Fall Symposia, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Computer Vision and Pattern Recognition, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997.
Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber. Going beyond lincar transformers with recurrent fast weight programmers. Neural information processing systems, 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Flcuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, 2020.
Khandelwal, He, Qi, and Jurafsky. Sharp nearby, fuzzy far away: How neural language models use context.
In Association for Compudational Linguistics, 2018.
14

---

Published in Transactions on Machine Learning Research (10/2024)
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv:2001.04451, 2020.
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, maxime gazeau, Himanshu Sahni, Satinder Singh, and Volodymyr Mnih. In-context reinforcement learning with algorithm distillation. In International Conference on Learning Representations, 2023.
Kuang-Huel Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman, Sergio Guadarrama, lan Fischer, Winnie Xu, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. In Neural Information Processing Systems, 2022.
Licong Lin, Yu Bai, and Song Mei. Transformers as decision makers: Provable in-context reinforcement learning via supervised pretraining. In International Conference Learning Representations, 2024.
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, and Feryal Bchbahani. Structured state space models for in-context reinforcement learning. Advances in Neural Information Processing Systems, 2024.
Michael Matthews, Michael Beukman, Benjamin Ellis, Mikayel Samvelyan, Matthow Jackson, Samuel Cow- ard, and Jakob Foerster. Craftax: lightning-fasi benchmark for open-ended reinforcement learning. International Conference on Machine Learning, 2024.
Andrew Kachites McCallum. Learning to use selective attention and short-term memory in sequential tasks. In International Conference on Simulation of Adaptive Behavior, 1996.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Neural Information Processing Systems, 2019.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforeement learning. In Inter- national Conference on Machine Learning, 2016.
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym: Bench- marking partially observable reinforcement learning. In International Conference on Learning Represen- tations, 2022.
Steven Morad, Chris Lu, Ryan Kortvelesy, Stephan Liwicki, Jakob Foerster, and Amanda Prorok. Revisiting recurrent reinforcement learning with memory monoids. arXiv, 2024.
Yann Ollivier, Corentin Tallec, and Guillaume Charpiat. Training recurrent networks online without back- tracking, orXiv:1507.00760, 2015.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, 2023.
Yangchen Pan, Adam White, and Martha White. Accelerated gradient temporal difference learning, In Association for the Advancement of Artificial Intelligence, 2017.
Emilio Parisotto and Ruslan Salakhutdinov. Efficient transformers in reinforcement learning using actor- learner distillation. In International Conference Learning Representations, 2021.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascamu, 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 International Conference Machine Learning, 2020.
Jurgis Pašukonis, Timothy Lillicrap, and Danijar Hafner. Evaluating long-term memory in 3D mazes. International Conference on Learning Representations, 2023.
15

---

Published in Transactions on Machine Learning Research (10/2024)
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, and Rui-Jie Zhu. Rwkv: Reinventing RNNs for the transformer era. arXiv:2305.1300 2023.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations, 2021.
Olivier Petit, Nicolas Thome, Clement Rambour, Loic Themyr, Toby Collins, and Lue Soler. U-Net trans- former: Self and cross attention for medical image segmentation. In International Workshop on Machine Learning in Medical Imaging, 2021.
Aleksei Petrenko, Zhehui Huang, Tushar Kumar, Gaurav Sukhatme, and Vladlen Koltun. Sample factory: Egocentric 3D control from pixels at 100000 FPS with asynchronous reinforcement learning. In Interna- tional Conference on Machine Learning, 2020.
Marco Pleines, Matthias Pallasch, Frank Zimmer, and Mike Preuss. Memory Gym: Partially observable challenges to memory-based agents. In International Conference on Learning Represenlations, 2023.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel Barth- maron, Mai Giménez, Yuyy Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. Transactions on Machine Learning Research, 2022.
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar. Mastering memory (asks with world models. The Twelfth International Conference on Learning Representalions, 2024.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. In International Conference on Learning Representations, 2014.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight program- mers. In International Conference on Machine Learning, 2021.
Schulman, Wolski, Dhariwal, Radford, and Klimov. Proximal policy optimization algorithms.
2017.
John Schulman, Sergey Levine, Picter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning, 2015.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In International Conference on Learning Representations, 2023.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference on Learning Representations, 2021.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: A unified understanding of transformer's attention via the lens of kernel. In Conference on Empirical Methods in Natural Language Processing and the International Joint Conference on Natural Language Processing, 2019.
Ashish Vaswani, Noam 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.
Oriol Vinyals, Timo Ewalds, Sergey Bartunov, Petko Georgiev, Alexander Sasha Vezhnevets, Michelle Yoo, Alireza Makhzani, Heinrich Küttler, John Agapiou, Julian Schrittwieser, et al. Starcraft ii: A new challenge for reinforcement learning. arxiv preprint Xiv: 708. 04 782, 2017.
16

---

Published in Transactions on Machine Learning Research (10/2024)
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. Xiviv preprinl or arXin:2006.00 04 768, 2020.
Eric W Weisstein. Fourier series. From MathWorld-A A Wolfram Web Resource.
Ronald J Williams and David Zipser. A learning algorithm for continually rumning fully recurrent neural networks. Newral Compulation, 1989.
Yuhuai Wu, Elman Mansimov, Roger B Grosse, Shun Liao, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation. In Neural Information Processing Systems, 2017.
Huasong Zhong, Jingyuan Chen, Chen Shen, Hanwang Zhang, Jianqiang Huang, and Xian-Sheng Hua. Self- adaptive neural module transformer for visual question answering. IEEE Transcttonns Multimedia, 2020.
Nicolas Zucchet, Robert Meier, Simon Schug, Asier Mujika, and João Sacramento. Online learning of long- range dependencies. Advances in Neurol Information Processing Systems, 2023.
17

---

Published in Transactions on Machine Learning Research (10/2024)
A Gated Linear Transformers (GaLiTe)
Algorithm 3 formalizes the self-attemtion mechanism introduced in GaLiTe. The algorithm introduces a hyperparameter, 7, and a few learnable parameters, Wa, W E [dddd and Wp W Wp, Rdxy, The hyperparameter 7) controls the size of the recurrent states, C and St, and of the key and the query vectors.
Algorithm 3 Gated Linear Transformer (GaLiTe) Self- Delf-Attention
Input: E Rd Ct-1 E xrds, St-1 Rnd
Hyperparameters:
W Parameters: WK, Wo, Wv, WB, W E Rdaxd and Wp P1o W P2o Pa E Rixd
if then
So f- 0, Co 0.
end if
{Calculate Key}
4: K, + f(rehu W p. xx) @ relu( (x)))
{Calculate Query}
q: + f(rehul W x, & relu relu(Waxi)) WQx:))
{Calculate Value}
6: V, + WyX,
[Generate Gating Vectors}
B1 - 0gl (W
7 Ao (W p. x1) & (W,x,))))
{Update Memory]
9: C. + (1- - (1- x @C.1 (3.ov,) @ (nok.)
10: St + (1- 70
[Calculate Attention Vector)
11: a (C.qi)/(sqqt)
Output: a, E Edd, C, E Rd xndn,
8, E Rod.
B Derivation of AGaLiTe
In this section, we walk through the derivations to approximate the GaLiTe self-attention mechanism. We first start with deriving approximation for the Kronecker Delta Function and then use these approximation results to derive the AGaLiTe's self-attention mechanism.
B.1 Approximation of Kronecker Delta Function
In this scction we derive an approximation of the Kronecker delta function. The Kronecker delta. function, 8mm is defined for integers and as if m and 0 if m & n.
We use a trigonometric identity that is used in computing Fourier series by relating the Kronecker delta function to an integral of a product of two cosine functions (Weisstein). is given by:
2m //²
1
8mn
cos(mz) cos(nx) dx.
(16)
18

---

Published in Transactions on Machine Learning Research (10/2024)
We the Trapezoidal rule to approximate the integral in Equation 16 The trapezoidal rule is numerical integration method that approxinates the integral of function by dividing the interval into sub-intervals and approximating the function in each sub-interval with a straight line connecting the endpoints. For a function f(x) that is integrable on the interval [a, b), it is given by:
S: f(x) (()) de 2
f(zk- 1) + f(xx) AT,
(17)
2
where Ar b-a r
3a a + kår, and r the number of sub-intervals used for the integral and it controls the degree of approximation. As + 00, the approximation becomes exact. Let 8mm be the Trapezoidal approximation of the integral defined in Equation 16
We then write 8mn as:
mm Σ
2Ti m
2 n
2 m
2T n
cos
cos
+
cos
COs
(18)
mm
Further, in the limit we have:
lom 8mn
8mn-
Next, we will simplify the above equation to combinc the two summations above into a single one:
8 mn = _
2T m cos
2 n
2ii m
2
cos
cos
cos
==
1=1
Adding and subtracting (cos(0) cos(O) (25mm cos( (2rr))):
2mi m
2Ti (mmin
r nn) ))
8mn
COS
cos
cos c2s(2 cos (2
2Ti m
2mi m ) cos(O) cos
cos
cos
r
(cos (0) cos (0) + cos (2mm) cos (2mm))
2mt m
2Tr n ) cos
22T m
22r
cos
cos
r
22t m
2Ti m )) cos(() cos
cos
(cos(0) cos(O) + cos (2Tm) cos(2TT))
2mi
2Tí m1 ))
cos
m
cos
(cos(0) cos(O) + cos(2mm) cos (osT))
r
i=0
Since and are integers:
2
2 m
2Ti n1 ))
2
8
Σ
cos
cos
(19)
1=0
In fact, in the limit, - oc, only the first term in the right hand side of Equation 19 matters in our approximation. Let
2
2mi m
2 11 ))
mn
cos
co8
(20)
i=0
19

---

Published in Transactions on Machine Learning Research (10/2024)
then have
2 --000 r
lim 7-00
8mn
lim 8mn lim
lim 6mn 0, 1-000
(21)
5-000
and consequently
lim 8m 8mn-
(22)
B.1.1 Using The Kronecker Delta Function to approximate GaLiTe
We start with the GaLiTe recurrent state update which we will then approximate using the Kronecker delta approximation introduced above. GaLiTe recurrent state update is expressed as follows:
C, - (1-))) o C--1 (p: (B) (
(23)
We use the approximation of the Kronecker delta function in Equation 20 to approximate the update in Equation 23 We will start by representing the recurrent state C₁ as sum of outer products. We will do by recursively expanding using the definition of C. Following the recursive definition have:
C. = ((1 - 34) -))) 2 C (3:00 @ v.) s (z))
ki)
=
+ (1- - @ @ ( (7-
(1- 2 k,-1) ((1 3.-1) @ (1- 2-1)) o CC-2)
= (18.0v)000000 + (1-))) (1- (1- 3 (Br- 0 (7-1 2 kk-1))
+ (11 @ (1-8--))
x- (11--)))
3 C,-2
Using the fact that (a & b) o (c (a o (bod) d) for arbitrary vectors a, b, c, d, have:
C. = (z0 + ((11 ((11 o @ (00-)))(1111)) (01-7)00-1 (k-1))
+ ((11-B))(1-B))) (1- @ ((1 ((1- ((-7)) )-1))) 2 C²-2
Recursively expanding it further and regrouping the terms similar to above:
= =( + ((L ((1-3) 3,) o G @v.-) ((1 101-711 7)) 2- ok-i) + ((1-3)) 3-1) 0 8 2 v ((1 (1- 7) o (1 7-1) @ n- 0k + C
Next, rewrite the recursive expansion single sunmation term of outer product. This is possible because upon following the recursive definition, we can see above that each term of the expansion could be written as an outer product of vectors. The vectors used to calculate the outer product at each timestep consists of element-wise vector product involving either the value or key, and their corresponding gating vectors. We show the expansion for only until 3, but the same steps could be followed until We define product operation II such that: I (a,) = if j i, and [T, = a, @ a .... ay. We introduce variables I; and m; represent. the left and right terms of the outer-product at given timestep.
20

---

Published in Transactions on Machine Learning Research (10/2024)
for 0, .t.
t. We then define C₁ in terms of these variables to simplify the equation above:
C Σ
m f-0
(24)
1 II (1 (1-3,)
(25)
j=++1
m,
II j=++1
01-k)
(26)
Next, we use the approximate Kronecker delta function in Equation 20 approximate the sum of outer products in Equation 24 We start by introducing the Kronecker delta function 6mn into the expression of C above by introducing double summation:
C
Σ 1=0
m
8ig 09 my
j=0
Then we replace the Replacing di, with رn we obtain an approximation C. of C2 as follows:
= C = 8. & mj
1= 1=0
2
2Tk
2nk [
0os
cos
Ø
my
1=0
2mk
2ntk
cos
my,
k= 1=0 ==0
where we first applied Equation 20 followed by rearranging the order of the summations. Let we = cos (2 then have:
2
C
cos (whz) cos (wkj) 1, & my.
k= 0=0
Because (c @ d) (ac) @ (bd) for scalars a, b and vectors c, d, this allows us to seperate the two cosine terms in the summation:
C
(cos (,) & (cos (waj) m,)
k=
Next, we will use the distributive property: (E a.) 3 (E = ET Ey @ by), where ag...., am and bo...., b, are arbitrary vectors. We will use to rewrite the above equation a single summation of outer products, where the summation is defined only over r. This is a key operation that allows us write the approximation to be written temporally iterative manner. Applying the distributive property from right
21

---

Published in Transactions on Machine Learning Research (10/2024)
to left into the equation above, and then using Equations 25 and 26 have:
C
(cos (wx+) 1,) @ (cos (wxj)) mj)-
k= == ==0
2
cos (wx+)1
& Σ cos (wx)) m:
Σ cos (wxi)
II (1 - 2
X cos (wx:t)
II (1- (1- 5=++1
(27)
-=++1
Next, we simplify the above equation and rewrite in a recurrent form. Let v and K2 be defined as:
cos (wxi) II (1 - 3,) oRo..
(28)
j=i+1
k
cos (wxi) II (1-2))
(29)
1=i+1
C, could then then written in terms of v and k as:
C c=
(30)
We regroup the terms and derive recursive relationship of ve and kt with respect to v and kg First, we separate the term in the summation when i i= t. Next, we factorize (1 - B)) from the summation term. Finally, following the definition in Equation 28 we replace the second term with vi-
V
cos (ossw
II
(1 (1l-))BOvv
1=0
$=1+1
cos(wit)33 V.
cos(wki)
II (1 8,),
j=it1
v, (1- Σ cos(wRi)
II (1 - (1-8))80...
1=0
j=++1
= cos(wn() 3 v, (1 3,) 1
(31)
Similarly, we apply the same operations to kj:
k
cos (wkz)
II (1-2))
0=
j=
H-1 cos(wie)
= cos(wet)))
@ k;
II (1 (1-2))omkk
1=0
j=++1
1-1
cos(wnt) k; + (1 7t) Σ cos(wRi)
II (1 75) 2 07
cos(wit) 3 k; + (1 7x)
k
(32)
22

---

Published in Transactions on Machine Learning Research (10/2024)
Using the recursive relationships in Equation 31 and 32 given we maintain recurrent states v and k for k = 0, 1, .. For we 2
can now present the final approximation. For
", and assuming vi
and kj are initialized as zeros, the recurrent updates to v and kž and further the approximation to C, are given by:
2 - Σ
C. Č
v
K
(33)
r km0
where, for = 0,1,2,.... . we have:
v cos(wxt)3 @ Ve + (1 (1- 36) ④ v k = cos(wal) ② k, + (1 - 71) @ k
(34)
(35)
Since lim 8mn = 6mn it follows that lim C, = Cr. Unlike Equation 23 Equation 34 and 35 define a recurrence over vectors instead of matrices, and if r < d, the recurrence is much more efficient in space than the recurrence in Equation 23 We leave to future work formally derive the approximation error. In Section we show the approximation error with a synthetic error under different values of r.
Lastly, since the current state C could be represented sum of outer products in non-recurrent manner, we can avoid explicitly calculating C₁ and instead calculate the attention output a as follows:
= E v/ ((k))
T qt)
(36)
2r(s, qt)
23

---

Published in Transactions on Machine Learning Research (10/2024)
C Approximate Gated Linear Transformer (AGaLiTe)
Algorithm shows the Approximate Gated Linear Transformer (AGaLiTe). We highlight changes from Algorithm 3 in blue. The algorithm maintains a set of vectors kp k k,_ E Rud,, Vi-1.-., vi-1 IRd,, and S,-1 E Rnd, as the recurrent state at given time-step controlled by modifying the hyperparameter r, which should ideally be set to small value. The key, query, and value vectors are calculated similarly to GaLiTe. The recurrent state update is modified to use the approximation in Equation 33 At each time step, the recurrent vectors are updated using element-wise vector multiplication and addition operations (lines 10-14). The operation on each recurrent vector could be exccuted parallel. The attention output calculated without ever explicitly calculating Č, (lines 16-18).
The number of vectors stored could be
Algorithm 4 Approximate Gated Linear Transformer (AGaLiTe) Self-Attention (Streaming Data)
Input: x E R2d kp-1.-.,,k;; k E Rudd, Vi-i--. Vi-1 E Rdh,, and S;-1 E RRdda
Hyperparameters: 7) and r.
Paramcters: WK. Wq, Wv, Wa, B, W. E Rdxx xd and Wp. W W Pg Ryxd
1: Assune So + 0, Co +.
{Calculate Key}
k, + -reell W pi xx) & relu W xx,))
{Calculate Query]
q. + [rrelu W x1) & relu W (xx)))
{Calculate Value]
4: v. + Wvx.
[Generate Gating Vectors}
5: B - 0g ( (x)
6: 7t + flo, W Ps xt) @ 0 (W,,,,W,,,, (W))
{Update Memory]
for i 0 to in parallel, do
wi (2mi)/r
V + V/-1 @1 (1- (p) + cos (wi() (Be O ve)
10:
k k (1 -)) + cos (wit) (rt @ k,)
11: end for
12: + (1 71)) 8:-1 + 7 ④ k.
{Calculate Attention Vector}
13: a E v [k] qi)
14: b 2r qi)
15: a, a/b
Output: a, E Rd, kf..... k E Rad, vl...., Vo. E Rdh, and s, Red,
D Results in Long Range Arena
We additionally evaluate on the ListOps and Text (IMDB) from the Long Range Arena TTy et al. 2021 to evaluate the architecture's ability to learn long-range dependencies in a supervised learning scenario. The performance AGaLiTe (1 = 8, = 1) is compared with the transformer's and linear transformer's in Table The exact hyperparameters for AGaLiTe are listed Table 4 We found that AGaLiTe outperforms the TTy et al., 2021]] in both
previously reported results of the transformer and linear transformer architecture of these tasks.
24

---

Published in Transactions on Machine Learning Research (10/2024)
Table 2: Results in Long Range Arena. Reported as mean = std error over seeds.
(a) ListOps Task
(b) Text (IMDB) Task
Model
Params
Score
Model
Params
Scorc
Linear Transformer
8.9M
16.13
Linear Transformer
3.4M
65.9
Transformer
8.9M
36.37
Transformer
3.4M
64.27
AGaLiTe
1.7M
39.33 + 0.34
AGaLiTe
1.7M
79.83 1.3
E
Effect of r on the Quality of Approximation in AGaLiTe
We empirically evaluate the effect of on the quality of the approximation of the current state matrix C.- Idcally, we want to set to small value as the space complexity of AGaLiTe is directly proportional to r. We consider synthetic example where the value and key kt at each time step are sampled randomly from a normal distribution. We set the embedding dimension d to 128 and randomly sample values and keys for 100 timesteps. Instead using vectors "/. and Br for gating at every timestep, we use a constant value c. We then compare the difference between the current state matrix C, computed using the exact method in Equation 5 with the current state matrix C, computed using the approximate method in Equation 33 at the 100th time-step. We use the Frobenius norm to measure the difference between the two matrices. We repeat the experiment for different values of and c. For each configuration, we report the mean error across 50 independent runs. Figure i shows the results of this experiment. We observe that the error in approximation decreases with increasing value of r. For most values of r and the approximation error is low. This useful since it allows us to set to a small value, thereby reducing the space complexity of the model. In fact, in the largest experiments described in this thesis, we set r 7. Interestingly, we observe periodic bands in the error plot. It is possible that this due to the periodicity of the cosine functions used in the attention mechanism. We leave further exploration around the thcoretical nature of the error in approximation for future work.
1.0
0.8
-300
Gating value (c)
0.6
200
0.4
100
0.2
0.0
24 Approximation hyperparameter (r)
8
16
32
Figure 7: Error in approximating the current state C₁ for different values and gating at t = 100 for randomly sampled values and keys.
25

---

Published in Transactions on Machine Learning Research (10/2024)
F Parallelization over an Input Sequence
Transformers are naturally designed for parallclism over a sequence of input data, the self-attention oper- ation does not have dependencies between different parts of the input sequence. It is essential to consider the parallelizaiilty of transformer architectures, when the input sequence presented in a batched fashion. Such a scenario is common in practice, as most existing actor-critic approaches such as PPO and A2C (Schulman ct al. 2017 Mnih al. 2016) estimate gradient updates to the actor and critic using batches of trajectories collected through aggnl-ennironment interactions. Furthermore, most modern hardware accelerators, such GPUs and TPUs, excel in handling parallclizable algorithms, and parallelization is vital for effectively training large models.
Extension of Algorithm B and to accommodate parallelization over a sequence of imputs is straightforward, depending on whether the computation has dependencics on the previous state or not. The majority of the computations both algorithms, which involve calculating keys, queries, values, gating vectors, and the attention vector, do not depend on the previous state and can be parallelized over the sequence. The only part of the algorithm that depends the previous state is the update of the current state. In Algorithm 3 this is done from lines 13-14, and in Algorithm from lines 10-15. The update of the current state in both algorithms is implemented as first order recurrence. This operation is parallelizable as such recurrences could be expressed as an associtiave binary operations (see Blelloch 1990). In our implementation, we used the associative_scom scon operation in Jax to parallelize GaLiTe and AGaLiTe over an input sequence.
26

---

Published in Transactions on Machine Learning Research (10/2024)
G Additional Experiment Details
G.1 T-Maze
Environment Description: The T-Maze environment considered in this paper is similar to the one pro- posed by Bakker 2. Figure ( shows two possible episodes in the T-Maze environment. At each timestep, the agent receives a 16-bit binary observation. The first two bits correspond to the cue signal which is either 01 or 10 at the first timestep of an episode, depending on whether the reward is located at the left or right turn at the intersoction, respoctively. The cue bits are zero all other timesteps. We consider the largest possible corridor length as 200. To encode the corridor information, the agent additionally receives 8-bit gray code encoding of its current location. The gray code encoding is zero at the beginning of an episode and is updated at cach timestep. To make the problem more challenging, we added 6 noisy distractor bits to the observation. The distractor bits are sampled uniformly at random at each timestep. The agent can take one of the four possible discrete actions at each timestep: up, down, left, or right. The agent receives a reward of -0.1 at each non-terminal timestep. At termination, the agent receives a reward of +4 for taking the correct turn and a reward of for taking an incorrect turn. The reward of +4 chosen to encourage the agent to take the correct turn at the intersection. The difficulty of this environment can be increased by increasing the corridor length. Increasing the corridor length requires the agent to remember the signal for longer number of timesteps. Since the agent's 's observations include distractor bits, the agent also needs to learn to ignore the distractor bits and focus on the cue signal.
+4
-1
-1
+4
Corridor Length
01
10
Episode A
Episode
Figure 8: The T-Maze environment. The agent has to remember a binary cue (denoted by green text), shown only at the beginning of the episode, order to take the correct turn at the intersection and receive positive reward. The figure shows two possible episodes and the optimal path an agent must take. The agent's current location is provided as gray code encoding in the observation, along with distractor signals. The corridor length could be varied to increase the difficulty of the problem.
Hyperparameters and Tuning Strategy: We include the architecture configuration for each of the 5 architectures in Table Our hyperparameter tuning strategy is as follows: We train 5 seeds per architecture for each corridor length in 120-200 and hyperparameter configuration for 5M steps. We identify the best hyperparameter configuration according to the best mean success rate in the last 100K steps across all corridor lengths.
A few additional dctails worth reporting for the purposes of reproducibility. We conducted all experiments using Python and implemented the agents using the Jax library (Bradbury el al. 2018)). We used the GTrXL implementation from the DIEngine library (engine Contributors] 2021). Each agent is trained using 16-core machine with 12GB RAM. The network weights are initialized using orthogonal initialization (Saxe et al. 2014)). single run using the slowest architecture takes around 20 hours to complete.
G.2 Partially Observable CartPole
Table 5 shows PPO hyperparameters used for CartPole experiments. We show additional results for the partially observable Cart environment when noise is added to the observation vector in figure
27

---

Published in Transactions on Machine Learning Research (10/2024)
Table 3: Hyperparameters and sweeps for the T-Maze experiments.
Hyperparameter
Value
Learning Rate
[0.001, 0.0001 0.0005, 0.00001, 0.00005)
Discount Factor (7)
0.99
Adv Estimation Coellicient. (A)
0.95
Entropy Coefficient
[0.1, 0.01, 0.001, 0.0001, 0.00001]
Value Loss Coefficient
0.5
Rollout Len
256
Num of Envs
8
Batch Size (Rollout Len x Num of Envs)
2048
Actor Layer Dimension
128
Critic Layer Dimension
128
Table 4: Architecture configuration for LSTM, GRU, GTrXL, GaLiTe, and AGaLiTe for T-Maze, Mystery-
Path experiments, and Craftax experiments.
Hyperparameter
LSTM
GRU
GTrXL
GaLiTe
AGaLiTe
Embedding Dimension (d)
600
680
128
128
128
Hidden Dimension
1200
1360
N/A
N/A
N/A
Num Heads
N/A
N/A
4
4
4
Head Dim (dr)
N/A
N/A
64
64
64
Num Layers (L)
1
4
4
4
Memory Size (M)
N/A
N/A
[128, 256]
N/A
N/A
Projection Hyperparameter (17
N/ A
N/A
N/A
[4,8]
Approximation Hyperparameter (r)
N N/A
N/A
N/A
N/A
1
Actor Layer Dimension
128
Critic Layer Dimension
128
Table 5: Hyperparameters and sweeps for the CartPole experiments.
Hyperparameter
Value
Learning Rate
[0.01,
0.001, 0.0001, 0.00001
Discount Factor (7)
0.99
Advantage Estimation Coefficient (A)
0.9
Entropy Coefficient
0.0
Value Loss Coefficient
1.0
Rollout Len
1024
Num of Eavs
1
Batch Size (Rollout Len Num of Envs)
1024
Number of Epochs
10
PPO Clip Ration
0.2
Max Gradient Norm
0.5
G.3 Mystery Path
Environment Description: Pleines et al. (2023) introduced the Mystery Path environment part of the Memory Gym benchmark, which aimed to test agents' abilitios to memorize many events over an episode. The Mystery Path is a 7 x 7 grid enviromment with pixel-based observations. At the beginning of each episode, the start position of the agent, the origin, is sampled from the grid's borders. Then, the target position sampled from the grid's borders on the opposite side of the origin. randomly generated path then connects both the origin and the goal. Figure 10p. shows an example of generated origin, goal, and path. The agent's observation, shown in Figure 10
is 64 64 RGB image containing the origin,
28

---

Published in Transactions on Machine Learning Research (10/2024)
Total Reward
Stateless CartPole (Without Noise)
AGaLiTe
400
GRU
300
200
GTrXL
100
2
Env Steps
1e6
Figure 9: Non-noisy Partially Observable CartPole
the target, and the agent. The agent gets a +1 reward when it reaches the goal and a 0.1 reward when visiting a new tile on the path to the goal. If the agent falls off the path, as in Figure [0p, a red cross appears as visual feedback, and the agent returns to the origin. The reward zero in all other timesteps. We consider two variants of this environment, MPGrid and MP. MPGrid has maximum episode length of 128, uses grid-like movements and possible actions (left, right, up and down). On the other hand, MP has maximum episode length of 512, has smoother movements, and larger action space that allows diagonal movements.
Path Goal
Feedback
Agent
Start
(b)
(a)
(c)
Figure 10: visualization of the Mystery Path environment.
Hyperparameters and Tuning Strategy: The architecture sizes used for Mystery Path experiments are kept same as in Table 4 however, we used actor and critic layer dimension of 256. We detail the hyperparam- eters used for the PPO algorithm that used for training the agents in the Mystery Path environment, in Table g We tune learning rate and entropy coefficient for the sweeps mentioned in Table G Our hyperparameter tuning stratogy is as follows: we train 3 seeds per architccture for each the hyperparameter configuration for 60M steps in the Mystery Path Grid environment. Finally, we identify the best hyperparameter configuration according to the best episodic reward in the last 1M training steps.
29

---

Published in Transactions on Machine Learning Research (10/2024)
Table Hyperparameters and sweeps for Mystery Path experiments.
Hyperparameter
Value
Learning Rate
00.0025, 0.00025, 0.000025)
Discount Factor (t)
0.99
Advantage Estimation Coefficient (A)
0.95
Entropy Coefficient
[0.03, 0.003, 0.0003, 0.00003]
Number of Epochs
3
Rollout Length
128
Sequence Length
128
Number of Env
128
Batch Size (Sequence Length x Number Env)
16384
Number of Mini Batches
PPO Clip Ratio
Max Gradient Norm
Value Function Coefficient
0.5
G.4 Memory Maze
The border indicates the object the agent must collect next
9x9
11x11
Objects are placed randomly around the maze
The border changes after the agent collects the matching object
13x13
15x15
Top down view of the maze layout
Figure 11: The Memory Maze environment. On the left, show possible maze layout for all four Memory Maze configurations. The maze layout is randomized each episode. On the right, show two sample observations that the agent receives. The agent's 's observation at each time-step is 64 x 64 RGB pixels and the action space is discrete. The border color of the observation image indicates the target object color which the agent needs to find to receive a reward. After collecting the object, the border color changes, indicating the next target object. The episode lengths are fixed depending on the Memory Maze configuration, with larger configurations having longer episodes.
Environment Description: The Memory Maze environment evaluates an agent's long-term memory capa- bilities in partially observable RL setting. Figure II illustrates this environment. The agent's 's observation at each time-step is image with 64 64 RGB pixels, and the action space is discrete. In each episode, the agent starts in randomly generated maze containing several objects of different colors. The agent's 's objective is to find the target object of specific color, indicated by the border color in the observation image. Upon successfully touching the correct object, the agent receives a +1 reward, and the next random object chosen as the new target. If the agent touches an object the wrong color, there no effect on the environment. The maze layout and object locations remain constant throughout the episode. Each episode lasts for a fixed amount time. Since the maze layout is randomized at each episode, the agent must learn to quickly remember the maze layout, the target object locations, and the paths leading them.
30

---

Published in Transactions on Machine Learning Research (10/2024)
Hyperparameters and Tuning Strategy: We include the details of the Memory Maze experiments. All of the experiments in that section were implemented using asynchronous PPO implementation from Sample Factory library (Petrenko et al. (2020)). We started with the default hyperparameters for the DMLab lab experiments in Schulman et al. (2015). and finetuned the learning rate and entropy coefficient. For each of LSTM, GTrXL and TG,LiTe, to tune the learning rale and entropy coefficient, we run a sweep for three seeds for 15M steps the Memory Maze 11 11 environment. We average the results for the last 1M steps across the three seeds and select the best hyperparameter according to total episodic reward. Using the best-identified hyperparameter,
generate the final results for 100M steps for each of the three seeds. We detail the hyperparameters along with the sweeps for the learning rate and entropy coefficient in Table 7 We include the architecture configuration for each of the architectures in Table
Table 7: Hyperparameters and sweeps for Memory Maze experiments.
Hyperparameter
Value
Learning Rate
0..025, 0.00025, 0.000025
Discount Factor (7)
0.99
Advantage Estimation Coefficient (A)
0.95
Entropy Coefficient
[0.03, 0.003, 0.0003]
Number of Epochs
Rollout Length
200
Sequence Length
100
Batch Size
3200
PPO Clip Ratio
0.1
PPO Clip Value
1
Max Gradient Norm
4
Value Function Coefficient
0.5
Number of Workers
32
Number of Envs per Worker
2
Table 8: Architecture configuration for GTrXL and AGaLiTe for Memory Maze experiments.
Hyperparameter
GTrXL
AGaLiTe
Embedding Dimension (d)
512
512
Num Heads
8
8
Head Dim (dr)
64
64
Num Layers (L)
Memory Size (M)
256
N/A
Projection Hyperparameter (7))
N/A
4
Approximation Hyperparameter (r)
N/A
7
G.5 Craftax
Hyperparameters and Tuning Strategy: The architecture sizes used for Craftax experiments are kept same as in Table We detail the hyperparameters used for the PPO algorithm that used for training the agents in the Craftax environment in Table g We tune learning rate and entropy coefficient for the sweeps mentioned in Table g Our hyperparameter tuning strategy as follows: we train seeds per architecture for cach the hyperparameter configuration for 100M steps in the Craftax symbolic environment. Finally, we identify the best hyperparameter configuration according to the best episodic reward in the last 1M training steps.
31

---

Published in Transactions on Machine Learning Research (10/2024)
Table 9: Hyperparameters and sweeps for Craftax experiments.
Hyperparameter
Value
Learning Rate
00.0002, 0.0003, 0.0004]
Discount Factor (7)
0.999
Advantage Estimation Coeflicient (A)
0.8
Entropy Coefficient
[0.0, 0.01, 0.001]
Number of Epochs
4
Rollout Length
128
Sequence Length
128
Number of Env
1024
Batch Size (Sequence Longth x Number of Env)
130944
Number of Mini Batches
8
PPO Clip Ratio
0.2
Max Gradient Norm
1.0
Value Function Coefficient
0.5
32

---

Published in Transactions on Machine Learning Research (10/2024)
H
Additional Learning Curves on Smaller Memory Maze Configurations
Memory Maze 9x9
Memory Maze 11x11
Oracle
60
Oracle
30
Total Reward
GTrXL-256 AGaLiTe
Total Reward
40
20
GTrXL-258
AGaLiTe
20
10
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Steps
x10
Steps
x10
Memory Maze 13x13
Memory Maze 15x15
Oracle
Oracle
60
Tobal Reward
Total Reward
60
40
GTrXL-256 AGaLITe
GTrXL-256
20
AGaLiTe
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Steps
x10
Steps
x10
Figure 12: Learning curves of GTrXL and AGaLiTe agents in the Memory Maze environment. The x-axis represents the number of environment steps, and the y-axis represents the total reward in episode. Each agent is trained with 3 different random seeds. The bold lines represent the mean return across the 3 seeds, and the blurred lines represent the individual seeds. Each point the average episodic reward over 1M environment steps. The dotted grey line represents the performance of an oracle agent that has access to the entire maze layout, Larget object locations and paths leading them.
Evaluating Impact of GTrXL's Context in Memory Maze
This experiment evaluates the impact of GTrXL's context length in the Memory Maze environment. We showed earlier that GTrXL's performance bottlenecked by the memory size in T-Maze. Our hypothesis is that a similar conclusion should hold in the Memory Maze environment. We expect that GTrXL with a larger memory size would outperform GTyXL with smaller memory size. We should also be able show that an AGaLiTe would outperform GTrXL with small memory size. To investigate this, we train two additional GTrXL agents with memory sizes of 64 and 128 in the Memory Maze 13 x 13 environment.
The learning curves of training the three memory sizes of GTrXL and AGaLiTe in the Memory Maze 13 x 13 environment is shown in Figure 13 Asymptotically, all four agents achieve similar performance. The individual learning curves, however, indicate that the GT+XT1-64 agent slower to converge than the GTYXL-122 and GTFXL-256 agents.
The results failed provide sufficient evidence to support our hypothesis. The performance obtained by the three agents does not appear to be different. This observation leads to the following speculation:
33

---

Published in Transactions on Machine Learning Research (10/2024)
the Memory Maze environment is too difficult for the agents to be able to utilize their long-term memory capabilities. The reward signal is sparse, which might make it difficult for the agent to learn long-term dependencies. It is also possible that learning long-term dependencies in navigation tasks is harder, in general, and longer training is necessary for the benefits of long-tern memory to show.
Memory Maze 13x13
30
GTrXL-256
GTrXL-128
Total Reward
AGaLiTe GTrXL-64
20
10
0
0
1
2
5
6
Steps
x00
Figure 13: Learning curves of GTrXL agents with diffcrent memory sizes in the Memory Maze 13 13 environment. The x-axis represents the number of environment steps, and the y-axis represents the total reward in an episode. Each agent is trained with 3 different random seeds. The bold lines represent the mean return across the 3 seeds, and the blurred lines represent the individual seeds. Each point is the average episodic reward over 1M environment steps.
34

---

Published in Transactions on Machine Learning Research (10/2024)
J Learning curves for various achievements in Craftax Symbolic
In Figure 14 plot the various achievements in the Craftax environement. Detailed description of these achievements could be found in Hafner (2021).
Gegg
Daps
Seps
2ep
Degt
Caps
Gegs
675
(16
1.00
600
8.50
077
8435
075
8.50
119
Sape
Dape
Dace
PALCON
625
615
075
1.00
8.50
115
8.5
-17
00
6.30
Kapm
30.00
30.
455
6.5
4.5
075
0.5
Figure 14: Learning curve for various achievements in Craftax Symbolic. Results are reported over 15 seeds + 95% bootstrapped CI. AGaLiTe; brown, GTFXL-128: blue.
35

---

Published in Transactions on Machine Learning Research (10/2024)
K Latency Measurements
In this section we provide additional empirical evidence of the computational efficiency of our proposed approach, by comparing the latency of forward pass using GTrXL and AGaLiTe. We measure the time required milliseconds (ms) to do a forward pass two scenarios: (1) processing single element in streaming sequence, (2) processing an entire sequence in parallel. We configure the architecture sizes of GTrXL and AGaLiTe according the values used by Parisotto et al. (2020): 12 layers, 8 heads, d, 64, d 256. We collected all data in single Google Cloud instance with NVIDIA A100 GPU, 12 CPUs and S0GB RAM.
First, compare the time required in milliseconds (ms) to do forward pass using single element in streaming sequence. We present the results of these comparisons in Figure 15a According to Dai et al. (2019), XL attention used in the GTrXL architecture has a limited context. The context length of XL attention, how far back in time the transformer architecture can remember, number of layers and M the memory size. We measure the impact of increasing the context length (varying M) of GTrXL (x-axis) on the latency do single forward pass (y-axis). AGaLiTe does not explicit hyper-parametts that allows controlling the context length, and the use of recurrent hidden state allows for a potentially unlimited context. Therefore, consider three AGaLiTe architectures with feature map hyper-parameter E [4,8, 16), and plot. it as straight line. We observe that the gap between GTT-XT, and AGaLi'le increases dramatically with increasing context length.
O(ML), where is the
Next, we measure the time required to do a forward pass over a batch, that is process an entire input sequence in parallel. We present the results of these comparisons Figure 15b We vary the length of the imput sequence (x-axis) and measure the time required to do a forward pass over the entire sequence (y-axis). We consider two GTrXL architectures with memory size M [128, 512]. We consider three AGaLiTe architectures with r 84, 8, 16]. We observe that the gap between GTrXL and AGaLiTe increases dramatically with increasing sequence length.
Single-Step Latency
Batch Latency
GTrXL
300
GTrXL (M 128)
AGaLiTe 4)
GTxXL (M 512)
30
AGaLiTe (r)88)
AGaLiTe (n= 4)
AGaLiTe (7=16)
AGaLiTe (n= 8)
Time (ms)
Time (ms)
200
AGaLITe (n=16)
20-
100
10
763
1536
3072
614 5144
12288
24576
491 1155
98304
196608
396216
786432
1572864
2
4
8
16
32
64
128
256
512
1024
2048
4096
8192
SequenceLength
Context Length
(b) Time (ui)) for processing an entire sequence in paral- lel.
(a) Time (ms) for processing a single element sequence.
Figure 15: Latency measurements for GTxXL and AGaLiTe. Each point is averaged 100 independent runs, and the shaded region is the standard error.
36


Reviews:
[{"review_id": "GwlymCLOmU", "reviewer": "Reviewer_E4UB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nS1.\tThe paper introduces three novel architectural innovations. \n\nS2.\tThe paper is clearly written, with an excellent motivation for the introduced techniques, and a detailed exposition of the methodology.\n\nS3.\tExperiments are well-formulated and provide compelling evidence in support of the claims of the paper in the context of reinforcement learning tasks.\n\nS4.\tThe approximation to the Kronecker delta function is innovative and well-motivated, the paper provides a proof of convergence, and the resultant architecture with a reduced memory footprint performs well. \n\nWeaknesses:\n\nW1.\tTo substantiate the claims, the paper needs to establish that the specified attributes of the existing architectures are indeed limitations and that they are overcome or mitigated by the proposed architecture. The paper successfully does this for partially observable RL problems, but it does not provide any results for other tasks. For example, Katharopoulos et al. (2020) report on image generation and speech recognition experiments. \n\nW2.\tInsufficient explanation for the name of the architecture and potential confusion with both the original linear transformer and other more recent works that adapt the recurrence relationships in other ways. \n\nW3.\tThe paper does not sufficiently discuss the relationship to important related work that presents techniques similar to some of the introduced methods.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper argues that the existing linear transformer architectures have important limitations – (1) the self attention mechanisms adds positive values to the recurrent state and therefore cannot delete previous information; (2) performance critically depends on the choice of the kernel function; (3) kernel functions often use element-wise or randomized expansive feature maps, which can result in limited memory capacity; (4) they have a high memory cost due to the maintenance of a matrix as a recurrent state. \n\nThe paper proposes a novel architecture (learning framework) that addresses these limitations by (i) introducing a gated structure that can focus on relationships with a much larger time gap; (ii) incorporating a learnable feature map in the self-attention mechanism (ensuring that it is amenable to sequential computation); and (iii) proposing an approximate version that eliminates the need to maintain a matrix as the recurrent state. \n\nThe paper provides an extensive set of experiments focusing on the partially-observable reinforcement learning task to demonstrate a performance improvement compared to the state-of-the-art transformer-based RL method. \n\nIn the appendix, the paper provides a proof that the limit of the approximating function is the Kronecker delta, as well as thorough details of the experiments and additional experimental results.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "(1)\tAlthough I believe it is important, weakness 1 is easily rectified by clearer identification in the introduction that the paper is focused on the RL setting and the claims apply to that setting. I don’t consider that this significantly diminishes the contribution of the paper. The alternative is to report the results of experiments addressing other tasks (in particular, tasks similar to those that have been explored in some of the papers that introduces the baseline techniques). Reinforcement learning is not explicitly mentioned in the first four paragraphs of the introduction. Yes, the first paragraph describes an RL setting, but the next paragraph switches to the more general task of “incremental state-construction”. The paper needs to explicitly specify that the identified weaknesses and rectifications are in the context of the RL setting (no experimental evidence is provided that the proposed techniques are beneficial for other tasks, so the claim is not supported in the general setting). For this reason, I would also suggest that the title is modified to include “Reinforcement Learning”, but this is only a suggestion. Mandatory: modify the text to more clearly specify the context of the claims. \n\n(2)\tI do not completely understand the choice of the name of the architecture. The linear transformer is already recurrent (as the paper clear specifies in Section 2.2 – “maintains a matrix … as a recurrent state”). How is a “recurrent linear transformer” different from a linear transformer? “Recurrent” is an unnecessary redundancy that doesn’t differentiate the proposed architecture in a clear way. Yes, the paper proposes a modification to the recurrence equations, but it doesn’t introduce recurrence, which is what the name implies. Beyond this, there are several works (discussed below) that modify the recurrence in other ways, or aim to introduce stronger recurrence relationships. Please either explain the choice of the name or consider a different, more informative and appropriate name. \n \n(3)\tThe discussion of the related work should be expanded and the connections and distinctions of the proposal compared to previously introduced methods made much clearer.\n\nThe paper cites Schlag et al. 2021, primarily in support of its motivations, but it does not provide a clear and thorough discussion of the relationship of the proposed method to the methods in that work. Nor does it cite or discuss the follow-up work [R1]. These two works introduce the DeltaNet, which does address the memory capacity concern, by allowing for both positive and negative updates and incorporating a learnable learning rate (although not element-wise). The related work in the submitted paper does explain that the key difference is that the proposed method includes an element-wise gate, although it is not acknowledged clearly that other techniques already address the positive addition and memory limitation. \n\nIn [R1], several extensions are presented that incorporate additional recurrence. The Delta RNN applies “an element-wise activation function f to the previous output of the fast network y(t−1) to obtain the recurrent query” – there is a separate set of weights for the recurrence. A Delta LSTM is also proposed. The Recurrent Delta Net introduces recurrence into the k, v, q update equations, as well as the scalar gate. \n\nThe paper should also cite and discuss [R2]. This is probably concurrent, independent work (the paper was arXived in Feb. 2024, but only published at an August 2024 conference). Now that it has been published, the paper should certainly be discussed in the related work section, probably with a statement that the research was conducted concurrently and independently. As the title implies, [R2] proposes the use of learnable kernels (“enabling the model to learn any quadratic function that is non-negative and has a single real root”). \n\n[R1] Irie, Kazuki, et al. \"Going beyond linear transformers with recurrent fast weight programmers.\" Advances in Neural Information Processing Systems 34 (2021): 7703-7717.\n\n[R2] Yaroslav Aksenov, Nikita Balagansky, Sofia Lo Cicero Vaina, Boris Shaposhnikov, Alexey Gorbatovski, and Daniil Gavrilov. 2024. Linear Transformers with Learnable Kernel Functions are Better In-Context Models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 9584–9597, Bangkok, Thailand. \n\nQuestions\n\n(Q1)\tIn Figure 2, the AReLiT behaviour seems considerably less stable than the GTrXL. Is this correct observation? If so, is there an explanation? The GTrXL also seems to still be improving at the selected endpoint. Does it reach/surpass the behaviour of AReLiT after more steps?\n\n(Q2)\tIs there a reason that a 2020 paper (Parisotto et al., 2020) is still the state-of-the-art baseline? Is there a need to incorporate other baselines that are not based on transformers?\n \n(Q3) “Note that we present the algorithm for processing a sequence in the case of the linear transformer. This is in contrast to Algorithm 1, which presents the algorithm for processing a sequence.” – I suspect this is a typo and it is the processing of a single input vector?", "broader_impact_concerns": "No concerns", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2uCG8nE5On", "reviewer": "Reviewer_cyG9", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n* The paper proposes a practical improvement on transformers that makes them more amendable for online RL problems.\n* The details of ReLiT and AReLiT are described in good detail and clarity.\n* The paper provides relevant details for the experiments with the ablation study providing evidence for the benefits of the proposed modification.\n\n**Weaknesses:**\n* The cases studied for the paper are limited to online only. It would be more convincing if the authors provided evidence that ReLiT and AReLiT can be useful for other ML problems. The paper mentions text classification in Section 5, but does not present any results on it. It would be good to have those results to see if the proposed method applies beyond the studied RL settings.\n* The experiments present evidence for problems that the baseline methods can solve. It would be more convincing if the authors showed cases where ReLiT's benefits of memory are critical. One such potential example could be extending T-Maze beyond the memory size of  GTrXL-256.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper introduces a new transformer architecture called Recurrent Linear Transformers (ReLiT), as well as Approximate Recurrent Linear Transformer (AReLit), to address shortcomings of transformer related to access of past information and expensive inference costs. The paper argues that the two shortcomings inhibit the application of transformers to RL settings, which motivates ReLiT and AReLiT and applies them to solve a number of online RL problems.\n\nThe paper first starts by providing an introduction to the application of transformers for RL in Section 1. Section outlines preliminaries of the transformer architecture, including the self-attention mechanism and the Linear Transformer. Section 3 introduces ReLiT with relevant details on how the Linear Transformer architecture is modified with a gating mechanism and learnable feature map. Section 4 describes AReLiT, which replaces the update of the recurrent state using a sum of cosine functions that approximate the Kronecker delta function. Section 5 shows an empirical evaluation on a series of online RL problems that require memory to be solved effectively. ReLiT and AReLiT generally perform fairly well in the studied examples providing competitive performance with better computational efficiency. Section 6 details an ablation of AReLiT showing that importance of different components of the proposed architecture. Section 7 describes related work and Section 8 provides a conlusion.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* It would be good if the paper can provide additional experiments to mitigate the weaknesses mentioned, as this would strengthen the general validity of the claims for the proposed architecture.\n* I think it would be good for the paper to provide a deeper discussion on applying transformers for offline RL settings, as well as pretraining plus finetuning [1] training settings, and motivate how the ReLiT could be useful there. On top of that, it would generally be helpful for the paper to provide a stronger motivation for why solving online RL problems is important since that is the primary focus of the current experiments.\n\nNice to Have:\n* I would some empirical numbers on the training time of ReLiT compared to the baselines in addition to the forward pass study in Appendix K.\n\n[1] Ghugare, Raj, et al. \"Searching for High-Value Molecules Using Reinforcement Learning and Transformers.\" The Twelfth International Conference on Learning Representations.", "broader_impact_concerns": "The paper could be strengthened with a broader impact section.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "I6f4L6ysDk", "reviewer": "Reviewer_i1t2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n* The writing is generally clear and easy to follow. I have checked most of the formulas in Section 3, and they look correct.\n\n### Weaknesses and other questions\n* The performance improvement in the challenging Crafter environment is\n* In the paragraph under Section 3, the authors mentioned that the recurrent equations in Algorithm 2 (line 5 and 6) add positive values to the recurrent state. From the equations alone, it seems not necessarily the case. Is it because in linear transformers, the function $\\phi$ always takes some special form?\n* Section 3.1 \"leading to numerical stability issues\": is there any evidence or reference to support this statement?\n* For the results in Figure 1, could the authors explain why the performance of linear transformers drops significantly at length=200?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper presents Recurrent Linear Transformer (ReLiT), a new architecture that aims to address the two limitations of linear transformers for reinforcement learning (RL). Specifically, ReLiT introduces a gating mechanism and a learnable feature map. The gating mechanism allows better control of the information flow within the context; the learnable feature map eliminates the choice of the kernel feature map, enhancing the expressiveness of the model. In addition, the authors present an approximate ReLiT as a practical implementation, which reduces the time and space complexity of processing a single element in a sequence. In a variety of environments, ReLiT has shown better performance than other baseline methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to the section above.", "broader_impact_concerns": "No", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors should be sure to provide a final, camera ready version that includes all of the modifications they have made in revision.
Decision: Accept (poster)