Paper: Recurrent Linear Transformers (ICLR-2024)

Reviewer comment (Theory):
The approach in recurrent linear transformers (Katharopoulos et. al 2020) need not add just positive values at each iteration. In particular, it relies on a notion of similarity through kernel functions, where the kernel itself is positive but the mapping $\phi$ need not be. While Katharopoulos et. al do use a mapping based on positive values, it can simply be extended to setups where $\phi$ can also map to negative values; and hence this does not seem to be a limitation of their framework.

Author response:
We agree that it is possible to have a non-positive feature map if the underlying kernel function produces non-negative similarity. We assumed that the feature map is positive, as done by Katharopoulos et. al (2020). We chose to follow this assumption for the following reasons:
A positive feature map is a simple way to produce positive kernel scores and has been a common assumption in previous extensions to the linear transformer paper (Peng et al., 2021; Schlag et al. 2021)
A positive feature map also ensures that the normalization vector $\mathbf{s}_t$ at each timestep $t$ is non-zero. This is necessary as calculating the attention vector $\mathbf{a}_t$ requires division by  $\mathbf{s}_t$.  $\mathbf{s}_t$ is defined as the sum of the query vectors. The assumption that the query vectors are positive is necessary, because otherwise the resultant sum $\mathbf{s}_t$ could become zero, causing the attention vector to be infinite.
We will modify the main text to expand on these points. An alternative feature map could be designed such that it (1) doesn’t have the above normalization issue and (2) generates positive similarity scores; we did not explore this in this work and consider it in future work.

Reviewer comment (Theory):
The motivation behind a learnable feature map for self attention is not clear. Why not just do a parameterized $\phi$ represented through a small neural network, as opposed to the complex scheme of computing outer products for queries and keys? What is the motivation behind doing so?

Author response:
The motivation behind a learnable feature map is to be able construct adaptable feature maps by learning them directly from the data. Our proposed feature map uses learnable parameters and an outer product operation to generate an expansive feature map. By introducing learnable parameters we are able to learn feature maps which are more suited to the data which they are trained on. This is in contrast to expansive feature maps in the existing literature (Peng et. al, 2021), which uses random vectors to generate an expansive feature map, and DPFP (Schlag et. al 2021) which assumes that the expansive features must be orthogonal. 

Using a neural network to generate a large feature map is interesting, but it will require more parameters than the proposed outer product-based mechanism. For instance, for an input vector of dimension $d$, consider that we want to produce an output feature of dimension $d\eta$. For concreteness lets assume $\eta$ is equal to 4 and $d$ is 64. Using a neural network to produce this feature would require $d^2\eta=16384$ parameters. On the other hand, our proposed outer-product-based rule requires $d\eta + d^2=4352$. Decoupling the $\eta$ hyperparameter from the $d^2$ allows us to increase the projection dimension to a large value without scaling quadratically with $d$.

Reviewer comment (Evaluation):
Given that an important benefit of the proposed approach is the savings obtained on time complexity, it would be nice to have some results highlighting the performance obtained by the different methods across wall clock time to really see the advantage of the method.

Author response:
The current manuscript has experiments that measure the wall clock time of our proposed approach when compared to transformers with softmax-based self-attention. We provide these results in Appendix J, Figure 12, where we compare the wall clock time/latency of forward pass using GTrXL and AReLiT. We find that in comparison to canonical self-attention, our proposed approach achieves significant savings in wall clock time with increased context or with increased sequence length. We also provide empirical evidence of the savings obtained by our approach by measuring the frames per second (FPS) in the Memory Maze experiments in Section 5.

Reviewer comment (Experiments):
There are also baselines that are missing in a number of experiments. None of the experiments benchmark against Katharopoulos et. al and there are also baselines missing in Figure 2.

Author response:
We agree with the reviewer that a comparison with Katharopoulos et al. is important. Fortunately, we have a comparison of our approach with Katharopoulos et. al (2020) augmented with a scalar gating mechanism (Peng et al. 2021) in Appendix K, Figure 13 (d). We performed this comparison in the T-Maze task with a corridor length of 200 and found that our proposed approach outperformed Katharopoulos et. al (2020).
As for Figure 2, we used GRU as our baseline for this diagnostic task since Morad et al. (2022) reported it to be the best-performing architecture (when comparing nearly a dozen other approaches) on partially observable classical control tasks, even when compared to transformer approaches.

Reviewer comment (Theory):
The authors claim that a positive feature map is required to ensure that the similarity scores produced by the underlying kernel function are positive. Technically a kernel function only requires mapping to some vector space where dot-product is defined, and thus defining a kernel function does not require positive feature maps. Could the authors clarify this?

