Gating is Weighting: Understanding Gated Linear Attention through In-context Learning

Weakness 1: The architecture considered in this paper appears to include only the GLA layer and no MLP layer. Consequently, a multi-layer GLA in section 3.2 would not fully align with a Transformer model. A discussion on the effects of the MLP layer would provide valuable insights.

Response: We appreciate the reviewer’s suggestion. The architecture in this paper includes the GLA layer without incorporating an MLP layer. Following your valuable feedback, we have added the following paragraph below Theorem 2 in Section 3.2:

Our theoretical results in Theorem 2 focus on multi-layer GLA without Multi-Layer Perceptron (MLP) layers to isolate and analyze the effects of the gating mechanism. However, MLP layers, a key component of standard Transformers, enable deeper feature transformations and non-linear interactions, potentially enhancing GLA's expressive power. Future work could explore the theoretical foundations of integrating MLPs into GLA and analyze the optimization landscape of general gated attention models, aligning them more closely with conventional Transformer architectures (Gu & Dao, 2023; Dao & Gu, 2024; Peng et al., 2024).

Please refer to the first paragraph (highlighted) after Corollary 1.

Weakness 2: In Section 5, the multitask setup appears to be simplified through the introduction of vectors $d$ as task boundaries and vectors $c$ as contextual features. Could you clarify the effect of each of $c$ and $d$ on optimal loss in this setting? Additionally, it would be insightful to evaluate the performance impact of including versus excluding each of $c$ and $d$ when training on real data.

Response: Thank you for the insightful question regarding the role of contextual features in Section 5. Here is our detailed response:

Theoretical and empirical perspective:

• As shown in Theorem 1, the weighting matrix $\Omega$ is highly complex, as it depends heavily on the non-linear gating function and the input embeddings. To ensure that the optimal weighting induced by gating remains predictable, we introduced contextual features $c$ and $d$.

• The green curves in Figure 1 illustrate that, without delimiters, a one-layer GLA cannot achieve optimal performance, as demonstrated by the theoretical predictions (black dashed curves). However, with delimiters, the optimal loss is achievable, as shown in Figure 1(a-c), where red and black curves overlap.

• Theorem 5 demonstrates that when $c$ and $d$'s are linearly independent, the choice of these contextual features does not affect the optimal loss, and any valid contextual features will result in the same loss. However, when $c$ and $d$ are not linearly independent, their influence becomes significant. In such cases, better optimization can be achieved by assigning relevant contextual features to related tasks and distinct contextual features to unrelated tasks.

Practical perspective on real data: Real data often includes inherent structures like verbs, objects, and nouns in language data. Proper embeddings allow models to categorize and distinguish these elements into different tasks effectively. However, in our setting, the input $z \in Z$ consists of random features $x \sim \mathcal{N}(0, I)$ without task-specific information. In such scenarios, without contextual features, the model struggles to distinguish between tasks based solely on the random input. As demonstrated empirically in Figure 1 (green curves), excluding contextual features leads to significantly degraded predictions by the GLA model. Additionally, contextual features have been utilized in prior works such as Wang et al. (2024), Asai et al. (2022), and Dun et al. (2023). These works demonstrate the practical significance of contextual features in handling multitask setups. We have expanded the discussion in the paper to include these insights and references.

Weakness 3: Minor comments

Response: Fixed. Thank you!

Weakness 1: The paper has strong theoretical contributions. But, more empirical studies and comparisons could strengthen the practical applicability of the theoretical contributions.

Response: We thank the reviewer for acknowledging the strong theoretical foundation of our work. Regarding the suggestion for additional empirical studies, in Appendix A.1, we have extended our experiments to multi-layer GLA models and verified that deeper models achieve better predictions. Furthermore, numerous prior works (e.g., Table 1 in Yang et al. (2024)) have successfully implemented GLA in real-world applications. While our study does not include real-application results, our theoretical contributions are both foundational and novel, and we believe our work provides significant contributions (as discussed in the General Response).

Weakness 2: There exists another line of work (which is not based on in-context learning) such as papers cited below. They propose an implicit self-attention formulation for GLA models (including Mamba and RWKV). Do you think there is a connection between your work and this line of work, and is it possible to apply in-context learning for model explainability and empirical studies?

Response: We appreciate the suggestion to explore connections with implicit self-attention frameworks.

