Gating is Weighting: Understanding Gated Linear Attention through In-context Learning
Authors: Yingcong Li
Davoud Ataee Tarzanagh
Ankit Singh Rawat
Maryam Fazel
Samet Oymak
Venue: COLM 2025 (2025)
Keywords: linear attention
gating
in-context learning
weighted gradient descent
optimization landscape
TL;DR: This work offers theoretical insights into the ICL capabilities of gated linear attention models, demonstrates how gating is crucial for achieving stronger data adaptivity, and characterizes the loss landscape of weighted preconditioned GD.

Abstract: Linear attention methods offer a compelling alternative to softmax attention due to their efficiency in recurrent decoding. Recent research has focused on enhancing standard linear attention by incorporating gating while retaining its computational benefits. Such Gated Linear Attention (GLA) architectures include highly competitive models such as Mamba and RWKV. In this work, we investigate the in-context learning capabilities of the GLA model and make the following contributions. We show that a multilayer GLA can implement a general class of Weighted Preconditioned Gradient Descent (WPGD) algorithms with data-dependent weights. These weights are induced by the gating mechanism and the input, enabling the model to control the contribution of individual tokens to prediction. To further understand the mechanics of this weighting, we introduce a novel data model with multitask prompts and characterize the optimization landscape of learning a WPGD algorithm. We identify mild conditions under which there exists a unique global minimum, up to scaling invariance, and the associated WPGD algorithm is unique as well. Finally, we translate these findings to explore the optimization landscape of GLA and shed light on how gating facilitates context-aware learning and when it is provably better than vanilla linear attention.

Published as a conference paper at COLM 2025
Gating is Weighting: Understanding Gated Linear Attention
through In-context Learning
Yingcong Li* University of Michigan yingcongêumich. edu
Davoud Ataee Tarzanagh" Samsung SDS Research America d. ttazzanaghssmssunggg. com
Ankit Singh Rawat Google Research NYC anki mokitssraattggogle.coo
Maryam Fazel University of Washington mf mfazel@uw.ee edu
Samet Oymak
University of Michigan oymakêumich. edu
Abstract
Linear attention methods offer compelling alternative to softmax atten- tion due to their efficiency in recurrent decoding. Recent research has focused on enhancing standard linear attention by incorporating gating while retaining its computational benefits. Such Gated Linear Attention (GLA) architectures include highly competitive models such as Mamba and RWKV. In this work, we investigate the in-context learning capabilities of the GLA model and make the following contributions. We show that a multilayer GLA can implement general class of Weighted Preconditioned Gradient Descent (WPGD) algorithms with data-dependent weights. These weights are induced by the gating mechanism and the input, enabling the model to control the contribution of individual tokens to prediction. To further understand the mechanics of this weighting, we introduce novel data model with multitask prompts and characterize the optimization land- scape of learning a WPGD algorithm. We identify mild conditions under which there exists a unique global minimum, up to scaling invariance, and the associated WPGD algorithm is unique well. Finally, we translate these findings to explore the optimization landscape of GLA and shed light on how gating facilitates context-aware learning and when it is provably better than vanilla linear attention.
1 Introduction
The Transformer (Vaswani, 2017) has become the de facto standard for language model- ing tasks. The key component of the Transformer is the self-attention mechanism, which computes softmax-based similarities between all token pairs. Despite its success, the self- attention mechanism has quadratic complexity with respect to sequence length, making computationally expensive for long sequences. To address this issue, a growing body of work has proposed near-linear time approaches to sequence modeling. The initial ap- proaches included linear attention and state-space models, both achieving O(1) inference complexity per generated token, thanks their recurrent form. While these initial architec- tures typically do not match softmax attention in performance, recent recurrent models such as Mamba (Gu & Dao, 2023; Dao & Gu, 2024), mLSTM (Beck et al., 2024), DeltaNet (Yang et al., 2024b;a), GLA Transformer (Yang et al., 2023), and RWKV-6 (Peng et al., 2024) achieve highly competitive results with the softmax Transformer. Notably, as highlighted in Yang et al. (2023), these architectures can be viewed as variants of gated linear attention (GLA), which incorporates a gating mechanism within the recurrence of linear attention. Addi- tionally, Behrouz et al. (2025a;b) unify those models associative memory modules that optimize internal objectives through iterative algorithms, revealing connections to online optimization and memory management dynamics.
'Equal contribution.

---

Published as a conference paper at COLM 2025
Given sequence of tokens (2) E RR+1 and the associated query, key, and value embeddings (qi, K v:) C RR with d + being the embedding dimension", the GLA recurrence is given by
S = G + F and o = S,g; iE (1,
(GLA)
Here, S E R ((++1)x(++1) represents the 2D state variable with So 0; o: E RR+1 represents the ith output token; and the gating variable G; := G(z;) E R(d+1)x(a+1) is applied to the state S through the Hadamard product @ and G represents the gating function. When G; is a matrix of all ones, (GLA) reduces to causal linear attention (Katharopoulos et al., 2020).
The central objective of this work is to enhance the mathematical understanding of the GLA mechanism. In-context learning (ICL), one of the most remarkable features of modern sequence models, provides powerful framework to achieve this aim. ICL refers to the ability of a sequence model to implicitly infer functional relationships from the demon- strations provided in its context window (Brown, 2020; Min et al., 2022). It is inherently related to the model's ability to emulate learning algorithms. Notably, ICL has been a major topic of empirical and theoretical interest in recent years. For example, Ali et al. (2024) show that Mamba embeds implicit attention matrices via data-controlled linear operators, while Sieber et al. (2024) unify SSMs, attention, and RNNs within dynamical-systems framework. More specifically, series of works have examined the approximation and optimization characteristics of linear attention, and have provably connected linear attention to the preconditioned gradient descent algorithm (Von Oswald et al., 2023; Ahn et al., 2024; Zhang et al., 2024). Given that the GLA recurrence in (GLA) has richer design space, this leads us to ask:
What learning algorithm does GLA emulate in ICL?
The (GLA) recurrence enables the sequence model to weight past informa- tion in a data-dependent manner through the gating mechanism ( Building on this observation, we demonstrate that GLA model can implement data-dependent Weighted Preconditioned Gradient Descent (WPGD) algorithm. Specifically, a one-step WPGD with scalar gating, where all entries of G, are identical, is described by the prediction:
Contributions:
g x P (y ④ w))
(1)
Here, X E Raxd is the input feature matrix; y E R" is the associated label vector, x e R4 Rddd is the preconditioning matrix; and
represents the test/ query input to predict; w E IR" weights the individual samples. When w is fixed, we drop "data-dependent" and simply refer to this algorithm as the WPGD algorithm. However, for GLA, w depends the data through recursive multiplication of the gating variables. Building on this formalism, we make the following specific contributions:
ICL capabilities of GLA ($3): Through constructive arguments, we demonstrate that multilayer GLA model can implement data-dependent WPGD iterations, with weights induced by the gating function. This construction sheds light on the role of causal masking and the expressivity distinctions between scalar- and vector-valued gating functions.
Landscape of one-step WPGD (64): The GLA WPGD connection motivates us to ask: How does WPGD weigh demonstrations in terms of their relevance to the query? To address this, we study the fundamental problem of learning an optimal WPGD algorithm: Given a tuple x, ~ D, with y E R being the label associated with the query, we investigate the population risk minimization:
GLA+WPGD
CAPGO :
min PER "ER"
LWPcu (P,co(P,co),
(2)
[
where LUPED P. (w))) = ED
The sequence length and embedding dimension are set to and d per the prompt definition
in (3).
2

---

Published as a conference paper at COLM 2025
As our primary mathematical contribution, we characterize the loss landscape under a general multitask data setting, where the tasks associated with the demonstrations (X,y) have varying degrees of correlation to the target task (x, We carefully analyze this loss landscape and show that, under mild conditions, there is a unique global minimum (P) up to scaling invariance, and the associated WPGD algorithm is also unique.
Loss landscape of one-layer GLA (55): The landscape is highly intricate due to the recursively multiplied gating variables. We show that learning the optimal GLA layer can be connected to solving (2) with a constraint induced by the choice of gating function and input space. Solidifying this connection, we introduce multitask prompt model under which we characterize the loss landscape of GLA and the influence of task correlations. Our analysis and experiments reveal insightful distinctions between linear attention, GLA with scalar gating, and GLA with vector-valued gating.
E W, where the restriction W is
Problem setup
Notations. IR4 is the d-dimensional real space, with R4 and R as its positive and strictly Bold letters, e-g., and A, represent vectors and matrices. The identity matrix of size n is denoted by In- The symbols and 0 denote the all-one and all-zero vectors or matrices of proper size, such as 1 E Rd and lxx E Rddxd The subscript is omitted when the dimension is clear from the context. The Gaussian distribution with mean a and covariance Σ is written as N ((u, Σ). The Hadamard product (element-wise multiplication) is denoted by C and Hadamard division (element-wise division) is denoted
positive orthants. The set [n] denotes (1,...,m)) n1.
by @. Given any di+1,...,, aj E Rd, we define dijj a;+1 C ... aj for < j, and a:ı = 1a-
The objective of this work is to develop a theoretical understanding of GLA through ICL. The optimization landscape of standard linear attention has been a topic of significant interest in the ICL literature (Ahn et al., 2024; Li et al., 2024b). Following these works, we consider the input prompt
Z = [2
Zn
Za+1]
x1
Xy
xp-1l
E (((+))x(4-1),
(3)
y1
yn
where tokens encode the input-label pairs
W:) C R2 x R.
We aim to enable ICL by training a sequence model F R(++1)x(++)) + R that predicts the label y := Yn+1 associated with the query x := xu+1- This model will utilize the demonstrations (x,, yi) to infer the mapping between x and y. Assuming that the data is distributed as (y, Z) ~ D, the ICL objective is defined as
((F) = Ep 1 (())
Linear attention and shared-task distribution. Central to our paper is the choice of the function class F. When F is a linear attention model, the prediction fF := F(Z) takes the form fF = + Wg W, Z ZWah, where W5, Wg. Wp E R(d-1)x(d-I) are attention parameters, and dd+ is the linear prediction head.
T
We assume that the in-context input-label pairs follow a shared-task distribution, where ß ~ (0, Eg), x; are i.i.d. with x; (0, Ex), and y; N T x2, where > 0 represents the noise level. Under this shared distribution, is shown (Von Oswald et al., 2023; Ahn et al., 2024; Zhang et al., 2024) that the optimal one-layer linear attention prediction coincides with the one-step optimal preconditioned gradient descent (PGD). In particular, considering the data distribution discussed above, we have
OF x TBB where B P y,
55)
3

---

Published as a conference paper at COLM 2025
and
p argmin E (y-x* x ) with x (xq Xn] and y = ly ynl (6)
PeRdxd
Linear attention and gating. Given the input prompt Z as in (3), let kg P;) (W. q 2i, W Zi, W zi) be the corresponding query, key, and value embeddings. The out- put of causal linear attention at time can be computed using (GLA) with G; 1. This recurrent form implies that linear attention has O(d)) cost, that is independent of sequence length n, to generate per-token. GLA follows the same structure as linear attention but with a gating mechanism (ie., G; # 1), which equips the model with the option to pass or supress the history. As discussed in Yang al. (2023), the different choices of the gating function correspond to different popular recurrent architectures such as Mamba (Gu & Dao, 2023), Mamba2 (Dao & Gu, 2024), RWKV (Peng et al., 2024), etc.
We will show that GLA can weigh the context window through gating, thus, its capabilities are linked to the WPGD algorithm described in (8). This will in turn facilitate GLA to effectively learn multitask prompt distributions described by y; ~ (V x o2) with B not necessarily identical.
What gradient methods can GLA emulate?
In this section, we investigate the ICL capabilities of GLA and show that under suitable instantiations of model weights, GLA can implement data data-dependent WPGD.
3.1 GLA data-dependent WPGD predictor
Data-Dependent WPGD. Given X and as defined in (6), consider the weighted least
squares objective L(B) Σ W (y: B x2) with weights = [W1, wn]" E IR". To optimize this, we use gradient descent (GD) starting from zero initialization, Bo with step size of 17 1/2. One step of standard GD is given by
B1 Bo - VLB))
Σ
w
x = X (w)
Given test/query feature the corresponding prediction is x T where B = BB Additionally, if we using PGD with preconditioning matrix P E Rddxd Bo 0, and = 1/2, then single iteration results in
g - where - = PX (w
Above is the basic sample-wise WPGD predictor which weights individual datapoints, in contrast to the PGD predictor as presented in (5). It turns out, vecfor-aalued gating can facilitate more general estimator which weights individual coordinates. To this aim, we introduce an extension as follows: Let Pt, P2 E dxd denote the preconditioning matrices, and Ω E RRa denote the vector-valued weighting matrix. Note that Ω is a weight matrix that enables coordinate-wise weighting. Throughout the paper, sample-wise and coordinate-wise weighting parameters, corresponding to scalar and vector gating strategies, respectively. Then given Ω can similarly define
w and Ω to denote
B (Pi, P2, 0) P2 0)
(8a)
one-step of (generalized) WPGD. Its corresponding prediction test query x is:
g x where ß == / (2,
(8b)
We note that by removing the weighting matrix Ω (8) reduces to standard PGD (cf. (5)) and setting Ω w1 reduces to sample-wise WPGD (cf. (7)). Li et al. (2024b) has demonstrated
4

---

Published as a conference paper at COLM 2025
that H3-like models implement one-step sample-wise WPGD, and they focus the shared- task distribution where B: = 8 for all É [n]. In contrast, our work considers a more general data setting where tasks within an in-context prompt are not necessarily identical.
We introduce the following model constructions under which we establish the equivalence between GLA (cf. (GLA)) and WPGD (cf. (8)), where the weighting matrix is induced by the input data and the gating function. Inspired by prior works (Von Oswald et al., 2023; Ahn et al., 2024), we consider the following restricted attention matrices:
Pk 0
0 0
Pg 0
0 0
0 Odxd
0 1
W
Wg
and W.
where Pk, Py E Rdxd Note that the (d + 1, 1)-th entry of W, is set to one for simplicity. More generally, this entry can take any nonzero value, e.g,, V E R. Parameterizing W with / produces the same output as in (9).
Theorem 1. Recall (GLA) with imput sequence Z [z] ... In zn+1]T defined in (3), the gating
G(z,) G E (dd 11), and outputs (a:) Consider model construction in (9) and take the last coordinate of the last token output denoted by (on+1)a+1 as prediction. Then, we have
(Z) (0n=1)ar1= = where B B-B (Ps, (ut Pq.
Here, f7 is one-step WPGD feature predictor defined in (Sa); Px and Pq correspond to attention weights following (9); and Ω [gim+1 grm+1] E IRuxd, where gim+1, E [n] is given by
gin+1 = gi+1 8142' 8n+1 E R², and G.
(10)
Here and throughout, we use * to fill the entries of the matrices that do not affect the final output. Based on the model construction in (9), these entries can be assigned any value.
Observe that, crucially, since g; is associated with 2,, z; influences the weighting of 2j for all j<i. í. We defer the proof of Theorem to the Appendix D.1. It is noticeable that only d of the total (d + 1)² entries in each gating matrix G, are useful due to the model construction presented in (9). However, if we relax the weight restriction, e.g., W. 10 [d u]
where u E Rd+1 then the weighting matrix Ω in Theorem is associated with all rows of the G, matrices. We defer to Eqn. (27) and the discussion in Appendix D.1.
Capabilities of multi-layer GLA. Ahn et al. (2024) demonstrated that, with appropriate construction, an L-layer linear attention model performs steps of PGD on the dataset (x; y:) provided within the prompt. In Appendix A, we extend this analysis to multi- layer GLA and characterize the algorithmic class it can emulate. Importantly, Ahn et al. (2024) does not account for causal masking, which is a fundamental aspect of multi-layer GLA due to its recurrent structure as described in (GLA). Our main result, Theorem 6 in Appendix A, establishes that an L-layer GLA implements L steps of WPGD, where the gradients are computed in recurrent form.
GLA with scalar gating. Theorem 1 establishes connection between one-layer GLA (cf. (GLA)) and one-step WPGD (cf. (8)), where the weighting in WPGD corresponds to the gating G(z,) G; in GLA, as detailed in Theorem 1. Now let us consider the widely used types of gating functions, such as G; = a;1241 (Yang et al., 2023; Katsch, 2023; Qin et al., 2024; Peng et al., 2024) or G; 77 (241)444)) (Dao & Gu, 2024; Beck et al., 2024; Peng al., 2021; Sun et al., 2024), where a; E Id+1 and y; E IR. In both cases, the gating matrices in (10)
take the form of
*
for some R, thus simplifying the predictor to sample-wise
WPGD, as given by
(Z) = BTx, with -= B
(11)
where = P.P.P. and w = 18in ... Sum i1] T R". In the remainder, we will mostly focus the one-layer GLA with scalar gating as presented in (11).
5

