Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix
Authors: Yingyu Liang
Jiangxuan Long
Zhenmei Shi
Zhao Song
Yufa Zhou
Venue: ICLR 2025 Poster (2025)
Keywords: Weights Pruning
Attention Approximation
Gradient Descent Optimization
Primary area: foundation or frontier models, including LLMs

Abstract: Large Language Models (LLMs) have shown immense potential in enhancing various aspects of our daily lives, from conversational AI to search and AI assistants. However, their growing capabilities come at the cost of extremely large model sizes, making deployment on edge devices challenging due to memory and computational constraints. This paper introduces a novel approach to LLM weight pruning that directly optimizes for approximating the attention matrix, a core component of transformer architectures. Unlike existing methods that focus on linear approximations, our approach accounts for the non-linear nature of the Softmax attention mechanism. We provide theoretical guarantees for the convergence of our Gradient Descent-based optimization method to a near-optimal pruning mask solution. Our empirical results demonstrate the effectiveness of our non-linear pruning approach in maintaining model performance while significantly reducing computational costs, which is beyond the current state-of-the-art methods, i.e., SparseGPT and Wanda, by a large margin. This work establishes a new theoretical foundation for pruning algorithm design in LLMs, potentially paving the way for more efficient LLM inference on resource-constrained devices.

Published as a conference paper at ICLR 2025
BEYOND LINEAR APPROXIMATIONS: A NOVEL PRUN- ING APPROACH FOR ATTENTION MATRIX
Yingyu Liang"
Jiangxuan Longl
Zhenmei Shi
Zhao Song
Yufa Zhou
ABSTRACT
Large Language Models (LLMs) have shown immense potential in enhancing var- tous aspects of our daily lives, from conversational AI to search and Al assistants. However, their growing capabilities come the cost of extremely large model sizes, making deployment on edge devices challenging due to memory and com- putational constraints. This paper introduces novel approach LLM weight pruning that directly optimizes for approximating the attention matrix, a core com- ponent of transformer architectures. Unlike existing methods that focus on linear approximations, our approach accounts for the non-linear nature of the Softmax attention mechanism. We provide theoretical guarantees for the convergence our Gradient Descent-based optimization method to near-optimal pruning mask solution. Our empirical results demonstrate the effectiveness of our non-linear pruning approach in maintaining model performance while significantly reducing computational costs, which is beyond the current state-of- of-thee-rt methods, i.c., theoreti-
SparseGPT and Wanda, by large margin. This work establishes cal foundation for pruning algorithm design in LLMs, potentially paving the way for more efficient LLM inference on resource-constrained devices.
INTRODUCTION
Large Language Models (LLMs) based on the transformer architecture Vaswwnn et al. 201 7), in-
cluding GPT-4o (OpenAl 2024a), Claude (Anthropic) 2024). and OpenAl's recent ol (OpenAl
2024b), have shown immense potential to enhance our daily lives. They revolutionize fields like conversational Al (Liu al., 2024). Al agents (Xi et al. 2023 Chen al. 2024b), search Al Ope- nAl 2024b). and AI assistants (Mahmood et al.) [2023] Zhang et al. 2023a Kuo et al.) 2024 Feng al. 2024). With their growing capabilities, LLMs are powerful tools shaping the future of tech- nology. However, the current state-of-the-and L.M weights number is extremely large. For instance, the smallest version of Llama 3.1 (Llama Team 2024) needs billion parameters, which takes more than 16GB GPU memory with half float precision and requires significant inference time. Deploying such models on edge devices such as smartphones becomes challenging due to their large memory and high computational cost. To reduce the LLM model size, many studies work pruning the LLMs model weights to relax the device memory constraint and minimize response latency. The classical pruning problem in LLMs can be formulated as follows. Given a weight matrix W E Rddd and some calibration data X E Rnxdd where is input token length, and d is feature dimension, the goal is to find matrix W under some sparse constraint such that || XW] being small under some norm function. The above formulation has been widely used in many state-of-the-arr pruning methods, such as SparseGPT (Frantar & Alistarh 2023) and Wanda (Sun al. 2024).
However, the current object functions only focus on the approximation of a linear function XW. Their optimal solutions do not have good approximation to the attention matrix (see Figure 2 for details). Note that the attention mechanism is the kernel module of the transformer architecture. The high-level insight of their bad performance is that the Softmax function is very sensitive to
yingyul @hku. hk. The University of Hong Kong.
yliangees. wisc. edu. University of
Wisconsin-
-isconsin-Madison. Madison.
t mangchianghsuannggme com. South China University of Technology.
zhmeishi@ca edu. University of Wisconsin-Madison.
$ magic. magic.linuxker. winuxkdeeggaaaaaaa. com. The Simons Institute for the Theory of Computing at UC, Berkeley.
yufazhcu@seas, upenn, edu. University of Pennsylvania.

---