Zimmerman et al. (2024) propose a framework demonstrating how various architectures, including GLA models like Mamba and RWKV, can be viewed under a unified implicit attention perspective. This aligns with our theoretical exploration of GLA’s data-dependent Weighted Preconditioned Gradient Descent (WPGD), as both approaches emphasize data-adaptive weights driven by gating mechanisms. We believe that the theoretical results in this paper, combined with the unified perspective of GLA as a variant of attention, offer a potential pathway for extending GLA’s analysis to the optimization landscapes of models like Mamba and RWKV and connecting them to WPGD.

Zong et al. (2024) leverage a gated cross-attention mechanism for robust multimodal fusion. Their approach emphasizes stable integration of heterogeneous data streams. While their task differs, the underlying gated mechanism aligns with GLA’s capacity to manage multi-task prompts by dynamically weighting inputs. This suggests that GLA’s gating mechanism can be repurposed for tasks beyond sequence modeling, including robust multimodal fusion. We have incorporated a summary of the above discussion into the Related Work section (Section 1.1)

Weakness 1: They claim that they establish ... which is far away from your claimed equivalence.

Response: We appreciate the reviewer pointing out the ambiguity in our use of the term “equivalence”.

• In Section 3.1, the equivalence refers to the fact that a one-layer GLA model under the specific construction implements one step of WPGD considering the weighting matrix $\Omega$ determined by the choice of gating function and input space. We have clarified this distinction in the revised manuscript to avoid misinterpretation.
• In Section 5.2, the equivalence refers to the optimization prediction for GLA and WPGD being identical under certain data assumptions. This result aligns with the findings presented in Theorem 1 of Ahn et al. (2023) and Proposition 1 of Li et al. (2024). However, our work extends these analyses by considering a more complex model architecture, i.e., GLA, and data settings, i.e., task-mixture ICL.

We kindly refer the reviewer to our response to Weakness 4 for further justification regarding the "strong" assumptions.

Weakness 2: In the preliminary section, you introduce ... involved token embedding matrix in Equation 20 needs more verification.

Response: To clarify, we address each point of this comment sequentially:

• The embedding matrix Eq. (4) is used in Sections 2 and 3, including Eq. (5), Theorem 1, and relevant portions of the text.
• In Section 3, we demonstrate that a one-layer GLA performs one step of WPGD with the weighting matrix $\Omega$ determined by the gating function and input embedding. However, we do not claim that their optimization landscapes are the same. Analyzing the optimization of GLA with the token embedding in Eq. (4) is challenging because the weight matrix $\Omega$ heavily depends on the embedding (as per Theorem 1), and the embedding itself incorporates random features $x \sim \mathcal{N}(0, I)$. To further investigate GLA's optimization performance, we introduce the embedding matrix in Eq. (20) with additional delimiters in Section 5. This ensures that the optimal data-dependent weighting is both analyzable and achievable. Furthermore, our empirical results show that without delimiters, the performance is non-smooth, non-optimal (as shown by the green curves in Figure 1), and cannot achieve the theoretically optimal performance (as indicated by the black dashed curves).
• Regarding practical relevance, delimiters can be interpreted as "task transition" identifiers. For instance, delimiters/prompts have been used in the mixture-of-prompt literature, such as Eq. (1) in Wang et al. (2024), Fig. 2 in Asai et al. (2022), and Section 2 in Dun et al. (2023), to separate different tasks. We appreciate the reviewer’s comment and have added this discussion to the paper.

Weakness 3: You claim that you prove ... how large the risk gap is and what magnitude it enjoys.

Response: Thank you for suggesting a clearer comparison of the ICL risks for GLA and LSA. Based on your valuable feedback, we have added Corollary 4 to our paper, explicitly demonstrating that GLA outperforms linear attention in terms of ICL risk. This corollary now quantifies the risk gap and specifies the assumptions under which the comparison holds.

Strength 2: If the results stated in the paper are correct, it would be exciting and reassuring to know that a very broad class of models, beside Transformers, is provably capable of ICL despite some findings in the literature (e.g. [1])

Response: Thank you for recognizing the strengths of our paper. Our framework builds on the observation that the Mamba architecture leverages a gated linear attention layer. Notably, our findings align with those of Jelassi et al. [1], which highlight the limitations of recurrent models in memory-intensive ICL tasks like associative recall. In contrast, we focus on characterizing algorithms expressible through gated linear attention, particularly for linear regression-type ICL tasks. Additionally, by distinguishing between "gated linear attention and linear attention" and "scalar-valued and vector-valued gating," our results shed light on how more advanced gating mechanisms can more effectively use the memory to enhance recall capabilities.

