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

We thank the reviewer for the positive assessment and detailed feedback.

I can spot no good reasons, really. … relevant to empirical work in the field.

We appreciate the acknowledgment that this work fits well within COLM's scope. The reviewer correctly notes this is a "more theoretical-leaning paper," and we believe it aligns strongly with COLM's Science of LMs track, specifically the areas of "complexity, training dynamics, grokking, learning theory for LMs." Our work makes several contributions that bridge theory and empirical understanding:

Learning Theory for LMs: We provide the first rigorous characterization of the optimization landscape for gated linear attention models, establishing conditions under which unique global minima exist (Theorems 2-3) and connecting these architectures to interpretable learning algorithms (WPGD).

Training Dynamics Insights: Our analysis of multi-layer GLA (Theorem 6) reveals how causal masking creates coupled gradient trajectories across layers, providing theoretical foundations for understanding training behavior in autoregressive models such as Mamba and RWKV.

Practical Impact: By characterizing the algorithmic capabilities of different gating strategies (scalar vs. vector as provided in Theorems 4 and 5), we provide theoretical insights that directly inform architectural choices and model design decisions for practitioners working with modern efficient sequence models.

Question to Authors

Below we address each question:

• Yes, $P^\star$ in Eq.(15) has the structure of a preconditioning matrix that improves the condition number. Specifically, $P^\star= \Sigma^{-1/2}[\frac{\gamma^\star}{M} \cdot \Sigma + I]^{-1}\Sigma^{-1/2}$ effectively whitens the data covariance $\Sigma$ and applies additional regularization through the $\gamma^\star$ term. This reduces the condition number compared to the identity preconditioner, particularly when the input covariance $\Sigma$ is ill-conditioned. The factor $\gamma^\star/M$ adaptively adjusts the regularization strength based on the task correlation structure and noise level.

• We will add a brief explanation: H3 (Fu et al. 2022) is structured state space model that uses diagonal state matrices and gating mechanisms, representing an efficient alternative to transformers for long sequences.

• We use "effective" to distinguish from the standard spectral gap (largest minus smallest eigenvalue) because our definition in Eq.(13) considers only nonzero eigenvalues: $\Delta_R := \lambda_{\max} - \lambda_{\min}$ where $\lambda_{\min}$ and $\lambda_{\max}$ are the nonzero smallest and largest eigenvalues. This is necessary since correlation matrices $R$ may have zero eigenvalues in degenerate cases.

• Previous work (e.g., Ahn et al. 2024) analyzed linear attention without causal masking (i.e., full attention where each token can attend to all positions). Our multi-layer analysis (Theorem 6) accounts for causal masking where token $i$ can only attend to positions $j \leq i$, which is essential for autoregressive language modeling and introduces coupling between the gradient trajectories across layers.

• Thank you for the detailed feedback on Figure 1. We will revise the figure with your suggestions: (i) Use different line styles and color-coding to distinguish LinAtt vs GLA theoretical bounds, (ii) Add opacity for overlapping lines, (iii) Share legend and y-axis, (iv) Clarify that GLA-vector is only needed in subplot (d) where Assumption 3 fails ($r_1 > r_2$), while in (a-c) scalar gating achieves the theoretical optimum.

• This occurs because when Assumption 3 fails ($r_1 > r_2$), scalar gating cannot achieve the optimal WPGD solution due to the monotonicity constraint $\omega \in \mathcal{W}$. In this case, the best achievable solution for scalar gating reduces to uniform weighting $\omega = \mathbf{1}$, which is equivalent to linear attention.

• Yes, Assumption 3 ($\mathbb{E}[\beta_i^\top\beta] \leq \mathbb{E}[\beta_j^\top\beta]$ for $i \leq j$) holds in subplots (a-c) where $(r_1, r_2) = (0,1), (0.2,0.8), (0.5,0.5)$ respectively, ensuring $r_1 \leq r_2$. We will add this verification to the caption.

• Fixed.

• These functions arise from spectral decompositions in the first-order optimality conditions and share structural similarities with secular equations. $h_2(\gamma)$ emerges from the eigendecomposition of $\Sigma$, while $h_1(\bar{\gamma})$ comes from the eigendecomposition of $R$ and captures weighted task correlations. Both have the characteristic rational form with poles at eigenvalues, similar to secular equations. Their fixed-point relationship $\gamma^* = h_1(h_2(\gamma^*)) + 1$ balances the spectral properties of data covariance and task correlation structures for optimal WPGD weighting.

We again thank the reviewer for the detailed questions and suggestions, and we will incorporate them into our final manuscript.

We thank the reviewer for the thoughtful evaluation and recognition of our novel theoretical contributions. We address each concern below:

Reason to Reject 1: The paper aims to understand GLA ... clarify its specific contribution relative to prior art.