Published as a conference paper at ICLR 2025
Softmax nk
X
x
ck
M
o
o
W
x
XT
d
d
n
X
X ck
M
ck
We
n
X
x
d
M
o
W
Figure 1: Comparison of our Attention Weights Pruning method and Linear Pruming method such as Wanda and SparseGPT. The top figure illustrates our proposed method of the attention matrix approximation, where pruning is applied directly to the fused attention weight matrix W, using only one pruning mask M. The bottom figure describes the Linear Pruning method of the linear function approximation, where pruning applied separately to the query weight matrix Wq and key weight matrix WK, using two different pruning masks Mq and Mr. respectively.
the large positive values of the input due to its exp scaling effect, while pruning mask based on linear approximation cannot capture this sensitivity. Thus, in this work, we directly compute the pruning mask weights to approximate the attention matrix, which is highly non-linear function, Softmax( XWX X) E R To the best of our knowledge, this paper is the first work studying attention weight pruning directly approximate the attention matrix. We provide a theoretical guarantee that optimization based on Gradient Descent (GD) on our loss function can converge to a good pruning mask solution (Theorem 1.3). Furthermore, we preliminarily verified the effectiveness of our method with empirical support (Section G). Our theoretical foundation may pave the way for more efficient LLM inference on resource-constrained devices.
Key background. We introduce some key backgrounds. We define the attention matrix in the self-attention mechanism as below:
Definition 1.1 (Attention Matrix). Let X Raaxd be the input. Given query and key weights matrix Wo, WK E Idddd we define W := WWW K. Then, we have the Softmax attention matrix being
Softmax(XWXX] = D- exx((XXTXX
where (1) D = entry-wisely, (3) diag() operation takes vector and outputs diagonal matrix with the entries of that vector, and (4) 1, denotes the length-n all ones vector.
1,), (2) exp denotes the exponential function and is applied
Further, introduce the problem setup of our Attention Weights Pruning. By selectively reducing the number of non-zero elements in the attention weight matrix W in Definition we can preserve model performance while lowering computational cost and GPU memory usage. Below, we formally define the Attention Weights Pruning problem and the corresponding loss function:
Definition 1.2 (Attention Weights Pruning). Let M E [0, 1]dxd x be the pruning mask. Let X, W be defined in Definition 1.1. Let A = exp(XXXX X and À = exp(X(M o W) where o is the Hadamard product. Let D == diag( A 1,) and D == A Let X E R4, be the regularization parameter. We define the Attention Weights Praning loss function to be
((M) = 1A0==1ID'AA--AA Al + MI
Thus, the Attention Weights Pruning optimization problem is mins L(M).
2

---

Published as a conference paper at ICLR 2025
Our contributions. This is the first work studying the Attention Weights Pruning problem, which an approximation problem to a non-linear function. We provide an algorithm for obtaining the near-optimal pruning mask based on Gradient Descent (GD) with convergence guarantee.
Theorem 1.3 (Main result, informal version of Theorem 4.1 For any 6 > O, our Algorithm !
can converge to the near-optimal pruning mask for the Attention Weights Pruning problem (Defini- tion 1.2 in (dd poly(n) /c) time with O + +) error, where is small term depending on intrinsic property of the data and weights.
In the above theorem, & can be arbitrarily small as & - when the regularization coefficient À O. Thus, our analysis shows that although the objective function is highly non-linear, the GD training can converge a near-optimal pruning mask solution.
Our experiments on synthetic clearly align and support our theoretical analysis (Section 6.. D. Fur- thermore, we evaluate our non-linear pruning method and show it beyond SparseGPT and Wanda by a large margin on real data (C4 Dataset Raffel et al. (2020) under real LLM (Llama 3.2-IB Meta (2024) in Section 6.2 and Section 6.3
Our contributions are as follows:
This is the first work that analyzes the weights pruning problem based on Softmax attention, which is non-linear function.
We provide the closed form of the gradient of Attention Weights Pruning loss function (Theorem 5.3). and Lipschitz of that gradient (Theorem 5.4
We provide Gradient Descent based Algorithm I to obtain the near-optimal pruning mask and its convergence guarantee (Theorem 4.1)
We conduct experiments verify the effectiveness of method (Section 6). showing our non-linear pruning method is beyond SparseGPT and Wanda by a large margin.
Roadmap. In Section we review the related work. Section introduccs key concepts and defini- tions essential for the subsequent sections. In Section we present our main result. Section 5 offers a technical overview of the methods employed. Experimental results are discussed in Section Finally, Section summarizes our findings and offers concluding remarks.
2
RELATED WORK
2.1 PRUNING AND COMPRESSION FOR LLMS
Model compression plays a critical role in improving the efficiency and deployment of large lan- guage models (LLLMs) (Zhu et al.) 2023) for its effectiveness in reducing computational overhead while preserving performance. Common compression techniques include quantization (Park et al. Xiao et al. 2023; Hooper et al. 2024), pruning (Chen et al., 2021 Hoefler al. 1021
2024
Hubara et al. [022 Jin et al. 2022 Frantar & Alistarr 2022 2023) Sun et al. 2024 Li et al.
2024a
Zandieh al.) 2024 Zhang et al., 2024c| Xia al. 2023 Ashkboos al. 2024 Chen
et al. 2025b). and knowledge distillation (Hsich et al. [003 Shridhar et al.] 2023 Jiang et al.
2023
Wang al. 2023). Specifically, pruning techniques have been developed extensively, such as unstructured pruning, which removes individual weights (Li et al. 2024a Sun et al. 2024), and structured pruning, which eliminates entire components like neurons or attention heads (Michel eet al. 2019. Ashkboos et al. 2024 Xia et al,, 2024). [Snn et al. (202)) proposed Wanda, novel unstructured pruning technique that uses weight-activation products to induce up to 50% sparsity in LLMs without retraining, achieving competitive results with significantly lower computational cost. SparseGPT (Frantar & Alistarh, 2023) introduced a one-shot pruning method that achieves up to 60% sparsity in large GPT-family models with minimal impact on performance. A follow- up work (Li et al. 2024a) improved the complexity analysis of SparseGPT, reducing the running time from (dd) to ((d2-53) enabling faster pruning LLMs. Together, these techniques tribute to more scalable and resource-efficieen LLMs, maintaining competitive performance while substantially reducing computational resources.
3

---

Published as a conference paper at ICLR 2025
2.2 ATTENTION ACCELERATION
The attention mechanism has faced criticism due to its quadratic time complexity with respect to context length (Vaswani et al. 2017). Addressing this criticism, a variety of approaches are em-
2023 Frantar & Alistarh 2023 Li al. 2024a). low-rank approximations (Razenshteyn et al. 2016
ployed, including sparse attention (Hubara et al. [2021] Kurtic al.
Li et al. H0 et al. 2022
Zeng & Lee 2024; Hu al. 2024d Li et al. 2025a), and kernel-based methods (Charikar al.
2020 Liu & Zenke 2020 Deng et al., 2025] Zandieh et al.] [2023] Liang et al.] 2024b), to reduce
computational overhead and improve scalability, Aggarwal Aggarraa & Alman (2022) enable the derivation of
a low-rank representation of the attention matrix, which accelerates both the training and inference
processes of single attention layer, tensor attention, and multi-layer transformer, achieving almost
linear time complexity (Alman & Song 2023 2024a b Liang et al. 2024g c Ke et al. 2025a Hu
al., 2024h). Other approaches like Mamba (Gu Dao 2023] Dao & Gu [2024). Linearizing
Transformers (Zhang et al.] 2024b Mercat et al. 2024). Hopfield Models (Hu et al., [0233 Wu et al.,
2024b, Hu et al.] [024cc] Xu et al.
2024a; Wu et 2024aa Hu et al. 2024a blt). and PolySketch-
Former KKachm et al.) 2023) focus architectural modifications and implementation optimizations to enhance performance. System-level optimizations such as FlashAttention (Dao et al. 2022 Dao
2023 Shah et al. 2024 and block-wise parallel decoding (Stern al.
2018) further improve efñ- ciency. Collectively, these innovations have significantly augmented transformer models' ability to
handle longer input sequences, unlocking broader applications across multiple sectors (Chen al.
2023
Peng et al. 20233] Ding et al., 2024 Ma et al. 2024 Xu et al. [2024d] An el al.
2024
Jin
al.
2024
Li
2024c
Liang et al. 2024d Shi al. 2024a).
3 PRELIMINARY
Notations. For any positive integer n, we use [n]] to denote set (1.2.....,)) ,n]. b E I" to denote the Hadamard product, i.c., the i-th entry of (a o 6) is a.b; for all i E [n]] For E Rmx* let A, E R" denote the i-th row and A., RT denote the j-th column of A, where E [m] and j E [n]] We use exp(A) to denote a matrix where exp(A)i, == .xp(A,,)) for matrix A Roxd We use ||A|| to denote the Frobenius norm of matrix A E Rnxdd i.e.,
For each a, b e R",
|||||e = / EiEfmi EjEiq |2 For symmetric matrix A E R A Z 0 means that A is positive semidefinite (PSD), i.e., for all R", we have Ax 2 0.
Attention weights pruning. We aim to determine a near-optimal pruning mask M for the attention weights in self-attention mechanism. Furthermore, we incorporate causal attention masking into our method to be more aligned with the current decoder-only LLM architecture, while our analysis can be applied to any general attention mask, e.g., block-wise attention mask. To formalize this, we provide the formal definition of causal attention mask and attention weights pruning in this section.
The causal attention mask ensures that each token in the sequence can attend only to itself and the preceding tokens. Here, we provide the formal definition of the causal attention mask:
Definition 3.1 (Causal attention mask. Liang et al. (204c) M, E {0, 1 1 where (Ma)k.j = 1 1/12 and (M.).J
We define the causal attention mask as O otherwise.
Now, we incorporate Attention Weights Pruning (see Definition 1.2 with causal attention mask M.-
Definition 3.2 (Attention Weights Pruning with Causal Attention Mask). Let M. E {0.1)ax be the causal attention mask defined in Definition 3. l
Let A = exp(XW XT) o Ma and Ā == eX o W) a Mc- Let D == diag(A (A 1,) and D = diag(A A 1n). We define Attention Weights Pruning with Causal Attention Mask loss function to be L(M) Catta (M) (rr
where Latta = D A Đ Ā1 and (Me = 2 al| |3.
In this paper, we always use pruning mask to refer M Raxd, which is our target, and use causal attention mask to refer M. IR x which is fixed mask for standard self-attention.
4

---

Published as a conference paper at ICLR 2025
Algorithm Gradient Descent for Pruning Mask (Theorem 4.1) Let X1, X2,..., : Xx E RRuxd be calibration dataset of size k. We iteratively GD method on this dataset.
1: procedure PRUNEMASKGD( X2,........ Xx E Raxd W E Rdddd Mc E {0, 1]dxd p
[0, 11, A [0, 1]. (0.).1)
D Theorem 4.1
2:
Initialize M E [1]axd
3: Initialize 7, T by Theorem 4.1
4:
for j= = k do
Iterate over the dataset
5:
"j exx((xxXX
6:
f < diag((u))))) ((u, Mc) 1n. (u, Mc)
7:
end for
8:
for
+
Iterate over T steps Iterate over the dataset
9:
for
j 1 + k do
10:
" - x o ))
f, diag((U, Mc) 1n) 4n)=((ū; Mc)
12:
+ fj - f
13:
Pıj + 9o
14:
P2. + diag(P1.j 1
15:
Pj + Pı. p2.j
16:
end for
17:
M M (n/k) (W (E)- xJ)) AM)
Gradient Descent
18:
end for
19:
m f vec(M)
> Flatten M into a vector Sort the elements of M
20:
Maorted + sort(m)
21:
T T Maortedllp Ilp d
b Get the p-th largest element
1
if Mi >
22:
Mi
> Set the top p entries to 1, others to 0
10
otherwise
23:
return M end procedure
4 MAIN RESULTS
We provide an Algorithm foo Attention Weights Pruning problem based on Gradient Descent (GD). We also prove the convergence for our GD algorithm Theorem 4.1
Theorem 4.1 (Main result, formal version of Theorem 1.3). Let M, X, W, W,D,A,, x, L, Latta be
defined in Desinition 3.2 Assume XXT BI and ming ( AA) 2 8 0. Furthermore,
Let μ min.,je(d [W 4. - 3 6. Let 12n 1.5 je(d [W.,1] || Ad/p. Then. for any > 0, provided < 1/L where L is the Lipschitz constant for V C(M) (see Theorem 5.4. GD (Algorithm) 7) with fixed step size and run for T = ( iterations results in the following guarantee,
X2
min Cattn (M +
M S (E e)/2.
n²
EC1
un?
|| Note that L(M) satisfies the (g( n2)),nz. p)-proxy PL inequality (Lemma (.5). Also, we have L(M) is non-negative and has L-Lipschitz gradients Theorem 5.4 Thus, we finish the proof using Theorem 5.2
Proof. Let (MM) 2Latta (M)+ 2x
Remark 4.2. The two assumptions in Theorem 4.1 are practical. The first assumption of the positive definite matrix is widely used theoretical deep learning analysis, e.g., & Liang (2018); Du al. 2019); Allen-Zhu et al. (2019b); Arora el al. (2019). The second assumption is natural, as D the pruned attention matrix, where each entry is
exp(X (M W)X JaJ 0. j= expeexXXXX (M W)))),
5

---

Published as a conference paper at ICLR 2025
which has a natural lower bound.
Our error upper bound in Theorem 4.1 is O + e), where can be arbitrarily small. For &, we can let it be small by choosing a proper X, i.e., the error term can be made arbitrarily small by choosing small A. However, if we choose very small X, the algorithm's run time gets larger as x 1/n & L & (A + other terms). Thus, although the objective function is highly non-linear, show that the Gradient Descent of our Algorithm can converge good solution of the Attention Weights Pruning problem.
After solving the optimization problem, obtain pruning mask with real-valued entries. In prac- tice, however, this pruning mask must be converted into a binary form, specifically M E {0, ] We define the pruning ratio / as the percentage of weights to be pruned. We apply this ratio by setting the pruning mask entries to zero for weights that fall below the p-th percentile and to one for those above. This ensures that only the specified proportion of weights is pruned.
5 TECHNIQUE OVERVIEW
In Section [.1 we introduce some useful tools from previous work. In Section 5.2 we derive the close form the gradient of Attention Weights Pruming. In Section [5.3] we calculate the Lipschitz constant of that gradient. In Section 5.4, we prove the PI. inequality for our loss function.
5.1 PREVIOUS TOOLS ON CONVERGENCE OF GD
To analyze the convergence behavior of GD for optimization problem (Definition 3.2 introduce the concept of g-proxy. E-optimal Malyak--ojasiewicz/PP inequality (Polyak 1963 Lo-
first
jasiewicz 1963 Karimi et al. 2016), under which GD will converge:
Definition 5.1 (g-proxy, f-optimal PL inequality, Definition 1.2 in Frei & Gu (2021). We say that function f TRP > IR. satisfies a g-proxy, E-optimal Sojyak-tojssieiicz inequality with parameters and μ O (in short, fsatisfies the (g, E, a, p)-PL inequality) if there exists a function IR + R and scalars & E R, p > 0 such that for all w E IR,,, [[w]| =p(g(w) E)
PL inequality is a powerful tool for studying non-convex optimization, and it has been used in recent studies provable guarantees for neural networks trained by gradient descent (Li & Liang.
2018 Allen-Zhh Zhu et al. 2019a b
c Frei et al. 2019; Cao & Gul 2020 Ji & Telgarsky] [019 Frei al. 2021 Shi et al. 2021 2024b). It provides proxy convexity property, although the objective function is non-convex. In detail, for a function with good smoothness property, we can find some proxy functions and show the convergence by utilizing these proxy functions.
Leveraging this PL inequality, Frei & Gu (202)) derives the following GD convergence guarantees.
Theorem 5.2 (Theorem 3.1 in Frei & Gu (2021). Suppose (w) satisfies the (g(.), E, a, p)-proxy
PL inequality for some function g(-) RD gradients. Then for any O, provided n < 1/L, GD with fixed step size n and run for T
>> IR. Assume that f is negativee and has L-Lipschirz
21 (jue 2220 (w) iterations results in the following guarantee, minkeT g(w(t)) +
The above theorem establishes that under the (g, E, a, PP inequality and Lipschitz continuity of the gradient, GD converges to point where the proxy function g(w) within c of &. To apply this result to our specific problem, we need to verify these conditions for our loss function L(M).
5.2 CLOSED FORM OF GRADIENT
As first step, compute the close form of the gradient C(M). The pruning mask M is inside a non- non-linear function Softmax, which complicates our calculation. We defer the proof to Section C
Theorem 5.3 (Closed form of gradient, informal version of Theorem D.5). Let L(M) be defined in Definition 3.2 Let p be defined in Definition D.1 Let X E Bxxd M [0. 1]dxd W E Rddd Then, we have
dC(M) W o (XT pX) + AM.
dM

---

Published as a conference paper at ICLR 2025
Based on Theorem 5.3 calculate the gradient of the pruning mask from Line 10 Line
13
our Algorithm
5.3 LIPSCHITZ OF GRADIENT
Having obtained the close form of gradient, we proceed investigate its Lipschitz continuity, We aim to show that the gradient V Var((M) is Lipschitz continuous with respect to M.
Theorem 5.4 (Lipschitz of the gradient, informal version of Theorem F.8). Let R be some fixed constant that satisfies R > 1. Let X E Raxd W Iddd We have || S R. and |[]] & R.
Let (MM) be defined Definition 3.2 For M, M Rdddd we have
IV M(M) - VM
(A 30dn R*) ||M M.
We defer the proof to Section Establishing the Lipschitz continuity the gradient satisfies one of the necessary conditions for applying Theorem 5.2 The above theorem implicates that the gradient for M is upper bounded, providing a way to choose step size.
5.4 PL INEQUALITY OF GRADIENT
Next, we need verify that loss function satisfies the PL inequality with appropriate parameters. To complete the verification of the conditions required for convergence, we demonstrate that L(M)
satisfies the PL inequality. We show that as C(M) satisfies the PL inequality in this lemma:
Lemma 5.5 (PL inequality, informal version of Lemma G.10 Let X, W, D, x, L, Cattn be
defined in Desinition 3.2 Assume that XXT ž 3I and JE (D- A)i.j 2 6 >>> 0. Also,
. Let μ = 2 min (IWas 8 8.
Leet = 12 max (WW || Ad/p.
have M[ Z (M(2Catn +
M E).
We defer the proof to Section G By confirming that L(M) satisfies the PI. inequality and that its gradient is Lipschitz continuous, we then apply Theorem 5.2 to conclude that GD will converge to solution within our desired error tolerance, and further prove Theorem 4.1
To prove the PL inequality, we also need the following two key Lemmas, which introduce our two assumptions our Theorem 4.1 X = BT and ming,je(m) (D-1 2 O.
Lemma 5.6 (Informal version of Lemma G.4). Let B E [ and X E Radd Assume that
XX ž =BI. Then, we have || BX F 2 B Ir.
Lemma 5.7 (Informal version of Lemma G.7 Let E RRxn and each row summation is zero, i.e. 1, = O,, Let B [0, 1]nx> and each row summation 1, i.e. 1, = 1n- Assume that min Br.j 2 8 > 0. Then, we can show [| o B diag((B 3 o B) 1, BlIr 2 5 |[
6 EXPERIMENT
6.1 EVALUATION ON SYNTHETIC DATA
We discuss the synthetic experiments conducted illustrate the effectiveness of our Algorithm
Method and evaluation. We implement our method following the pseudocode in Algorithm T us- ing NumPy and JAX for acceleration. We evaluate our method on unstructured sparsity, meaning that zeros can occur anywhere within the attention weight matrix W. Specifically, we use Defini- tion 3.2 as our loss function, optimizing over the pruning mask J using gradient descent based on the closed-form expression derived in Theorem 5.3 To accelerate convergence, leverage mentum into the optimization process and fix the momentum parameter 0.9. After obtaining the optimal pruning mask, convert M to binary pruning mask to prune W. maintaining sparsity

---

Published as a conference paper at ICLR 2025
16, O.,.A 0.04,
122
Ours
Wanda
1.0
SparseGPT
Rolative
64123
256
512
1024
0.4
0.5
0.6
0.7
0.8
6.
64, p 0.5, 128,
64, p 0.S, 0.04,
128,
1.4
Ours
Wanda
SparseGPT
1.4
8.0
1.0
66 128 238
1024
0.4
0.5
0.6
0.7
0.8
64
(a) Varying A.
(b) Varying
(c) Varying
Figure 2: The comparison among our Algorithm| Wanda, and SparseGPT. The y-axis a relative - All: where D-1 A is original attention matrix and Đ-1 Á is approximated attention matrix based on three methods. We always use d = 64. We use k = 16 for the first row and & = 64 for the second row. The r-axis is (a) regularization coefficient for the left column; (b) input sequence length n for the middle column; (c) pruning ratio for the right column.
error, which is defined as [| |
the desired pruning ratio We use the relative error our evaluation metric, which is defined as || - A D-////D-AAA/ where D, A, D. A are defined in Definition 3.2
Baselines. We compare our method with two linear pruning approaches, namely Wanda (Sun et al,, 2024) and SparseGPT (Frantar & Alistarh 2023). Wanda is pruning method that removes weights with the smallest magnitudes multiplied by the corresponding input activations, achieving sparsity without requiring retraining or weight updates. SparseGPT is second-order pruning method that utilizes the Hessian matrix to prune a portion of the weight matrix while simultaneously updating the remaining parameters. We implement Wanda and SparseGPT as described in their respective papers. Notably, since the settings of SparseGPT and Wanda are linear, we do not prune the fused weight matrix W directly; instead, we prune Wo and WK separately (see Figure T..
Data. In order to assess the efficacy of different methods in approximating the attention matrix, we construct the data via carefully defined generating process. Specifically, we create multiple independent random Gaussian matrices G E Raxd, where each entry of G' drawn from normal distribution, i.e., G. NO0 1) for i, E [d]. Then, we perform singular value decomposition (SVD) on matrix G, i.e., U, V = SVD(G). We retain the first four singular values in and set others to zero, constraining the rank to four. Our We and W K then constructed as U V The weight matrix W used in our setting is formed by taking the product W WoW K. For X E Rx we generate it full-rank Gaussian random matrix.
64 is kept constant across all figures, and we simulate two datasets of size & 16 and k 64. We set the input sequence length n = 128 for experiments (a) and (c) in Figure 2 The pruning ratio p = 0.5 set for experiments (a) and (b) in Figure 2] For our method, the regularization coefficient À == X/n where we abuse the notation to denote À the same used in Definition 3.2 and À here is the parameter we really control in experiments. is set as 0.04 for experiments (b) and (c) in Figure 2 (intuition drawn from experiment (a)). The total number of epochs is set T 100. The step size is set as 0.1/A because Theorem 4.1 indicatss that is inversely proportional to À with some constant, i.e., n x 1/L 1/(A + other terms).
Setup. our experiments, the weight matrix dimension
8

---

Published as a conference paper at ICLR 2025
0.7, 0.05, d 2048
0.005, 512, d 2048
Wanda SparseGPT
0.34
0.68
6.00
0.64
2/5 211 2
Figure 3: The comparison among our algorithm, Wanda, and SparseGPT on Llama 3.2-1B,
show that our Algorithm D outperforms Wanda and SparseGPT with large margin, which supports theoretical analysis in Theorem 4.1 In the following, we will discuss each setting in detail.
Results. Overall, the results Figure 2
Relation with regularization coefficient X. The leftmost column of Figure investigates the impact of the regularization coefficient À on relative error. As increases from very small values, the relative error initially decreases sharply for our algorithm, reaching minimum before gradually rising again, which forms a U shape curve. This behavior indicates that there is an optimal À where our algorithm achieves its best performance around 2 4 The U-shape curve phenomena are well- known in most syper--arameter choosing, e.g., regularization coefficient.
Relation with input sequence length n. The center column of Figure Seareeexpppe explores how the relative error changes with respect to the input sequence length n. As increases, the relative error for all three methods grows, though at different rates. Our method demonstrates slower increase, maintaining significant margin over both Wanda and SparseGPT, particularly for larger values of n. Wanda, while showing better performance than SparseGPT for larger sequence lengths, becomes comparable SparseGPT
is relatively small.
Relation with pruming ratio p. The rightmost column of Figure 2 illustrates the relationship between the relative enror and the pruning ratio p for the three methods under comparison: algorithm, Wanda, and SparseGPT. As the pruning ratio P increases, all methods exhibit a rise in relative error, indicating degradation in approximation accuracy. However, our algorithm consistently outper- forms both Wanda and SparseGPT across the range of p. with a lower relative error. SparseGPT and Wanda follow a similar trend, closely tracking each other.
6.2 EXPERIMENT REAL DATASET AND LLMS
In this subsection, we discuss the experiments conducted on the real dataset and LLMs to illustrate effectiveness of our method.
Method and evaluation. We evaluate our method on unstructured sparsity, meaning that zeros can occur anywhere within the attention weight matrices Wq and WK. Specifically, we use the loss function defined below:
[MM. MR) = 5- Al + (1 ||
1
where D and A are the original attention matrix, Mq, Mr are pruning masks, and
Ā : cxp(X WQ) [x)) D diag(A A 1,).
After obtaining the optimal pruning mask, we convert the pruning mask to a binary pruning mask to prune WQ and WK. maintaining sparsity the desired pruning ratio p. We the relative error as our evaluation metric.
Baselines. The pruning is performed the last attention layer of the pretrained model Llama 3.2- IB Metal (024). The baselines are Wanda and SparseGPT, the same as Section 6.2
Data. To simulate real-world large language models, utilize the Colossal Clean Crawled Corpus (C4 Dataset) Raffel et al. (2020), which also used as the calibration dataset in our baselines Wanda and SparseGPT. Additionally, with primary focus pruning the attention matrix, extract the

---

Published as a conference paper at ICLR 2025
imput hidden states corresponding the target attention matrix from the pretrained Llama 3.2 model using a customized hook function and use these as our input X.
Setup. The weight matrix dimension is d 2048 in Llama 3.2-1B. For all the experiments, set as 0.05 and y as 0.005. For the varying n experiment, we set the pruning ratio as 0.7. For the varying P experiment,
set the input sequence length n as 512.
Results. Overall, the results in siow that our algorithm outperforms Wanda and SparseGPT in real-world LLMs, which supports theoretical analysis in Theorem [4. and enhances pre- liminary experiment in Section 6 In the following, we will discuss each setting in detail.
Relation with input sequence length The left column of Figure 3 shows that algorithm out- performs the baselines in different sequence lengths continuously.
Relation with pruning ratio The right column of Figure 3 illustrates that the pruning ratio p increases, all methods exhibit a rise in relative error. But our algorithm consistently outperforms both Wanda and SparseGPT across the range of with much lower increasing rate.
Assumptions verification. Notice that Theorem 4.1 relies on two assumptions: XXT > 37 and ming,jejm] A-j,) 2 8 > O. We verify these assumptions using the C4 dataset, obtaining , m 0.034 and 0.0025, thereby demonstrating the practicality our theoretical framework.
6.3 EXPERIMENT ON END-TO-END PERPLEXITY
We present the experiment on end-to-end per- plexity in this section.
Table 1: Comparison of different methods based on Perplexity (PPL).
Baselines. As our method focuses on pruning the attention matrix, it can be seamlessly com- bined with approaches that perform linear prun- ing on the MLP, such Wanda and SParseGPT. Therefore, we conduct three groups of experi- ments: (1) using dense MLP without pruning, (2) applying SparseGPT to prune the MLP layer, and (3) using Wanda to prune the MLP layer. Subsequently, we apply Wanda, SparseGPT, and our method to prune the attention weights.
Method
PPL
Dense MLP Dense Attn
12.487
Dense MLP + SparseGPT Attn
14.269
Dense MLP + Wanda Attn
14.912
Dense MLP Our Attn
13.885
Wanda MLP + Wanda Attn
30.426
Wanda MLP + SparseGPT Attn
26.074
Wanda MLP + Our Attn
24.427
SparseGPT ML. + Wanda Attn
45.435
SparseGPT MI.P + SparseGPT Attn
36.641
Method and data. We use the same method and data as Section 6.2
SparseGPT MLP + Our Attn
34.946
Setup. We use Llama 3.2 IB as the target model to prune, in which the hidden state dimension is 2048. We the pruming ratio as 0.5 for MLP and 0.5 for Attention Layer, when pruning MLP and attention weights at the same time, we have pruning ratio of the whole model as 0.5. We use 8 sentences the C4 training dataset Raffel et (2020) as calibration data, and we use 128 sentences in the C4 validation dataset Raffel et al. (2020) to evaluate perplexity. We set 0.005 as the learning rate 77 and 0.05 as the regularization coefficient X.
Results. As can see in Table I our Attention pruning methods always outperform Wanda Attention pruning and SparseGPT Attention pruning by a large margin when combined with Cense/Wanda/SparseGGi
MLP pruning methods. These empirical results support theoretical analysis that our pruming method can converge. Furthermore, our method is broadly applicable in the real-world case and can be combined with many other pruning methods.
CONCLUSION
This paper introduces a novel approach to LLM weight pruning that directly optimizes for approxi- mating the attention matrix. We provide theoretical guarantees for the convergence of our Gradient buscent-based algorithm to near-optimal pruning mask solution. Experimental results demon- strated the method's effectiveness in maintaining model performance while reducing computational costs. This work establishes new theoretical foundation for pruning algorithm design in LLMs, potentially enabling more efficient inference resource-constrained devices.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
Research is partially supported by the National Science Foundation (NSF) Grants 2023239-DMS, CCF-2046710, and Air Force Grant FA9550-18-1-0166.
REFERENCES
Amol Aggarwal and Josh Alman. Optimal-degres polynomial approximations for exponentials and gaussian kernel density estimation. In Proceedings of the 37th Computational Complexity Con- ference, pp. 1-23, 2022.
Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances neural information processing systems, 32, 2019a.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, 2019b.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. In NeurIPS, 2019c.
Josh Alman and Zhao Song. Fast attention requires bounded entries. Advances in Neural Information Processing Systems, 36, 2023.
Josh Alman and Zhao Song. The fine-grained complexity of gradient computation for training large language models. arXiv preprint arXivv2402.04497. 2024a.
Josh Alman and Zhao Song. How to capture higher-order correlations? generalizing matrix soft- max attention to kronecker computation. In The Twelfth International Conference on Learning Representations, 2024b.
Josh Alman and Zhao Song. Fast rope attention: Combining the polynomial method and fast fourier transform. manuscript, 2024c.
Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, and Lingpeng Kong.
Training-free free long-contex. scaling of large language models. arXiv preprint a7X 2024.
Anthropic. The claude 3 model family: Opus, sonnet, haiku, 2024. https://www-cd: anthropic wwtp.///eeeabbleeeaeaeooe orc.aaaww0s.cccaaacccaeaooa0ee1. Card_Claude__0 .pd
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang. Fine-grained analysis of op- timization and generalization for overparameterized two-layer neural networks. In International Conference on Machine Learning. pp. 322-332. PMLR, 2019.
Saleh Ashkboos, Maximilian Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns. In The Twelfth International Conference on Learning Representations, 2024.
Rishi Bommasani, Drew A Hudson. Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportu- nities and risks of foundation models. arXiv preprint arXiv: 200.::2000.025 2021.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Bee Kamar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, Harsha Nori, Hamid Palangi. Mareo Tulio Ribeiro, and Yi Zhang. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint IrXiv:2303.12712 2023.
Yuan Cao and Quanquan Gu. Generalization error bounds of gradient descent for learning over- parameterized deep relu networks. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 34, pp. 3349-3356, 2020.
11

---

Published as a conference paper at ICLR 2025
Ya-Ting Ting Chang, Zhibo Hu, Xiaoyu Li, Shuiqiao Yang. Jiaojiao Jiang, and Nan Sun. Dihan: A novel dynamic hierarchical graph attention network for fake news detection. In Proceedings of the 33rd ACM International Conference on Information and Knowledge Management, pp. 197-206, 2024.
Moses Charikar, Michael Kapralov, Navid Nouri, and Paris Siminelakis. Kernel density estimation through density constrained near neighbor search. In 2020 IEEE 61 Annual Symposium Foundations of Computer Science (FOCS), pp. 172-183. IEEE, 2020.
Bo Chen, Xiaoyu Li, Yingyu Liang, Jiangxuan Long, Zhenmei Shi, and Zhao Song. Circuit com- plexity bounds for rope-based transformer architecture. arXiv e-prints, pp. arXiv-2411, 2024a.
Bo Chen, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Bypassing the exponential dependency: Looped transformers efficiently learn in-context by multi-step gradient descent. In International Conference on Artificial Intelligence and Statistics, 2025a
Bo Chen, Yingyu Liang, Zhizhou Sha, Zhenmei Shi, and Zhao Song. Hsr-enhanced sparse attention acceleration. Conference on Parsimony and Learning. PMLR, 2025b.
Tianyi Chen, Bo Ji, Tianyu Ding, Biyi Fang, Guanyi Wang, Zhihui Zhu, Luming Liang, Yixin Shi, Sheng Yi, and Xiao Tu. Only train once: A one-shot neural network training and pruning framework. Advances in Neural Information Processing Systems, 34: 34:19637-19651, 2021.
Weize Chen, Ziming You, Ran Li, Yitong Guan, Chen Qian, Chenyang Zhao, Cheng Yang, Ruobing Xie, Zhiyuan Liu, and Maosong Sun. Internet of agents: Weaving a web of heterogeneous agents for collaborative intelligence. arXiv preprint ar 2rXiv:2407.07011, 2024b.
Yifang Chen, Jiayan Huo, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Fast gradient computation for rope attention in almost linear time. arXiv preprint arXiv:2412.17316 2024c.
Yifang Chen, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. The computational limits of state-space models and mamba via the lens of circuit complexity. In Conference on Parsimony and Learning. PMLR, 2025c.
Yifang Chen, Xiaoyu Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. Universal approximation of visual autoregressive transformers. arXiv preprint orXiv:2502.06167, 2025d.
Yukang Chen, Shengju Qian, Haotian Tang, Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longlora: Efficient fine-tuning of long-context large language models. arXiv preprint 2tf.::2309113007 2023.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:220..02311, 2022.
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finet finetunndd language mod- els. arXiv preprint arXiv:2210.1144 14416, 2022.
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.0m 08699,, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2106/ 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
Mehmet F Demirel, Shengchao Liu, Siddhant Garg, Zhenmei Shi, and Yingyu Liang. Attentive walk-aggregating graph neural networks. Transactions on Machine Learning Research, 2022.
Yichuan Deng, Zhao Song, Yitan Wang, and Yuanyuan Yang. A nearly optimal size coreset algo- rithm with nearly linear time. arXiv preprint 0rXiv:22/0.083:0 2022.
12

---

Published as a conference paper at ICLR 2025
Yichuan Deng, Zhao Song, Shenghao Xie, and Chiwun Yang. Unmasking transformers: A theoret- ical approach to data recovery via attention weights. arXiv preprint arXiv:2310.1246.. 2023.
Yichuan Deng, Jiangxuan Long, Zhao Song, Zifan Wang, and Han Zhang. Streaming kernel pea algorithm with small space. In Conference on Parsimony and Learning. PMLR, 2025.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Papers), pp. 4171-4186. Association for Computational Linguistics, 2019.
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond 2 million tokens. arXiv preprint 4rXiv:2402.13753. 2024.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In ICLR. arXiv preprint arXiv: 1rXiv:1810.0200 1810.02044 2019.
Tao Feng, Chuanyang Jin, Jingyu Liu, Kunlun Zhu, Haoqin Tu, Zirui Cheng, Guanyu Lin, and Jiaxuan You. How far are we from agi. arXiv preprint arXiv:2405. 2024.
Elias Frantar and Dan Alistarh. Optimal brain compression: framework for accurate post-training quantization and pruning. Advances in Neural Information Processing Systems, 35:4475-4488, 2022.
Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot. In International Conference on Machine Learning, pp. 10323-10337. PMLR, 2023.
Spencer Frei and Quanquan Gu. Proxy convexity: A unified framework for the analysis of neural networks trained by gradient descent, Advances Neural Information Processing Systems, 34: 7937-7949, 2021.
Spencer Frei, Yuan Cao, and Quanquan Gu. Aggorithm-dependen generalization bounds for over- parameterized deep residual networks. Advances in neural information processing systems, 32, 2019.
Spencer Frei, Yuan Cao, and Quanquan Gu. Provable generalization of sgd-trained neural networks of any width in the presence of adversarial label noise. In International Conference on Machine Learning, pp. 3427-3438. PMLR, 2021.
Peng Gao, Jiaming Han, Renrui Zhang, Ziyi Lin, Shijie Geng, Aojun Zhou, Wei Zhang, Pan Lu, Conghui He, Xiangyu Yue, al. Llama-adapter v2: Parameter-efficient visual instruction model. arXiv preprint arXiv:2304. 15010, 2023a.
Tianyu Gao, Adam Fisch, and Danqi Chen. Making pre-trained language models better few-shot learners. Pruceedings of the 59th Annual Meeting of the Association for Computational Lin- guistics and the 11th International Joint Conference Natural Language Processing, 2021.
Yeqi Gao, Sridhar Mahadevan, and Zhao Song. An over-parameterized exponential regression. arXiv preprint arXiv:2303.16550 2023b.
Yeqi Gao, Zhao Song, Weixin Wang, and Junze Yin. A fast optimization view: Reformulating single layer attention in Ilm based tensor and svm trick, and solving it in matrix multiplication time. arXiv preprint wrXiv:2309.07411 2023c.
Yeqi Gao, Zhao Song, and Shenghao Xie. In-context learning for attention scheme: from single softmax regression to multiple softmax regression via a tensor trick. arXiv preprint 004.::2307.02419, 2023d.
Albert Gu and Tri Dao. Mamba: Linear- sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 2023.
13

---

Published as a conference paper at ICLR 2025
Torsten Hoefler, Dan Alistarh, Tal Ben-Nun, Nikoli Dryden, and Alexandra Peste. Sparsity in deep learning: Pruning and growth for efficient inference and training in neural networks. Journal of Machine Learning Research, 22(241):1-124, 2021.
Coleman Hooper, Schoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length Ilm inference with kv cache quantization. arXiv preprint arXiv:2401. 18079, 2024.
Cheng- Yu Hsich, Chun-Liang Li, Chih-kuan Ych, Hootan Nakhost, Yasuhisa Fujii, Alex Ratner, Ranjay Krishna, Chen- Yu Lee, and Tomas Pfister. Distilling step-by-step! outperforming larger language models with less training data and smaller model sizes. In Findings of the Association for Computational Linguistics: ACL. 2023, pp. 5003-8017,2000 2023.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. LoRA: Low-rank adaptation of large language models. In International Con- ference Learning Representations, 2022.
Jerry Yao-Chieh Hu, Donglin Yang, Dennis Wu, Chenwei Xu, Bo- Yu Chen, and Han Liu. On sparse modern hopfield model. In Thirty-sevnnth Conference on Neural Information Processing Systems (NeurIPS), 2023.
Jerry Yao-Chieh Hu, Pei-Hsuan Chang, Haozheng Luo, Hong- Yu Chen, Weijian Li, Wei-Po Wang, and Han Liu. Outlier-efficient hopfield layers for large transformer-based models. In Forty-first International Conference on Machine Learning (ICML), 2024a.
Jerry Yao-Chieh Hu, Bo-Yu Yu Chen, Dennis Wu, Feng Ruan, and Han Liu. Nonparametric modern hopfield models. arXiv preprint arXiv:2404.03990, 2024b.
Jerry Yao-Chich Hu, Thomas Lin, Zhao Song, and Han Liu. On computational limits of modern hopfield models: A fine-griind complexity analysis. In Forty-first International Conference Machine Learning (ICML), 2024c.
Jerry Yao-Chieh Hu, Maojiang Su, En-Jui Kuo, Zhao Song, and Han Liu. Computational limits of low-rank adaptation (lora) for transformer based models. arXiv preprint arXiv:2406.03136, 2024d.
Jerry Yao-Chieh Hu, Wei-Po Wang, Ammar Gilani, Chenyang Li, Zhao Song, and Han Liu. Funda- mental limits of prompt tuning transformers: Universality. capacity and efficiency. arXiv preprint arXiv:2411.1652. 2024e.
Jerry Yao-Chieh Hu, Dennis Wu, and Han Liu. Provably optimal memory capacity for modern hop- dense associative memories. In Advances
field models: Tight analysis for transformer-coppan in Neural Information Processing Systems (NeurIPS), volume 37, 2024f.
Jerry Yao-Chieh Hu, Weimin Wu, Yi-Chen Lee, Yu-Chao Huang. Minshuo Chen, and Han Liu. On statistical rates of conditional diffusion transformers: Approximation, estimation and minimax optimality. arXiv preprint arXiv:241 17522, 2024g.
Jerry Yao Chihh Hu, Weimin Wu, Zhuoru Li, Sophia Pi, Zhao Song, and Han Liu. On statistical rates and provably efficient criteria of latent diffusion transformers (dits). Advances in Neural Information Processing Systems, 38, 2024h.
Itay Hubara, Brian Chmiel, Moshe Island, Ron Banner, Joseph Naor, and Daniel Soudry. Accel- erated sparse neural training: A provable and efficient method to find transposable masks. Advances in neural information processing systems, 1199119991111,, 2021.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent achieve arbi- trarily small test emrrr with shallow relu networks. arXiv preprint arXiv: 12292, 2019.
Yuxin Jiang, Chunkit Chan, Mingyang Chen, and Wei Wang. Lion: Adversarial distillation of proprietary large language models. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
14

---

Published as a conference paper at ICLR 2025
Hongye Jin, Xiaotian Han, Jingfeng Yang, Zhimeng Jiang, Zirui Liu, Chia- Yuan Chang. Huiyuan Chen, and Xia Hu. Lim maybe longlm: Self-extend Ilm context window without tuning. arXiv preprint ar arXiv:2401.01322 2024.
Tian Jin, Michael Carbin, Dan Roy, Jonathan Frankle, and Gintare Karolina Dziugaite. Pruning's effect generalization through the lens of training and regularization. Advances in Neural In- formation Processing Systems, 35:37947-37961, 2022.
Praneeth Kacham, Vahab Mirrokni, and Peilin Zhong. Polysketchformer: Fast transformers via sketches for polynomial kernels. arXiv preprini arXiv:2310.0163 2023.
Hamed Karimi, Julie Nutini, and Mark Schmidt. Linear convergence of gradient and proximal- gradient methods under the polyak-tojasiewicz condition. In Machine Learning and Knowl- edge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings. Part I 16, pp. 795-8195-111 Springer, 2016.
Yekun Ke, Xiaoyu Li, Zhao Song, and Tianyi Zhou. Faster sampling algorithms for polytopes with small treewidth. In 2024 IEEE International Conference on Big Data (BigData). IEEE, 2024.
Yekun Ke, Xiaoyu Li, Yingyu Liang, Zhizhou Sha, Zhenmei Shi, and Zhao Song. On computational limits and provably efficient criteria of visual autoregressive models: A fine-grained complexity analysis. arXiv preprint arXiv:2501.044: 2025a.
Yekun Ke, Yingyu Liang, Zhenmei Shi, Zhao Song, and Chiwun Yang. Curse of attention: kernel- based perspective for why transformers fail to generalize on time series forecasting and beyond. In Conference on Parsimony and Learning. PMLR, 2025b.
Omar Khattab, Keshav Santhanam, Xiang Lisa Li, David Hall, Percy Liang, Christopher Potts, and Matei Zaharia. Demonstrate- weaee-searrhh--eeeeee preeccc::r: Composing retrieval and language models for snowledge-intensiig nip. arXiv preprint arXiv:2212. 14024, 2022.
Tzu-Sheng Kuo, Aaron Lee Halfaker, Zirui Cheng. Jiwoo Kim, Meng- Beig-Hsin Wu, Tongshuang Wu, Kenneth Holstein, and Haiyi Zhu. Wikibench: Community-drive data curation for ai evaluation on wikipedia. In Proceedings of the CHll Conference on Human Factors in Computing Systems, pp. 1-24, 2024.
Eldar Kurtic, Denis Kuznedelev, Elias Frantar, Michael Goin, and Dan Alistarh. Sparse finetuning for inference acceleration of large language models. arXiv preprint arXiv:2310.06997, 2023.
Brian Lester, Rami A1-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Pro- cessing. Association for Computational Linguistics, 2021.
Chenyang Li, Yingyu Liang, Zhenmei Shi, and Zhao Song. When rank approximation problem in truly subquadratic time? In International Conference on Artificial Intelligence and Statistics, 2025a.
solve the weighted low
Chenyang Li, Yingyu Liang, Zhenmei Shi, Zhao Song, and Tianyi Zhou. Fourier circuits in neural networks and transformers: A case study of modular arithmetic with multiple inputs. In Interna- tional Conference on Artificial Intelligence and Statistics, 2025b.
Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101. arX1v:2101.00190, 2021.
Xiaoyu Li. Yingyu Liang, Zhenmei Shi, and Zhao Song. A tighter complexity analysis of sparsegpt. arXiv preprint ar S1V::2488.2255555 2024a.
Xiaoyu Li. Yingyu Liang. Zhenmei Shi, Zhao Song, and Junwei Yu. Fast john ellipsoid computation with differential privacy optimization. arXiv preprint arXiv:2408.06 06399, 2024b.
Xiaoyu Li. Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Fine- grainer attention i/o complexity: Comprehensive analysis for backward passes. arXiv preprint arXiv:2410.09397, 2024c.
15

---

Published as a conference paper at ICLR 2025
Xiaoyu Li, Jiangxuan Long, Zhao Song, and Tianyi Zhou. Fast second-order method for neural network under small treewidth setting. In 2024 IEEE International Conference on Big Data (BigData). IEEE, 2024d.
Xiaoyu Li, Zhao Song, and Junwei Yu. Quantum speedups for approximating the john ellipsoid. arXiv preprinl arXiv:2408. /4018, 2024e.
Xiaoyu Li, Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, and Zhen Zhuang. Neural al- gorithmic reasoning for hypergraphs with looped transformers. arXiv preprint arXiv:2501. 1222221110088 10688, 2025c.
Xiaoyu Li, Yingyu Liang, Zhenmei Shi, Zhao Song, Wei Wang, and Jiahao Zhang. On the com- putational capability of graph neural networks: A circuit complexity bound perspective. arXiv preprint arXiv:2501.06/44 2025d.
Yuanzhi Li and Yingyu Liang. Learning overparameterized neural networks via stochastic gradient descent on structured data. Advances in neural information processing systems, 31, 201 8.
Yuanzhi Li, Yingyu Liang. and Andrej Risteski. Recovery guarantee of weighted low-rank ap- proximation via alternating minimization. In International Conference on Machine Learning, pp- 2358-2367, PMLR, 2016.
Yingyu Liang Heshan Liu, Zhenmei Shi, Zhao Song, and Junze Yin. Conv-basis: new paradigm for efficient attention inference and gradient computation in transformers. arXiv preprint arXiv:2405.05219, 2024a.
Yingyu Liang, Zhizhou Sha, Zhenmei Shi, and Zhao Song. Differential privacy mechanisms in neural tangent kernel regression. arXiv preprint arXiv:2407.1362 2024b.
Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, and Yufa Zhou. Multi-layer transformers gradient can be approximated in almost linear time, arXiv preprint vrXiv:2408.13233 2024c.
Yingyu Liang Zhenmei Shi, Zhao Song. and Chiwun Yang. Toward infinite-long prefix in trans- former. arXiv preprint arXiv:2406. 14036, 2024d.
Yingyu Liang Zhenmei Shi, Zhao Song, and Yufa Zhou. Unraveling the smoothness properties of diffusion models: A gaussian mixture perspective. arXiv preprint arXiv:2405.16418. 2024e.
Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Differential privacy of cross-attention with provable guarantee. arXiv preprint arXiv:2407.
2024f.
Yingyu Liang, Zhenmei Shi, Zhao Song, and Yufa Zhou. Tensor attention training: Provably effi- cient learning of higher-order transformers. arXiv preprint 0rXiv:2405.16411 2024g.
Yingyu Liang, Zhizhou Sha, Zhenmei Shi, Zhao Song, and Yufa Zhou. Looped relu mips may be all you need as practical programmable computers. In International Conference on Artificial Intelligence and Statistics, 2025.
Na Liu, Liangyu Chen, Xiaoyu Tian, Wei Zou, Kaijiang Chen, and Ming Cui. From Ilm to conversa- tional agent: A memory enhanced architecture with fine-tuning of large language models. arXiv preprint ar XXiv:2401.02777 2024.
Tianlin Liu and Friedemann Zenke. Finding trainable sparse networks through neural tangent trans- fer. In International Conference on Machine Learning, pp. 6336-6347. PMLR, 2020.
AI @ Meta Llama Team. The llama 3 herd of models. arXiv preprint arXiv:2407.21788. 2024.
Stanislaw Lojasiewicz. A topological property of real analytic subsets. Coll. du CNRS, Les équations aux dérivées partielles, 117(87-89):2, 1963.
Xuezhe Ma, Xiaomeng Yang. Wenhan Xiong, Beidi Chen, Lili Yu, Hao Zhang, Jonathan May. Luke Zettlemoyer, Omer Levy, and Chunting Zhou. Megalodon: Efficient 1lm pretraining and inference with unlimited context length. arXiv preprint 2444:22004.8880,,, 2024.
16

---

Published as a conference paper at ICLR 2025
Amama Mahmood, Junxiang Wang, Bingsheng Yao, Dakuo Wang, and Chien-Ming Huang. Llm- powered conversational voice assistants: Interaction patterns, opportunities, challenges, and de- sign guidelines. arXiv preprint 2..1::2309.18879 2023.
Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. arXiv preprint arXiv:240S. 06640, 2024.
Meta.
and vision with meta. com/blog/ 1lama 3-2 2-conecct 2024 -iii-vvisinn -dgee-dgeeee-- mobile-devices/ 2024. Accessed:
Llama 3.2:
Revolutionizing
edge ai
open,
customizable models.
https:
2024-11-21.
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in neural information processing systems, 32, 2019.
Swaroop Mishra, Daniel Khashabi, Chitta Baral, and Hannaneh Hajishirzi. Cross-task generalization via natural language crowdsourcing instructions. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics, 2022.
Maxwell Nye, Anders Johan Andreassen, Gur AriGuy, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint 1rXiv:21/2.00//4 14, 2021.
OpenAl. Hello gpt-4o. https://openaai.com////eee////// 1aptll1l-ppppp/p//1-111 40/ 2024a. Accessed:
May 14.
OpenAl. Introducing openai ol-preview.
htTp:://opeaa..comm2eees
introducing-opente oop--o-eniiiooo -ppeeiie---preveeww/ 2024b. Accessed: September 12.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 2022.
Gunho Park, Minsub Kim, Sungjae Lee, Jeonghoon Kim, Beomseok Kwon, Se Jung Kwon, Byeong- wook Kim, Youngjoo Lee, Dongsoo Lee, et al. Lut-gemm: Quantized matrix multiplication based on luts for efficient inference in large-scale generative language models. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprint 3xXiv:339000077 2023.
Boris T Polyak. Gradient methods for the minimisation of functionals. USSR Computational Math- ematics and Mathematical Physics, 3(4):644-878 1963.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1-67 2020.
Ilya Razenshteyn, Zhao Song, and David P Woodruff. Weighted low rank approximations with provable guarantees. In Proceedings of the forty-eighth annual ACM symposium on Theory of Computing, pp. 250-263, 2016.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2024.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. arXiv preprint arXiv:2407.08668, 2024.
Xuan Shen, Zhao Song, Yufa Zhou, Bo Chen, Yanyu Li, Yifan Gong, Kai Zhang, Hao Tan, Jason Kuen, Henghui Ding, Zhihao Shu, Wei Niu, Pu Zhao, Yanzhi Wang, and Jiuxiang Gu. Lazydit: Lazy learning for the acceleration diffusion transformers. arXiv preprint 2412:24111.1444, 2024a.
17

---

Published as a conference paper at ICLR 2025
Xuan Shen, Zhao Song, Yufa Zhou, Bo Chen, Jing Liu, Ruiyi Zhang, Ryan A. Rossi, Hao Tan, Tong Yu, Xiang Chen, Yufan Zhou, Tong Sun, Pu Zhao, Yanzhi Wang. and Jiuxiang Gu. Numerical pruning for efficient autoregressive models. arXiv preprint mrXiv22412.12441, 2024b.
Zhenmei Shi, Junyi Wei, and Yingyu Liang. A theoretical analysis on feature learning in neural networks: Emergence from inputs and advantage over fixed features. In International Conference on Learning Representations, 2021.
Zhenmei Shi, Jiefeng Chen, Kunyang Li, Jayaram Raghuram, Xi Wu, Yingyu Liang, and Somesh Jha. The trade-off between universality and label efficiency of representations from contrastive learning. In The Eleventh International Conference Learning Representations, 2023.
Zhenmei Shi, Yifei Ming, Xuan-Phi Nguyen, Yingyu Liang, and Shafiq Joty. Discovering the gems in early layers: Accelerating long-context Ilms with 1000x input token reduction. arXiv preprini arXiv:2409.17422, 2024a.
Zhenmei Shi, Junyi Wei, and Yingyu Liang. Provable guarantees for neural networks via gradient feature learning. Advances in Neural Information Processing Systems, 36, 2024b.
Kumar Shridhar, Alessandro Stolfo, and Mrinmaya Sachan. Distilling reasoning capabilities into smaller language models. In The 61st Annual Meeting Of The Association For Computational Linguistics, 2023.
Anshumali Shrivastava, Zhao Song, and Zhaozhuo Xu. A theoretical analysis of nearest neighbor search on approximate near neighbor graph. arXiv preprint arXiv:2303.0m 06210, 2023.
Ritwik Sinha, Zhao Song, and Tianyi Zhou. A mathematical abstraction for balancing the trade-off between creativity and reality in large language models. arXiv preprint Stt::2306.022999 2023.
Zhao Song and Chiwun Yang. An automatic learning rate schedule algorithm for achieving faster convergence and steeper descent. arXiv preprint arXiv:2310. 1291, 2023.
Zhao Song, Lichen Zhang, and Ruizhe Zhang. Training multi-layer over-parametrized neural net- work in subquadratic time. In Innovations Theoretical Computer Science (ITCS), pp. 93:1- 93:15, 2024.
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. Advances in Neural Information Processing Systems, 31, 2018.
Mingjie Sun, Zhuang Liu, Anna Bair, and Zico Kolter. simple and effective pruning approach for large language models. In The Twelfth International Conference on Learning Representations, 2024.
Tianxiang Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tuning for language-modellas--e
In International Conference on Machine Learning. PMLR, 2022.
Xiaoyu Tan, Shaojie Shi, Xihe Qiu, Chao Qu, Zhenting Qi, Yinghui Xu, and Yuan Qi. Self-criticimm: Aligning large language models with their understanding of helpfulness, honesty, and harmless- ness. In Proceedings of the 2023 conference empirical methods in natural language process- ing: industry truck, pp. 650-662, 2023.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 13771 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordy- intsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning. PMLR, 2023.
18

---

Published as a conference paper at ICLR 2025
Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannanch Hajishirzi. Self-instruct: Aligning language models with self-generated instructions. In The 6l Annual Meeting Of The Association For Computational Linguistics, 2023.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, and prompting elicits reasoning in large language models. Advances
Denny Zhou. in neural information processing systems, 35:24824-24837, 2022.
Jerry Wei, Le Hou, Andrew Kyle Lampinen, Xiangning Chen, Da Huang, Yi Tay, Xinyun Chen, Yifeng Lu, Denny Zhou, Tengyu Ma, and Quoc Le. Symbol tuning improves in-context learn- ing in language models. In The 2023 Conference on Empirical Methods in Natural Language Processing, 2023.
Dennis Wu, Jerry Yao-Chieh Hu, Teng- Yun Hsiao, and Han Liu. Uniform memory retrieval with larger capacity for modern hopfield models. In Forty-first International Conference on Machine Learning (ICML), 2024a.
Dennis Wu, Jerry Yao-Chieh Hu, Weijian Li, Bo- Bo-Yu Chen, and Han Liu. STanhop: Sparse tan- dem hopfield model for memory-enhanced time series prediction. In The Twelfth International Conference on Learning Representations (ICLR), 2024b.
Weimin Wu, Maojiang Su, Jerry Yao-Chieh Hu, Zhao Song, and Han Liu. Transformers are deep op- timizers: Provable in-context learning for deep model training. arXiv preprint 11 2024c.
Zhiheng Xi, Wenxiang Chen, Xin Guo, Wei Hc, Yiwen Ding, Boyang Hong, Ming Zhang, Junzhe Wang, Senjie Jin, Enyu Zhou, et The rise and potential of large language model based agents: A survey. arXiv preprint arXiv:2309.07880
2023.
Haojun Xia, Zhen Zheng, Yuchao Li, Donglin Zhuang, Zhongzhu Zhou, Xiafci Qiu, Yong Li, Wei Lin, and Shuaiwen Leon Song. Flash-llm: Enabling cost-effective and highly-efficient large gen- erative model inference with unstructured sparsity. Proceedings of the VLDB Endowment, 17(2): 211 211,224, 2023.
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning, In The Twelfth International Conference on Learning Representations, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference Machine Learning, pp. 38087-38099, PMLR, 2023.
Sicong Xie, Qunwei Li, Weidi Xu, Kaiming Shen, Shaohu Chen, and Wenliang Zhong. Denoising time cycle modeling for recommendation. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development Information Retrieval, pp. 1950-1955, 2022.
Chenwei Xu, Yu-Chao Huang, Jerry Yao-Chieh Hu, Weijian Li, Ammar Gilani, Hsi-Sheng Goan, and Han Liu. Bishop: Bi-directional cellular learning for tabular data with generalized sparse modern hopfield model. In Forty-first International Conference on Machine Learning (ICML), 2024a.
Jun Xu, Weidi Xu, Mengshu Sun, Taifeng Wang, and Wei Chu. Extracting trigger-sharing events via an event matrix. In Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 1189-1201, 2022.
Weidi Xu, Jingwei Wang, Lele Xie, Jianshan He, Hongting Zhou, Taifeng Wang. Xiaopei Wan, Jingdong Chen, Chao Qu, and Wei Chu. Logiemp: A neuro-symbolic approach for encoding first- order logic constraints. In The Twelfth International Conference on Learning Representations, 2024b.
Zhuoyan Xu, Zhenmei Shi, Junyi Wei, Yin Li, and Yingyu Liang. Improving foundation models for few-shot learning via multitask finetuning. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.
19

---

Published as a conference paper at ICLR 2025
Zhuoyan Xu, Zhenmei Shi, and Yingyu Liang. Do large language models have compositional abil- ity? an investigation into limitations and scalability. In First Conference on Language Modeling, 2024c.
Zhuoyan Xu, Zhenmei Shi, Junyi Wei, Fangzhou Mu, Yin Li, and Yingyu Liang. Towards few- shot adaptation of foundation models via multitask finetuning, arXiv preprini arXiv:2402.1500 2024d.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao, and Karthik R Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kemne density estimation. In ICML. arXiv preprint arXiv:2302.02551, 2023.
Amir Zandich, Majid Daliri, and Insu Han. Qjl: 1-bit quantized transform for kv cache quantiza- tion with zero overhead. arXiv preprint arXiv:2406.0344 2024.
Yuchen Zeng and Kangwook Lee. The expressive power of low-rank adaptation. In The Twelfth International Conference on Learning Representations, 2024.
Jiahao Zhang. Graph unlearning with efficient partial retraining. In Companion Proceedings of the ACM on Web Conference 2024, pp.
2024.
Jiahao Zhang, Rui Xue, Wenqi Fan, Xin Xu, Qing Li, Jian Pei, and Xiaorui Liu. Linear-time graph neural networks for scalable recommendations. In Proceedings of the ACM on Web Conference 2024, pp. 3533-3544, 2024a.
Jieyu Zhang, Ranjay Krishna, Ahmed H Awadallah, and Chi Wang. Ecoassistant: Using Ilm assistant more affordably and accurately, arXiv preprint arXiv:2310.0304 2023a.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the por- cupine: Expressive linear attentions with softmax mimiery. arXiv preprint arXiv:240.00337,, 2024b.
Renrui Zhang, Jiaming Han, Aojun Zhou, Xiangfei Hu, Shilin Yan, Pan Lu, Hongsheng Li, Peng Gao, and Yu Qiao. Llama-adapter: Efficient tine-tuning of language models with zero-init atten- tion. arXiv preprint arXiv:2303. 16/99, 2023b.
Yifan Zhang, Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, and Andrew Chi- Chih Yao. Tensor product attention is all you need. arXiv preprint arXiv:2501.06425, 2025.
Yuxin Zhang, Lirui Zhao, Mingbao Lin, Sun Yunyun, Yiwu Yao, Xingjia Han, Jared Tanner, Shiwei Liu, and Rongrong Ji. Dynamic sparse no training: Training-free fine-tuning for sparse 1lms. In The Twelfth International Conference on Learning Representations, 2024c.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Sameer Singh. Calibrate before use: Improving few-shot performance of language models. In International Conference on Machine Learning. PMLR, 2021.
Huaixiu Steven Zheng, Swaroop Mishra, Xinyun Chen, Heng-Tze Cheng, Ed Chi, Quoe V Le, and Denny Zhou. Step-back prompting enables reasoning via abstraction in large language models. In The Twelfth International Conference on Learning Representations, 2024.
Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, LILI YU, Susan Zhang, Gargi Ghosh, Mike Lewis. Luke Zettlemoyer and Omer Levy. LIMA: Less is more for alignment. In Thirty-seventh Conference on Neural Information Process- ing Systems, 2023.
Xunyu Zhu, Jian Li, Yong Liu, Can Ma, and Weiping Wang. A survey on model compression for large language models. arXiv preprint 2xXx::2308..7633, 2023.
20

---

Published as a conference paper at ICLR 2025
Appendix
CONTENTS
Introduction
Related Work
2.1
Pruning and Compression for LLMS
Attention Acceleration
3
Preliminary
Main Results
5
Technique Overview
5.1
Previous Tools on Convergence of GD
Closed Form of Gradient
5.3
Lipschitz of Gradient
5.4
PL Inequality of Gradient
6
Experiment
6.1
Evaluation on Synthetic Data
6.2
Experiment on Real Dataset and LLMS
6.3
Experiment on End-to-end Perplexity
10
7
Conclusion
10
More Related Work
23
Preliminary
23
Notations
23
B.2 Facts
24
Gradient Calculation
26
Definitions]
26
C.2
Gradient for Each Row of X(M W)X
................................
28
C,3
Gradient for Each Row of u(M)
29
C.4
Gradient for Each Entry of a(M)
29
C,5
Gradient for Each Entry of f(M)
30
C.6
Gradient for Each Entry of c(M)
31
Gradient for Latia M
32
C.8
Gradient for Leeg M
32
21

---

Published as a conference paper at ICLR 2025
C.9 Gradient for L(M)
33
Matrix Form
34
D.I Matrix Form of B(M)
34
D.2 Matrix Form of aM L(M)
37
Bounds for Basic F Functions
38
Basic Assumptions
38
Bounds for Basie Functions
38
E.3
Bounds for Gradient of ((M)
40
I Lipschitz of Gradient
41
Useful Facts
F.2
Lipschitz of f(M)
42
F3
Lipschitz of c(M)
42
Lipschitz of f(M) c(M)
42
F55
Lipschitz of diag((f(M) o c(M) 1.
43
F6
Lipschitz of cliag( o cl (M) 1(M)
44
F.7
Lipschitz of Gradient
Convergence of Gradient Descent
45
G.1
Helpful Statements
45
G.2
Lower Bound on Frobenius Norm
46
Sandwich Lower Bound on Frobenius Norm
47
4
Lower Bound on Hadamard Product Between Two Matrices
48
G.5
Final Bound
48
G.6
Pl. Inequality
49
22

---

Published as a conference paper at ICLR 2025
Roadmap. The appendix is organized as follows. In Section A we reviewed more literature lated to our paper. In Section [B we give the preliminary of our paper. In Section C we provide a detailed gradient analysis of the loss function. In Section D we provide details about how we inte- grate the gradient of loss function into matrix form. In Section
H we bound some basic functions to we provide proof for the Lipschitz property of the gradient of the loss we provide proof of convergence for GD.
be used later. In Section F function. In Section G
A MORE RELATED WORK
Large Language Models. Transformer-base neural networks (Vaswani et al. (017)) have rapidly emerged as the dominant architecture for natural language processing in machine learning. When expanded to billions of parameters and trained on vast, diverse datasets, these systems are typically
termed large language models (LLMs) or foundation models (Bommasani et al.] 2021). Promi-
nent LLM examples encompass BERT (Devlin et al., 2019). PaLM (Chowdhery et al. 2022)) ,
Llama Touvron et al. 2023
Meta 2024), and GPT4o OppeAAl, 2024a), which display adaptable
competencies (Bubeck et al.
2023 across numerous downstream applications. To enhance LLMs
for domain-specificc uses, researchers have created multiple adaptation approaches. These include:
adapter modules (Hu et al. [022]] Gao et al.
2023a Zhang et al. 2023b Shi et al.) [2023]; calibration
mechanisms (Zhou et al. [203] Zhao et al.
2021): multitask refinement (Gao et 20.] Von Os-
wald et al., 2023] Xu et al. [202d]]]]] 2023); along with prompt engineering (Lester al.) 2021
scratchpad approaches (Nye et al. 2021). instruction optimization Chung et al. [022] Li & Liang
2021 Mishra et al. (2022), ; symbolic adaptation (Wei et al,, 2023; Xu et al. 2024b) 2022).
black-
box adjustments (Sun et al., 2022), human-aligned reinforcement learning (Ouyang et al. 2022).
and structured reasoning techniques (Khattab et al.) 2022 Yao et al. 2023 Wei et al. 2022 Zheng
ct al., 2024). Contemporary investigations cover tensor architecture innovations Alman & Song
{2024b); Liang et al. (2024g); Sanford et al. (2024); Zhang et al. (202)), efficiency enhancements
Alman & Song (202c);; Chen et al. (2024c (2055)) Hu et al. (2024a) b h 2023); Ke et al. (2024);
Liang et al. (2024a))) Li et al. (2024b d daa); Liang et al. (2024d); Shi et al. (2024a); [Shen et al.
(2024a b): Song et al. (2024); y Wu al. (2024a b); Xu al. (2024a). plus ancillary studies Chen
et al. (2024a) Demirel et al. (2022): Chang et al. (2024): Deng et a (2022) Gao et al. (2023b db Li
al.
(2025b
Liang et al. (2024f) Li al. (2024e): Shrivastava al. (2023); Sinha et al. (202 }
Song & Yang (2023); Tan et al. (2023); Xie et al. (202), Xu et al. (202cc) Zhang (2024); Zhang
al.
(0244)) ; Chen et (2025d):
: Li al. (2025c), Ke al. (225b): Liang et al. (2025) Chen
et al.
2025a c): Gao et al.
2023c Hu et al. 22024g Wu al. (002cc).
B PRELIMINARY
In Section B.1 introduce some notations use in this paper. In Section B.2 we provide some basic facts.
B. NOTATIONS
For any positive integer n, we use [n]] to denote set (.. nn.. y E R", we use (r, to denote the inner product between ær y. i.e. (x. Er Tilh. For each a, b E R", we use a IR" to denote the Hadamard product, i.c. the i-th entry of (a is a,b, for all i e [n]. We use e, to denote a vector where only i-th coordinate is 1, and other entries are 0. We use 1, to denote a length-n vector where all the entries are ones. We use || to denote the l, norm of a vector & E IR", i.e. || == Em_ |az], ||œ]|2 == (E 1 x2) 22/1/2 and || == max4g(n] [ms]- For A E Rax= let A, E Ra denote the i-th row and A E 1 denote the j-th column of A, where i E [m] and j E [n). For square matrix A, we use tr[A] to denote the trace of A, i.c., tr[4] Σn Az.. For two matrices X, Y E Rm the standard inner product between matrices is defined by (X, XX,Y)=
For two vectors x E R and
tr|x YY. We use exp(A) to denote a matrix where exp(A)s,, == exp( As,)) for a matrix A E R For k > n, for any matrix A E R*x" we use |||]]] to denote the spectral norm of A, i.e. ||||| := supzer= || || 12- We use ||A|| to denote the fox norm of matrix A Baxd ie.. ||A|| := max,e |Ar We use ||A]] to denote the Frobenius norm of
a matrix E gaxd i.e. || = / Vie(M) Eje(g Ar, 12. For a symmetric matrix A E Raxx,
23

---

Published as a conference paper at ICLR 2025
A > 0 (positive semidefinite (PSD)), if for all E R", have x Axr 2 O. We Amin ( (A) and Amax (A) to denote the minimum and the maximum eigenvalue of the square matrix A, respectively. Let A Rxxd We use vec(A) denote length nd vector. We stack rows of into a column vector, i.e. (A) == [a] a.....,a. ".] where a, is the i-th row of A, or simply vec( (A) (( Ai for any [m].] [d]
B.2 FACTS
Fact B.1 (Indexing). Suppose we have matrices U E Raxm V E Rmxd
We define
X
axd
RXm mxd
Then, we have the following:
Indexing for one row: X, = V E Rd, iee X. = =U, V. Voror [n].
Indexing for one column: X.,, = UV E R" for E [d].
Fact B.2. We have
Part 1.Suppose we have vectors E R", D E R For i E [n],, we define
It = U/
Then we have the following:
diag(u) nx1 axn
uou
diag(r) vt
ax1
xXn
Part 2.Suppose we have matrix W E Raxn vector E R". For i E [n], define
X.j= = W.juj-
Then have the following:
X = W diag(u)
Fact B.3 (Calculus). We have
Part 1. (Scalar calculus) For any E R. function f: IR + R, we have
4 nf (t) dfit)
de
Part 2. (Vector calculus) For any Z. E E t R. have
d(woy)
y + dy de æ. (Product rule of vector Hadamard product)
de
d(w.y) dt
<dr y) + (x, de). (Product rule of inner product)
de. E;.
Part 3. (Matrix calculus) For any X,Y E puxm Z E Ra E R. which is independent of Z, function f:R + nxd ffunctions f(.)) IR gndd we have
d(XoY) dx Y + dy X. (Product rule of matrix Hadamard product) dt
dexp(/(0) exp(f(t)
d/it) de
where exp(-) applied entry-wise.
d(ZZ)
dx Z.
dt
dt
d(2xT) dt
7
24

---

Published as a conference paper at ICLR 2025
d Σ fi(t)
df.0) EE=
dx
ci
dX
axl
WER" X Fact B.4 (Basic algebra). Let u E CR,,,,,, Then, we have
Rand Y BRxd and Z E R
(u, (u, u) u == vJu
uon = r diag(u)u = diag(v)u
(u, = (u u, 1m)
(uo w) = (u w, v) = (w u. u)
u wwooww) = u diag(v)w
(XoY) =xT o OYT
Xoere Xijete]
diag(u) diag(v) (wx" ) ) Z
xY==
= -Eeed X
XiyYis (XoY)s
EjEMn uo A..j = uo EjE(n) Ans
|x = tr[XX
[=[XT]==r]p X]
|daa(()]ll= ||u||2
Fact B.5 (Norm bounds). For a E R. RRR X,Y Raxd Z E Rdxm have
|aX||r
| = |1]1x||| (absolute homogeneity).
|X+Y|| = |||||
F + ||Y]e (triangle inequality).
(X, S || || (Canchy-Sehwarz inequality).
|XT||= = |X||
[Xulla = || ||ull2
|X o 1| < || [[.
For any E [n). j E [d], we have [Xisl & |||||
|X]|| s X F S / x where k is the rank of X.
|| 7|FF < |||| ||
Fact B.6. For matrices A, B E mxx, we have
|| 1 = |||| + B + +(AA B).
Proof. We can show
| + BI = + B) (A + B)]
=tr[ATA A A B + B A BT
25

---

Published as a conference paper at ICLR 2025
tr(A A| tr(B' B] + 2tr[A' B)
-|A4|] + || + B)
= -| + ||]| + B)
where the first step follows from tr[A4' A] = || for matrix A E gmxx, the second step follows from the basic algebra, the third follows from (r[X"Y] = tr[XYT] for matrices X, E RR**, the fourth step follows from tr(AA 'A] ||A|| for matrix A mx and the last step follows from definition of inner product matrices.
Lemma B.7. Let M E R Let X E R=x" be independent of M. We have
d(MoX) Xizeze]
dMaj
Proof. We can show
d(M X)
M dx
Xo dM dM
dM1
dM3
dM
X aM,j
o (e,e,)
= ]
where the first step, the second step, and the third step follow from Fact B.3 the fourth step follows from Fact B.4
C GRADIENT CALCULATION
C.I DEFINITIONS
In this section, we introduce some definitions we used to compute dC(M) exponential function.
First, we introduce the
d.V.
Definition C.1 (Exponential function u, ū). If the following conditions hold
. Let X E Rnad
Let W Raxd
Met [0, 1]dxd
Let [m].
We define E Ruxx as follows
u u: exxxxxXXXXX"
We define ü(M) E Rnxn as follows
= expXXXM T)
We define ih row of ūl (M) as follows
ū(M) (M
Then, we introduce the sum function.
Definition C.2 (Sum function of softmax a, ã). If the following conditions hold
Let M [0, 1jxd
26

---

Published as a conference paper at ICLR 2025
Let Mc [0, be the causal attention mask defined in Definition 3.1
Let u, ū(M) be defined Definition C.1
Le [n].
define a as follows
= (u o M.) 1n.
We define ((M) IR" as follows
((M) (ū(M) Mo) 1n.
We define in-th entry of a(M) follows
i(M) ((ā(M) Mc)ia 1n)
Then, we introduce the Softmax probability function.
Definition C.3 (Softmax probability function f. f). If the following conditions hold
Let M [0, 1]dxd.
Let Ma {0, 1jmx be the causal attention mask defined in Definition 3.1
Let u, ū(M) defined as Definition C.1
Let ä(M) be defined as Definition C.2
Let io E [n].
We define f E Bax" for each j E [n] as follows
== diag(a) Mc).
We define F(M) guxn for each j E [n|] as follows
J(M) == diag ding(((M) M.). We define ig-th row of F(M) as follows
J(M) == a(M) (ū(M) o Mc).
define the entry in in-th row, jo-th column of f(M) as follows
J(M) a(M), (@(M) Mc)so.jo
Then, we introduce the one-unit loss function.
Definition C.4 (One unit loss function c). If the following conditions hold
Let f, f be defined in Definition C.3
Let M [0, 1jddd
Let io, E [n]]
We define c(M) E R=x* as follows
c(M) = J(M)-
define ig row of e(M) as follows
c(M) = F(M) fio
We define jo-th column of c(M) as follows
c(M)..jo = {MM)).jo - fn.j0 define the entry in in-th row, jo-th column of ((M) as follows
c(M) M J(M) ip:5o fro
27

---

Published as a conference paper at ICLR 2025
Then, introduce the reconstruction error.
Definition C.5 (Reconstruction Error Lusra) If the following conditions hold
Let M [0, 1].d
Let c(M) be defined in Definition C.4
We define (M) E IR as follows
Caun = (l
1
Σ
Σ
2 i0=
Then, introduce the regularization term.
Definition C.6 (Regularization Term Lreg). h If the following conditions hold
M [0,10x 1jaxd.
We define Crog (M) I as follows
1 Creg == > 1 l
Finally, we introduce the overall loss function.
Definition C.7 (Overall loss function L). If the following conditions hold
Let [0, 1]dxd.
Let Cattn(M) be defined in Definition C.5
Let Creg(M) be defined in Definition C.6
Let X E T., be the regularization parameter.
We define L(M) as follows
L(M) = Lattn(M) + (M)
C.2 GRADIENT FOR EACH Row OF X(M o W)XT
We introduce the Lemma of gradient for each row of X(M W)XT. Lemma C.8. Let i1 [d]. [d], io [ni]. we have
d(x (M W)X )m
Whij Xinoh Xnsi scalar scalar nxl
dMM,
Proof. We can simplify the derivative expression
d(X(M o W)X )io (X o W)),, dM
dM
(X o W)T Xio
dMiis
ddMoW) X
x
dMr
where the first and second step follows from Fact B.1 the third step follows from Fact B.3 We further compute Eq.
d(M W)T
dM WT
dM
dM,
28

---

Published as a conference paper at ICLR 2025
dM o WW d(M Dinc
(3)
where the first follows from Fact B.4 the second step follows from for any matrix X, X (X Jas. the third step follows from Fact B.7 and the fourth step follows from for any matrix X, Xig )aa Finally, we have
a(X W
XW Che Xio
dMins
Whis (Xer, (e, Xio)
Wiy X. X Xio.i
where the first step follows from Eq. (2) and Eq. from basic algebra.
33, and the second step and the third step follow
We introduce the Lemma of the gradient for each row of ū(M).
C.3 GRADIENT FOR EACH Row or ū(M)
Lemma C.9. If the following conditions hold:
Let ü(M) I) be defined in Definition C.1
Ler e [d], fi [d], to E [n], we have
dū(M) Ka
((M)
o WWis Xindy (esi)
aM
ax1
scalar scular nx
Proof. We have
dū(MM) io dM
dexp(XX(M W)X )io
dM
nx1
nx1
d(X(M o W)X )io dM
exp X (M o WM X
nxd
dxd
dxl
nx1
d(X(M W)X )
T(M)
dM
nxl
nx1
a(M), o (Wino Xig.n Xa.ji
axl
scalar
scalar
nxl
where the first step follows from Definition C.I the second step follows from Fact B.3 the third step follows from Definition C.1
and the fourth step follows from Lemma C.8
C.4 GRADIENT FOR EACH ENTRY OF ((M)
We introduce the Lemma of gradient for each entry of ã ((M).
Lemma C.10. If the following conditions hold:
29

---

Published as a conference paper at ICLR 2025
Let ü(M) be defined in Definition
C.1
Let (MM) be defined in Definition C.2
Let e [d] ji E [d] to E [nd], we have
dā (T(M)
o (Mc)ss Wiohh X X1.1)
dM
scalar
Proof. We have
dā(M)
d((ū(M) o Mc)in 1n)
dM
dM
scalar
scalar
d(ī(M) Molia 1n)
dMi
dia(Mia o (Mc) 1m)
(M(M)io (Waish Xiom X..i) o (Melio: 1n)
(m(M)., (Ma)io: Wh Xin., X.ji)
where the first step follows from Definition C.2 the second step follows from the product rule of inner product in Fact B.3] the third step follows from the product rule of Hadamard product in Fact B.3 the fourth step follows from Lemma
C.9
and the last step follows from Fact B.4
C.5 GRADIENT FOR EACH ENTRY OF J(M)
We introduce the Lemma of the gradient for each entry of J(M).
Lemma C.11. If the the following conditions hold:
Let ü(M) be defined in Definition C.1
Let ((M) be defined in Definition C.2
Let F(M) be defined in Definition C.3
*
Let E [d], E [d], to E [n]] jo E [n]. we have
df(M) J(M)in.to Wish Xioni Xjoch
W Xio Xa
dMish
Proof. We have
df(M)
dā(() (m(M) Ma)inco
dM
dM
dā(M)
d(T(M) Mc)in 5(M)
(ū(M) M Mc)
dM
dM
(4)
where the first step follows from Definition C3 and the second step follows from Fact B.3
In the following part, compute the two terms separately.
For the first term above, we have
dā (o) (ū(M) o Ma)
d.
30

---

Published as a conference paper at ICLR 2025
(T(M) Mclio.io (-1)@ dā(A))
(ū(M) o io (M(M) i. (Me))o: W Wh. in Xass))(AMM))
(ā(M) (o ((M) to) (M()) k(M),s (Mc)ios W Wi Nin Xt
i(M) (f(M) W Wi X..i)
(5)
where the first step follows from Fact B.3 the second step follows from Lemma C.10 the third step follows from basic algehra, and the fourth step follows from Definition
For the second term above, we have
d(ū(M)) Me)
dMinch
dï M Yo (Mo)hunda (M)
dMi
du((M) 10
iMooo
) (M)
aM
= ((Mo) ((M) i0 to 0 (M) W in:h Xi Xoos
f(M) Wi..j X Xin.i X Xjo
(6)
where the first step and the second step follow from basic algebra, the third step follows from Lemma C.9 and the fourth step follows from Definition
C.3
So, we have
af(M)ioJo
dā(M) (ū(M) Me)io.to + d(ū(M) Mo)in a(M)
dMing
d.M..hh
dM
f(M)
Winh ii X - J(M) (f(M) With Xin.i Xa.i))
where the first step follows from Eq. and the second step follows from Eq. 5 and Eq. (6).
C.6 GRADIENT FOR EACH ENTRY OF C(M)
We introduce the Lemma of gradient for each entry of c(M).
Lemma C.12. If the following conditions hold:
Let F(M), f be defined Definition C.3
Let c(M) be defined in Definition C.4
Let i E [d], E [d], io E [rd], jo E [n]. we have
de(M) io Winch Xin Xjo.i J(M (FM)
Whnh Xiosi X.
dM
Proof. We have
de( dc(M)
d(F(M)
io. Sto.jo)
dM
dMinn
df(MM)
dMin
= F(M) Win. Xio F(M))..5 (M))
Wii Xinsi
where the first step follows from Definition C.4 the second step follows from Fact B.3 and the third step follows from Lemma C.11
31

---

Published as a conference paper at ICLR 2025
C.7 GRADIENT FOR (M)
We introduce the Lemma of the gradient for Lattn (M).
Lemma C.13. If the the following conditions hold:
Let ((M be defined in Definition C.3
Let c(M) be defined in Definition C.4
Let Catt be defined in Definition C.5
Let 11 e [d E [d]. io [n], jo E [n], we have
dCaut
Bi(M) B2(M)
aM
t0
where we have definitions:
BB. (MM) Wins: Xio Xjo
B2(M) -c M) (M) io J( No (f(M)io Wiys Xiosi X
Proof. We have
dLatto(M)
ddllc(M)1
dM,
2 dM
1 1d
in= 1d
2
dMisdi
1
M(c(M)io.5o)
2
dMiinh
10
dc(M)
cMM
d.M.
where the first step follows from Definition C.5 the second step follows from the definition of Frobenius's norm of the matrix, the third step follows from Fact B.3 and the fourth step follows from Fact B.3
Following Lemma C.12 we have
de( M io.to
of M iM Jo
dM
Jo=1
c(M) M (F(M) W Winn Xty X F(M too (F(M) Wonch Xi Xe
An 1
((M)io
( {(M) jo Wi. Xio Xjo.s ((M) J (MM) (FM Wh X Xig X
(M) + B2(M)
1o=1
where the second step follows from basic algebra.
C.8 GRADIENT FOR Creg(M)
We introduce the Lemma of the gradient for freg (M).
Lemma C.14. If the following conditions hold:
32

---

Published as a conference paper at ICLR 2025
Let Lreg (M) be defined in Definition C.6
Let E [d], E [d], we have
dLreg(M)
B3(M)
dMi
where we have the definition:
Ba(M) AM
Proof. We have
dCreg(MM (M)
- dlI
dMiy
dMMi
d
> dM
Σ
Σ
M
Mo
1.50
= J
(3(M)
where the first step follows from Definition C.6 the second step follows from the definition of Frohenius's norm of the matrix, and the third step follows from Fact B.3
C.9 GRADIENT FOR L(M)
We introduce the Lemma of the gradient for L(M).
Lemma C.15. If the following conditions hold:
Let ū(M) be defined in Definition C.1
Let ä(M) be defined in Definition C.2
Let F(M) be defined in Definition C.3
Let (atn(M) be defined in Definition C.5
Let Lreg(M) be defined in Definition C.6
Let L(M) be defined in Definition C.7
Leti E [d], E [d], to E [ni], jo E [ni], we have
dC(M)
n
n
Σ
Σ
(M + B2(M)) + Ba(M)
d.M..
to 15
where we have definitions:
B.(M) c(M) (MM Wn.ji Xio.i X jo.j1
B2(M) ccMM F(M (f() Waig Xio.s Xa.i
Bs(M) AM
Proof.
dC(M)
dLatn(M) + Lreg((M)
d.M...h
dMi..
(Ma
(M)
-M,
d.
33

---

Published as a conference paper at ICLR 2025
(M) B2(M)) B3(M)
where the first step follows from Definition C.7
the second step follows from Fact B.3 and the third
step follows from Lemma C.13 and Lemma
D MATRIX FORM
D.1 MATRIX FORM OF ((M)
Given the matrix form, we define p to simplify the caleulation.
Definition D.1. the following conditions hold
Let X Rexd
Let M [0, 1]xdd
Let W Raxd
Let c() be defined Definition C.4
Let F(M) be defined in Definition C
We define pı as follows
= c(M) F(M)
We define the jo-th column of pı as follows
(P1)a.jo = (c(M) f(M))
define P2 as follows
p2: :2 diag(P -1))f(M)
We define the ig-th row of P2 follows
(p2)s, == (p1) i(M) = F(MM), F(M)
We define p follows
P P2 P = [MM diag((c(M)) F(M)) 1.)5(M)
We introduce the matrix view of B1 (M) and its summation.
Lemma D.2 (Matrix view of Bj (M)). If we have the below conditions,
Let (M,,in, 11,11 = (MM Wins Xiodi Xjodv which defined in Lemma C.15
We define (i(M) Iddxd For all i1 i e [d], let C1(11,1) denote the (i1, (1) -th entry of (MMM). We define
Then, show that
Part 1. For in, jo E [n]
(i(M = c(M) F(M)
to. (W (X)
Part 2.
Σ
C(MM W (XT PiX)
34

---

dxd Published as a conference paper at ICLR 2025
Proof. Part 1. We have
Ci(i1,51) c(M) io
io Wisch Xxoiis Xjovi
e(MMio.t to do (M to
X )as (W (Wo,tis dis (X
scalar
scalar
= c(MM (M) dodio
(diag( X Waii dis
(X (X50)
scalar
dx1
scalar
where the first step follows from the definition of C1, the second step follows from Fact B.1 and the third step follows from Fact B.2
Following from Fact B.1 can get ýi -th column of Ci
Ci (+,,)))= = c(M) (MMoooooo Joo..o J(M diag( (Xio)
We.h (Xjo)i scalar
scalar
dxd
dxl
=c(MM f (MM No (Xag(Xio)) W diag (Xx)))
scalar
dxd
axd
dxd
where the second step follows from Fact B.2
Following from Fact B.1 we can get (M
C1(M) = F(M) diag(Xio)
W diag(X;,)
scaler
dxd
dxd
ccM) io (M to (W (x))
scolar
where the second step follows from Fact B.4
Part 2. We further compute the summation of (MM).
Ci(M)
c (M) Mi (MM Joo (W Xix
10 Jo
io=
ja
scalar
W JMM) X x(o))
W
cM No. F(M)
7 a
W Σ Σ do 10=1
where the first step follows from Eq. (7). the second step follows from basic algebra, the third step follows from Fact B.4 and the fourth step follows from Definition D.I
Then following from Fact B.2 we have
7 W.
7
Σ
Σ
x==1
W
(pi)aso
J==1
W (XX)
35

---

Published as a conference paper at ICLR 2025
We introduce the matrix view of (2(M) and its summation.
Lemma D.3 (Matrix view of (2)M)). If we have the below conditions,
Let (1.1) = -clMM)o J(M) (f(M) Wiis Xiosi Xa be defined in Lemma C.I5
We define (MM) E Rddxd For all 81 e [d], ler C2(11, denote the (i1.31) entry of C2(M). We define C2( (2,M,,,,))
Then, we can show that
Part J. For in, jo [ni]
C2(M) = - io. (M No.so diag( Xio) W
diag(X (X f(M)
axd
dxd
axd
Part 2.
C2(M) = -W (XT p2X)
10=1 1
Proof. Part 1. We have
-C2(i1.h1) c(M) Jo ((M) (f(M) to: W Xx.si)
= ccMM).... (MM J(M)) Xaji Wins Xioni
1XR
-c(M)in F (MM) Jo.jo
F(M) Xx,5, (Xio)i (W (Wajiii
1XR
ux1
= =c(MM No.jo F(M) Xx,j, (diag( (1 W Wi )i
1xn
ax1
where the first step follows from the definition of C2. the second step, the third step, and the fourth step follow from Fact B.4
Following from Fact B.I
we can get ji-th column C2
(2) diag( Xi
W..,, e(M) ii J(M) X
dxd
dx1
1xn
= coMM Jo iMM) diag(X W..j F(M) To X..i dxd
dxl
1xn
nxl
= =( Jo F(M) diag(Xi,) W (X F(M))an)
sealar diag( X ) (W diag(X X [(M))
dxd
= (MM
dxd
dxd
where the second step and the fourth step follows from Fact B.4 and the third step follows from F.
Following from Fact B.I we can get Cz.
No (MM c(M) F(M Jo diag (Xio) W diag(X diag(XF((M)m.))
(8)
dxd dxd
dxd
Part We further compute the summation C2
C2(M)
Σ Σ clM F(M Mo.io diag(Xis) W diag(X F(M)
=1 i0=1 100
dxd.
36

---

Published as a conference paper at ICLR 2025
(M) (MM) io ((Xi,
(M) ((M)
[[
W)
W
(X (M) x) Σ ((p2)
jo=1
where the first step follows from Eq.
8 the second step and the third step follow from Fact B.4
Following from Fact B.2. we have
Wo Σ (X2 ((M) T X) Σ ((r))oo50
1==1
Jo=1
Wo Σ (X (X (MM)[(()))) X X)1 (pi)
10=1
= W (X" diag(pi 1x) F(M)x)
W W (X) P2X)
where the third step follows from Definition D.I
We introduce the matrix view of (y
Lemma D.4 (Matrix view of B2(M)). If the following conditions hold
Let Pa(M,én,Ji = AM be defined in Lemma C.15
We define (M) E RRdxd, For all i1, J1 E [d], Cz(11,h) denote the (11: entry of Ca(M). We define Ca(i1.51)
We can show that
Cs(M) AM.
Proof. The proof is straightforward. By the definition (M)) for all 4,5 E [d], the (1 )- entry of (M( is given by Ci = B3 (M, 51) AM Thus, the entire matrix (2(M) has entries that correspond to those of AM. Therefore, we can conclude that Ca(M) AM as required.
D.2 MATRIX FORM OF d aM L(M)
We introduce the matrix form of the overall loss function.
Theorem D.5 (Close form of gradient, formal version of Theorem 5.3 If the following conditions hold
Let L(M) be defined in Definition C.7
Let be defined in Definition D.1
Let X RRaxd
Let M [0, 1]dx
Let W e Raxd
We can show that
dC(M)
WW (XT pX) + AM.
dM
37

---

Published as a conference paper at ICLR 2025
Proof. We have
dC(M)
Σ E + + (MM
dM
1o=1 jo=1
= -Wo (X p1 W o (x" p2X) + AM
= W (X (pi X2) + AM
W (XT pX) + A.
where the first step follows from Lemma C.15 the second step follows from Lemma D.2
Lemma [33] and Lemma D.4 the third step follows from basic algebra, and the fourth step fol- lows from Definition D.I
E BOUNDS FOR BASIC FUNCTIONS
E. BASIC ASSUMPTIONS
Here, introduce our bounded parameters assumption.
Assumption E.1 (Bounded parameters), We assume the following conditions
. Let R be some fixed constant satisfies R. > 1.
Let X E Endd W Rdxd We have | p S and |W]| S R.
Here, we present the lemma of bounds for M and Mc-
Lemma E.2 (Bounds for M and Mc). Let M E 11,1d and Me {0,1)axn be the causal attention mask defined Definition3.1 For M, we have
| <d
For Mo. we have
|Mell S
Proof. This Lemma simply follows from the definition of the Frobenius norm, given that the max value of each entry in M and Me is 1.
E.2 BOUNDS FOR BASIC FUNCTIONS
We first introduce the lemma of bounds for basic functions.
Lemma E.3. Under Assumption E. for all [ni. [m], in [d], jn [d], have the following bounds
Part
[F(M))le $ vn
Part
[c(M)] 2
Part
[[(((M) o F(M)) 22/n
Part 4.
I(MM)oo5ol 1
38

---

Published as a conference paper at ICLR 2025
Part 5.
Wast S
Part 6.
[ S R
Parrt
i
Part
io Wins Xio Xjosil S R2
Part
If(M) (MM) Which Xinsi R2
Part 10.
o 1,) S 2n
Proof. Proof of Part Each entry in f(M) present probability, thus for [n], jo E [n]. have
S io 1
For any in-th row of J(M), following from the definition of Softmax function, we know
Σ io. 1. 10=1
So have
Σ F(M) Jo. 1
30
which follows from F(M) ioo5o I (F(M)innto ?. Then, we can show
IF(M))l==
F(M S Vn
1o=1j==1
Proof of Part 2. Following from Part 1., we can show
[[(M)) S vn
and
I[fle S Vn.
Then have
(e(M)]l==113) = [I((M) - 111
< IIf((M] + ||
< 2//7
where the first step follows from Definition C.4 the second step follows the triangle inequality.
Proof of Part 3. We have ( S 7(M)oo. I 1, so we have
[[((M) F( < [|((M)]
39

---

Published as a conference paper at ICLR 2025
$ 2nn
where the second step follows from Part 2..
Proof of Part 4. See Proof of Part 1..
Proof of Part 5. The proof simply follows from Assumption E. and Fact B.5
Proof of Part 6. The proof simply follows from Assumption E. and Fact B.5
Proof of Part See Proof of Part
Proof of Part 8. The proof simply follows from Part 4., Part 5., Part 6. and Part 7.,
Proof of Part The proof simply follows from Part 4., Part 5., Part 6. and Part 7..
Proof of Part 10. We have
d o f(M)) I.)le = [I(((M) F(M)) 1.112
|1, [alll f(
/n 2/n
2n
where the first step follows from Fact B.4 the second step follows from Fact B.5 the third step follows from Part 3., and the last step follows from simple algebra.
E.3 BOUNDS FOR GRADIENT OF f(M)
We introduce the lemma of bounds for the gradient of f(M).
Lemma E.4. the following conditions hold
Let F(M) be defined in Definition C.3
Assumption E. holds.
Then we have
d vec(F((M)
l|F S 2dn d vee(M)
Proof. We have
df(M
Doodo i
d.
We Wanch Xind Xjoss F(M) No (f() Windh Xinis Xasii!
(f(M) (M)io Wan. Xind Xiosil (F(M) (f(M)io Waysh Xinai Xa
2R
For d vecl F(M) verl ())
we can show
d dvvc(F(M))
d2
d vvc(F((M)) d veol M)
1|
d vec(MM)
k==1 1==1
$ 2ndR3
40

---

Published as a conference paper at ICLR 2025
F LIPSCHITZ OF GRADIENT
F.1
USEFUL FACTS
Here, we introduce the fact of the mean value theorem for matrix function.
Fact F.1 (Mean value theorem for matrix function, Fact C.6 in Liang et al. (2024e) If the following conditions hold
Let Let E C C 1ddxd where C an open convex domain.
Let g(x): : C TRxx be differentiable matrix function on C.
Let
|| S for all x C.
We have
ll((Y) - g(X)IlF RJjY - X||r.
Proof. For the convenience of proof, we define x and y as follows:
veo(X) and y vec(Y).
h(x) == vec(g(X)) and h(y) == vec(g(Y)).
N (a) denotes a matrix which its (i, j)th term is dh(a);
do;
Assume we have 1-variable function ((c) = f(x + ely x)), we can apply Mean Value Theorem:
f(y) - f(x) = >(1) -(0) (f) 0) vf(r + x)) (v -x) x)
(9)
where [0, 1]. Let G(c) (h(v) h(=))()))))))))(()))) M(C), have
llg(Y) - 9(x)11 G(y) G(x)
-G(x + t(y - 2)) (y - x))
(h'(z + t(y - xx) h(y) - h(z)) (y x)
Pxn²
n2xl
$ [1 - t(y - =))l (1(()) - h(x) 12 |ly - 212
where the second step follows from Eq. 9 the third step follows from the chain rule, the fourth step follows from the Schwar-Schwartz inequality.
By definition matrix Frobenius norm and vector l2 norm, we have
llg(Y) g(X)||= [1h()) (x))ll
and
||Y - X 1-111~a0 xxl2
so we can show
ll g(x))lx S RIIY - X||F
which follows from
ver[g( X))
d vX) |F < for all C.
Here, we introduce the fact of Lipschitz for the product of functions.
Fact E2 (Lipschitz for product of functions, Fact H.2 in Deng et al. 2023 ) Under following
conditions
Let (f(())) be sequence functions with the domain and range.
41

---

Published as a conference paper at ICLR 2025
For each i [n], we have
-(( is bounded: Yx, [[(()))le & R2 with R; 1.
(z) Lipschitz continuous: Vz, y, llf(()) f((v)lF L;| yll
Then have
f(()) II f(v)lle 2 max [L,} (II Ra) 1 vl1
se|n
F.2 LIPSCHITZ F(M)
We introduce the lemma about Lipschitz of F(M).
Lemma F. (Lipschitz of F(M)). Under the following conditions
Assumption E.1 holds.
Let F(M) be defined as Definition C.3
For M, M E Baxd, we have
[F(M) F(Mill S 2dn RM M/I
Proof. We have
(I(M) J(M)] < [V](M) ||M M||F
< 2dn
||M M
|F
where the first step follows from Fact R.I the second step follows from Lemma E.4
F.3 LIPSCUITZ OF c(M)
We introduce the lemma about Lipschitz of c(M). Lemma E44 (Lipschitz of c(M)). Under the following conditions
Assumption E.1 holds.
Let c(M) be defined Definition C.4
For M, M E Rdxd we have
lc(M) S 2dnR2 M|
Proof. We have
[c((M) - o(())] - \cc(M)] |M - M
||/ || M/
< 2dnRt ||M - ||M-
where the first step follows from Fact the second step follows from Lemma C.12 the third step follows from Lemma E.4
F.4 LIPSCHITZ oF F(M) c(M)
We introduce the lemma about Lipschitz of F(M) c(M).
Lemma F.5 (Lipschitz of f(M) c(M)). Under the following conditions
42

---

Published as a conference paper at ICLR 2025
Assumption E.1 holds.
Let c(M) be defined Definition C.4
Let (M) be defined as Definition C.3
.
For M, M. E Rdddd have
If(M) o F(M) co 6dn - Mll
Proof. We have
LHS < (IF(M) o c(M) - (MM o c(M) F(M) o c(M)ll
s [lf(M)))F [|((M] c(M) | [lf(M) ) F(M)Ile
$ vn. 2 s Vn 2dnnR" - M
(1]( F(M))
f(M]II
where the first step follows from triangle inequality, the second step follows from Fact F.5 the third step follows from Lemma
E.3 the fourth step follows from Lemma F.4 and Lemma F3
So we have
(F(M) e(M) J(M) c(M)le S 6dn 2 M
F.5 LIPSCHITZ OF diag(([(M) c 1..)
We introduce the lemma about Lipschitz of diag((F(M) o c(M)) 1,).
Lemma F.6 (Lipschitz of diag((((M) ((M) 1,)). If the following conditions hold
Assumption E. holds.
Let c(M) be defined as Definition C.4
Let F(M) be defined as Definition C.3
For M,M E Raxd, we have
ocM)) ||diag( [ 1x) - diag((f(M) o c(M)) 1.)lFF S 6dn 2dn2R [|M - F
Proof. We have
LHS
Il((M) o c(M)) 1, - (F(M) c(M) 1
II(F((M) o c(M) -(((M) o c( c(M)) 1,1122
(I((M) c(M) j(M) (M) || 1211
= c(M) F(M) )ob))a c(M)]
(10)
where the first step follows from Fact B.4 the second step follows from basic algebra, the third step follows from Fact B.5 and the fourth step follows from | = vn.
Then have
(1(M) c(M) F(M) o c(M) <
(11)
which follows from Fact B.5
Following Eq. 10 Eq. and Lemma F.5 we have
LHS < Vn. 6dn 3/2 M M 6dn,
R R M]
43

---

Published as a conference paper at ICLR 2025
F.6 LIPSCHITZ OF diag((f(M) c(M) 1.)
We introduce the lemma about Lipschitz of diag(( o c(M)) (.M)
Lemma F.7 (Lipschitz of diag( c(M)) (M)) If the following conditions hold
Assumption
E.I holds.
Let c(M) be defined as Definition C.4
Let F(M) be defined as Definition C.3
For M, M,M E Baxd, have
2/2 ||diag( e c(M)) LM) ( 1n)f(M)ll = 24dh RM M M
Proof. Following Fact F2 we have
|| diag ((M() o c(M) 1.)) (M) diagg( oc (M)]l
2 max Rd, 6dn (²) (vn 2n)]| M
= 2/2 M M
where we have the upper bound in Lemma E.3 the Lipschitz of diag(( c(M)) and f(M) in
Lemma [33 and Lemma F.6
F.7 LIPSCHITZ OF GRADIENT
We introduce the lemma about Lipschitz of the gradient.
Theorem E8 (Lipschitz of the gradient, formal version of Theorem 5.4 We can show Var L(M)
is L-Lipschitz
the following conditions hold
Assumption E. holds.
Let c(M) be defined Definition C.4
Let F(M) be defined as Definition C.3
*
For M, M E Raxd, we have
[V V (A IIM MF
Proof. We have
[VMC(M) - VM()T
(X F(M) diag(c(M) o F(M) 11))(MM
-c(M f(M) + diag(c(A) F(M 1,) (,) + XM - XMIL
S o (c(M((c ( F(M)
diag
((M)((M)
-c(M) f(M) + diag(c(M) J(M) 1,,) 1,.) (MM) + |IA(- (MMM-Mii MiIle
(12)
where the first step follows from Theorem D.5 and the second step follows the triangle inequality. Now, proof these two terms separately.
For the first term, we have
|| (x (XT (c(M) F(M) - diag(c(M)) F(M) 1.M -c((T J(M) + diag(c(M) ((M))1))((MM) (M))
44

---

Published as a conference paper at ICLR 2025
S [|W] (c(M) f(M) diag(c(M) f(M) LnM)(M)
- -(() o f(M) + diag(c ((M) (M)
|| || |lc f(M) diag(c(M)) f(M) 1I.
-((M - [(M) + d [(M) 1,)J(M)ll
S 1| -1x (lc((M) f(M) - c() F(M)lle
+ dIiig(()))M) o f(M) (M) diag ((lg((((M))) f(M) (m)f(())
R² (Il((M) o F(M) c(M) o J(M)IlF
+ || diag(e(MM) f(M) fAV) diag(c(M) f(M) (M)
where the first step and the second step follow from Fact B.5 the third step follows from triangle inequality, and the fourth step follows from Assumption E.1
Then we have
(e.((c((M o ((M) c(M) f(M)lr
+ || diM) f(M) (MM) diag(c(M) F(M) (M)F())
< R (24dn R M M1FF F 6/ M11F)
< R (30dn R M Ml)r)
30dn 7/2 12a |M - M||F
(13)
where the first step follows from Lemma F.5 and Lemma F7 the second step follows from 2 For the second term, we have
JM( M)lE = MII Ml
(14)
which follows from Fact B.5
Finally, we have
V MCM)
VM ||r < (A + 30dn 1IM
M/
which follows from Eq.
12 Eq.
13 and Eq. 14].
G CONVERGENCE OF GRADIENT DESCENT
G.1 HELPFUL STATEMENTS
Here, present useful facts that to prove convergence result. Fact G.1. We can show that for a, b E R
Part
+
02 + +
V2
Part 2. Suppose |a] >
Mal 16 2 Vla
V
Proof. Proof of Part 1. Square both side of the inequality in Part 1., have
LHS = a² +b²
and
RHS a² + 2/a|| -18| + 62
2
45

---

Published as a conference paper at ICLR 2025
So just need to prove which is hold because for any E R, x² 20. O.
+
[6] + b
LHS RHS a² + 6²
2
a² 12
2/a| 16|
2
(1a) [b))
2
20
Proof of Part Square both side of the inequality in Part we have
LHS |a| - 16]
and
RHS la| + |6] 2/ 22/a//1 |a|b
So just need to prove
LHS RHS |a| - |6|
|a|
18| + 2 |a|bb
2 |a| |
2|b
2 [B| / a
M
which is hold because |a| > [b] and [b] 2 O.
G.2 LOWER BOUND ON FROBENTUS NORM
In this section, we present the lemma for the lower bound on the Frobenius norm. Lemma G.2. If the following conditions hold
Let B Rdddd
Let M [0, 1]axd
Let X E [0, 1] some constant.
Suppose that || B R.
Then, show
Part J.
|B + AMI 2 |B + + 2RAd
Part
1
B + AM
B|| + >||MIa)) 2RAd
/2
Proof. Proof of Part 1. We can show that
| + AMI = | + X2 || 2 + +(BB,AM)
> | + X2 M B|I
||M
2 B + MI 2RAd
(15)
where the first step follows from Fact B.6 the second step follows from Fact B.5 the third step follows from the upper bound of B3| F and |
46

---

Published as a conference paper at ICLR 2025
Proof of Part 2. Taking the square root on both sides, get
B + AMIlF 2 VIBIE + AiI 2RAd
B + M1 v2RAd
B| + )))
V2RAd
where the first step follows from Eq. 15) the second step follows from Part 2. of Fact G.1 and third step follows from Part 1. of Fact G.
G.3 SANDWICH LOWER BOUND ON FROBENIUS NORM
Here, we introduce sandwich trace fact.
Fact G.3. A = BI, then tr(B" AB) 2 Br B).
Proof. As A ž 3I, we have 31 ž Multiplying both sides by B on the left and B on the right (noting that these operations preserve the positive semidefiniteness) we have
B ( 31) N
Taking the trace and utilizing the property that the trace of a positive semidefinite matrix is non- negative, we have
tr T AB ---BBBB B 20
which simplifies to
Lr Lr(BT AB] - BIr(B B N 0.
This concludes the proof.
We establish sandwich lower bound on the Frobenius norm.
Lemma G.4 (Formal version of Lemma 5.6). If following conditions hold
Let B E RAXD and X E Baxd
Assume that XXT 3I.
Then, we have
|X BX| BIBBF
Proof. We can show that
||X = TXX
T T BT x]
2 2r BB X]
= -3BBBB X B]
2 8d Ir[B" B]
82 |BI
where the first step, the third step, and the fifth step follow from Fact B.4 the second step and the fourth step follows from Fact G.3 and X 2 31.
Taking the square root of both sides, we finish the proof.
47

---

Published as a conference paper at ICLR 2025
G.4 LOWER BOUND ON HADAMARD PRODUCT BETWEEN TWo MATRICES
We present the lemma for the lower bound on the Hadamard product between two matrices in this section.
Lemma G.5. If the following conditions hold
Let B, W Rdxd.
Then, have
max
[|,5] ||]| 2 I|W Bll 2 min (Wil [B].
<Je[d
iJe(d]
Proof. The proof directly follows from the definition of the Frobenius norm.
G.5 FINAL BOUND
We introduce some useful lemmas that we use to prove the final bound. Lemma G.6. If the following conditions hold
Let b TR and (b, I,.) 0.
Let f 1) and (ʃ. In) 1.
Then we have
- (b,)) flla 0|/6112-
Proof. Note that (b. I,,) = 0 so that and 1, are orthogonal with each other. Then, we have
[1(b (b. f),,) fll2 8l|b 1)1nll2
= 1611 + f1..111113
8116112
where the second step is from the Pythagorean theorem.
We present our final bound for proving the PL inequality.
Lemma G.7 (Formal version of Lemma 5.7p If the following conditions hold
Let RRxn and each row summation is zero, i.e., B 1, = On.-
Let F(M) E [0, 1]mxn and each summation is Le., F(M) 1, 1n
Assume that jinj F(M) 2 8>00 O.
Then, we can show
| F(M) diag((B F(M)) 1, 1.) 2 8- |B1F
Proof. For any [n]. let B; Rn be the i-th row of B, and we have (B,, 1n) by the first condition.
For any i E [n], let f(M), E R be the i-th row of f(M), and have (f(M), 1x) by the second condition and f(M),, [8, 1] by the third condition.
By Lemma G.6 for any E [n], we have
[IB - (Bi, F(M). F(M),l 2 B.]B.
Then, have
|Bo J(M) diag( J(M) 1.)/(M)
48

---

Published as a conference paper at ICLR 2025
Σ
[(B;
2
Σ
5²1 1B:12
62| B
G.6 PI. INEQUALITY
Here, we present the bound for unit loss function. Lemma G.8. If the following conditions hold
Let c(M) be defined in Definition C.4
We have
|ccM)]|F S 2Vn.
Proof. We have
|c((M]] S [If(M)]l + I
22vn,
where the first step follows from Definition C.4 and triangle inequality, the second step follows of + 2 I (I1 +... xn)² when r; 2 for any E [n].
We present the lemma to prove the PL inequality. Lemma G.9. If the following conditions hold
Let f(M) be defined in Definition C33
Let e(M) be defined in Definition C.4
We have
|| |dia((c((M) F(M)) 11) s Vn.
Proof. We have
|| diag((c((M) ((M) (M))( max(((((M)) J(M))) 1./ -1/(0
iEn
S Vn.
where the first step is by Frobenius norm definition and the second step follows from (F(M), 1n) 1 and c(M); [-1, 11"d for any i E [rs].
Finally, we can show the lemma for PL inequality.
Lemma G.10 (PL inequality, formal version of 5.5). If the following conditions hold,
Let M 1jddd
Let X [0, 1] some constant.
Assume that XXT = BT.
Assume that min J(M),, 2 8> 0.
49

---

Published as a conference paper at ICLR 2025
Let LMM be defined in Definition C.7
Furthermore,
Let a =
Let μ = 2 jel { 8.
Lett 12 max. sel4 ([a,) ||x|| Ad/p.
We have
2A² V ju(ll(AMM13 + 1
|M E).
Proof. We have f(M) 1n = 1n and f f-1, 1, 1n by Definition by Definition C.4 Thus, we have c(M) 1, O,,
C.3 Note that c(M) F(M)
On the other hand, we have
I|W (X (c(M) o F(M) diag((c(M) o F(M)) diag((((M)))
((M(()))
< max
1.jE/d] [/a,]) -- ||X (c(M) o J(M) - diag((c(M)) F(M)) LM)
[War]] ||X|| -lo([l/())))))) - F(M) diag((c(M) F(M)) 1m)F(M)ll
max
[[aa]] (0%/)(--1XX) -IX1 ((((MM) J(M) (I o J(M) 1M)((M)
max JJEM
$ max
[ || (I F(M)) ( (m)f(l)
[JE[4]
S max [Wa,]) } -|XI (2V7+vm)
max A.Eld
{ } ||X|| 3/7
where the first and fourth steps follow Lemma G.5 the second step follows from Frobenius norm property, the third step follows from triangle inequality, the fifth step follows from Lemma G.8 and Lemma G.9
Let 2. We have the following
V M L(M)
= || o (x (c(M) o f(M) diag((c(M) - (.)))))))))((()))+)) AMI
2 |[ (XT
J(M) diag((e(M) f(M)) (.))(())))I 11IM- 0n
2 02 |X (c(M) o F(M) diag((c(M) F(M)) -1))(
02 0s [c((M) J(M) diag((c(M) 0 ((M)) 1..) +
2 a2 0s as [lc 1 a1
2X2 M - E)
= "(((
where the second step follows from Lemma G.2 and i1 бş/ max ( 3 | Ad. the third step follows from Lemma G.5 and 0:2 min,,eei ( the fourth step follows from Lemma G.4 and 0: = 8, the fifth step follows from Lemma G.7 and 0t ö, and the last step
follows from μ = 20² Qg 04 and { 201/14
50


Reviews:
[{"review_id": "jwlQ7P5PCH", "reviewer": "Reviewer_opkW", "summary": "This paper proposes a novel approach to pruning QK projection matrices by directly optimizing for an approximation of the attention matrix, bypassing traditional linear approximations applied before the Softmax operation. By addressing the non-linear nature of the attention mechanism, the authors introduce a gradient descent-based optimization method with theoretical guarantees on convergence, aiming to produce a near-optimal pruning mask. Preliminary empirical results, although limited, demonstrate that this approach can maintain model performance while significantly reducing computational costs.", "questions": "- I assume the cost of pruning itself is very small given the closed-form gradient, but it would be helpful to see an analysis of this cost. Is this approach lightweight enough for large models?\n- The assumptions regarding the Lipschitz continuity and PL inequality of the gradient are central to the theoretical guarantees. Are there empirical checks or practical insights into how these assumptions hold up in real LLMs?\n- The paper uses synthetic data by generating $W_Q$ and $W_K$ through SVD on Gaussian random matrices. A bit more explanation for this choice—or testing on real model weights—would clarify whether the setup could translate to practical applications.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper offers a robust theoretical analysis, including convergence guarantees for the gradient descent algorithm, which enhances confidence in the method’s foundational soundness.\n- The proposed pruning method consistently outperforms baseline approaches in terms of relative error, effectively recovering the attention matrix activations. This suggests its potential for maintaining model performance while reducing computation.", "weaknesses": "- The empirical analysis relies heavily on synthetic data, specifically a full-rank Gaussian random matrix, with matrix dimensions (64x128) that are small relative to the scale of practical LLMs. This somewhat limits confidence in the method’s real-world applicability.\n- I understand the main contributions are theoretical, but I think the paper would benefit from experiments with actual LLM activations and weights, even small-scale LLMs. Results from real models would give a stronger sense of how effective this method might be in practice. \n- One big limitation is that the approach only targets $W_Q$ and $W_K$ matrices in the attention mechanism, while the majority of parameters in transformers come from the MLP layers. It would be interesting to see if this technique could be extended to cover the MLP parameters as well, which could have a more substantial impact on reducing model size.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EZyCuKiXFN", "reviewer": "Reviewer_WRvG", "summary": "The paper proposes to prune the projection weight matrices of the Softmax attention conditioned on the similarity between the pruned and the original attention matrices. The proposed approach leverages gradient decent to solve the constraint opimitzation problem and has theoretical gurantees for near-optimality with polynormal time.", "questions": "1. How to revert the pruned fused matrix back to the query and key projection matrices respectively?\n2. How does the proposed attention pruning affect length extrapolation ability?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Considering the similarity between attention matrices during pruning of the projection weight is novel and interesting.\n2. The assumptions in theoretical analyses are mild and sensible.", "weaknesses": "1. The applicability of the approach is limited. The proposed pruning approach needs to fuse the query and the key projection matrices into a single matrix through matrix multiplication. This means that it will eliminate the inference efficiency benefits of popular LLMs, such as Llama 3, which use Grouped Query Attention that shares the key-value pairs between queries. This is because after the proposed fushion, the original input X needs to be saved during decoding, while for GQA, only much smaller key-value vectors are needed which has much smaller memory I/O footprint.\n2. Lacks of empirical results on pruning real-world LLM. The low error rates on synthetic data may not translate to real world dowstream performances.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yj2JWLADnX", "reviewer": "Reviewer_N1Ar", "summary": "The paper introduces a novel approach to pruning weights in large language models (LLMs) to improve computational efficiency, particularly for deployment on resource-constrained devices. This method focuses on optimizing the approximation of the attention matrix, a critical component in transformer architectures, by accounting for the non-linear properties of the Softmax attention mechanism. The authors present a Gradient Descent-based optimization technique, offering theoretical guarantees for its convergence toward a near-optimal pruning solution. Preliminary results indicate that this approach maintains model performance while significantly reducing computational costs. This work provides a new theoretical basis for designing pruning algorithms for LLMs, potentially enabling more efficient inference on devices with limited resources.", "questions": "Conside the weakness", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The presented technique enhances convergence guarantees for non-convex functions, crucial for optimization challenges in areas like deep learning, through the application of the Polyak-Łojasiewicz inequality. It showcases mathematical rigor with detailed derivations and validations of key conditions, supported by established theorems, notably Theorem 5.2 from Frei & Gu (2021). Practical implications include improved computational efficiency via attention weight pruning without notable performance trade-offs, assured by the method’s convergence mechanisms. Additionally, its framework offers potential applicability across various non-convex optimization scenarios, expanding its usefulness in the optimization toolkit.\n2.  The experimental results suggest that the proposed algorithm outperforms existing methods under certain conditions, the above weaknesses highlight areas where the study could be improved.", "weaknesses": "1. The convergence guarantees depend on strict conditions like the boundedness and positive definiteness of matrices (e.g., $\\( X^T X \\succeq \\beta I \\))$. These conditions may not always hold in real-world applications, limiting the practical relevance of the results. Additionally, the analysis does not consider the stochasticity and noise common in real-world data and optimization, which could affect gradient descent’s performance.\n\n2. The experimental section lacks clarity on whether the tested input sequence lengths cover the full range required for real-world corpus. Given that modern llm often deals with lengthy texts like entire documents or extended conversations, the study does not explore the algorithm’s performance in such scenarios, raising concerns about its practical utility.", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Mv5ZW1dxsi", "reviewer": "Reviewer_QLv4", "summary": "This paper proposes a method for unstructured pruning of the query and key projection matrices in attention. The paper provides theoretical guarantees for the convergence of the pruning method. It also provides some experiments on synthetic data demonstrating that the pruning method outperforms other pruning methods.", "questions": "How does the convergence rate in Theorem 4.1 compare to the rates one can achieve with naive baselines?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "*Original problem.* The paper highlights an interesting, niche problem in the pruning literature, which, to my knowledge, has not been analyzed theoretically. \n\n*Clear presentation.* The paper is well-written and the structure is easy to follow. \n\n*Extensive theoretical analysis.* The authors provide an extensive theoretical analysis on the convergence of their pruning algorithm.", "weaknesses": "*Limited potential for speedup on large scale models.* I have two concerns around the potential for this method to provide significant speedups for large language models. (1) It is typically difficult to leverage fine-grained, unstructured sparsity for speedups on GPUs because standard vector and matrix instructions are not compatible with sparsity (though, recent changes to Nvidia GPUs could begin to mitigate this problem [Nvidia Blog](https://developer.nvidia.com/blog/accelerating-inference-with-sparsity-using-ampere-and-tensorrt/)). (2) The technique is applicable to a small fraction of the parameters in the full model. In large models, the MLPs dominate the total parameter and FLOP counts, but this method is only applicable to the Q and K projections, which account for only 9.4% of parameters in Llama-405B (excluding embeddings, [source](https://huggingface.co/meta-llama/Llama-3.1-405B/tree/main)). So even if the method could achieve 100% sparsity, the relative reduction in FLOPs/memory consumption would be limited. \n\n*Experiments limited to synthetic data.* Because experiments are only performed on synthetic data and a single layer, it remains unclear how the performance of real language models would degrade with this method. Without experiments on real language models it is hard to assess the improvement of this method relative to prior work (Wanda and SparseGPT).", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)