MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map

Summary

Proposes a unifying view of recent linear attention/SSM/linear RNN methods and compares these models to Softmax Attention. Proposes MetaLA to address the shortcomings of the prior methods in approximating Softmax attention. Performs experiments on MQAR, language modeling, image classification and LRA.

Strengths

- Addresses an important and relevant topic of analyzing efficient sequence models that attempt to replace Softmax Attention - Systematically comparing and unifying the formulations of the various linear attentions and linear RNNs/SSMs is useful - The tables classifying similarities and differences of the models are nice - The theoretical results in the appendix appear to be sound - The experiments performed suggest promising results

Weaknesses

- The general presentation of the "optimal linear approximation" and theoretical analysis is dense and confusing and obscures the contributions. - I would recommend moving the Proposition statements to the main text (the proofs can remain in the appendix). E.g. Currently I get to line 218 and am told that the key value is being dropped based on theoretical analysis performed in Appendix A2, but I would rather already have an idea of this result before getting to this point. Similarly, the "self-augmentation" change seems to come out of nowhere. - Once you have moved these Proposition statements to the Section 4, I would recommend rewriting/reorganizing Section 4 framed around these proposition statements - The MQAR analysis is weak. I would like to see more difficult versions of this task (larger vocab sizes, more retrievals, longer contexts) and an attention baseline to better understand the limitations and differences between the different methods. - I think more downstream tasks should be included that require the recall/retrieval abilities that are known to affect all these fixed state size methods (e.g. https://arxiv.org/abs/2402.01032, https://arxiv.org/abs/2402.04248, https://arxiv.org/abs/2402.18510v3 to name a few, more generally these and other works should be cited) would make the proposed approach and value of the proposed framework more convincing. This would help the reader understand how much progress as actually been made toward approximating Softmax attention with more efficient methods. Even simple needle in the haystack, passkey retrieval, or phonebook retrieval (as in the Griffin paper) would help give a better sense of this. - In the conclusion and discussion, some potential weaknesses and questions regarding the differences between Softmax attention and other methods are mentioned and proposed as future work (model capabilities, insufficient training, or eval issues). However, I would suggest some of these questions should be better explored in this work since that would help provide a better sense of how useful the proposed framework and method is.

Questions

My main questions and issues are listed above. - Line 287: typo, "Value" is misspelled in the text

Rating

6

Confidence

4

Soundness

2

Presentation

1

Contribution

2

Limitations

Sufficient

Summary

The paper presents a theoretical analysis of existing linear attention methods such as LinFormer, SSM, and LinRNN. Building on this analysis, the authors propose a unified framework that combines the strengths of these methods. Utilizing this framework, authors develop a novel linear attention model called MetaLA. The key innovations of MetaLA include the elimination of the Key matrix, the introduction of dynamic decay for achieving dynamic memory and static approximation, and the integration of self-augmentation with short convolution. These enhancements aim to improve the approximation accuracy of softmax attention and outperform existing linear attention methods.

Strengths

- The paper presents a comprehensive unification of various linear attention models (LinFormer, SSM, LinRNN), offering a deeper understanding of their underlying mechanisms and differences. - The proposed MetaLA model is theoretically grounded based on the proposed framework. - The authors provide empirical evaluation of the MetaLA model on NLP tasks, demonstrating its effectiveness and robustness.

Weaknesses

- The only result for long sequences provided by the authors is on the LRA benchmark. This is not sufficient to justify the scalability of the method for real LLMs trained on sequences longer than 1k tokens since LRA is an artificial and non-representative benchmark. - Authors provide empirical results mainly for classification tasks. It is not clear how their model performs in the text generation scenario. - The authors neither mention nor present any experiments with low-precision formats like FP16 and BF16. Many linear attention methods fail when trained in these popular formats. - The authors do not provide any measurements or comparisons of computational efficiency and memory requirements of their method compared to softmax attention and other linear methods. - Since the authors do not provide code for the MetaLA model, it is hard to assess the practicality and ease of implementation.

Questions

- How does MetaLA perform on tasks with longer sequences compared to standard softmax attention and other linear attention mechanisms? For example, on NarrativeQA or other tasks from LongBench. - Can the authors provide more information on the setting of hyperparameters for their method, specifically whether the decay factor should be selected by the user? - Another linear attention work [1], which the authors did not cite, proposes evaluating the concentration ability of linear attention using entropy and spectral gap metrics. What is the authors' opinion on these concentration metrics in the context of the MetaLA framework analysis? [1] Linear Log-Normal Attention with Unbiased Concentration

Rating

8

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Summary

They proposed MetaLA, which solved problems of previous attention alternatives (LinRNN, SSM, LinFormer). They start to build their MetaLA by deriving general form of linear alternative of softmax attention. 1. They remove the K matrice redundant parameters and achieve better training efficiency 2. Add self-augmentation. This allows the input value to immediately affect the output depending on the query to prevent forgetting current token information. 3. Adding a convolutional layer at input X. This enhances the local interaction further, motivating Mamba and Griffin.

Strengths

Strong build-up toward their method from well-curated baselines. Their method shows high performance empirically.

Weaknesses

1. The experiment is not done with various scales. It is hard to know the effect of the scaling model and training dataset with MetaLA. 2. The performance of MetaLA is mainly shown with only benchmark scores, and there are no latency reports. Therefore, it is hard to know if this performance is Pareto optimal in the trade-off between latency and performance.

Questions

1. Is approximating softmax attention the optimal solution for sequential modeling? This paper aims to approximate the various characteristics of the computational and model aspects very well. However, I wonder if softmax attention is not the optimal solution for sequential modeling. Is there any justification for it? 2. Can you provide the code?

Rating

7

Confidence

2

Soundness

3

Presentation

4

Contribution

3

Limitations

It feels minor improvement on top of previous linear attention alternative studies.

Summary

This paper discusses the development and evaluation of linear complexity models for Transformers, which aim to replace the conventional softmax attention mechanism. The authors first unify existing models including LinFormer, SSM, and LinRNN into the framework of Linear Attention. Then they establish three (actually four) criteria for optimal linear attention: (0) linear complexity, (1) dynamic memory capability, (2) static approximation ability, and (3) minimal parameter usage. This paper shows that the three existing models mentioned above do not meet all these conditions and therefore perform suboptimally. In response, they introduce a new model, Meta Linear Attention (MetaLA), designed to fulfill these criteria. MetaLA is tested across various tasks including the Multi-Query Associative Recall (MQAR) task, language modeling, image classification, and the Long-Range Arena (LRA) benchmark, where it outperforms existing linear models, demonstrating its effectiveness.

Strengths

1, This work unifies three widely-applied linear-complexity substituitions (LinFormer, SSM, LinRNN) of Transformer into the same framework, including a recurrent form and a parallel form, which enables higher-level studies and analysis. 2, This work proposes three criteria for evaluating a linear-complexity model (memory, approximation, parameter), which brings insights and targets for future developments of linear-complexity models. 3, Combining 1 and 2, this work evaluates whether each of the three existings models satisfies each criteria, and shows that every model does not fulfill all requirements. Based on this insight, they propose a MetaLA Transformer model that satisfies all three criteria and performs well in multiple tasks.

Weaknesses

(After rebuttal: The authors carefully addressed my main concerns.) -------------- Please correct me if there are any mistakes. 1, The main concern originates from the three criteria of "**optimality**" for linear attention (LA). In general, "optimality" indicates a situation that any possible method/result cannot be substantially better. However, it seems like these three criteria are all *necessary* conditions. As far as I noticed, this work has not stated whether these criteria are *sufficient* for optimality. As a consequence, even if MetaLA satisfies all three criteria, it is not proper to state that MetaLA is *optimal*. 2, The presentations could be improved, especially in Section 2. It is important to define or denote the dimensionality of not only each matrix/vector (which the authors have done) but also each function and operation (which they have not).

Questions

The authors are encouraged to disagree or discuss on my questions/comments. 1, On Page 2 Line 67, what is the mask matrix $M$ like? From the context, I guess $M$ is a lower triangular matrix with all valid elements $=1$? 2, On Page 3 Line 78 (top), why does the LinFormer consume $O(1)$ time and memory complexity? Also, would you please elaborate on the "Chunkwise algorithm" mensioned in Line 80? 3, On Page 5 Line 147, why this softmax-generated expressive attention map requires infinite states? 4, On Page 5 Line 177, what is the condition (or consequence) of "good enough" inputs? 5, On Appendix A 2.3 (and also A 2.4), it is claimed that (Q, $\Lambda_t$) can meet C0, C1, C2 with fewest parameters. Is this concept of "fewest" restricted in a certain domain of model selection? In other words, do you have a lower-bound proof showing that any model utilizing fewer parameters than (Q, $\Lambda_t$) will fail?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

Limitations are adequately discussed in the Section of Conclusion and Discussion.

Thanks for the replies

The reviewer thanks the authors for your thorough and informative explanations. I am not an expert in this area, and therefore I am not able to advocate for this paper. However, I will increase my confidence from 2 to 3 as a reflection of what I have learned from your response.