---

Published as a conference paper at COLM 2025
4 Optimization landscape of WPGD
In this section, we explore the problem of learning the optimal sample-wise WPGD algo- rithm described in (11), key step leading to our analysis of GLA. The problem is as follows. Recap from (6) that we are given the tuple y,, matrix, y E R" is the label vector, x E R is the query, and y E R is its associated label. The goal is to use X, y to predict given x via the one-step WPGD prediction y = x B, with B as in (11). The algorithm learning problem is given by (2) which minimizes the WPGD risk
D, where X E Rexd x is the input
EDl(y x PX( 2 y))
Prior research (Mahankali et al., 2023; Li et al., 2024b; Ahn et al., 2024) has studied the problem of learning PGD when input-label pairs follow an i..d. distribution. It is worth noting that while Li et al. (2024b) establishes a connection between H3-like models and (11) similar to ours, their work assumes that the optimal consists of all ones and does not specifically explore the optimization landscape of w when in-context samples are not generated from the same task vector B. Departing from this, we introduce a realistic model where each input-label pair is allowed to come from a distinct task.
Definition 1 (Correlated Task Model). Suppose B; E Rd (0, I) are jointly Gaussian for all i E [11 + 1]. Define
:: B:=Bn+I+ B = B Bn] R = E BB and r: 1 E BBp
(12)
Note that in (12), we have R E Ro*=" and r E R", with normalization ensuring that the entries of R and lie in the range 1, 1], corresponding to correlation coefficients. Further, due to the joint Gaussian nature of the tasks, for any i, E In + 1), the residual TijB; is independent of Bj.
Definition 2 (Multitask Distribution). Let (Bi) be drawn according to the correlated task model of Definition 1, (x:) E IRd be iid. with x; ~ (O,E), and yi ~ Nxx B,, 02) for all ie [n +++1 1].
Definition 3. Let the eigen decompositions of E and R be denoted by E = Udia and R = Ediag E where s = [$1 sal E TRC and A [A [A1.... ... An) T E ER" Let Smin and max denole the smallest and largest eigenvalues of Σ, respectively. Further, let Amin and Amax denote the smallest and largest nonzero eigenroluuss of R. Define the effective spectral gap of L and R, respectively, as
Ar = Smax - Smin, and AR := Amax Amin-
(13)
Assumption 1. The correlation vector r from (12) lies in the range of E, i.e., Ea for some nonzero [a] an] E R".
Assumption essentially ensures that (representing the correlations between in-context tasks) can be expressed value. Note that since r lies in the range of the eigenvector matrix E, it also lies in the range of R defined in (12). This guarantees that the correlation structure is non-degenerate, meaning that all elements of are influenced by meaningful correlations. Assumption avoids trivial cases where there are correlations between tasks. By requiring at least one nonzero element in a, the assumption ensures that the tasks are interrelated.
linear transformation of vector with at least one nonzero
The following theorem characterizes the stationary points () of the WPGD objective (2). Theorem 2. Consider independent linear data as described in Definition Suppose Assumption on the correlation vector holds. Define the functions h1 R R and h2 : 10) - IR, as
Aja7
a?
h1( =
(14a)
(1
2/7)2
(1 1
d
h1(7)
1+ M
Σ
(14b)
M )
+ s
6

---

Published as a conference paper at COLM 2025
where {sili=1 and {A: the eigenvalues of I and R, respectively; {a: are given in Assumption 1; and M 0² + +
The risk function LuPGU SPPGO(P.CO (P..)) in (2) has stationary point (p=,,"") , up to rescaling, defined
P-=-- 2
(M (M.I+I E and e (12(7)) - R+
(15)
where 7 is a fixed point of composite function ht (1((())) + 1.
Theorem 2 characterizes the stationary points (P*, w²), which exist up to re-scaling. This result presents the first landscape analysis of GLA for the joint learning of (P,w), while also exploring the stationary points (P*, w²). In the following, we provide mild conditions on effective spectral gaps of R and E under which unique (global) minimum (P*, w²) exists.
Theorem 3 (Uniqueness of the WPGD Predictor). Consider independent linear data as given in Definition Suppose Assumption 1 on the correlation vector holds, and
AΣ AR < M + Smins
(16)
where Dy and Ag denote the effective spectral gaps of I and R, respectively, as given in (13); Smin is the smallest eigenvalue of I; and M = =² + Si-
T1. The function hy (h2(7)) + 1 is contraction mapping and admits a unique fixed point y = y- T2. The loss Lurco (P, ()) has unique (global) minima (P", w²), up to re-scaling, given by (15).
Theorem 3 establishes mild conditions under which a unique (global) minimum (P", w²) exists, up scaling invariance, and guarantees the uniqueness of the associated WPGD algorithm. It provides the first global landscape analysis for GLA and generalizes prior work (Li et al., 2024b; Ahn al., 2024) on the global landscape by extending the optimization properties of linear attention to the more complex GLA with joint (P,w) w)) optimization.
Remark 1. An interesting observation about the optimal gating parameter w² is its connec- tion to the correlation matrix R, which captures the task correlations in a multitask learning setting. Specifically, the optimal gating given in (15) highlights how w² depends directly on both the task correlation matrix and the vector r, which encodes the correlations between the tasks and the target task.
Remark 2. Condition (16) provides a sufficient condition for the uniqueness of a fixed point. This implies that whenever Ar AR < M + Smin, the mapping h (12(2)) + 1 is a contraction, ensuring the existence of unique fixed point. However, there may be cases where the mapping hy (l2(7)) + does not satisfy Condition (16), yet a unique fixed point (and unique global minimum) still exists. This is because the Banach Fixed-Point Theorem does not provide a necessary condition.
Corollary 1. Suppose Σ = Then, Ag = 0, satisfying Condition (16), and we have h2 h2(r*)) d1o2 which yields p = I and w² = (R+ (d + 02 + 11) Thus, the optimal risk Lipco
defined in (2) is given by
Cipc = -2 d.r R (d (d+ (1+))))))) r.
(17)
5 Optimization landscape of GLA
In this section, we analyze the loss landscape for training a one-layer GLA model and explore the scenarios under which GLA can reach the optimal WPGD risk.
5.1 Multi-task prompt model
Following Definitions 1 and 2, we consider the multi-task prompt setting with K correlated tasks (Ba) and one query task B. For each correlated task [K], prompt of length
7

---

Published as a conference paper at COLM 2025
Livit GLA
Livi GLA
Like GLA
Livitt
GLA
Testritt
0.8
GL4-40
Jestrok
0.8
CLA-90
Test rotr
0.9
GLA-80
0.8
GLA no
0.6
0.6
0.9
0.9
(a) (21,12) = (0,1)
(b) (2) ) (0.2,088)
(c) (11,12) = 0.55, (d) (f1,12 = (0.8, 0.2)
Figure 1: We consider four different types of model training: LinAtt (blue solid): Standard linear attention training. GLA (red solid): GLA training using prompts with delimiters (see (19)) and scalar gating. GLA-wo (green solid): GLA training using prompts without delimiters and with scalar gating. GLA-vector (cyan solid): GLA training using prompts with delimiters and vector gating. The blue and black dashed curves represent the optimal linear attention and WPGD risks from (24) and (17), respectively, as the number of in-context examples и increases. Implementation details are provided in Appendix B.
fg drawn, consisting of ITD input-label pairs {(* Vikk M where yik ~ N X x2) to
obtain sequence 2
Kik
Kngak Vns.
Let E lk.-
These sequences (Zx) K
k-1 along with the query token zn+1 [x,+1 OJT, are concatenated to form a single prompt Z. Recall (GLA), and let fala(Z) (Z) denote the GLA prediction as defined
in Theorem 1. Additionally, consider the model construction in (9), where Pq. Ps E Rdxd are trainable parameters. The GLA optimization problem is described as follows:
CLLA :=
min PEPpER GEG
LOLA( Pq, G),
(18)
where (Ps G) (BV
Here, G(-) represents the gating function and G denotes the function search space, which is determined by the chosen gating strategies (cf. Table in Yang et al. (2023)).
Note that 1) the task vectors ( are not explicitly shown in the prompt, 2) in-context
features Xi are randomly drawn, and 3) the gating function is applied to the tokens ( (Zx) (==)
Given the above three evidences, the implicit weighting induced by the GLA model varies across different prompts, and it prevents the GLA from learning the optimal weighting. To
address this, we introduce delimiters to mark the boundary of each task. Let (de){_ be the delimiters that determine stop of the tasks. Specifically, the final prompt is given by
Z = T
d
z dx 2n+1]
(19)
Additionally, to decouple the influence of gating and data, we envision that each token is |x; Vi C; IJ where c; # 0 E R' is the contextual features with p being its dimension ... ,2x E IPP be K + contextual feature vectors. and (xi,y) Vi) are the data features. Let co....
For task prompts Zg: Contextual features are set to fixed vector ēσ # 0.
For delimiters de: Data features (e.g., yi) are set to zero so that d: = 10 JT.
Explicit delimiters have been utilized in addressing real-world problems (Wang et al., 2024a; Asai al., 2022; Dun et al., 2023) due their ability to improve efficiency and enhance generalization, particularly in task-mixture multi-document scenarios. To further motivate the introduction of (dx) k- we present in Figure the results of GLA training with and without delimiters, represented by the red and green curves, respectively. The black dashed curves indicate the optimal WPGD loss, Lipso under different scenarios. Notably, training GLA without delimiters (the green solid curve) performs strictly worse. In contrast, training with delimiters can achieve optimal performance under certain conditions (see Figures la, 1b, and 1c). Theorem 4, presented in the next section, provides a theoretical explanation for these observations, including the misalignment observed in Figure 1d.
8

---

Published as a conference paper at COLM 2025
5.2 Loss landscape of one-layer GLA
Given the input tokens with extended dimension, to ensure that GLA still implements WPGD in Theorem 1, we propose the following model construction.
Wx 00. 0
Wg and W. 0
We 0
Wg
(20)
Here, Wgg E R p p () and Wkg, R(+11)x(+11) are constructed via (9). The main idea is to set the last p rows and columns of attention matrices to zeros, ensuring that the delimiters do not affect the final prediction.
Assumption 2. Contextual feature vectors 2o. ,ēk are linearly independent, and activation function (z): R + [0, 1] is continuous, satisfying (0-o0) and op =
Assumption 3. The correlation between context tasks (Bx and query task f satisfies EiA Bj] = 0 and B B) < B) all S sisisk,
Given context examples ((+,,)) (x Vik 1 define the concatenated data (X,y):
X =
X
X E Raad and y b
E ER.
(21)
Based on the assumptions above, we are able to establish the equivalence between optimiz- ing one-layer GLA and optimizing one-step WPGD predictor under scalar gating.
Theorem 4 (Scalar Gating). Suppose Assumption 2 holds. Recap the function Lweeo (P, ()) from (2) with dataset y) defined in (21). Consider (GLA) with input prompt defined in (19), the model constructions described in (20), and scalar gating G(z) = 4(w, 2) (a+p+1)x(d=p=x)) where
wg is trainable parameter. Then, the optimal risk Lāl defined in (18) obeys
COLA = =L
where
=
min PERAXU NEW
LWPGD (P,
(22)
Here, ::= = {arlm ... w R" o5 wi w/ S SL
Additionally, suppose Assumption holds and n1; = ny, for all i, E [K]. Let Lipen be the optimal WPGD risk (cf. (2)). Then aa satisfies
CoLA
(23)
Assumption 2 ensures that any in W can be achieved by an appropriate choice of gating parameters. Furthermore, Assumption 3 guarantees that the optimal choice of w under the WPGD objective lies within the search space W. The proof is provided in Appendix F.1.
In Figure 1, we conduct model training to validate our findings. Consider the setting where K = 2 and let 12) [E Ed B) /d). In Figures 1a, 1b, and 1c, Assump- tion 3 holds, and the GLA results (shown in solid red) align with the optimal WPGD risk (represented by the dashed black), validating (23). However, in Figure 1d, since F1 > 72, Assumption does not hold, and as a result, the optimal GLA loss L obtained from (22) is lower than the optimal WPGD loss - Further details are deferred to Appendix B.
Loss landscape of vector gating. Till now, our discussion has focused on the scalar gating setting. It is important to highlight that, even in the scalar-weighting context, analyzing the WPGD problem remains non-trivial due the joint optimization over (P, ()).). However, as demonstrated in Theorem 4, scalar gating can only express weightings within the set W. If Assumption 3 does not hold, Lā cannot achieve the optimal WPGD loss (see the misalign- ment between red solid curve, presenting Lå and black dashed curve, presenting LiRDD in Figure 1d). We argue that vector gating overcomes this limitation by applying distinct weighting mechanisms across different dimensions, facilitating stronger expressivity.
Theorem 5 (Vector Gating). Suppose Assumption 2 holds. Consider (GLA) with input prompt Z. from (19), the model constructions from (20) but with W., [0 p 1)xd 0(11p1))xp)

---

Published as a conference paper at COLM 2025
and vector gating G(z) falAZ) op4ph, where h R&+p+1 is a linear prediction head. Here, u, W and h are trainable parameters. Let LaLAND denote its optimal risk, and Lúpco be defined as in (2). Then, CaLAV LMPGD
Recap Problem (18) with prediction defined as
In Theorem 4, the equivalence between Laa and Cipco is established only when both Assumptions and 3 are satisfied. In contrast, Theorem 5 demonstrates that applying vector gating requires only Assumption 2 to establish GLA LaPaD Specifically, under the bounded activation model of Assumption 2, scalar gating is unable to express non- monotonic weighting schemes. For instance, suppose there are two tasks: Even if Task is more relevant to the query, Assumption 2 will assign higher (or equal) weight to examples in Task 2 resulting in sub-optimal prediction. Theorem 5 shows that vector gating can avoid such bottlenecks by potentially encoding tasks in distinct subspaces. To verify these intuitions, in Figure 1d, we train GLA model with vector gating and results are presented in cyan curve, which outperform the scalar gating results (red solid) and align with the optimal WPGD loss (black dashed).
Loss landscape of one-layer linear attention. Given the fact that linear attention is equiv- alent to (GLA) when implemented with all ones gating, that is, G; = 1, we derive the following corollary: Consider training standard single-layer linear attention, i.e., by setting
G = in (GLA), and let FATT (Z) be its prediction. Let LATT be the corresponding optimal risk following (18).
Corollary 2. Consider a single-layer linear attention following model construction in (9) and let linear data as given in Definition 2. Let R and r be the corresponding correlation matrix and vector as defined in Definition 1. Suppose = Then, the optimal risk obeys
d(1 r)
CATT = min LIPGD (P 1) = =d n(d 02 + +1 R1
(24)
PeRdxd
In the Figure 1, blue solid curves represent the linear attention results and blue dashed are the theory curves following (24). The two curves are aligned in all the subfigures, which validate our Corollary 2. More implementation details are deferred to Appendix B.
Discussion
Our analysis is currently limited to GLA models under specific weight construction as- sumptions outlined in Section 3, which may not capture the full expressivity of practical GLA implementations. The theoretical framework we developed requires certain structural constraints the attention matrices (as specified in Equation 9) to establish the connection with WPGD algorithms.
Several important directions for future research include:
1. Extending our analysis to more general GLA architectures without the weight con- struction constraints, which would better reflect deployed models such as Mamba (Gu & Dao, 2023; Dao & Gu, 2024) and RWKV (Peng et al., 2024).
2. Investigating when delimiters are necessary for effective learning in multi-tasi settings. Our experiments show that without delimiters, GLA models perform sub- optimally, but deeper theoretical characterization of this phenomenon is needed.
3. Exploring the GLA landscape where gating functions depend on input features in more complex ways than our current formulation allows.
4. Analyzing the effects of incorporating MLP layers between GLA layers, which could enhance the model's expressive power beyond what we have characterized.
Addressing these limitations would further bridge the gap between our theoretical under- standing and the practical performance of state-of-the-art GLA-bbsed sequence models.
10

---

Published as a conference paper at COLM 2025
Acknowledgements
Y.L. and S.O. were supported in part by the NSF grants CCF2046816, CCH-2403075, CCF- 2008020, the Office of Naval Research grant N0001 42401224, and by gifts awards from Open Philanthropy, OpenAl, Amazon Research, and Google Research. M.F. was supported in part by awards NSF TRIPODS II 2023166, CCF 2007036, 2212261, and CCF 2312775.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. Adoances in Neural Information Processing Systems, 36, 2024.
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Represenlations, 2023.
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arviv preprint Xio:2403.00559, 2024.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-hhroughput tradeoff. aiv preprint arXivv2402.11668, 2024.
Akari Asai, Mohammadreza Salehi, Matthew E Peters, and Hannaneh Hajishirzi. Attempt: efrfieter-efficient multi-task tuning via attentional mixtures of soft prompts. arXiv preprint ar xrXiv:2205.11961, 2022.
Han Bao, Ryuichiro Hataya, and Ryo Karakida. Self-attention networks localize when qk<eigenspectrum concentrates. Xiv preprint arXio:2402.02099 2024.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prud- nikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Extended long short-term memory. arXiv preprint ar arXio:2405.04517, 2024
Ali Behrouz, Zeman Li, Praneeth Kacham, Majid Daliri, Yuan Deng, Peilin Zhong, Meisam Razaviyayn, and Vahab Mirrokni. Atlas: Learning to optimally memorize the context at test time. arviv preprint 2:Xi::2500.23735, 2025a.
Ali Behrouz, Meisam Razaviyayn, Peilin Zhong, and Vahab Mirrokni. It's all connected: A journey through test-time memorization, attentional bias, retention, and online optimiza- tion. Xrv preprint ar arXio:2504.13173, 2025b.
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005.14166 2020.
Sitan Chen and Yuanzhi Li. Provably learning a multi-head attention layer. arXiv preprint ar trXiv:2402.04085 2024.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi- head softmax attention for in-context learning: Emergence, convergence, and optimality. ar Xiv preprint ar Xrx::22402.1944,, 2024.
Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts function lipschitzness. Xrv preprint ar arXio:2402.1163,, 2024.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arviv preprint arXiv:2405.2100 2024.
Soham De, Samuel Smith, Anushan Fernando, Aleksandar Botev, George Fristian-Murany
Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. Xiv preprint ar 2022229402 2024.
11

---

Published as a conference paper at COLM 2025
Yichuan Deng, Zhao Song, Shenghao Xie, and Chiwun Yang. Unmasking transform- ers: A theoretical approach to data recovery via attention weights. arXiv preprint #rXiv:2310.12462 2023.
Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention. ar arv preprint ar aXXto:3310.12680, 2023.
Nan Ding, Tomer Levinboim, Jialin Wu, Sebastian Goodman, and Radu Soricut. Causallm is not optimal for in-context learning. arv preprint ar arXiv:2308.06912 2023.
Chen Dun, Mirian Hipolito Garcia, Guoqing Zheng, Ahmed Hassan Awadallah, Anastasios Kyrillidis, and Robert Sim. Sweeping heterogeneity with smart mops: Mixture prompts for Ilm task adaptation. arXiv preprint arXiv:2310.02842, 2023.
Tolga Ergen, Behnam Neyshabur, and Harsh Mehta. Convexifying transformers: Improving optimization and understanding of transformer networks. wwxXi:22211.100:5 2022.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint ar Xiv:2212. 14052, 2022.
Deqing Fu, Tianqi Chen, Robin Jia, and Vatsal Sharan. Transformers learn higher-order optimization methods for in-context learning: A study with linear models. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning, 2023.
Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systemss, 35:30583-30599 2022.
Khashayar Gatmiry, Nikunj Saunshi, Sashank Reddi, Stefanie Jegelka, and Sanjiv Kumar. Can looped transformers learn to implement multi-step gradient descent for in-context learning? In Proceedings of the 41st International Conference on Machine Learning, pp. 15130- 15152, 2024.
Riccardo Grazzi, Julien Siems, Simon Schrodi, Thomas Brox, and Frank Hutter. Is mamba capable of in-context learning? arviv preprint ar XrXio:2402.03177, 2024.
Albert Gu and Tri Dao. Mamba: senear-time sequence modeling with selective state spaces. arXiv preprint arXit:2312.0075: 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with
structured state spaces. arv preprint xxXio:2111.00366, 2021.
Ruiquan Huang, Yingbin Liang, and Jing Yang. Non-asymptotic convergence of training
transformers for next-token prediction. arXiv preprint arXio:2409.17333 2024.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. Xrviv preprint ar arXi::2310.05249 2023.
M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint arxo:2240..1351, 2024.
Samy Jelassi, Michael Eli Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying, arxiv preprint 4rXiv:2402.01032, 2024.
Hong Jun Jeon, Jason Lee, Qi Lei, and Benjamin Van Roy. An information-theerret analysis of in-context learning. ar Xiv preprini ar arXto:2401.1553 2024.
12

---

Published as a conference paper at COLM 2025
Aaron Alvarado Kristanto Julistiono, Davoud Ataee Tarzanagh, and Navid Azizan. Opti- mizing attention with mirror descent: Generalized max-margin token selection. arXiv preprint ar xrX::2240.14581, 2024.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rons: Fast autoregressive transformers with linear attention. In International conference on machine learning, pp. 5156-5165. PMLR, 2020.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. Xrviv preprint arXiv:2311.019277 2023.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint ar Xiv:2302.06011, 2023a.
Yingcong Li, Muhammed Emrullah Tidiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pp. 19594. PMLR, 2023b.
Yingcong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with self-attention. In International Conference on Artificial Intelligence and Statistics, pp. 685-693. PMLR, 2024a.
Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. arXiv preprint ar erXiv:2407.10055, 2024b.
Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. arXiv preprint ar drXiv:2402.8811, 2024.
Arvind Mahankali, Tatsunori B Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. arv preprint ar Xix::2307 231/:23607.0 2023.
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains. arXiv preprint arXiv:2402.0416/ 2024.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? arXiv preprint ar Xro::22202.12837, 2022.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 26670-26698. PMLR, 2023.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In International Conference on Machine Learning, 2023.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oy- mak, Kangwook Lee, and Dimitris Paaailiiopoulos. Can mamba learn how learn? comparative study on in-context learning tasks. arxiv preprint ar fit:2202004248 2024.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing rnns for the transformer era. ar Xiv preprint ar Xrv:22005.13045 2023.
Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Teddy Ferdinan, Haowen Hou, Przemyslaw Kazienko, et al. Eagle and finch: RWKY with matrix-valued states and dynamic recurrence. arviv preprint ar arXi::2404.05892 2024.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. Xiv preprint arXit:2103.021443 2021.
13

---

Published as a conference paper at COLM 2025
Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. Hgrm2: Gated linear mns with state expansion. ar preprint Kio:2404.07904 2024.
Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pi- lanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In International Conference Machine Learning, pp. 19050-19088. PMLR, 2022.
Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison H Zhou. Implicit regularization of gradient flow on one-layer softmax attention. arXiv preprint ar Kkt:22403008699 2024.
Lingfeng Shen, Aayush Mishra, and Daniel Khashabi. Position: Do pretrained transformers learn in-context by gradient descent. In Proceedings of the 41st International Conference on Machine Learning, volume 235, pp. 44712-4474h 2024.
Jerome Sieber, Carmen Amo Alonso, Alexandre Didier, Melanie N Zeilinger, and Antonio Orvieto. Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks. arXiv preprint 1rXi::2405.15311, 2024.
Haoyuan Sun, Ali Jadbabaie, and Navid Azizan, In-context learning of polynomial kemnll regression in transformers with glu layers. arXiv preprint ar Xrio:2501.18117, 2025.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. Xiv preprint orXiv:2307.0862) 2023.
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models. Xiv preprint ar XrX::2205.00254, 2024.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Trans- formers as support vector machines. arXiv preprint arXiv:2308..168 168998 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Advances in Neural Information Processing Systems, 36, 2024.
Christos Thrampoulidis. drXio:2402.18551, 2024.
Implicit bias of next-token prediction. aiv preprint
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du Joma: Demys- tifying multilayer transformers via joint dynamics of mlp and attention. ar preprint a/Xiv=2310.00535, 2023.
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis. Implicit bias and fast convergence rates for self-attention. arXiv preprint arXio:2402.05738, 2024a.
Bhavya Vasudeva, Deqing Fu, Tianyi Zhou, Elliott Kau, Youqi Huang, and Vatsal Sha- ran. Simplicity bias of transformers to learn low sensitivity functions. arXio preprint ar xxXi:2243.00625, 2024b.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pp. 35151-35174. PMLR, 2023.
Johannes von Oswald, Eyvind Niklasson, Maximilian Schlegel, Seijin Kobayashi, Nicolas Zucchet, Nino Scherrer, Nolan Miller, Mark Sandler, Max Vladymyrov, Razvan Pas- canu, et al. Uncovering mesa-optimization algorithms in transformers. Xiv preprint XXii:2309.05855 2023.
14