Author response:
We agree that it is possible to have a non-positive feature map if the underlying kernel function produces non-negative similarity. We assumed that the feature map is positive, as done by Katharopoulos et. al (2020). We chose to follow this assumption for the following reasons:
A positive feature map is a simple way to produce positive kernel scores and has been a common assumption in previous extensions to the linear transformer paper (Peng et al., 2021; Schlag et al. 2021)
A positive feature map also ensures that the normalization vector $\mathbf{s}_t$ at each timestep $t$ is non-zero. This is necessary as calculating the attention vector $\mathbf{a}_t$ requires division by  $\mathbf{s}_t$.  $\mathbf{s}_t$ is defined as the sum of the query vectors. The assumption that the query vectors are positive is necessary, because otherwise the resultant sum $\mathbf{s}_t$ could become zero, causing the attention vector to be infinite.
We will modify the main text to expand on these points. An alternative feature map could be designed such that it (1) doesn’t have the above normalization issue and (2) generates positive similarity scores; we did not explore this in this work and consider it in future work.

Reviewer comment (Evaluation):
Direct comparison in terms of number of learnable parameters (additional parameters introduced in Section 3.1 and 3.2 seems a lot), execution times or memory footprint (see Table 3 and Figure 2 of [2]) are needed.

Author response:
> Section 3.1 and 3.2 seem to introduce a large amount of new learnable parameters. 

In fact, the additional learnable parameters we introduce are of the same size as the weight vectors used to generate the key and query vectors; thus they only marginally increase the number of learnable parameters. For instance, in Appendix J Figure 13 (d), we compare the performance of AReLiT with Linear Transformers. With the same configuration, we observe that Linear Transformer uses 1,718,928 parameters, whereas AReLiT uses 1,803,072 parameters (less than 5% increase). We will include these numbers in the camera ready version of the paper. 

> Direct comparison in terms of number of learnable parameters

In Appendix J, Figure 13 (d) we provide comparison of AReLiT with Linear Transformer ([1], [2]). As mentioned in our previous comment, the introduced additional parameters only marginally increase the total number of parameters. 


> Direct comparison in terms of ... execution times or memory footprint (see Table 3 and Figure 2 of [2]) with respect to direct competitors [1,2] ... are needed (also, the tradeoff w.r.t performances).