We acknowledge the importance of positioning our work relative to existing online learning frameworks, particularly DeltaNet [1]. Recent work [Behrouz et al. 2025a] and [Behrouz et al. 2025b] unify sequence models as associative memory modules that optimize internal objectives through iterative algorithms, revealing connections to online optimization and memory management dynamics. Both DeltaNet and GLA implement variants of this online framework but with different online learning-retaining objectives (MSE vs dot product-based) [Behrouz et al. 2025b]. This yields for (Gated) DeltaNet $S_t = \alpha_t (I - \beta_t k_t k_t^\top) S_{t-1} + \beta_t v_t k_t^\top$ with similarity-based forgetting (controlled by $\beta_t k_t k_t^\top$), while GLA employs $S_t = G_t \odot S_{t-1} + v_t k_t^\top$ with data-dependent gating ($G_t$). From the online learning perspective, both balance acquiring new information versus retaining past information, but GLA uses element-wise gating $G_t$ for retention while (Gated) DeltaNet combines scalar retention with content-based removal through key similarity (controlled by $\beta_t k_t k_t^\top$); please see Table 1 and Sec. 4 in [Behrouz et al. 2025b] for further details.

This online learning lens illuminates our gating↔weighting connection: the unrolling of optimization steps corresponds to context weighting. While we focus on in-context learning capabilities rather than online optimization per se, this connection helps contextualize our findings. We provide the first rigorous optimization landscape analysis for GLA under a heterogeneous/multitask context model, revealing how GLA attends and selects relevant context. We prove a unique global minimum (Thm. 3), classify all stationary points (Thm. 2), and specify conditions for optimality—insights absent from DeltaNet and related work. Our comparison of scalar vs. vector gating (Thms. 4–5) further uncovers expressivity trade-offs, offering principled guidance for architectural design.

We will discuss the connection to [1], online learning frameworks [Behrouz et al. 2025b], and related work in our revised manuscript.

Reason to Reject 2: The theoretical claims are validated ... real-world applications is substantial.

Synthetic tasks have been an important proxy to explain and improve the mechanics and capabilities of more complex models. For instance [Park et al. 2024] and [Grazzi et al. 2024] entirely focus on comparing Mamba and Transformer on synthetic tasks. While we acknowledge the limitation of synthetic experiments, this conscious methodological choice aligns with established theoretical ICL literature [Ahn et al., 2024; Von Oswald et al., 2023; Mahankali et al., 2023]. Following your valuable suggestion, we expand experiments to higher-dimensional settings. Results follow Figure 1(b) with $(r_1,r_2)=(0.2,0.8)$ but $d=100$ instead of $d=10$:

Model Performance refers to empirical test risk achieved by training models; Theory Prediction refers to theoretical optimal risks: $L_{ATT}^\star$ from Corollary 3 (Eq. (25)) for Linear Attention and $L_{WPGD}^\star$ from Corollary 2 (Eq. (24)) for GLA.

These results demonstrate that theoretical predictions remain highly accurate in higher-dimensional settings ($d=100$) and across varying context lengths up to $n=1000$, validating that our WPGD framework captures fundamental properties extending beyond our initial low-dimensional experiments.

We thank the reviewer for recognizing our theoretical contributions. We address the comments below.

While I really appreciate that theoretical … guidelines can be drawn.

Existing works on the ICL–GD connection (Mahankali et al., 2023; Li et al., 2024b; Ahn et al., 2024) have focused on efficient architectures such as linear attention or state-space models. In contrast, real-world LLMs incorporate additional components (e.g., softmax attention, layer normalization, MLPs) that enable more sophisticated learning and lead to expected performance gaps (Shen et al., 2024). GLA or Mamba models helped bridge the gap from linear attention to transformers. Our framework enhances earlier ICL–GD settings through a richer heterogeneous context model and develops the theoretical foundations that shed light on these practically relevant architectures. Thus, while we build on the ICL–GD connection, the results and implications go further.

Additionally, by formalizing GLA as a data-dependent weighting scheme, we show that GLA-type architectures can provably attend and select the relevant subset of the tokens. By characterizing scalar vs. vector gating, task correlation effects on GLA expressivity, and optimization landscape properties, our work also sheds light on architectural principles. For instance, Figure 1 demonstrates clear separations between vector gating, scalar with delimiters, and scalar w/o delimiters in line with our theory.

We will elaborate on these points and expand the discussion on the gap between Linear Attention-ICL (GLA-ICL in our case) and LLM-ICL raised by Shen et al. (2024) in the revision. For details on attention-weight construction and prompt design, please see our response to Q1.

Q1: The equivalence between GLA and WPGD… with practical setups?

Our attention weight construction and input prompt structure build on earlier theoretical studies linking ICL and preconditioned GD (Mahankali et al., 2023; Li et al., 2024b; Ahn et al., 2024). Although our theory targets specific attention weight settings, Figure 1 shows that WPGD-like dynamics still emerge in trained GLA models even when attention weights are unconstrained (see lines 665–666). For further details, please see our response to Reviewer qjca, Q1.

We agree that exploring more practical prompting setup is valuable. For example, features and labels can be released sequentially in the prompt (e.g., $x_1,y_1,x_2,y_2,...$). Then, let $v_{2i}$ and $k_{2i}$ represent feature $x_i$, and $v_{2i+1}$ and $k_{2i+1}$ the label $y_i$. In this setup, the state update at time step $2i+1$ becomes:

$$S_{2i+1} = G_{2i+1} \odot S_{2i} + v_{2i+1}k_{2i+1}^\top = G_{2i+1} \odot G_{2i} \odot S_{2i-1} + G_{2i} \odot v_{2i}k_{2i}^\top + v_{2i+1}k_{2i+1}^\top.$$

This reduces to our Eq. (GLA) if (i) features and labels are embedded in orthogonal subspaces, and (ii) $G_{2i}$ is the all-one matrix. Finally, for more complex dataset models, a multilayer GLA model could emulate more sophisticated learners such as lasso or random forest regression (extrapolating from Garg et al. 2022, Park et al. 2024).

We will discuss these practical variations and future directions in the revision.

Q2: The problem setting, including the correlated … learning tasks?

Our correlated task model and multitask distribution reflect the reality that real-world prompts often include examples from multiple related tasks, not just a single one. Production systems frequently encounter heterogeneous task distributions within a single context, such as document processing with multiple sub-tasks, conversational AI handling related queries, and educational platforms with correlated problems. Recent works (Wang et al., 2024a; Dun et al., 2023; Asai et al., 2022; Lines 279–281) have documented these scenarios and proposed methods for multitask prompts. Our framework provides the first theoretical analysis of how GLA's data-dependent weighting addresses such scenarios. We will expand on this discussion in the revision.

We thank the reviewer for these questions. We address your concerns below, noting some details require verification.

Q1: Eqn. (9) involves specific forms of Wk,Wq,Wv, … still recovers WPGD-like behavior or achieves similar optimality?

We clarify that our empirical experiments (see Lines 665-666) impose no constraints on the model weights. Indeed, despite training fully parameterized GLA models, the empirical results (solid curves in Fig. 1) align closely with our theoretical predictions (dashed curves in Fig. 1), indicating that fully parameterized GLA models implicitly exhibit WPGD-like dynamics.

From a theoretical perspective, we emphasize that the restricted structure of the attention matrices is not arbitrary but is strongly supported by prior literature. For instance

• Proposition 1 in Von Oswald et al. (2023), Section 2.4 in Lu et al. (2024), and Appendix B in Wei et al. (2024) all utilize similar structural assumptions about attention weights.

• Theorem 1 in Ahn et al. (2024) and Proposition 1 in Li et al. (2024) demonstrate equivalence between predictions from optimized single-layer linear attention models, with and without such structural constraints.

• Theorem 4.1 in Zhang et al. (2024) and Eq. (2) in Huang et al. (2023) show that initializing attention weights according to similar structural constraints ensures that these forms persist throughout training, with zero entries remaining zero and weights converging to forms consistent with our Eq. (9).

Our study builds upon these foundations by extending the analysis to GLA models featuring nonlinear, data-dependent gating mechanisms and multitask prompting, which introduce additional complexities. Even within the structural framework provided by Eq. (9), the resulting optimization landscape is intricate and nontrivial, as thoroughly detailed in Section 4 of our paper.

Q2: While Theorem 3 establishes the uniqueness of the global minimum for the WPGD objective,... are guaranteed to converge to this optimum?

We thank the reviewer for this insightful question. Although Theorem 3 establishes that the WPGD population risk has a uniquely attainable global minimizer with a contraction property, our analysis is limited to the optimization landscape of the WPGD problem itself. Extending these guarantees to standard training procedures for GLA introduces additional challenges, as outlined below.

For the WPGD optimization directly, Theorems 2 and 3 do provide strong convergence foundations: Theorem 2 characterizes all stationary points (up to rescaling), and Theorem 3 establishes a unique global minimum (up to rescaling), together suggesting that standard optimization algorithms on WPGD population risk $L_{WPGD}(P, \omega)$ should converge to the optimum. However, the main difficulty arises because standard training for GLA optimizes the parameters $(W_k, W_q, W_v, \text{gating functions})$ of a GLA module, which implicitly define the WPGD parameters $(P, \omega)$ through a highly coupled and non-convex reparameterization. This joint optimization over attention matrices and gating mechanisms adds another layer of non-convexity absent from existing analyses of transformer optimization and the ICL–GD connection (e.g., Ahn et al., 2024; Mahankali et al., 2023). Furthermore, any rigorous convergence statement would be optimizer-specific, requiring detailed analysis of the chosen optimization method such as stochastic GD, step-size conditions, and the smoothness of the composed GLA objective.

A systematic convergence analysis would ideally proceed by modeling the training dynamics over the full joint parameter space of attention matrices and gating mechanisms, and by examining how feasible regions defined by different gating schemes (cf. Theorems 4 and 5) interact with specific optimization algorithms. Theorems 4-5 suggest a pathway where sufficiently expressive GLA parameterizations can inherit WPGD convergence properties. Empirically (see Fig. 1), we observe that when our theoretical conditions hold and the gating class is sufficiently expressive, standard optimizers do converge to the WPGD optimum, suggesting that the favorable landscape properties may translate to practical GLA training.

However, formalizing these empirical observations would require overcoming the theoretical challenges outlined above, representing substantial future research directions beyond the current understanding of attention mechanism optimization.