MetaLA: Unified Optimal Linear Approximation to Softmax Attention Map

Thanks for your insightful feedback and your time in reading our paper.

Weakness1: More long sequence validation besides the LRA benchmark. This is not sufficient to justify the scalability of the method for real LLMs.

A: Thank you for the suggestion. In addition to the LRA benchmark, we address longer token sequences in LLM pre-training. Specifically, we pre-train a MetaLA-based language model with token counts ranging from 15B to 300B and a sequence length of 2048. To demonstrate MetaLA’s scalability, we also extended the 3B scale model experiment to 300B tokens on the CommonSense Benchmark as follows, where MetaLA showed strong performance.

Weakness3: About the low-precision training of FP16 and BF16.

A: Our model is trained using fp16, and most linear models, such as [2,3,4], can be trained with half-precision, including bf16 and fp16. The numerical precision in our training shares the same characteristics as theirs.

Weakness4: Computational efficiency and memory requirements compared to softmax attention and other linear methods.

A: We report the throughput and memory footprint of 1.3B models in the general response (please refer to the PDF). The findings indicate: (1) Our model demonstrates good linearity, maintaining processing speed and memory efficiency with increasing sequence length, unlike the Transformer, which experiences a sharp drop in token processing speed as sequence length increases. (2) Our model matches the computational efficiency of linear models like GLA [2] in both latency and memory, and is significantly faster than Mamba [5], which also has linear complexity.

Weakness5: Code Link.

A: https://anonymous.4open.science/r/MetaLA-1BB3/README.md

Weakness2 and Q1: It is not clear how their model performs in the text generation scenario. 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.

A: We have added some additional experiments based on your suggestions. We validate MetaLA in NarrativeQA, a natural language processing task that challenges LLMs to comprehend and answer questions about long-form narratives. The results show that MetaLA performs well on this long-sequence tasks.

Q2: About the hyperparameters.

A: Our pre-training hyperparameter settings were based on the experience summarized by predecessors [6]. The majority of the settings are identical to those used in Pythia. Our hyperparameter settings are provided in the code link as a YML file, based on the parameter passing format used in gpt-neox [7].

Q3: Relathionship between the model in paper "Linear Log-Normal Attention with Unbiased Concentration" and the general linear attention model in this work.

A: This paper and our work are complementary. In lines 633-645, we discuss that the approximation of softmax attention can be divided into two parts: value approximation and functional approximation. Our work primarily focuses on functional approximation, while the paper 'Linear Log-Normal Attention with Unbiased Concentration' that you mentioned belongs to value approximation.

The paper 'Linear Log-Normal Attention with Unbiased Concentration' is a commendable work, offering extensive mathematical derivations and identifying the Log-Normal method as suitable for linear models to approximate the softmax attention map from the perspectives of entropy and spectral angles. This method has demonstrated strong performance in language models, particularly on the GLUE benchmark, and highlighted the linear complexity advantage of Linear Attention in long-sequence modeling. Our paper, on the other hand, unifies the token mixing mechanisms of Linear Attention, Linear RNN, and State Space Model, linking them to the softmax attention map from the viewpoint of attention maps. We analyze under what circumstances these models can functionally approximate softmax attention.

Thus, both papers offer complementary insights into the theoretical analysis of the softmax attention map from different perspectives, indicating great potential for their combination. We are genuinely inspired by this work and will cite the paper in our own.

[1] Scrolls: Standardized Comparison over Long Language Sequences, In EMNLP 2022.

[2] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.

[3] RWKV: Reinventing RNNs for the Transformer Era, In EMNLP 2023.

[4] TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer, In Arxiv 2023.

[5] Mamba: Linear-time Sequence Modeling with Selective State Spaces, In COLM 2024.

[6] Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling, In ICML 2023.

[7] https://github.com/EleutherAI/gpt-neox

[8] Llama 2: Open foundation and fine-tuned chat models, In Arxiv 2023.

[9] HGRN2: Gated Linear RNNs with State Expansion, In COLM 2024.

[10] You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet, In Arxiv 2024.

Thank you for the rebuttal. The authors have addressed my concerns regarding long sequences and numerical formats. Additionally, they provided further experimental results, comparing their method to the LLaMa2 and MAMBA models. Overall I think it's a good paper and following the additional information provided by the authors, I have decided to increase my score.