---

Published as a conference paper at COLM 2025
Ruochen Wang, Sohyun An, Minhao Cheng, Tianyi Zhou, Sung Ju Hwang, and Cho-Jui Hsieh. One prompt is not enough: Automated construction of a mixture-of-expert prompts. arXio preprint ar axXio:2407.00256 2024a.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. ar Xiv preprint arXiv:2406606889, 2024b.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter Bartlett. How many pretraining tasks are needed for in-context learning of linear regres- sion? Xiv preprint ar Xrii:2310008891,, 2023.
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2022.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. ar preprint XrXi::2312.06655, 2023.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule, Xiv preprint ar Xiv:2412.066664 2024a.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491-115522, 2024b.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25-4 2024.
Yize Zhao and Christos Thrampoulidis. Geometry of concepts in next-token prediction: Neural-collapse meets semantics. In The Second Conference on Parsimony and Learning (Recent Spotlight Track), 2025. URL
www.eeewne/ffoum/// /ttt//ofummdddtt1 weettt61woooog woQ84.
Yize Zhao, Tina Behnia, Vala Vakilian, and Christos Thrampoulidis. Implicit geometry of next-token prediction: From language sparsity patterns to model representations. ar Xiv preprint ar Xrx::2248.1155
2024.
Junhao Zheng, Shengjie Qiu, and Qianli Ma. Learn or recall? revisiting incremental learning with pre-trained language models. ar preprint axXio:2312.07887, 2023.
Itamar Zimerman, Ameen Ali, and Lior Wolf. A unified implicit attention formulation for gated-linear recurrent sequence models. arXiv preprint arXivo2405.16504, 2024.
Chang Zong, Jian Shao, Weiming Lu, and Yueting Zhuang. Stock movement predic- tion with multimodal stable fusion via gated cross-attention mechanism. arXiv preprint drXiv:2406.06554, 2024.
15