The computational complexity of Linear Transformer is of the same order as that of Linear Transformer (Appendix A, Table 1). AReLiT, however, is more computationally efficient than Linear Transformer due to the introduced approximation. In terms of number of operations in T-Maze, we observed that ReLiT uses 18 times more memory to store the recurrent state, and 3.13 times more operations than AReLiT for a single attention head with similar configuration. In the camera ready version of the paper, we will include comparison of AReLiT wrt to Linear Transformer/ReLiT  in terms of wall-clock time (as demonstrated in Appendix J Figure 12. 

> Direct comparison in … with respect to direct competitors [3]

Our experiments focus on RL. We don't understand how RWKV could be applied in RL. Extending RWKV to RL could be a whole paper in and of itself. Does the reviewer have suggested baselines that would be applicable in RL?

Reviewer comment (Experiments):
The authors proposed **ReLit** and **AReLit** in order to solve the issue of positive value addition in Linear Transformers recurrence (that could grow arbitrarily large). It could be interesting to better investigate this issue (also empirically -- or at least make the reader understand where the Linear Transformer fails) and how the proposed model solves it.

Author response:
> Q2. The authors proposed ReLit and AReLit in order to solve the issue of positive value addition in Linear Transformers recurrence (that could grow arbitrarily large). It could be interesting to better investigate this issue (also empirically -- or at least make the reader understand where the Linear Transformer fails) and how the proposed model solve it.

It is tricky to report negative results in a paper. Obviously, our claim did not come out of nowhere. We did run vanilla Linear Transformers in T-Maze, however, due to the issue of positive value addition, the gradients quickly blew up and became NaNs after a few iterations. We can include this result in the camera ready version of the paper, by plotting the magnitude of the gradients during the initial steps of training.  

Finally, in Appendix J, Figure 13 (d) we compare the performance of Linear Transformer with scalar gating mechanism with AReLiT and found it to be worse.

Reviewer comment (Novelty):
The relation with Peng et al. [2] should be better highlighted in the paper text, given that a (scalar) gating mechanism was already introduced in such paper.

Author response:
> The relation with Peng et al. [2] should be better highlighted in the paper text, given that a (scalar) gating mechanism was already introduced in such paper.

We agree with the reviewer that a comparison with Peng et al. is important as their work relates to our gating mechanism. Fortunately, we already have such a comparison as part of the ablation experiments in Appendix K, figure 13 (a), where we evaluated the impact of each of our proposed contributions, including gating mechanisms. We changed paragraph 4 of section 5, diagnostic MDP, to highlight these comparisons. To summarize the relation to Peng et al., we found that  the learning decay parameters for each element of $\mathbf{C}$ (gating) are better than a scalar decay used in the Linear Transformer (Peng et al. 2021) and that our proposed approach outperforms Linear Transformer with scalar gating presented by Peng et al. (2021). We will update the paper to (1) more clearly highlight that this comparison is in the appendix, and (2) expand the discussion of Peng et al. [2].

Reviewer comment (Theory):
At the end of page 2, the paper describes the time complexity of the self-attention module in canonical transformers to be O(N d^2), which does not seem correct. I tried calculating the complexity on my own and also verified with a few other sources, too, and found the time complexity to be O(N^2 d + N d^2). I also found it hard to understand why Linear Transformer's self-attention complexity was independent of N.

Author response:
I see why you thought the complexity would be O(N^2 d + N d^2), because you assumed that the computational complexity is calculated for processing an entire sequence of length N. We, however, report the computational complexity of self-attention for processing a single element in a sequence.  For a matrix $\mathbf{X}$, the inference cost is the cost of processing a single element $\mathbf{X}[i]$. For canonical self-attention, the cost of processing a single element depends linearly on the input sequence length N. The following is the computational complexity of the steps performed by canonical self-attention: 
1. Calculate the key and value vectors for N timesteps, and query for a single timestep: O(N d^2)
2. Calculate the dot product of query of a single timestep with key vectors of previous timesteps, apply softmax, and multiply the resultant with the value vectors for N timesteps: O(N d)
This results in a overall complexity of O(N d^2 + N d) = O(N d^2)

Reviewer comment (Novelty):
Authors could revise the related works section by adding discussion on the aforementioned related directions.

Author response:
We thank the reviewer for highlighting these great papers from the literature. We will update the related work section to cover previous approaches that address increasing the context length of transformers.

Reviewer comment (Experiments):
Authors could look into ways of highlighting the usefulness of the gating mechanism described in ReLit.

Author response:
We have already done this Appendix K, Figure 13 (a), where we compare the gating mechanism proposed in ReLiT with existing gating mechanisms proposed for the Linear Transformer architecture.  We find that learning decay parameters for each element of $\mathbf{C}$ (gating) is better than a scalar decay used in the Linear Transformer (Peng et al. 2021). We will update the main text of the paper to more clearly point the reader to this result in the appendix

Reviewer comment (Experiments):
How do different kernel functions affect the working of ReLiT or linear transformers in general? Can the authors please demonstrate an ablation study?

Author response:
We have already done this Appendix K, Figure 13 (b), where we compare the kernel function proposed in ReLiT with feature maps used in the existing literature. We found that our expansive feature map such as ours outperforms element-wise maps like ELU+1 and deterministic feature maps like DPFP (Schlag et al. 2021). We will update the main text of the paper to more clearly point the reader to this result in the appendix.

Reviewer comment (Theory):
Can authors please provide a step-by-step derivation of how time and space complexities are derived for different transformer variants?

Author response:
Similar to the complexities reported for canonical transformers, we report the computational cost of processing a single element using the Linear Transformer architecture. The Linear Transformer algorithm presented in Algorithm 2 requires only a fixed dimension hidden state to generate the output at a given timestep. As such, the computational complexity of self-attention for processing a single element is independent of the input sequence length N.

Reviewer comment (Theory):
The motivation part can be more precise and some of your motivation is not really important, from my perspective. In the third paragraph of the introduction section, you propose three motivations, while I think your methods do not really address the second one (the kernel function one). The authors claimed that the Transformers of linear Transformers have a manually-designed kernel function (softmax / Gaussian), but in your ReLiT model, although there are some learnable parameters in your kernel functions, the kernel function class is still manually-designed and fixed. Moreover, the kernel function class in ReLiT is not rich enough to make Transformers learn the 'globally optimal' kernel. You can only expect that the Transformers learn the 'optimal' kernel in the kernel class you specify.

Author response:
> The motivation for the work is unclear

The motivation behind  each of the three proposed contributions are as follows:
1. An outer-product based gating mechanism to control the flow of past information in the Linear Transformer self-attention: Gating mechanisms allow for a more fine grained control of past information that needs to be deleted and new information that gets added. This is important to handle long contexts as the network can selectively retain past information.
2. An outer-product based parameterized expansive kernel feature map: Using outer-product 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 as it correlates to the memory capacity (see Schlag et. 2021). Using a parameterized feature map allows us to learn the feature map directly from the data.
3. Approximation of recurrent state update: approximation of recurrent state is essential to reduce the computational complexity of applying the recurrent state update. 

Taken together, these three components are essential for constructing a transformer architecture that can be updated online in a streaming fashion, which requires it to be computationally efficient without, of course, losing expressiveness. The outer product-based gating mechanism gives the architecture flexibility to attend to information from the past, the expansive kernel map gives us flexibility on the representation for learning it from data, and the approximation gives us computational efficiency. The results in the main text and in the appendix (highlighted below) provide empirical evidence of the benefits of these 3 components.

> The motivation part can be more precise and some of your motivation is not really important, from my perspective….while I think your methods do not really address the second one (the kernel function one).

As always motivation can be sharpened. Existing kernel functions which use element-wise (Katharopoulos et al., 2020) or randomized expansive (Peng et al., 2021) feature maps have been shown to have limited memory capacity (Schlag et al., 2021). Schlag et al. (2021) showed that expansive feature maps outperform element-wise feature maps such as ELU+1. They also showed that deterministic feature maps have higher memory capacity than randomized ones. Following the results in Schlag et al. (2021) our motivation was to design a feature map that is deterministic and generates an expansive representation of the input. We calculate an expansive representation of the input vector using outer products. As an added benefit, our proposed feature map is parameterized to enable learning more expressive expansive representations of the input. 

> The authors claimed that the Transformers of linear Transformers have a manually-designed kernel function (softmax / Gaussian), but in your ReLiT model, although there are some learnable parameters in your kernel functions, the kernel function class is still manually-designed and fixed.

We did not mean to imply that manually designed kernel functions are bad, more so we should have emphasized that our kernel functions are better because our proposed kernel function uses learnable parameters to learn an expansive feature map, instead of using a static function. It is deterministic similar to DPFP (Schlag et al., 2021), but the learnable parameters allow it to learn a more fine tuned feature map from the training data. 

> Moreover, the kernel function class in ReLiT is not rich enough to make Transformers learn the 'globally optimal' kernel. You can only expect that the Transformers learn the 'optimal' kernel in the kernel class you specify.

We are not sure what you exactly mean by a ‘globally optimal kernel’? Can you clarify what you mean?

Reviewer comment (Experiments):
I am not sure how the authors think. My suggestion is, if you want to claim that the learnable parametric kernel is important and the kernel class you specify is good, you may need to do some ablation experiments. For example, what happen is you replace the softmax or fixed kernel function with this learnable kernel in the standard /linear Transformers?

Author response:
> ...if you want to claim that the learnable parametric kernel is important and the kernel class you specify is good, you may need to do some ablation experiments. For example, what happen is you replace the softmax or fixed kernel function with this learnable kernel in the standard /linear Transformers?

We actually included some of these ablations in the appendix of the submitted paper.  To highlight the advantage of our proposed feature map compared to existing feature maps, we present an ablation conducted in corridor length 200 in the T-Maze task. In Appendix K, Figure 13 (b), we show that replacing the proposed feature map in AReLiT with previously used feature maps such as ELU and DPFP caused degradation in performance.

Reviewer comment (Theory):
In the last paragraph in the section 2, you claim 'The Linear Transformer’s self-attention has a context-independent inference cost, unlike the canonical self-attention mechanism', which I feel it kind of unfair. Although the inference cost for each token in linear Transformer is independent of the context length, the inference cost of a token matrix (a sequence) does depend on that, right? In sec 2.1, you showed the costs for standard TF for inferring a token matrix while in sec2.2, you showed the costs for one token. I believe you may modify your claim here.

Author response:
> Although the inference cost for each token in linear Transformer is independent of the context length, the inference cost of a token matrix (a sequence) does depend on that, right?

Yes, in fact the computational cost of linear transformer and its derivative approaches are independent of the context length only when processing a single token, and not a token matrix.

> In sec 2.1, you showed the costs for standard TF for inferring a token matrix while in sec2.2, you showed the costs for one token. I believe you may modify your claim here.

We do not need to modify our claim because all the computational costs presented in our paper regarding our approach, the canonical transformer and linear transformers are in terms of the inference cost.   Let’s be clear about what we mean by inference cost: the cost of processing a single element in a sequence. We report all complexities under this definition of inference cost. In Sec 2.1, we report the computational cost of processing a single element using the canonical transformer, which depends linearly on the input sequence length N, as canonical self-attention requires past input to be provided to generate the output representation at a given timestep. In Sec 2.2, we report the computational cost of processing a single element using ReLiT, which in contrast to canonical transformers, doesn’t depend on having the full sequence.