Thanks for your insightful feedback. We will outline your suggestions and questions, followed by our detailed responses. We hope our answers address your concerns.

Weakness1: 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...

A: We sincerely apologize for any confusion in the organization of this section. We carefully considered and discussed your suggestions, and thank you for your kind advice. We find ourselves in a bit of a dilemma. On one hand, we strive to maintain the completeness and rigor of the paper by including the proposition statement. On the other hand, we have to carefully consider space limitations, which may increase the difficulty for readers. We hope you can understand the balance we are striving to achieve. While we have not yet identified a perfect solution, we will make every effort to enhance the clarity of our presentation while taking your concerns into account. We greatly appreciate your understanding and patience.

Weakness2: The MQAR analysis is weak. The results of more difficult versions of this task and an attention baseline will be conducive to understanding.

A: Based on your suggestion, we evaluate the performance of several models on more challenging settings. The attention baseline benefits from global modeling capabilities, achieving optimal results under both conditions discussed in the paper (>99.0). The additional experiments show that MetaLA outperforms Mamba (does not converge under the same training conditions), and there is still a significant gap compared to transformers:

For the results of longer sequences such as 1024, since it takes a long time for the linear model to converge, it is difficult for us to get the results during the rebuttal period. We sincerely ask for your understanding. As a supplement, we test another long sequence recall benchmark NIAH. Details are in the response to weakness3.

Weakness3: More downstream tasks that require the recall/retrieval abilities should be included.

A: According to your suggestions, we evaluate several models on MAD [1], a collection of six synthetic tasks predictive of scaling laws, including recall, memorization, and compression. As shown in the following table, MetaLA achieves the best results across various linear complexity models.

Three distinct versions of the recall task and the selective copy task serve to underscore the effectiveness of our model in the domain of information retrieval. And we will cite the relevant articles. We also present experimental results on the Needle in a Haystack (NIAH) task, which is designed to evaluate the in-context retrieval capabilities of LLMs:

Retrieval ability in long texts is a significant challenge for linear models, as all current linear models lack good solutions to this problem. Nonetheless, MetaLA has achieved satisfactory results in comparisons among linear models. Compared to Transformer models, this performance is still insufficient. This is precisely the issue we will address next, following the unification of linear model forms.

Weakness4: Some of these questions should be better explored in this work .

A: Thank you for your valuable suggestions. During the rebuttal period, we conducted validation for MQAR, MAD, and long-text tasks based on your feedback. Additionally, we supplemented pre-training from 370M to 3B model scales on 300B tokens. We conducted more experiments around the existing work in this paper, verifying the advantages and shortcomings of MetaLA. Regarding other directions mentioned in the discussion, such as the issue of sufficient training [3] and benchmark evaluation, these are based on different fields of foundational models. We fully acknowledge the research value of these directions, which are essential for improving linear foundational models and will be part of our future work.

[1] Mechanistic Design and Scaling of Hybrid Architectures, In Arxiv 2024.

[2] HGRN2: Gated Linear RNNs with State Expansion, In COLM 2024.

[3] Never Train from Scratch: FAIR COMPARISON OF LONGSEQUENCE MODELS REQUIRES DATA-DRIVEN PRIORS In Arxiv 2024

[4] Mamba: Linear-time Sequence Modeling with Selective State Spaces, In COLM 2024.

[5] Llama 2: Open foundation and fine-tuned chat models, In Arxiv 2023.

[6] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.

Thank you very much for your thoughtful feedback and for raising your score. We sincerely appreciate your valuable insights, and we will make sure to thoroughly discuss the gap with global attention in recall-intensive tasks as you suggested. Additionally, we will include a softmax attention baseline for the MAD tasks in the final version to enhance the comparison. Your comments have been instrumental in improving our work, and we are grateful for your time and consideration.

Thank you for your insightful feedback. Your questions are crucial and something we have been thinking about.

Weakness1: The experiment is not done with various scales. It is hard to know the effect of the scaling model and training dataset with MetaLA.

A: Based on your suggestions, we conduct additional experiments. We further evaluate our model ranging in size from 380M to 3B, trained with 300B tokens, on the CommonSense Reasoning benchmark. Below are experimental results:

The experimental results corroborate the scalability and performance of MetaLA.

Weakness2: 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.

A: Thanks for your suggestion, we report the latency/throughput results of 1.3 B models in the general response to all reviewers, please refer to the PDF file. The report indicates that:

• Our model exhibits good linearity, maintaining processing speed with increasing sequence length, in contrast to the Transformer, which shows a sharp drop in token processing speed as sequence length increases.
• Our model is computationally as efficient as linear models like Gated Linear Attention [1], while being significantly faster than the Mamba model [2], which also has linear complexity.

Question1: 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?

A: You raise a very profound and interesting question, one that we have been thinking about and we would love to discuss with you.

The AI (performance) perspective: Although there is no strict theoretical proof of optimality for softmax attention, its excellent performance has been empirically validated by many large models, making it undoubtedly the foundation of existing large models.

The neuroscience (bio-plausibility) perspective: It is almost certain that softmax attention is NOT the optimal solution for sequence modeling. To achieve AGI (Artificial General Intelligence), there are two constraints we cannot bypass.

• Power and Latency: In terms of power and latency, the human brain operates very stable with bounded power consumption and trascient reaction. In contrast to Transformer models, if we assume the sequence input and output are interactions over time with the environment, then it faces a severe issue. Over time, its power consumption and latency will grow, a limitation brought by quadratic complexity.

• Physical Space Limitation: As we pointed out in the paper, softmax attention is an infinite-state mechanism, leading to an increasing demand for physical space (for storage) as the model runs. This contradicts the human brain's ability to maintain a large and efficient memory capacity despite its finite size.

It can be seen that the perspective from which we view softmax attention determines whether it is optimal, and there are trade-offs involved. Our approximation strives to restore its performance optimality (which has been empirically validated), but for aspects where softmax is not optimal, we still possess unique advantages.

Question2: Can you provide the code?

A: We provide the code link here: https://anonymous.4open.science/r/MetaLA-1BB3/README.md

[1] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.

[2] Mamba: Linear-time Sequence Modeling with Selective State Spaces, In Arxiv 2023.

[3] Llama 2: Open Foundation and Fine-tuned Chat Models, In Arxiv 2023.

[4] You Only Scan Once: Efficient Multi-dimension Sequential Modeling with LightNet, In Arxiv 2024.

Thanks for your insightful feedback and your time in reading our paper.

Weakness 1: Concerns about the optimality for linear attention.

A: Your concern is valid. There is no definitive evidence to prove that MetaLA is the optimal approximation of self-attention. Please allow me to explain why we chose to use this term.

• First, we repeatedly state in the paper that we are discussing the necessary conditions for optimality. In fact, these necessary conditions highlight an important issue, namely, what kind of models can approximate softmax attention.

• Second, we use the term "optimal" to draw the community's attention to the issue of how to approximate softmax attention. We believe that better approximating softmax attention is one of the potential ways to address cutting-edge issues in linear models, such as retrieval and in-context learning.

Weakness 2: The presentation in Section 2.

A: Thank you for your suggestion. We will carefully check and optimize these expressions.

Q1: Explanation of the mask matrix $M$.

A: Your understanding is correct. For the matrix $M$, the lower triangular part (including the diagonal) is set to 1 to maintain model causality, ensuring that the output at time $t$ depends only on inputs before $t$. As for the upper triangular part (excluding the diagonal),

• Linear attention: The upper triangular elements are set to 0, as no exponential operation is involved.

• Softmax attention: The upper triangular elements are set to $-\infty$, ensuring that their exponentiated values become 0, thereby not affecting the normalization of values before $t$.

Q2: About the complexity of linear attention and the chunk-wise algorithm.

A: As shown in Line 78, during inference, the time and memory complexity per token is $\mathcal{O}(1)$, meaning that computational and memory costs remain constant. From the parallel formula Eq.5 in Line74, we can express the inference in a recurrent form (operating separately on the numerator and denominator):

$S_{t} = \sum_{i=1}^{t}\phi^{T}(k_{i})v_{i} = S_{t-1} + \phi^{T}(k_{t})v_{t},$

