Latte: Latent Attention for Linear Time Transformers

Weaknesses:

1. • The first point is incorrect. The definition of a latent variable is: “An unobserved quantity during train and test time”. [1] Provides a good description of latent variable models and [2] provides a full derivation of examples like GMM which are latent variable models that are unrelated to causal inference. Latent variable models and causality are two orthogonal concepts.
   • There is a general misunderstanding here. We now added the graphical representation of our model in Figure 8 (page 16). Note that by $p(s|t)$ we do not mean $p(x_s|x_t)$ which perhaps might be a misunderstanding from the reviewer.
   • All discrete latent variable models can be interpreted as clustering models. Naturally, our method performs a form of clustering in token embedding space, with the latent token states corresponding to cluster centres. We assumed readers would be familiar with the well-recognised parallel between discrete latent variable modelling and clustering. However, we are happy to explain this point further if this would be useful to a reader.
2. • This point is also incorrect. As mentioned previously, latent variables and causal inference are two independent notions. See [1] and [2].
3. • We mention that latent states represent global concepts on line 175. Latent states are associated with general concepts such as colours or shapes.
   • Consider computing the attention between two tokens “red” and “green”. In standard attention, this similarity comparison goes directly via the token embeddings of these two concepts. However, one might argue that the reason “red” and “green” should be similar is because they are both colours. Therefore, in Latte, we know that “red” and “green” are similar since “red” is similar to the latent concept/token “colour”, and “green” is also similar to “colour”.
4. • This point is incorrect. Check the newly added Figure 8 from page 16. As mentioned previously our model consists of 2 discrete random variables s,t which denote positions and a discrete latent variable l. Under our graphical model, the assumption is correct.
5. • We compare it to standard attention and we consider it fair because Flash Attention is a Hardware optimisation of the standard attention. The time complexity of Flash Attention is still quadratic.
   • One can also provide hardware optimisation for Latte as in [3]. We are working on such a CUDA kernel, but it is not the main objective of our work and we feel would be a separate paper in itself. For example, Flash Attention is based on a careful and well-researched design decision, which is non-trivial. There is no in principle difficulty in creating a hardware-optimised version of Latte, but this will require some research and experimentation.
   • We now added Figure 11 on page 19 which has a comparison with Flash Attention. The figure shows that FlashAttention is still quadratic.
6. • This is wrong because we weigh the two attentions. Their sum is always one. The full attention is weighted by $p(l=0|t)$ and the latte attention is weighted by $p(l>0|t)$. Our attention scores always add up to 1 and are valid probabilities.
7. • Overall we strongly disagree with the reviewer. We suspect the reviewer’s belief in the lack of mathematical soundness points is based on a fundamental misunderstanding of the generally accepted definitions of latent variables and causality.
   • We agree that Latte can be viewed as clustering and we can mention this in the paper. Indeed, this is a strength of latte and formed the basic motivation and intuition for the approach; we strongly disagree that our description based on using latent variable models is wrong and that we need to re-write the paper around clustering. However, we are happy to explain to readers the intimate relationship between clustering and discrete latent variable models.

Questions:

1. We showed this in Figure 5.b
2. It is a fixed hyperparameter that we choose empirically. We have 128 latent states, information present in Tables 6 and 7 in the Appendix.
3. We modify the batch size as we increase the sequence length so that the number of tokens stays constant. We will improve our caption to explain this.
4. We do not say that there is a collapse of latent states in other linear attention models since other linear attention methods do not use the concept of latent variables. We only say that is generally possible in latent variable models to have latent collapse where only a few states are used. Then we show that this phenomenon does not occur in our model. We can explain this better if it would benefit the reader.

[1] https://medium.com/@manasmahanta10/latent-variable-models-demystified-7f1342698985
[2] Christopher M. Bishop. 2006. Pattern Recognition and Machine Learning
[3] Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y., “Gated Linear Attention Transformers with Hardware-Efficient Training”

Weaknesses:

• We now introduced experiments in Table 5 showing the results of the pre-trained model on natural language harness (MMLU, Hellaswagm, ARC).
• Compared to T2R and Hedgehog our method has a few advantages:
  • Unlike T2R and SUPRA we have local attention which helps improve results as shown by experiments in Table 1 and Table 2.
  • We do a weighted combination of sliding window attention and our Latte attention. Hence we can freeze all the parameters of the network relating to sliding window attention (and the MLPs of the transformed) and adapt only the parameters of Latte attention. This means that we can readily take an off-the-shelf model and readily merge this with our Latte model (resulting in the Latte-Macchiato approach). This reduces the memory required by the optimizer and the training speed.
  • Hedgehog requires quadratic training cost since it needs to compute the full attention in its training objective, hence it cannot be trained on long sequences. One can apply the methodology in HedgeHog to our model as well and this would make for interesting future work. We see these as independent contributions.
• We thank the reviewer for suggesting these related works, some of which we were unaware of and appeared since we began our research. We will cite them.
• We see our work as part of the major effort to make linear sequence models as performant as standard quadratic attention transformers. While the problem is not completely solved we tackle some issues by introducing a hybrid model which uses a valid probability distribution for attention.

2. • No, we do not use Yarn. Methods like Yarn or Xpos can extend the context of a pre-trained transformer, but they maintain its quadratic complexity. Conversely, our model extends the context at only linear scaling complexity.

Weaknesses:

• Our work is novel because, as far as we know, it is the only causal linear attention model that has the intuitive interpretation of defining latent states (corresponding to token clusters) whilst retaining a correctly normalised attention distribution.
• Furthermore, our model has natural formulations for both the bidirectional and causal cases and provides an extension to hybrid models that integrate standard sliding window attention with linear attention, retaining correctly normalised attention throughout. As far as we know, this is the only model with such features. Efficient attention cannot be used to model the causal case nor can it be combined with sliding window attention.
• The Q and K matrices are related to attention because our model effectively computes attention between tokens of the sequence and latent (token) states. An intuitive example is to think about each latent state as a general concept like shape or colour. Consider computing the attention between two tokens “red” and “green”. In standard attention, this similarity comparison goes directly via the token embeddings of these two concepts. However, one might argue that the reason “red” and “green” should be similar is because they are both colours. Therefore, in Latte, we know that “red” and “green” are similar since “red” is similar to the latent concept/token “colour”, and “green” is also similar to “colour”. This is the essential intuition behind Latte — that similarity between tokens can be more succinctly expressed by similarly between tokens and learned latent tokens.

2. • There are in-depth comparisons with other sequence models in Table 2 and, at the reviewer’s suggestion, we further introduced GLA and Lightning Attention, showing that Latte-Macchiato outperforms those approaches.
   • As a humble university research department, we do not have the compute necessary to run large-scale experiments and the experiments we have presented are already at the limit of our computational resources. We provided comparisons of Latte with other linear models showing the potential of our idea. We hope that others with larger computational resources can be inspired by these results and explore the scaling of Latte to larger models. Latte is, we believe, a remarkably simple and easily implementable drop-in replacement for standard attention. The simplicity of Latte means that it can also be easily integrated with standard sliding window attention, as we have shown with Latte-Macchiato.
   • We understand the reviewer's concern that our distillation results don’t necessarily provide evidence of scaling in terms of pre-training. However, we strongly believe that our results show evidence of Latte-Macchiato working at scale. Our distilled model is a 2.7B parameter Gemma model extended with Latte-Macchiato and we show in Table 4, that it outperforms the original model on autoregressive language modelling. We believe this provides clear evidence of Latte’s scaling capabilities. Naturally, we would like to apply the same approach to much larger base models, but this is currently beyond our computational resources.