---

Published as a conference paper at COLM 2025
A Capabilities of multi-layer GLA
Ahn et al. (2024) demonstrated that, with appropriate construction, an L-layer linear atten- tion model performs L-step preconditioned GD on the dataset Vi) provided within the prompt. In this work, we study multi-layer GLA and analyze the associated algorithm class it can emulate. It is worth mentioning that Ahn et al. (2024) does not consider causal masking which is integral to multilayer GLA due to its recurrent nature described in (GLA). Our analysis will capture the impact of gating and causal mask through и separate GD trajectories that are coupled.
Multi-layer GLA. For any input prompt E &R(n ()x(d 1), let GLA [01 02 ... 0++1] E ((++1)x(++I) denote its GLA output, as defined in (GLA). We consider L-layer GLA model follows: For each layer é E [L], let Z, denote its input, where we set Z = Z, and let WK W Wo denote the key, query, and value matrices for layer E, respectively. After applying residual connection, the input of the (-th layer is given by
: = Ze + GLAA-1 (Z
(25)
Here, GLAa(-) denotes the output of the f-th GLA layer, associated with the attention matrices (Wa.r Wa,6, we)) for all & [L]. In the following, we focus on the (d + 1)-th entry of each token's output at each layer (after the residual connection), denoted by 0ig == [Ze
GLAy 11 for all E [n++]] and E [L].
Theorem 6. Consider an L-layer GLA defined in (25), where WE Wa. in the l'th layer param- eterized by PK,S, Pa.. E Rddxd, for all E E [L] following (9). Let the gating be a function of the
features, e.g., G; = -G(x;), and define Ω as in Theorem 1. Additionally, denote the masking as
M; = 1; 0
E RRnxn, and let Bo = Bi = 0 for all E [m]. Recall that 0i, for all i E [n + 1] and
o
E [L] stands for the last entry of the i'th token output at the l'th layer. We have
For all i S n, Oife yi x7 Biee, where Bi, Bi,--1 P... / ( ie 0 gim+1);
-x -T Be, where Be (1 (1- a Pge (V nt 2 Bat1), and at
On 12 x P.P.PPT. k,
y;= = diag(XB) we define Here, letting B; = |B,, ... BAd X, = XPx. C 2 and yı
Vi x M; (ye y).
We defer the proof of Theorem 6 to the Appendix D.2. Theorem 6 states that L-layer GLA (25) implements L steps of WPGD with gradient in recurrent form and additional weight decay ay. To recap, given data (X,y) and prediction B, the gradient with respect to the squared loss takes the form X (xB- -y), up to some constant c. In comparison, Pq.e (Vi, O gis+1) similarly acts as a gradient but incorporates layer-wise feature preconditioners (P Psçe), data weighting (O), and causality (gin+1, Mj). Here, M; represents causal masking, ensuring that at time i, only inputs from j S i are used for prediction. Notably, the recurrent structure of GLA allows the gating mechanism apply context-dependent weighting strategies.
To simplify the theorem statement, we assume that the gating function depends only on the input feature, e.g., G, G(x;), ensuring that the corresponding data -deeeddennd weighting is uniform across all layers. This assumption is included solely for clarity in the theorem statement, and the gating function varies across layers or depends on the entire (d + 1)- dimensional token rather than just the feature x;, each layer has its own gating O and X, = X @ Dr. Note that our inclusion of the additional term & captures the influence of the last token's 's output on the next layer's prediction. Based on the above L-layer GLA result, we have the following corollary for multi-layer linear attention network with causal mask in each layer.
16

---

Published as a conference paper at COLM 2025
Corollary 3. Consider an L-layer linear attention model with causal mask and residual connection in each layer. Let Ethh layer follows the weight construction in (9) with Wa,t, W q, parameterized by
PK.C- q as in Theorem 6 and define p::= == Pa Kr ker l [L]. Let Bo O. Then, the (d + 1 yth entry of the last token output of each layer satisfies:
On+1E xxx where
Br = (1 az ar) PP (ye y).
Here, we define at = T Pex and yt follows the same definition as in Theorem 6.
Note that Corollary 3 generalizes (Ahn et al., 2024, Lemma 1) by extending it to causal attention. When considering full attention (i.e., without applying the causal mask), we have y3- = XB,, and the expression X (y: - y) = X (XB, y) corresponds to the standard gradient of the squared loss. Furthermore, if we disregard the impact of the last token (e.g., by incorporating the matrix M as Eq. (3) of Ahn et al. (2024)), then a₂ = 0. These results are also consistent with Ding al. (2023), which demonstrate that causal masking limits convergence by introducing sequence biases, akin to online gradient descent with non-decaying step sizes.
Our theoretical results in Theorem 6 focus on multi-layer GLA without Multi-Layer Per- ceptron (MLP) layers isolate and analyze the effects of the gating mechanism. However, MLP layers, a key component of standard Transformers, facilitate further nonlinear feature transformations and 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).
B Experimental setup
B.1 Implementation detail
Data generation. Consider ICL problem with input in the form of multi-task prompt as described in Section 5.1. In the experiments, we set 2, dimensions d = 10 and p = 5, uniform context length nı #12 = ń where we have = Kř, and vary ň from 0 to 50. Let (21,12) [H B) /d, Eig ß] d) denote the correlations between in-context tasks B1 and query task B.. We generate task vectors as follows:
B1,B2 (0, and B N r2B2, (1 - (2)
Input features are randomly sampled (Lxz) ~ (O,, and we have V/x B xi (o 0), k E {1,2} and yut1 B xnt1. Additionally, delimiters Co ,kg are randomly sampled from N 1p)
Implementation setting. We train 1-layer linear attention and GLA models for solving multi-prompt ICL problem as described in Section 5.1. For GLA model, we consider
sigmoid-tyrn gating function given by scalar gating: G(z) = ()+++++))(( or
vector gating: G(z) q (Wg (Wg)) where p(z) (1 + 2) is the activation function. Note that although the theoretical results are based on the model constructions (c.f. (9) and (20)), we do not restrict the attention weights our implementation. We train each model for 10000 iterations with batch size 256 and Adam optimizer with learning rate 10-3, Similar to the previous work (Li et al., 2024b), since our study focuses on the optimization landscape, ICL problems using linear attention/ /LA models are non-convex, and experiments are implemented via gradient descent, we repeat 10 model trainings from different model initialization and data sampling (e.g,, different choice of delimiters) and results are presented as the minimal test risk among those 10 trails. Results presented have been normalized by d.
17

---

Published as a conference paper at COLM 2025
100
1-layer
2-layer
3-layer
Test risk
10-1
10
0
10
20
30
40
50
n
Figure 2: Multi-layer GLA experiments with (r1,12) = (0,1).
Experimental results. Based the experimental setting, can obtain the correlation matrix and vector following Definition
a 1 0
R
and r r= [n r21
1
Then dotted curves display our theoretical results derive using I = and R, above. Specifically, in Figure 1, black dashed curves represent Li following (17) and blues dashed curves represent La following (24). We consider scenarios where (2,)2) e {(0,1), 0.8), (0.5, 0.5), (0.8, 0.2)) and results are presented in Figures (la), (Ib), (1c) and (1d), respectively.
GLA-wo achieves the worst performance among all the methods. We claim that it is due to the randomness of input tokens as discussed in Section 5.1. Thanks to the introduction of delimiters as described in (19), data and gating is decoupled and a task-dependent weighting is learnt. Hence, GLA is able to achieve comparable performance to the optimal one (Capgy, red dashed). Note that GLA-wo performs even worse than LinAtt. It comes from the fact the weighting induced by GLA-wo varies over different input prompts and it can not implement all ones weight.
The alignments between LinAtt (blue solid) and blue dashed curves validate our Corol- lary 2. In Figures 1a, 1b and 1c, the alignments between GLA (red solid) and LWOGD (black dashed) verify our Theorem 4, specifically, Equation 23. While in 1c and 1d, GLA achieves the same performance as LinAtt. It is due to the fact that GLA can not weight the history higher than its present. Then the equal-weighting, e-g., w is the optimal weighting given such constraint. What's s more, the alignment between GLA-vector (cyan curves) and red dashed in Figure 1d validates our vector gating theorem in Theorem 5.
B.2 Multi-layer experiments
In this section, present additional experiments multi-layer GLA models. We adopt the same experimental setup as described in Figure la and Appendix B, with parameters setting to (f1,12) r2) (0, 1). The results displayed in Figure 2, where the blue, red, and green curves correspond to the performance of one-, two-, and three-layer GLA models, respectively, with the y-axis presented in log-scale. According to Theorem 6, L-layer GLA performs L steps of WPGD, suggesting that deeper models should yield improved predictive performance. The experimental findings in Figure align with the theoretical predictions of Theorem 6.
C Related work
Efficient sequence models. Recent sequence model proposals such as RetNet (Sun et al., 2023), Mamba (Gu & Dao, 2023), xL.STM (Beck et al., 2024), GLA Transformer (Yang al.,
18

---

Published as a conference paper at COLM 2025
2023), RWKV-6 (Peng et al., 2024) admit efficient recurrent forms while being increasingly competitive with the transformer architecture with softmax-attention. However, we have a rather limited theoretical understanding of these architectures, especially, when it comes to their optimization landscape and ICL, capabilittess Park et al. (2024); Grazzi et al. (2024) demonstrate that Mamba is effective and competitive with transformer of similar size in various ICL tasks, whereas Arora et al. (2024); Jelassi et al. (2024) establish theoretical and empirical shortcomings of recurrent models for solving recall tasks. It is worth mentioning that, GLA models also connect to state-space models and linear RNNs (De et al., 2024; Orvieto et al., 2023; Gu et al., 2021; Fu et al., 2022), as they could be viewed as time- varying SSMs (Dao & Gu, 2024; Sieber et al., 2024). Finally, GLA models are also closely related to implicit self-attention frameworks. For example, the work by Zimerman et al. (2024) on unified implicit attention highlights how models such as Mamba (Gu & Dao, 2023) and RWKV (Peng et al., 2023) can be viewed under a shared attention mechanism. Additionally, Zong et al. (2024) leverage gated tross-attention for robust multimodal fusion, demonstrating another practical application of gated mechanisms. Both approaches align with GLA's data lataddependent gating, suggesting its potential for explainability and stable fusion tasks.
Theory of in-context learning. The theoretical aspects of ICL has been studied by growing body of works during the past few years (Xie et al., 2022; von Oswald et al., 2023; Gatmiry et al., 2024; Li et al., 2023b; Collins al., 2024; Wu et al., 2023; Fu et al., 2023; Lin & Lee, 2024; Akyürek et al., 2023; Zhang et al., 2024). A subset of these follow the setting of Garg et al. (2022) which investigates the ICL ability of transformers by focusing on prompts where each example is labeled by task function from a specific function class, such as linear models. Akyürek et al. (2023) focuses on linear regression and provide transformer construction that can perform a single step of GD based on in-context examples. Similarly, Von Oswald et al. (2023) provide a construction of weights in linear attention-only transformers that can replicate GD steps for linear regression task on in-context examples. Notably, they observe similarities between their constructed networks and those resulting from training on ICL prompts for linear regression tasks. Building on these, Zhang et al. (2024); Mahankali et al. (2023); Ahn et al. (2024) focus on the loss landscape of ICL for linear attention models. For a single-layer model trained on in-context prompts for random linear regression tasks, Mahankali et al. (2023); Ahn et al. (2024) show that the resulting model performs a single preconditioned GD step on in-context examples in test prompt, aligning with the findings of Von Oswald et al. (2023). In contrast, real-world LLMs incorporate additional components (e.g., softmax attention, causal masking, MLPs) that enable more sophisticated learning and lead to expected performance gaps (Shen et al., 2024). GLA models helped bridge the gap from linear atfention to transformers. More recent work (Ding et al., 2023) analyzes the challenges of causal masking in causal language models Mc) showing that their suboptimal convergence dynamics closely resemble those of online gradient descent with non-decaying step sizes. Additionally, Li et al. (2024b) analyzes the landscape of the H3 architecture, an SSM, under the same dataset model. They filter. However, their
show that H3 can implement WPGD thanks to its convolutional/SSI) WPGD theory is restricted to the trivial setting of equal weights, relying on the standard prompt model with IID examples and shared tasks. In contrast, we propose novel multitask datasets and prompt models where nontrivial weighting is provably optimal. This allows us to characterize the loss landscape of WPGD and explore advanced GLA models, linking them to data data-dependent WPGD algorithms.
Optimization landscape of attention mechanisms. The optimization behavior of attention mechanisms has become rapidly growing area of research (Deora et al., 2023; Huang et al., 2023; Tian et al., 2023; Fu et al., 2023; Li et al., 2024a; Tarzanagh et al., 2024; 2023; Deng et al., 2023; Makkuva et al., 2024; Jeon et al., 2024; Zheng et al., 2023; Collins et al., 2024;
Chen & Li, 2024; Li et al., 2023a; lldiz et al., 2024; Vasudeva et al., 2024b; Bao et al., 2024; Chen et al., 2024; Huang et al., 2024; Wang et al., 2024b; Sun et al., 2025). Particularly relevant to our work are studies investigating convex relaxation approaches to optimize attention models (Sahiner et al., 2022; Ergen et al., 2022), gradient gradient-based optimization methods in vision transformers (Jelassi et al., 2022), optimization dynamics in prompt- attention (Oymak et al., 2023; Tarzanagh et al., 2024), implicit bias of gradient descent
19

---

Published as a conference paper at COLM 2025
for attention optimization (Tarzanagh et al., 2024; 2023; Julistiono et al., 2024; Li et al., 2024a; Thrampoulidis, 2024; Zhao et al., 2024; Vasudeva et al., 2024a; Sheen et al., 2024), and optimization geometry analysis of next-token prediction (Zhao & Thrampoulidis, 2025). While these works mainly focus on standard attention, our work provides the first optimization landscape analysis of gated attention, establishing conditions for unique global minimum and providing comprehensive insights into the implications of data- dependent weighting in the optimization dynamics encountered during training.
D
GLA + WPGD
D.1 Proof of Theroem
Proof. Recap the problem settings from Section where in-context samples are given by
Z
[1 Zn Znz 11
x1
Xn Xn+1 0
y
yn
and let the value, key and query embeddings at time be
=
Wo Zz,
k; = W T Zi, and qi = Wg T Zz-
Then we can rewrite the GLA output (cf. (GLA)) as follows:
0:
Siqi and
S G S 1 1 n
= ....G. Orik + + G.OM-IN v.kk
where we define
Gj G O Gj+2 G₁, j<i, i, and === (
Consider the prediction based on the last token, then we obtain
R+1
Ou+1 Sp+14a41 and Sn+1
Σ Gjn+1 Oryk
Construction
Recall the model construction from (9) where
Ps
Pg
Oaxd
W
W
0 and W.
(26)
0
1
Then, given each token [x] Vil E [m], single-layer GLA returns
P x
[P_ x
v;
k;
and qi
Vi
0
and we obtain
O Odxd
0
v;k =
isn, n, and Pa = ((++1)x(d+1)
yix,
Pk
0
Therefore, since only d entries in v;k matrix are nonzero, given as the Hadamard product, only the corresponding d entries in all G; matrices are useful. Based on this observation, let
G
and Gju
g
g
where 8ja 1gj11 1 8j12" 2 gi E RR for i and gi: 1
20

---

Published as a conference paper at COLM 2025
Combing all together, and letting y, follow the definitions in (6), we obtain
0n+1 14+19n+1
G Gjatt ① rykk]
au+1
Oaxd 1 1jj Pe C g ja++1
P x
g
0
0
pq x
where
Ω g1 g2m++1 Then taking the last entry of 0x+1 as final prediction, we get
gmn+1] E IRa
(Z) = Pg (XPx (U) y
which completes the proof of Theorem 1.
Based the construction given in (26), only d elements of G; matrices are useful. One might ask about the effect of other entries of G,. Therefore, in the following, we introduce other model construction showing that different row of G implements WPGD with different weighting. Similarly, let Ws, Wa be the same as (26) but with W, constructed by
Construction 2:
W. (d+1)xd
where [i4 U2 u E Rd+1
Then, the value embeddings have the form of vi yau, which gives
v;k = u [yix, Pk 0
Next, let
(g))T
(B))T
T
(8%)) T
G
and
Gja
{g)
d++1
( ddt1
where g Rdd corresponds the -' row of G, and g the output
8/+1
g/:1
g Then we get
E "1ypx Pk O in T
0
x Pq (XPx 01
E u2y,x
Pk
o
(g
[PT x
Pg (XPx Ω2)
Ox+1
0
Ej= u //+1)
Pt 3 (g T
x P. (XP; 2 Dan1)
where
N =1; 1 2
g E Re is d1
(27)
Therefore, consider (1+11)dimenssonee output on 1 Each entry implements a 1-step WPGD with same preconditioners Pt. P, and different weighting matrices Q's. The weighting matrix of i'th entry is determined by the i'th row of all gating matrices. Note that if consider the last entry of 0n+1 as prediction, returns the same result as Construction above, where only last rows of G 's are useful.
Additionally, suppose that the final prediction is given after linear head that is, fELA (Z) h 0 and let h [/t h2 ... h4+1]" E IR Then
(z) h 0x41 = x
XXP 2 ")
(28)
21

---

Published as a conference paper at COLM 2025
where
d11
d11
Ω
Σ
0
hall
[8im11
82m
g
RRxd
(29)
Then, single-layer GLA still returns one-step WPGD with updated weighting matrix.
D.2 Proof of Theorem 6
In this section, we first prove the following theorem, which differs from Theorem 6 only in that W. parameterized by pg instead of -...- Theorem can then be directly obtained by setting qq -.
Theorem 7. Consider an L-layer GLA with residual connections, where each layer follows the weight construction specified in (9). Let Wx and Wq. in the E'th layer parameterized by Pkąr, q. E Raxd for & E [L]. Let the gating be a function of the features, e.g., G; G(x;), and define O as in
Theorem 1. Additionally, denote the masking as M; 1; 0 all i E [x]]
E Rux" and let Bo = Bi,0- = for
0
Then, the last entry of the i'th token output at the E'th layer (0i,e) is:
For all S n, Oife = yi x Bier where Bi. = Big 1t Pat (Vie 2 gim 11),
On-18 = -x BBr where Be (1 + a-aaBBB-11 Puge (Vnt O Bn+1), and at = x PTPePP
Here, letting By BB B X1 XPx. O D and yt diag(XB) we define
Vi x Mi y).
Proof. Suppose that the input prompt for l'th layer is
Z: g
In
(++++)(((+1))
(30)
Yn
Øu+1
and let ki,f, fi,t, Figg be their corresponding key, query and value embeddings. Recap the model construction from (9) and follow the same analysis in Appendix D.1. Let E [n + 1] be the corresponding states. We have
s
Σ F1
of
Oaxd
0
E y/z
PR. 2 (8f
where
G G(R,)
*
and G =
[gg)T
*
I:
0
Additionally, recap that we have Mj
and @ denotes Hadamard division. Let
0
0
n =
[g
g 1
Then, defining
x [1 82 ... 2n] E nxd and
y = [71 122 ... ynl E R",
and letting 0,00 i [n 1] be their outputs, get:
22

---

Published as a conference paper at COLM 2025
For i S n,
Σ 9/ S
o 0
9iPPitt ke O 888mt1 3 gi
(XPrs 2 Mig 2 & Bin+1r
and
Oaxd
p g
01 S/qi
( M 2 8ix+1
x Pq ((PP:
M, 2 8(m+1)
(31)
For i n n 1,
Σ
P R O = (XPx, o o + Pk
and
Oaxd
p
1
((EPr) ( Juti p Ee )
T
RT & Pa.e (XPrs 9 P (ookt1))
8 P.i + P (
(32)
From (31) and (32), the first d entries of all the (d + 1)-dieension outputs (e.g, 0j E [n+ + 1) are zero. Given the multi-layer GLA as formulated in (25) where residual connections are applied, we have that the first d dimension of all the input tokens remain the same. That is, ã = x, Additionally, since gating only depends on the feature then all the layers return the same gatings (e.g, G G,).
Note that if the gating function varies across layers or depends on the entire token rather than just the feature x;, each layer will have its own gating, leading to G # G,. In this theorem, assume identical gating matrices across layers for simplicity and clarity. However, extended version of this theorem, without this assumption, can be directly derived.
Therefore, we obtain
x Pge (x. M,y 2 Bin m)) is n
0,,-
1 P q PT x (x Bmn)
0
0,
where we define X, == XPxx6 @ Ω.
Since we focus on the last/ (d + 1)tt entry of each token output, we have
og = x: q (x] MD 2 Bimt1), isn, n,
(33)
0, 1 = Fn 1x
P.
K +x
Pap (x 8m11).
It remains to get the yi for each layer's 's input. Let inputs of (L - 1)'th 'tth and H'th layer be
X1
Xn
x
Xn x
Z
1/1
yh
1-1
and 2; =
y
y
y
Define
1-
luf
f-1
yn 11 E R²
23

---

Published as a conference paper at COLM 2025
For i I n, from (33) and (25), we have
y = y 11 + q, { M,y 11 2 gin+1
= y y 1 1 x qe
{ M,y e 2 8im
Suppose that y/ 1// x; B. Then
y' = - |x]
xn B = ---da diag
(34)
and
1/i x; B (y: x; p 2) x; pue [x] M; diag (XB )) 2 8im+1)
xx B x B x p (x, M2 diag XB )) 2 Bin-1).
Hence,
v = yi - where B 1- p.. { My diag XB Ø gim+1
1-P
(35)
For = n 1, from (33) and (25), have
yf = yf yf xx P. qe (xTs² 2 @n+1).
Setting at x S.PPa pqeePee Fe and recalling ye from (34), we get
y= = (1 (1)y
1-1 + +x
diag ()Bx 8n+1). (x7 TT
08)(111)))))))
Additionally, let let yf = -x B² Then,
Bd 1 x pae (x. (y- diag(XD ()) 2 8n+1).
Hence,
y' x where B² = (1+ 1 1 P-Pas
Pal (x. (y diag(XB) ()) @ gn+1
(36)
Combining (35) and (36) together completes the proof.
E Optimization landscape of WPGD
We first provide the following Lemma.
Lemma 1 (Existence of Fixed Point). Consider the functions h1 (7) and h2(7) defined in (14a) and (14b), respectively. The composite function h (h2(7)) + has at least one fixed point y 2 i.e, there exists y 2 1 such that hy (hz(y")) +1=
Proof. We prove existence using continuity and boundedness arguments along with the intermediate value theorem. Note that h1 (T) and l12(7) are continuous functions. Thus, their composition f(r) = h1 (h2(7)) + continuous for all y 2 1.
We examine the behavior of f()) at the boundaries of its domain: At 1, we have h2(1)
co >> 0, where Co is a positive constant. Therefore, f(1 h (h2 (1)) +1 (1())) + 1 0, which is finite and positive since h is positive for all non-negative inputs. As 7 + co, (2(r)) + Coo > O, where Cos is a positive constant.
For sufficiently large 7,
establish that f(y) y. To see this, note that we have shown that h2(7) + Cop as 7 + co, where Coo is a fixed positive constant. Further, for any fixed positive value v, the function hy (7) is bounded, i.e., 1 (7) maxisisa A, which follows
24

---

Published as a conference paper at COLM 2025
from the fact that (7) is weighted average of the A; values. Therefore, there exists constant K such that f(r) & K for all y 2 1. Hence, follows that for all y > K, we have S 5 <<
y, which is continuous on (1) since both f(r) = - - h1 (Co) > 0 and for some '1 K, we ((n)~? < 0. Since f is continuous and changes sign from positive to negative on the interval [1, m., by the Intermediate Value Theorem, there exists at least one
Now, consider the function and y are continuous. We have
have f(11)
E (1) such that f(r*) = This implies f(r*) - 0, or equivalently, f(y*) =-
Therefore, there exists at least one fixed point 7 Z 1 for the function h1 (h2(2)) 1.
E.1 Proof of Theorem
Proof. Recapping the objective from (2) and following Definitions 1 and 2, we have
((P,w) = E (y (y-x*) P
E [ 2E yx T PX(w + E ( (x(())))))
Lety = x + and y; x B: + Ei, for all i [r], where 5,5i (O, o2) are i.d. Then,
Ely²1 = E (x B + E) = tr (E)
and
E yx PX(w
= E
(B' T x + ()x T P (x B: 2i)
B xx w;xx,' B: E
ΣPE Σ w/E [BB
tr (2PP) w
Here, the last equality comes from the fact that since B3 ri;B] is independent of Bj for
i [n + 1] following Definition 1, we have E [B = rin+tla and _n1 W B
returns w I.
Hence,
E (x PX(
=
x
(x; B: + +z))xx
Σ
H
W (x, (x, B: + (i) pTx
Σ w/ (x. B: + +: 1)
tr
P ΣPE
Σ wiw (x) Bi+ x1) (x]] 21) 1
+
p L'PEE
where
H
w/ (x, B: + 2 )0mm71 1) tr P EPE
Σ (x B.B. xi ) ) )
tr
P EPE
Σ
-|wull? tr P EP E xx xxx + +o2I))
-I|eo1?3 trr (IPTEP) (tr (Z) (tr
(I)) + 2tr (I²PTEP))
25

---

Published as a conference paper at COLM 2025
and
Σ (x) B: ())x (x) Bj + {i) 1)
wiwxxxx B:B; xxx] 1)
tr
P EPE
= tr P EPEE
Σ
(+)
tr ( EP
Σ
Wiw,, E B,1
(*)
tr (E'P T R 1)..
Combining all together and letting M tr (Z) + 0², we obtain
(() - M - 2tr
(22)) TO
+ ||ll tr EP EP) + (1|wl1 + +w Ro)tr (I'PTIP))
(37)
For simplicity, and without loss of generality, let
p= VEPVE.
(38)
Then, we obtain
[(P,w) = M - 2tr (IP) wT
(39)
+ M |||l tr ( (TT)) + (11wo12, +w RR))) (EPT) (IPTP).
Further, the gradients can be written
2M | ll. P 2 2(11w01 + - Rw)EP, pp -2 rI
(40)
VoL(P,w) = -tr (IP) r 2Mtr
(p p + 2tr [IP P) (In + R)w.
(41)
Using the first-order optimality condition, and setting wP) = 0 and C w) = 0, we obtain
-1 Tw
p
MM |1011 I ([looll,, + w (.)RR
Tr
||wll + +w Rw
-1
1 I
(42a)
M|WwI
M M|wll
w
( M
1+ ---1-- -1
M ||
where
TT T Rw
=
||00| 2
+ 1.
Further,
w ((Mer (p (PTP)= P
tr (EPTP) I ( R) tr (EP) r
(42b)
tr (IP)
tr (IPTT
1+ (t Mtr (pT
tr ΣP P) R
Mtr (pT P) + tr (EPTP)
Let
- Y I+
M
26

---

Published as a conference paper at COLM 2025
Then, get
tr (EPTT P) Mtr (P P) t (PPTP)
1 M tr tr [EP P)
P P)
tr [m7)
1+ M
tr [Σn'"
-1
$
d
s
1 M
Σ
(M + ys;)2
(M + 7s:)2
h2
Here, the last equality follows from eigen decomposition Σ Sa] E R
Udiag(s)U* with s =
Now, plugging P defined in (42a) within w given in (42b), we obtain
tr (IP) w Mtr (PT + tr (EPTP) (h2(7) R R+ R+
(43)
Using the above formulae for w, we rewrite 7 w Rw/ ||0|| 122
+1as as
7 T (h2(T)) 1) (((2(()) +1) -1r
1r
(12(T)) + 1)
2r
Ana
"
(44)
Σ
Σ (1+
I (1+
=: h1 (b2(T)),
where the second equality follows from Assumption where Ea with R", and the fact that R Ediag(d)E denotes the eigen decomposition of R, with = App E R
[a1,-... "......anpllpp fn)
follows from Lemma that there exists 7" such that by (h2(?*)) + y From (42a) and (43), we obtain
c(), (?* 1+ )
and
M
(45)
cr, E) (h2(7*) + 1)-1
for some C(r, w. Σ)
oTr and M
tr(IP)
Mtr(P ir +r
Now, using the our definition p= VEP we obtain
P(r*) = C(r,w,E) =ー 7-*
E +1)
Σ and
M
x(>*) = c)r, (12(7")- R R+1)
This completes the proof.
E.2 Proof Theorem
We first provide the following Lemma.
27

---

Published as a conference paper at COLM 2025
Lemma 2. Consider the functions h1 (f) and lt2(r) defined in (14a) and (14b), respectively, where
A; 2 0, a; + 0 for at least one i, s; > 0, and M 0² + +E - > Suppose Åy AR < M + Smin where Dg and DR denote the effective spectral gaps of Σ and R, respectively, as given in (13); and
Smin is the smallest eigenvalue of Σ. We have that
al1 ah1
< 1.
dr al12
(M + Smin) )2
Proof. Let
s7
A(r) 1+ M C(r)
B(7) = Σ The derivatives of B(7) and C(r) are
C(r) = Σ
(M + rs;) 2
(M + ysi) 2'
B(r)
s
B'(r) 2 Σ
C((r) -2 Σ
(M 73/)3 3r
(M + 7s;) )
1-1
The gradient of h2 A(7) -1 is
alt2
2 M A(T)B(7)) )²
(C()))())
(46)
dy
It can be seen that
2
2
s
S M
Σ
A(7)
(M + ys,)) ?
(M + ysi) 2
which implies that
2
d
-2
SM² 2
Σ
(47a)
A())
M + +y,)
Further, we have
d
-2 Σ
(M + 7
3
(M + 7s;)2
i=1
i=1
d
o7
d
s
2
(M + 7s;)
(M + 73
1=1
(47b)
d
Tij
ΣΣ
(M + yss (M 7s;)3
5753 - si)
M Σ
(M + 75)) (M + 79;)33
where
Tij (M + 18)s + 6/3 7s,)
- (M + -5)) + 7s;)s;
(47c)
5778 [M. (07 5 - 25;8;) + (5:37 + -;5? -s7s;)
Thus, substituting (47a) and (47b) into (46), we obtain
dhz2 δy
d
$7
-2
d
37.7 52)2
< M- M
Σ
(48)
(M
7s,)2
(M + 7 (M + 75,)3
28

---

Published as a conference paper at COLM 2025
Next, we derive di 09
Let
D(9) = Σ 1a
"1
af
E(9) = L
(1 + FA,)2'
(1 + 9A:}2
1=1
have
11
D'(7) - Σ 1=1
E'(7) ( Σ
(1 + 9A,33
(1 + 92,)3'
1-1
The derivative of It with respect to y is given by
ah1
(E())(
(49)
E($)
Substituting into (49), we get
-2
1
d
(50a)
E(7)
(1 + 9A,}3
and
11
11
af
( = Σ
Σ
(1 pA,)3
7A12
n
a²A;
Σ
(1+ +7;;)
(1 +1+791 FA,)
1=1
T
(50b)
= -E
(1 (1
+1
-
22 (1} + 21/1,)
ΣΣ
(1+ 12,13
Here,
T.j = (UNT(((++A)+++(()A))}{}
-,,a7 (1+ = (14,)
(50c)
=dfaß (17 (A,) + (1 + 7A,)13 - 1, + - A,A(1++
Hence, substituting (50a) and (50b) into (49) gives
ah d9
Σ
a
(A; 1222 (1 (1+ 31 + 92,13"
(51)
(1+
(1)
TA2)
2
Now, for the combined derivative, we have
ah2
dlt
s?
2
6,57(9; sj}2 (M + 79:)3 (M + ys;)
<
M + 7s,) 2
1
a,a 21,2
(1 72i2
(1 1-11
29

---

Published as a conference paper at COLM 2025
Note that M + ys; and 1 + 7A, are nonnegative for all j,j. Hence,
al12 dy
dh1
(M++ + 77;) d 57 (M (M
2
(M +ys))s
1-1
d
$/ A1 (M++5)) (M + 7s;)
(M + 7s:)² (M + 78))
a2 [ 911)
-2
11,A)3
d d2a7 A2 (1 + 9.,) (1 (1+ +9,)
(1+ 9/,)3 (1+ 11,)))
where
(s: 5/)
(Ai 1 12) iješ (1+ (1: + +7,1)
1 = max
A2 = max
(52)
ijes (M + 7sj) (M + 7s:)
Here, S {i [n] A; # 0} [n].
Finally, setting 7 h2(r), since y 2 1 and by our assumption Ar AR < M + Smin, we obtain
(12((()) k2(r)l = 2h2 ah1
4
- A
A2 5
1.
(M Smin )2
where ár and Ar are the spectral gaps of E and R; and Smin is the smallest eigenvalue of Σ; and M = =2 -2
Proof of Theorem follows from Lemma that there exists 7" 2 1 such that ft (h2(r*)) + 1 7 Further, Lemma 2 establishes that the mapping (h2 (2)) + is a contraction map- ping, since it satisfies Jält (hz / < 1. Therefore, by the Banach fixed-point theorem, there exists unique fixed-point solution, denoted by y*, satisfying Y it (h2(?*)) + This completes the proof of statement T1..
Next, we establish T2.. First, from Theorem 2, the stationary point (P*,w*), up to rescaling, is defined as
Y ( (MM-++)) Σー (==++)>'=+ and ~n = -hh((")) R R+ M
P P-
where 7 is a fixed point of composite function h (12(77) +
Given the coercive nature of the loss function L(P, .)) global minimum exists. Since the global minimum must satisfy first-order optimality conditions, it must lie within the set of stationary points characterized by Theorem 2.
Now, consider arbitrary global minimizers (f, Define scaling factors a, B > 0 and write (f, @) = (aP*, Substituting this scaling into the loss function gives:
L (aP", fw=") = M - -22Btr (I²p*)
pt TT r Ma $= 1| (IP"TEP')
+ a² =2 + wt Ro*) (I²prTH (I²P*TIP*) IPP)
(53)
Differentiating this expression with respect to a and / and setting the derivatives equal to zero, obtain the equation
2A + 2aB(MCB + (C+ 0,
(54a)
where we define
A = tr (I²pt) TT, B:= tr (IP=TT (
C:= |
D wt Rw" E tr II'p* EP-)
(54b)
30

---

Published as a conference paper at COLM 2025
Equation (54) implies unique constraint the product aB, ensuring no independent rescaling beyond a fixed ratio can further minimize the loss. Hence, no distinct minimizers exist other than scaling the original Theorem is indeed a unique minimizer up to rescaling.
pair. Thus, the stationary point (P", ,²,"") from
E.3 Proof Corollary
Proof. Since by assumption Σ it follows from (14b) that
-1
1
h2(r") 1 (d
0² + y + +12
(d + 02 +
+
1
d +1 +2
Substituting this into (15) gives
P =1 L, and w²= = R (d+ (d+
Now, recall that the objective function is given by
((w) = M 22r (2²P) To r M++M M|ll tr (EP EP) + (||wo1 + -o Rw)tt
and, by assumption, M = 0² + d.
Substituting p I and ~² (R + (d od 02 +1)1 r into the objective (37), and using E= I, we get:
L(w*) = (d² T ~ + (02 +d) 110 + +d (1101]2 + w (o=)
The expression simplifies
[(w") (o² - r R ( (d+ (d+)))))))) r+ (d2+))/(00"11 +d ( w Roo)
Next, we compute || and wn T Ro By definition, we have
||00112= (R (d+ od +1)1 2
and
T Rw" =x (R+ (d+ R ((R+((+00+1))) I r.
Thus,
(d + 02 +1]|[00112 + T "R= -r (R+ (R+(++++111))
(R+ + -1
=r- R (
(2+ r.
Substituting this result back into the objective function gives
L(w*) = (o2 +d) - -d.r" TT R (d o2 +1)1 r.
31

---

[K]. Published as a conference paper at COLM 2025
F Loss landscape of 1-layer GLA
E1 Proof of Theorem 4
Proof. We first prove that under Assumption 2, La min PeRiil LAPGO w)) where W is the search space of weighting vector w E R" defined as
W:
{
[ithh
wx]nk] E R"
0 S wi S w 1, su
Step 1: Define a set W = {wom
Wa]
E RD
0 s w S wj S1, 1A < siss S " and *
have W W. Given scalar gating G;
following (11), the weighting vector
18:1
returns
[81m 11
8un
Since GLA with scalar gating valued in [0, 1] following Assumption 2, that is, E [0, we have 8e++1 S 8ja+1 for 1 S S Therefore, any weighting vector w implemented by GLA gating should be inside W.
Step 2: Next, we will show that
o' EEW where wn = arg min LNPEO (P, P.EW
Define the weighting vector w [w] "f E R" where we have on = (woft)
(k) E R", For any w & W, there exist (ij,,) with j j- such that (K)
Way
wt
< sw)h Given gradient in (41), we have that
C
= =
(PT)) (Mtr
(K) + (IPTP) (1a -w)))
Here, (42a) gives that (optimal) P * and therefore, we have that (Mtr (PT P) M)) 0 and V w V uif Therefore either increasing
w (if 20) 0) or decreasing a0
(8) (if V oo 0) will reduce the loss. This results in showing that the optimal weighting vector wt satisfies with = w/th for any i E [mg] and k E [K]. Hence, w² W.
Step 3: Finally, we will show that any E W can be obtained via the GLA gating. Let [wi IT W be any vector in W and assume that wk < 1 without loss of generality. Then such sample weighting can be achieved via the gating
We WEe
WK
WIRK
ng
WK.K
Let wt = Se and let W' be in the form of wag
O411
Rd+p+1
W
Then, it remains to show that there exists w. satisfying
11,
k= 0, W k
( 0x)
The linear independence assumption in Assumption implies that the problem is feasible, which completes the proof of (22).
32

---

Published as a conference paper at COLM 2025
Proof of (23): Recap the optimal weighting from (15) which takes the form of
(h2(T)) R R+ 1)
where h2(7") > 0. Since Assumption 3 holds and #1
= "k == n, R is block
diagonal matrix, where each block is an all-ones matrix. That is
nxn
1 n
1
Therefore, we get
11((7)) h2(2*) - A
(h2(7")-8 R)
R.
Since Rr = Ar, then
e h2(r*)
Therefore, the optimal weighting (up to scalar) is inside the set W. Combining it with (22) completes the proof.
F2 Proof of Theorem 5
Proof. Following the similar proof of Theorem 4, and letting W
{
[[111] 1T ... wx1 [ E RR) we obtain
min PdRded WCW
LWPCO (Ppco(P,,)) (P,w)=
min W (P, (w)
PCRaxd WC
Therefore, it remains to show that any E W can be implemented via some gating function. Let [1]1]] [T]]] be arbitrary weighting in W. Theorem 4 has shown that if W1 2
WK. GLA with scalar function can implement such increasing weighting.
Now, inspired from Appendix D, Construction and (29) that all dimensions in the output implement individual WPGD, the weighting be composition of up to d different weights (u is (d + 1)-dieensi... Therefore, for any wT, ,WK E IR, we can get K separate weighting:
Φ1 = w wi[11, ITT
ng.
w2 (w2 w1) [0,, 1 ng.
03 = (w3 w2) 10.1, O 1
1 T nx
OK = (wo wx-1) 10 O o 0 Ong-1 Ang
Recap from Appendix D and (29), and consider the construction W, ((4+p+1)xd u (2+p+1)xp]] Assumption 2 implies that K S <
From (28) and (29), let i'th dimension implements the weighting w; for i E [K]. Specifically, let i'th row of each gating matrix G; implement weighting 10,, ... n 1 ... I (which is feasible due to Theorem 4) and set uilt; w; wi-1 with wo = O. Then the composed
weighting following (29) returns I wy which completes the proof.
33

---

Published as a conference paper at COLM 2025
E3 Proof of Corollary 2
Proof. Recap from (42a) that given Σ = / and w = 1,
+on++++++ p (d + + + (n+ ( -1
1
r
1 = cl.
n(d+ 02 + + RT
Then taking it back to the loss function (c.f. (37)) obtains
L(P", 1 = 1) = d + 0² 2ed1 + (d + ) + (n + 1 RI)cad
d + 0² cdl Tr. r.
It completes the proof.
34


Reviews:
[{"review_id": "Le6QR2tDDN", "reviewer": "Reviewer_Pxhm", "summary": "This paper investigates how Gated Linear Attention (GLA) models operate through the lens of In-Context Learning (ICL). It demonstrates that multi-layer GLA can execute a class of algorithms called Weighted Preconditioned Gradient Descent (WPGD), with weights that adapt dynamically based on the input data and gating mechanisms. To better illustrate how this adaptive weighting occurs, the authors introduce a multitask data model and analyzed how these WPGD-based models are optimized, proving under certain conditions the existence of a unique global minimum. Additionally, they apply these findings to single-layer GLA models, further examining their optimization landscape.\nThe paper also compares linear attention, scalar-gated GLA, and vector-gated GLA, highlighting the differences and explaining why gated attention models can offer superior performance compared to conventional linear attention methods.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "## Important and timely topic: \n\nEfficient new LLM architectures are crucial in current research. Although GLA and similar models have shown good results in practice, their theoretical understanding is still limited. This paper contributes significantly to the therectical understanding of the gating mechanism in GLA from the ICL perspective. \n\n## Solid theoretical framework and analysis\n\nThis paper makes a good case for connecting how GLA models work with WPGD algorithms and linking the gating mechanism and the date-dependent weighting mechasism. The interpretation of the results to ICL ablility is insightful. The optimization landscape analysis is also novel to me and adds to the theretical value of this work.", "reasons_to_reject": "While I really appreciate that theoretical contribution of this work, my biggest concern is its practical relevance. It follows the line of work that understands ICL as doing gradient descent. However, the connection requires rather strong conditions and hand-constructed transformer weights. There are other works arguing this does not relate well to how LLMs actually do ICL [1]. \nGiven the active research field of efficient new LLM architectures, it would be better if more practical insights/guidelines can be drawn. \n\n\n[1] Shen at al. Position: Do pretrained Transformers Learn In-Context by Gradient Descent?, ICML 2024.", "questions_to_authors": "* The equivalence between GLA and WPGD (Theorem 1) and the subsequent optimization landscape analysis rely on specific constructions for attention weights (Eq. 9) and input prompts. How strong are these setups and do the authors hypothesize that similar WPGD-like dynamics could emerge implicitly in GLA models trained with practical setups?\n\n* The problem setting, including the correlated task model (Definition 1), multitask distribution (Definition 2), and various assumptions on the problem seem too technical. Could the authors discuss how they relate to real-world in-context learning tasks? \n\n* The paper positions GLA as a framework encompassing models like Mamba and RWKV. Could the authors elaborate further on how the specific WPGD mechanism and the data-dependent weighting in this work can be helpful to better understand these models (e.g., their particular gating functions, state update rules, or recurrence mechanisms)?", "extra_scores": {}}, {"review_id": "bQJdBeaJZv", "reviewer": "Reviewer_qjca", "summary": "This paper provides a theoretical framework linking Gated Linear Attention (GLA) to Weighted Preconditioned Gradient Descent (WPGD). By constructing a GLA model with structured Q/K/V matrices, the authors show that GLA implements a class of WPGD, where gating acts as a weighting mechanism over in-context examples. They prove that the WPGD objective has a unique global minimum and show that GLA with scalar or vector gating can recover this optimum under certain conditions. \n\n**Quality.** This paper introduces precise and rigorous theoretical setup and derivations. It proves connections between the optimization landscapes of GLA and WPGD. Theoretical results are complemented with empirical validations.\n\n**Clarity.** The writing is clear and well-organized. Definitions and assumptions are clearly introduced, and contributions are outlined in relation to prior work throughout the paper.\n\n**Originality.** The work establishes a new theoretical interpretation of GLA as implementing a general class of WPGD algorithms, expanding prior results limited to linear attention or shared-task prompts in ICL. \n\n**Significance.** The paper addresses an important problem in understanding how gated attention architectures (e.g., Mamba, RWKV) perform ICL. By offering architectural insights, the work can shed light on how future gated models can be designed. However, the paper does assume a special attention weight structure, which might limit broader applications.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper addresses an important question of understanding the theoretical foundations of ICL in gated attention, which is relevant to recent architectures like Mamba and RWKV.\n2. The paper rigorously establishes a novel connection between GLA and WPGD, offering a principled interpretation of how gated mechanisms enable ICL.\n3. The paper adopts a multitask prompt model that allows task vectors to be correlated and not identical, extending beyond the shared-task assumption commonly made in prior ICL analyses.", "reasons_to_reject": "1. The theoretical results rely on a structured parameterization of $W_k, W_q, W_v$ (Eq. (9)), where many entries are fixed to zero to facilitate analytical tractability. However, real-world models are trained over fully parameterized weight spaces. It is unclear whether the GLA landscape results could be extended to this more general case.  \n2. The paper provides the conditions for uniqueness of the global minimum (Theorem 3). It would be more informative to include convergence guarantees for training GLA models to recover this optimum.", "questions_to_authors": "1. Eqn. (9) involves specific forms of $W_k, W_q, W_v$, which may not directly reflect how real-world architectures are trained. Can the authors empirically evaluate whether a fully parameterized GLA, trained end-to-end without these structural constraints, still recovers WPGD-like behavior or achieves similar optimality?\n2. While Theorem 3 establishes the uniqueness of the global minimum for the WPGD objective, can the authors discuss theoretical guarantees or insights on whether standard training procedures for GLA are guaranteed to converge to this optimum?", "extra_scores": {}}, {"review_id": "wK6e4AJcnt", "reviewer": "Reviewer_XH8t", "summary": "This paper investigates Gated Linear Attention (GLA) mechanisms, prevalent in modern efficient sequence models like Mamba and RWKV, through the lens of in-context learning (ICL). The core contribution is establishing a theoretical connection between multi-layer GLA and data-dependent Weighted Preconditioned Gradient Descent (WPGD) algorithms. The authors show that the gating mechanism induces data-dependent weights ω, allowing the model to modulate the influence of context tokens.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The paper establishes a novel and insightful connection between GLA architectures and a specific class of optimization algorithms (WPGD). This provides a valuable theoretical framework for understanding how gating contributes to the learning capabilities of models like Mamba and RWKV, going beyond standard linear attention.\n\n2. The paper sheds light on the functional role of the gating mechanism in GLA, linking it directly to data-dependent weighting. The distinction and comparison between scalar and vector gating provide valuable insights into architectural choices and their implications for model expressivity.", "reasons_to_reject": "1. The paper aims to understand GLA via its connection to gradient-based learning algorithms (WPGD). However, there is a significant body of existing work connecting linear attention and related recurrent models (like SSMs) to online learning. For instance, models like DeltaNet [1] frame state transitions explicitly as gradient steps minimizing an L2 loss. A more thorough discussion and comparison with these existing frameworks linking recurrence/attention to online optimization would strengthen the paper's positioning and clarify its specific contribution relative to prior art.\n\n2. The theoretical claims are validated primarily through experiments on low-dimensional synthetic data (e.g., d=10 feature dimension) using simplified 1-to-3 layer models. While useful for illustrating the core concepts, these experiments lack scale and complexity. It remains unclear how well the specific WPGD connection and optimization landscape results translate to practical, large-scale GLA-based language models operating on high-dimensional text data. The gap between the theory/simplified experiments and real-world applications is substantial.\n\n3. While GLA architectures are used in language models, the paper's core theoretical and experimental focus is on demonstrating ICL capabilities for linear regression tasks using synthetic data. The direct connection and implications for the complexities of *language modeling* (e.g., generative capabilities, linguistic structure understanding, large vocabulary handling) are not explicitly explored. The work feels more aligned with theoretical machine learning or representation learning venues than a conference specifically focused on language modeling. The link to language modeling is currently only asserted via the architectures studied rather than demonstrated through tasks or analysis relevant to language itself.\n\n[1]: Yang, Songlin, Jan Kautz, and Ali Hatamizadeh. \"Gated Delta Networks: Improving Mamba2 with Delta Rule.\" arXiv preprint arXiv:2412.06464 (2024).", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FbNFaMjvAq", "reviewer": "Reviewer_Xg8p", "summary": "This paper contributes to the theoretical understanding of in-context learning (ICL) as optimization. Specifically, it analyzes gated linear attention showing its relationship to a form of gradient descent on the squared loss (an established construction for the study of ICL). \n\nThe theoretical results developed are nontrivial yet clearly explained, their relevance well justified, and overall they provide valuable insight into how and why in-context learning works in GLA models. The work fits in a specific niche, using established constructions and definitions, and yet to the best of my assessment distinguishes itself and goes beyond the known results in that niche.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The questions and corresponding results are timely, and relevant, in the context of recent work in this direction.\n- A numerical simulation study confirms the theoretical findings.\n- The presentation is good, clear, and understandable to readers who are new to this line of work.", "reasons_to_reject": "I can spot no good reasons, really. I guess this might be a more theoretical-leaning paper low on the \"Empiricism, Data, and Evaluation\" dimension of COLM, but this is a weak reason to reject, as the paper still is relevant to empirical work in the field.\n\nThere are a few minor points in the presentation that I list below.", "questions_to_authors": "- I get that P fills the position that a preconditioning matrix would take in a gradient descent algorithm, but is P actually in some meaningful way a \"preconditioner\", does it reduce a condition number and/or improve convergence in some understandable sense? The form of P* in 15 suggests that maybe yes; could you please discuss/clarify? \n- H3 models are referenced without defining or explaining what they are. \n- In def 3, why do we call it an \"effective\" spectral gap, why not just a spectral gap? What does \"effective\" additionally denote here?\n- The discussion on line 182 about multi-layer models was a bit unclear: it seems from context that previous work does not consider the casual masking scenario (does this mean they consider fully unmasked \"encoder-style\", or what?), and I assume it's implied that the main result here does account for casual masking, but this is not said on this line, could you clarify?\n- Figure 1 is important and needs a bit of clarification. First, when skimming but before reaching the explanation for figure (d), it wasn't clear to me why GLA-vector would not be included in a-c also. I later understood (i think) that the fact that the red line (GLA) matches its dotted line is a sign that the theoretical result is reached and thus the vector formulation cannot improve anyway. However the fact that the dashed line for LinAtt and for GLA are not styled differently, this makes 1(d) more confusing. (in 1(d) it seems the GLA line matches the LinAtt theoretical bound, but it doesn't match the GLA theoretical bound, and we need the vector model for this. Is this right? It would be better if the LinAtt and GLA theoretical lines were styled differently (maybe dashed but color-coded). I recommend using opacity and line thickness in this figure to make it more readable when lines overlap a lot. The legend and the y axis can be shared to save space.\n- Still on Fig 1, Is it just a coincidence or is there a reason why  in 1d GLA matches so closely LinAtt(&LinAtt optimal)?\n- Still on Fig 1, is it right to infer that for a, b, and c Assumption 3 holds? Can this be checked/said explicitly? The caption could clarify why GLA-vector is only relevant in (d).\n- The Vaswani citation on line 520 is wrong, missing the other authors and also inconsistent with the rest, showing only first initial instead of full first name. This suggests that all references should be double-checked and published versions cited in lieu of arxiv ones whenever the case.\n- Is there some useful intuition for the form of $h_1$ and $h_2$? They look a bit like secular equations that appear in eigenvalue update problems but not quite.", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper explores the gated linear attention mechanism (popularized by models like Mamba) through a largely theoretical lens. The reviewers are unanimously positive about the paper, and were engaged in a lively discussion with the authors. it appears to be a solid contribution on an important topic. The primary concerns raised were about the strong (and perhaps unrealistic) assumptions that were made in the theoretical components in order to provide provable guarantees. However, for theoretical work, this is often the case, and it is nonetheless valuable to start somewhere. Overall, this work would be nice to see at COLM.
Decision: Accept