$n_{t} = \sum_{i=1}^{t}\phi^{T}(k_{i}) = n_{t-1} + \phi(k_{t}),$

$o_{t} = \frac{\phi(q_{t})S_{t}}{\phi(q_{t})n^{T}_{t}}.$

In this iterative form, each operation only needs to maintain a fixed state size $S_{t}\in\mathcal{R}^{d \times d}, n_{t}\in\mathcal{R}^d$, resulting in a storage complexity of $\mathcal{O}(1)$ (occupying $d^{2}+d$ space, where $d$ is a constant). Furthermore, each token computation involves $d^{2}+d$ multiplications and additions, ensuring the time complexity is also $\mathcal{O}(1)$.

About the chunk-wise algorithm. The Chunk-Wise algorithm combines the low complexity of recurrent operations with the high throughput of parallel processing during training. It achieves this by using serial processing between chunks and parallel processing within chunks, ensuring (1) full utilization of GPU computational units within a chunk and (2) reduced redundant calculations of historical information between chunks.

For Linear Attention, normalization by the denominator is typically omitted. We modify the earlier formulas for further discussion (a more general form is provided in Section 4.2 of [1]):

$S_{t} = S_{t-1} + k_{t}^{T}v_{t},$

$o_{t} = q_{t}S_{t},$

During training, if using a parallel form of computation, i.e., $O = (QK^{T}\odot M)V$, these matrices must be computed from left to right, resulting in a computational complexity of $\mathcal{O}(L^{2})$. To reduce the time complexity of the linear model during training, researchers use a chunk-wise algorithm, which consists of two parts: intra-chunk and inter-chunk. We divide a sequence of length $L$ into $\frac{L}{C}$ chunks (with padding if not divisible), each chunk of size $C$. The computation method is as follows(The first term on the right side of the equation is the inter-chunk, and the second term is the intra-chunk.):

$O_{[i]} = Q_{[i]}S_{[i-1]} + (Q_{[i]}K_{[i]}^{T}\odot M)V_{[i]},$

$S_{[i]} = S_{[i-1]} + K_{[i]}^{T}V_{[i]}.$

Thus, each chunk will only generate a parallel-like computation, where the resulting $C \times C$ matrix is a submatrix of the attention map. So the computational cost is $\frac{L}{C} [(Cd^{2}+C^{2}d)+Cd^{2}]$. If $d$ and $C$ are constants, the time complexity of the computation becomes $\mathcal{O}(L)$.

Q3: Why self-attention requires infinite states.

A: The computation for the $t$-th output of self-attention during inference is given by:

$o_{t} = \text{softmax}(q_{t}K_{t}^{T})V_{t} = \frac{\sum_{i=1}^{t}\langle q_{t}, k_{i}\rangle v_{i}}{\sum_{j=1}^{t}\langle q_{t},k_{j}\rangle}$

where $K_{t}^{T} = [k_{1}, \cdots, k_{t}]$ and $V_{t}^{T} = [v_{1}, \cdots, v_{t}]$.

In the Transformer model, a KV-Cache is maintained as a state, updated after processing each input to avoid redundant computations. Upon receiving the $(t+1)$ -th input, the state is updated by concatenating $k_t+1$ and $v_t+1$ to $K_t$ and $V_t$, respectively, which increases memory requirements over time.

Q4: "good enough" in line 177.

A: In the theoretical analysis, we solved for the $Q,K,\Lambda$ matrices given any attention scores distribution. This implies a prerequisite assumption: the parameter functions $f_{q/k/\alpha}$ can project the input $X$ into the corresponding matrix space. This prerequisite imposes constraints on $f_{q/k/\alpha}$ (expressive enough) and $X$ (good enough).

Q5: The explanation of "fewest parameters"

A: The minimum parameters refer to the smallest parameter group that achieves optimal approximation. If the model dimension $d, d_k, d_v$ are fixed, the minimum parameters equates to the smallest group. Please refer to line 185.

[1] Gated Linear Attention Transformers with Hardware-Efficient Training, In ICML2024.

Thank you for your thoughtful feedback and for increasing your confidence in our paper. We genuinely appreciate your time and effort in reviewing our paper and responses. Your support and understanding are invaluable to us. If you have any further questions or need additional clarification, please let us know. Thank you once again for your help.