3. • No. Originally, we only compared it with standard causal attention (without Flash Attention) because Flash Attention is a hardware optimisation of standard attention. The time complexity of Flash Attention is still quadratic and therefore will be slower than linear attention beyond a certain sequence length. The crossing point is naturally dependent on hardware optimisation. Our work is not yet hardware-optimised; we are working on a CUDA kernel, but it is not the main goal of this paper and (like Flash Attention) would be a separate research contribution. However as requested, we have now included a comparison with Flash Attention in Figure 11, Appendix C and show that for long sequences Latte is faster.
   • All linear methods have linear slopes in terms of time versus sequence length, albeit with potentially different slopes. This means that all linear methods (including Latte) will eventually outperform any quadratic attention (even those that are hardware-optimised). For this reason, we do not feel it useful to plot the scaling of all other linear attention mechanisms we considered.
4. • This is a very good suggestion. We now incorporated MQAR [Arora] synthetic dataset in our experiments and show in Appendix D that, for our settings, the model performs better than other linear models. We use a window of size 128 for full attention, which is smaller than the entire context. Whilst our method is, we believe, state-of-the-art for linear scaling methods, finding a linear scaling attention approach that preserves excellent performance on needle-in-a-haystack problems remains an elusive goal for the research community.

Weaknesses

1. • An advantage of our method over Mamba and other Linear Attention models is that it preserves the probabilistic properties of attention. This has intuitive appeal and practical benefits such as correct normalisation over sequences and the fact that we preserve some interactions between tokens through our latent variables.
   • The reviewer is correct that Latte-Macchiato (the extension of Latte) is dependent on SWA. However, we believe this to be a key advantage of our approach as we combine both linear and SWA attention with correct normalization such that their summation remains a valid probability distribution over tokens. We believe this is important to avoid overestimating tokens in the window over tokens outside of the window. We thank the reviewer for raising this point and we will make sure to amend this in our updated paper.

• Our method does not have slower theoretical inference — it is still constant next token prediction. Our model also supports hardware optimisation — however, hardware optimisation is not the main aspect of our paper and will be done in future extensions.

• We thank the reviewer for raising this point as we believe that our mathematically sound integration of SWA with Latte is a key advantage over models such as Luna. Indeed, models such as Luna do not support the combination of local standard attention and linear attention, whereas in Latte we combine linear attention with SWA whilst ensuring a valid probability distribution to attend over token indices.
• Luna has never been applied to autoregressive language modelling. In Table 2 we perform several of our own comparisons with other linear models and show that Latte achieves state-of-the-art performance on autoregressive language modelling. We also added a new experiment in Appendix D showing that our model is better at retrieval tasks than other linear models.

• We would like to stress that our model does not require Jax and is agnostic to whichever machine learning library one might prefer. We chose Jax for simplicity given the pre-existence of the “scan” function, although we stress again that this could be implemented in PyTorch if the user so wished.

• While such a comparison would be interesting, checking all possible model configurations is beyond the scope of our paper. In our paper, we performed a thorough investigation of the benefits of different components and believe that the novelty of our work comes from how we preserve a valid probability distribution over token indices while combining standard causal attention and linear Latte attention. We do not know of any other linear sequence models that can be combined with standard attention while preserving a probability distribution.

5. • Thanks, we fixed the typo.

Questions:

• No. Originally, we only compared it with standard causal attention (without Flash Attention) because Flash Attention is a hardware optimisation of standard attention. The time complexity of Flash Attention is still quadratic and therefore will be slower than linear attention beyond a certain sequence length. The crossing point is naturally dependent on hardware optimisation. Our work is not yet hardware-optimised; we are working on a CUDA kernel, but it is not the main goal of this paper and (like Flash Attention) would be a separate research contribution. However as requested, we have now included a comparison with Flash Attention in Figure 11, Appendix C and show that for long sequences Latte is faster.

2. • Yes, we find that Latte also extrapolates. Since Latte Macchiato combines standard sliding window attention and our latte attention, the excellent extrapolation properties of Latte Macchiato derive from both SWA and the Latte mechanism itself.