Weakness 1: Presentation is not self-contained, and it’s difficult to follow the arguments without prior reading of several referenced papers,... as well as attempting to independently draw missing analogies and reductions between them and the reviewed paper.

Response: We appreciate the reviewer’s feedback on the presentation of our paper. In response, we have revised Section 3 to provide a clearer explanation of the connection between GLA and WPGD, thereby better motivating our contributions. Additionally, we have outlined the novelty and key contributions of our work in General Response. We welcome further suggestions from the reviewer if there are any remaining points that require clarification.

Weakness 2: In the papers (Ahn et al. (2024), Li et al. (2024)) mentioned above, their authors research Preconditioned Gradient Descent. ... in case the authors use them interchangeably.

Response: Thank you for highlighting the confusion regarding terminology. We used the term "Projected Gradient Descent" because the $P_1$ and $P_2$ matrices in Eq. (7) can be viewed as projection matrices. However, we agree with the reviewer that "Preconditioned Gradient Descent" is more accurate and aligns better with the terminology used in the referenced literature. We have updated the paper to adopt this terminology consistently throughout.

Weakness 3: If I understand correctly, the equation 7 states that ... on priors $P_1$, $P_2$, and $\Omega$, are not provided.

Response: Thank you for your question. We have enhanced the exposition for clarity. In Section 3.1, we now begin with a standard “weighted least-squares objective” to first derive the scalar-weighted PGD algorithm. This corresponds to scalar-gated linear attention. We then generalize this to the vector-weighted PGD estimator $\hat{\beta} = \beta_1^{\text{gd}}(P_1, P_2, \Omega)$, as defined in Eq. (7), with preconditioning matrices $P_1,P_2$ and weighting matrix $\Omega$.

Our Theorem 1 establishes the mapping between this vector-weighted PGD and the attention weights in Eq. (8) and weights induced by vector-valued gating. To reduce ambiguity, we have revised Theorem 1 to include additional explanations that explicitly clarify the connection between GLA and WPGD, making it easier to understand.

Weakness 4: Moreover, Theorem 1 states that ... It would be helpful to explicitly demonstrate and prove it as in e.g., Lemma 1 of Ahn et al. (2024).

Response: Thank you for your comment. Our Theorem 2, which extends to multilayer architectures, builds directly upon and generalizes Lemma 1 of Ahn et al. (2024). It addresses additional complexities introduced by causal masking and non-linear gating mechanisms, which are crucial and widely used in practical applications but were not considered in previous theoretical work.

Additionally, we recently came across the work by Ding et al. (2024), which provides a theoretical analysis of causal masking in in-context learning. Their findings demonstrate that causal language models (causalLM) exhibit suboptimal convergence dynamics akin to those of online gradient descent with non-decaying step sizes. This behavior limits their ability to reach optimal solutions, even with an increasing number of in-context examples. We have cited this work below Theorem 2 to further underscore the challenges posed by gated and causally-masked architectures, which justify the extensions and contributions.

To further clarify, we have reorganized Theorem 1 and Section 3.1 to explicitly clarify that the forward pass of a one-layer GLA corresponds to a one step of WPGD.

Questions: Also, there might be some typos or minor mistakes that I would consider clarifying or correcting.

Response: Fixed. Thank you!

Dear authors,

Thank you for your response and sorry for the late answer.

I acknowledge that revision of the paper made it more clear. However, I remain skeptical due to the following:

$\beta_1^{\text{gd}}(P_1, P_2, \Omega) := P_2 \big( X P_1 \odot \Omega \big)^\top y$

is not a preconditioned gradient descent as it should have only one preconditioner matrix at the leftmost side of the expression as in the formula in line 178. For reference, see definitions of preconditioned gradient descent in e.g. https://www.cs.cornell.edu/courses/cs4787/2019sp/notes/lecture8.pdf, https://www.cs.princeton.edu/courses/archive/fall18/cos597G/lecnotes/lecture5.pdf, and https://www.cs.princeton.edu/%7Earora/TheoryDL.pdf (2.4.1). It seems that you creatively introduced an additional parameter $P_1$ to the PGD formula so it could align well with your derivation for the GLA output.

It stands as a major problem for me, because I believe it makes the core claim of the paper in its current state that the GLA implements WPGD unproved.

