Recurrent Linear Transformers

Dear reviewer gTNz, please see the general response to all reviewers about experiments and our focus on the RL setting. Thank you for carefully reading our paper and giving thoughtful comments and suggestions. We will respond to the key points of misunderstanding and answer your questions below in a point by point fashion.

The authors claim that a positive feature map is required… Could the authors clarify this?

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.

The motivation behind a learnable feature map for self attention is not clear. … What is the motivation behind doing so?

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.

Why not just do a parameterized represented through a small neural network, as opposed to the complex scheme of computing outer products for queries and keys?

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$.

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.

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.

Dear reviewer EZRU, please see the general response to all reviewers about experiments and our focus on the RL setting. Thank you for carefully reading our paper and giving thoughtful comments and suggestions. We will respond to the key points of misunderstanding and answer your questions below in a point by point fashion.

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].

Thank you for carefully reading our paper and giving thoughtful comments and suggestions. We will respond to the key points of misunderstanding and answer your questions below in a point by point fashion.

The paper would really benefit from the addition of a comparison with linear transformers

We have already done this in Appendix K, Figure 13 (d), where we compared ReLiT with linear transformers and showed the effect of different kernel feature maps, gating mechanisms, and approximation approaches. We will update the main text of the paper to more clearly point the reader to this result in the appendix.

Authors could revise the related works section by adding discussion on the aforementioned related directions.

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.

Authors could look into ways of highlighting the usefulness of the gating mechanism described in ReLit.

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

How do different kernel functions affect the working of ReLiT or linear transformers in general? Can the authors please demonstrate an ablation study?

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.

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 found the time complexity to be O(N^2 d + N d^2).

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)

I also found it hard to understand why Linear Transformer's self-attention complexity was independent of N.

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.

Thank you for carefully reading our paper and giving thoughtful comments and suggestions. We will respond to the key points of confusion and answer your questions below in a point by point fashion.

In addition, please see the general response to all reviewers about experiments and our focus on the RL setting.

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?

I think the tunable kernel function may not be the main reason for the experimental success of ReLiT on RL problems. The main advantage of your model is the low inference costs introduced by the low-rank approximation, and this learnable kernel is kind of far from your main contribution.

We agree that a major advantage of our method is the computation win from the low-rank approximation. However, we do not agree the learned kernel is not important as our feature map ablation experiment in Appendix K, Figure 13 (b) shows that the proposed feature map outperforms deterministic expansive feature map approach DPFP and element-wise feature map ELU+1.