And, anticipating a possible argument that $P_1$ could be treated as a parameter of underlying regression algorithm rather than a second "preconditioner", I note that in such a case 1) the underlying algorithm would likely no longer be an ordinary linear regression; 2) As an optimizable parameter in the corresponding GLA layer, it would also be optimizable in a regression, and the gradient would have been taken w.r.t both $\beta$ and $P_1$ further complicating matters.

Additionally, I would like to acknowledge that SSMs such as Mamba and RNNs (RWKV) are fundamentally different from Linear Transformers in their core sequence mixing mechanism, despite the similarities in gating mechanisms which were discussed in the GLA paper [1]. One class of models cannot be readily re-parametrized to represent another. Your derivations and proofs are valid only in case of the eponymous Gated Linear Attention class of models from the paper [1]. Therefore, I strongly suggest to remove the mentions of various recurrent and state space models which could mislead the reader into believing your results cover those models too.

Weakness 1: The attention matrices are restricted in certain form (eq (8), line 201-203). ... It's not sure whether GLA implements WPGD when attention matrices don't have these restricted forms.

Response: Thank you for pointing this out. We would like to emphasize that the restricted form/construction of attention matrices in Eq. (8) is not an arbitrary choice but is well-supported by prior literature. Several existing works have adopted similar constructions:

1. Proposition 1 in Von Oswald et al. (2023), Section 2.4 in Lu et al. (2024), and Appendix B in We et al. (2024) all make similar assumptions about attention weights.
2. Theorem 1 in Ahn et al. (2024) and Proposition 1 in Li et al. (2024) demonstrate that optimizing single-layer linear attention models with or without such restrictions yields equivalent predictions.
3. Theorem 4.1 in Zhang et al. (2024) and Eq (2) in Huang et al. (2023) show that, when attention weights are initialized following similar constraints, their structural form persists throughout training, as zero entries remain zero and weights converge to forms consistent with Eq. (8).

We note that our study focuses on the GLA model with a nonlinear gating function and task mixtures where sequences include multiple tasks. These settings introduce significantly more complexity compared to prior work, even under the constrained attention weight forms discussed in Eq. (8). Hence, our contributions extend beyond existing findings on single-task ICL and are non-trivial.

Weakness 2: The token embeddings also have restricted forms (line 426). It is not clear whether GLA can learn the contextual features if the token embeddings are learnable parameters.

Response: In both this work and previous studies on linear attention architectures, $z$ in $Z$ is referred to as both input tokens and embeddings due to the linearity of the model. If we consider a learnable linear embedding matrix $W_e$, the prediction for linear attention can be expressed as:

where $W_{q,k,v}'=W_eW_{q,k,v}$. Thus, the embedding matrix $W_e$ can be absorbed into the attention weights, yielding equivalent optimization results. This is why our framework does not consider learnable token embeddings without loss of generality.

Additionally, in our setting, the contextual features in Eq. (20) can be arbitrary random vectors once they are linearly independent (as per Assumption B). Therefore, this setup is broad and highly general.

Question 1: Line 486 says that Assumption B ensures that any $\omega$ in $W$ can be achieved by an appropriate choice of gating parameters. ... If the number of tasks $K$ is larger than the dimension of $w_g$ (the trainable parameter of the gating function), the above statement seems to be wrong.

Response: You are correct that if the number of tasks $K$ exceeds the dimension of $w_g$, not all $\omega \in \mathcal W$ can be represented via the gating parameters. However, Assumption B explicitly requires that the delimiters be linearly independent, which already implies that $K <\dim(w_g)$, as a larger $K$ would violate the linear independence condition and render Assumption B invalid.

Question 2: Theorem 2 states that an $L$-layer GLA implements $L$ steps of WPGD. The question is: when $L$ is large enough, does the $L$-layer GLA find a better predictor than the one-layer GLA? Can Theorem 2 demonstrate the advantage of deeper models?

Response: It is well-established that additional steps of gradient descent (with appropriately chosen step sizes) generally result in reduced loss. In an $L$-layer GLA, each layer corresponds to a step of GD as outlined in Theorem 2, meaning the $L$-layer GLA effectively performs $L$ steps of GD. Consequently, it achieves improved predictions compared to a single-layer GLA for $L > 1$.

Question 3: Are there more experiments of multi-layer GLA?

Response: In response to the reviewer’s request for additional experiments, we have added results in Appendix A.1 of the revised submission. Due to time constraints, we replicated the setting from Fig. 1(a) to demonstrate the improvements provided by deeper models. The additional experimental results again verify that deeper models yield better predictions. We are also considering conducting further experiments in different settings for the final submission.