Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix

## 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 preliminary empirical results demonstrate the effectiveness of this approach in maintaining model performance while significantly reducing computational costs.
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.

## 1 Introduction

Large Language Models (LLMs) based on the transformer architecture [], including GPT-4o [], Claude [], and OpenAI’s recent o1 [], have shown immense potential to enhance our daily lives.
They revolutionize fields like conversational AI [], AI agents [, ], search AI [], and AI assistants [, , , ].
With their growing capabilities, LLMs are powerful tools shaping the future of technology.
However, the current state-of-the-art LLM weights number is extremely large. For instance, the smallest version of Llama 3.1 [] needs 8 billion parameters, which takes more than 16GB GPU memory with half float precision and requires significant inference time.
Due to large memory and high computational cost, deploying such models on edge devices such as smartphones becomes challenging.

To reduce the LLM model size, many studies work on 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\in\mathbb{R}^{d\times d}$ and some calibration data $X\in\mathbb{R}^{n\times d}$, where $n$ is input token length and $d$ is feature dimension, the goal is to find a matrix $\widetilde{W}$ under some sparse constraint such that
$\|XW-X\widetilde{W}\|$
being small under some norm function.
The above formulation has been widely used in many state-of-the-art pruning methods, such as SparseGPT [] and Wanda [].

However, the current object functions only focus on the approximation of a linear function $XW$.
Their optimal solutions do not have a 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 $\mathsf{Softmax}$ function is very sensitive to 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 on weights to approximate the attention matrix, which is a highly non-linear function, $\mathsf{Softmax}(XWX^{\top})\in\mathbb{R}^{n\times n}$.
To the best of our knowledge, this paper is the first work studying attention weight pruning to 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 6).
Our theoretical foundation may pave the way for more efficient LLM inference on resource-constrained devices.

In the following, we introduce some key backgrounds and our contributions in detail.

Figure: Figure 1: Comparison of our Attention Weights Pruning method and Linear Pruning 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 is applied separately to the query weight matrix $W_{Q}$ and key weight matrix $W_{K}$, using two different pruning masks $M_{Q}$ and $M_{K}$, respectively.
Refer to caption: x1.png

### 1.1 Key Background

We define the attention matrix in self-attention mechanism as below:

###### Definition 1.1 (Attention Matrix) .

Let $X\in\mathbb{R}^{n\times d}$ be the input.
Given query and key weights matrix $W_{Q},W_{K}\in\mathbb{R}^{d\times d}$, we define $W:=W_{Q}W_{K}^{\top}$.
Then, we have the Softmax attention matrix being

$$ $\displaystyle\mathsf{Softmax}(XWX^{\top})=D^{-1}\exp(XWX^{\top}),$ $$

where (1) $D:=\operatorname{diag}(\exp(XWX^{\top})\cdot{\bf 1}_{n})$,
(2) $\exp$ denotes the exponential function and is applied entry-wisely,
(3) $\operatorname{diag}()$ operation takes a vector and outputs a diagonal matrix with the entries of that vector,
and (4) ${\bf 1}_{n}$ denotes the length-$n$ all ones vector.

Further, we 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 1.1, 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\in[0,1]^{d\times d}$ be the pruning mask.
Let $X,W$ be defined in Definition 1.1.
Let $A:=\exp(XWX^{\top})$ and $\widetilde{A}:=\exp(X(M\circ W)X^{\top})$, where $\circ$ is the Hadamard product.
Let $D:=\operatorname{diag}(A\cdot{\bf 1}_{n})$ and $\widetilde{D}:=\operatorname{diag}(\widetilde{A}\cdot{\bf 1}_{n})$.
Let $\lambda\in\mathbb{R}_{+}$ be the regularization parameter.
We define Attention Weights Pruning loss function to be

$$ $\displaystyle\mathcal{L}(M):=\frac{1}{2}\|D^{-1}A-\widetilde{D}^{-1}\widetilde {A}\|_{F}^{2}+\frac{1}{2}\lambda\|M\|_{F}^{2}.$ $$

Thus, the Attention Weights Pruning optimization problem is $\min_{M}~{}\mathcal{L}(M)$.

###### Definition 1.1 (Attention Matrix) .

###### Definition 1.2 (Attention Weights Pruning) .

### 1.2 Our Contributions

This is the first work studying the Attention Weights Pruning problem, which is an approximation problem to a non-linear function.
We provide an algorithm to obtain 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 $\epsilon>0$, our Algorithm 1 can converge to the near-optimal pruning mask for the Attention Weights Pruning problem (Definition 1.2) in $O(d\operatorname{poly}(n)/\epsilon)$ time with $O(\xi+\epsilon)$ error, where $\xi$ is a small term depending on intrinsic property of the data and weights.

In the above theorem, $\xi$ can be arbitrarily small as $\xi\rightarrow 0$ when the regularization coefficient $\lambda\rightarrow 0$.
So our analysis shows that although the objective function is highly non-linear, the GD training can converge to a near-optimal pruning mask solution, supported by our experiments in Section 6.

Our contributions are as follows:

- •
This is the first work that analyzes the weights pruning problem based on $\mathsf{Softmax}$ attention, which is a 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 1 to obtain the near-optimal pruning mask and its convergence guarantee (Theorem 4.1).
- •
We conduct preliminary experiments to verify the effectiveness of our method (Section 6).

###### Theorem 1.3 (Main result, informal version of Theorem 4.1 ) .

#### Roadmap.

Our paper is organized as follows.
In Section 2, we review the related work.
Section 3 introduces key concepts and definitions essential for the subsequent sections.
In Section 4, we present our main result.
Section 5 offers a technical overview of the methods employed.
Experimental results are discussed in Section 6.
Finally, Section 7 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 language models (LLMs) [] for its effectiveness in reducing computational overhead while preserving performance. Common compression techniques include quantization [, , ], pruning [, , , , , , , , , , , , ], and knowledge distillation [, , , ]. Specifically, pruning techniques have been developed extensively, such as unstructured pruning, which removes individual weights [, ], and structured pruning, which eliminates entire components like neurons or attention heads [, , ]. [] proposed Wanda, a 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 [] 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 [] improved the complexity analysis of SparseGPT, reducing the running time from $O(d^{3})$ to $O(d^{2.53})$, enabling faster pruning on LLMs.
These techniques together contribute to more scalable and resource-efficient LLMs, maintaining competitive performance while having substantial reductions in computational resources.

### 2.2 Attention Acceleration

Attention mechanism has faced criticism due to its quadratic time complexity with respect to context length []. Addressing this criticism, a variety of approaches are employed, including sparse attention [, , , ], low-rank approximations [, , , , ], and kernel-based methods [, , , , ], to reduce computational overhead and improve scalability. [] 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 [, , , , ].
Other approaches like Mamba [, ], Linearizing Transformers [, ], Hopfield Models [, , , , , , , ], and PolySketchFormer [] focus on architectural modifications and implementation optimizations to enhance performance. System-level optimizations such as FlashAttention [, , ] and block-wise parallel decoding [] further improve efficiency. Collectively, these innovations have significantly augmented transformer models’ ability to handle longer input sequences, unlocking broader applications across multiple sectors [, , , , , , , , , ].

## 3 Preliminary

In this section, we introduce some basic concepts and key definitions.
In Section 3.1, we introduce some basic notations we use in this paper.
In Section 3.2, we provide the definition of attention weights pruning.

### 3.1 Notations

For any positive integer $n$, we use $[n]$ to denote set $\{1,2,\cdots,n\}$.
For each $a,b\in\mathbb{R}^{n}$, we use $a\circ b\in\mathbb{R}^{n}$ to denote the Hadamard product, i.e., the $i$-th entry of $(a\circ b)$ is $a_{i}b_{i}$ for all $i\in[n]$.
For $A\in\mathbb{R}^{m\times n}$, let $A_{i}\in\mathbb{R}^{n}$ denote the $i$-th row and $A_{*,j}\in\mathbb{R}^{m}$ denote the $j$-th column of $A$, where $i\in[m]$ and $j\in[n]$.
We use $\exp(A)$ to denote a matrix where $\exp(A)_{i,j}:=\exp(A_{i,j})$ for a matrix $A\in\mathbb{R}^{n\times d}$.
We use $\|A\|_{F}$ to denote the Frobenius norm of a matrix $A\in\mathbb{R}^{n\times d}$, i.e., $\|A\|_{F}:=\sqrt{\sum_{i\in[n]}\sum_{j\in[d]}|A_{i,j}|^{2}}$.
For a symmetric matrix $A\in\mathbb{R}^{n\times n}$, $A\succeq 0$ means that $A$ is positive semidefinite (PSD), i.e., for all $x\in\mathbb{R}^{n}$, we have $x^{\top}Ax\geq 0$.

### 3.2 Attention Weights Pruning

We aim to determine a near-optimal pruning mask $M$ for the attention weights in a self-attention mechanism. Furthermore, we incorporate causal attention masking(^1^11In this paper, we always use pruning mask to refer $M\in\mathbb{R}^{d\times d}$, which is our target, and use causal attention mask to refer $M_{c}\in\mathbb{R}^{n\times n}$, which is a fixed mask for standard self-attention.) 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 preceding tokens. Here, we provide the formal definition for the causal attention mask:

###### Definition 3.1 (Causal attention mask, [ 57 ] ) .

We define the causal attention mask as $M_{c}\in\{0,1\}^{n\times n}$, where $(M_{c})_{i,j}=1$ if $i\geq j$ and $(M_{c})_{i,j}=0$ otherwise.

Now, we incorporate Attention Weights Pruning (see Definition 1.2) with causal attention mask $M_{c}$.

###### Definition 3.2 (Attention Weights Pruning with Causal Attention Mask) .

Let $M_{c}\in\{0,1\}^{n\times n}$ be the causal attention mask defined in Definition 3.1. Let $A:=\exp(XWX^{\top})\circ M_{c}$ and $\widetilde{A}:=\exp(X(M\circ W)X^{\top})\circ M_{c}$. Let $D:=\operatorname{diag}(A\cdot{\bf 1}_{n})$ and $\widetilde{D}:=\operatorname{diag}(\widetilde{A}\cdot{\bf 1}_{n})$.
We define Attention Weights Pruning with Causal Attention Mask loss function to be

$$ $\displaystyle\mathcal{L}(M):=\mathcal{L}_{\mathrm{attn}}(M)+\mathcal{L}_{ \mathrm{reg}}(M)$ $$

where

$$ $\displaystyle\mathcal{L}_{\mathrm{attn}}(M):=\frac{1}{2}\|D^{-1}A-\widetilde{D }^{-1}\widetilde{A}\|_{F}^{2}~{}~{}\mathrm{and}~{}~{}\mathcal{L}_{\mathrm{reg} }(M):=\frac{1}{2}\lambda\|M\|_{F}^{2}.$ $$

Figure: Algorithm 1 Gradient Descent for Pruning Mask (Theorem 4.1). Let $X_{1},X_{2},\dots,X_{k}\in\mathbb{R}^{n\times d}$ be our calibration dataset of size $k$. We iteratively run our GD method on this dataset.

###### Definition 3.1 (Causal attention mask, [ 57 ] ) .

###### Definition 3.2 (Attention Weights Pruning with Causal Attention Mask) .

## 4 Main Results

In this section, we provide our main results. We provide an Algorithm 1 for Attention Weights Pruning problem based on Gradient Descent (GD). We also prove the convergence for our GD algorithm in Theorem 4.1.

###### Theorem 4.1 (Main result, formal version of Theorem 1.3 ) .

Let $M,~{}X,~{}W,\widetilde{D},\widetilde{A}~{}\lambda,~{}\mathcal{L},~{}\mathcal{L
}_{\mathrm{attn}}$ be defined in Definition 3.2. Assume $XX^{\top}\succeq\beta I$ and $\min_{i,j\in[n]}(\widetilde{D}^{-1}\widetilde{A})_{i,j}\geq\delta>0$. Furthermore,

- •
Let $\mu=2\min_{i,j\in[d]}\{|W_{i,j}|\}\cdot\beta\cdot\delta$.
- •
Let $\xi=12n^{-1.5}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot\lambda d/\mu$.

Then, for any $\epsilon>0$, provided $\eta<1/L$ where $L$ is the Lipschitz constant for $\nabla_{M}\mathcal{L}(M)$ (see Theorem 5.4), GD (Algorithm 1) with fixed step size $\eta$ and run for $T=4\mathcal{L}(M^{(0)})/(\eta\mu\epsilon n^{2})$ iterations results in the following guarantee,

$$ $\displaystyle\frac{1}{n^{2}}\min_{t<T}\mathcal{L}_{\mathrm{attn}}(M^{(t)})+ \frac{\lambda^{2}}{\mu n^{2}}\|M^{(t)}\|_{F}^{2}\leq(\xi+\epsilon)/2.$ $$

###### Proof.

Let $g(M)=2\mathcal{L}_{\mathrm{attn}}(M)+\frac{2\lambda^{2}}{\mu}\|M\|_{F}^{2}$. Note that $\mathcal{L}(M)$ satisfies the $(g(M),n^{2}\xi,2,\mu)$-proxy PL inequality (Lemma 5.5).
Also, we have $\mathcal{L}(M)$ is non-negative and has $L$-Lipschitz gradients Theorem 5.4.
Thus, we finish the proof using Theorem 5.2.
∎

###### Remark 4.2 .

The two assumptions in Theorem 4.1 are practical.
The first assumption of the positive definite matrix is widely used in theoretical deep learning analysis, e.g., [, , , ].
The second assumption is natural, as $\widetilde{D}^{-1}\widetilde{A}$ is the pruned attention matrix, where each entry is

$$ $\displaystyle\frac{\exp(XWX^{\top})_{i,j}}{\sum_{j=1}^{n}\exp(XWX^{\top})_{i,j }}>0,$ $$

which has a natural lower bound.

Our error upper bound in Theorem 4.1 is $O(\xi+\epsilon)$, where $\epsilon$ can be arbitrarily small.
For $\xi$, we can let it be small by choosing a proper $\lambda$, i.e., the $\xi$ error term can be made arbitrarily small by choosing small $\lambda$.
However, if we choose a very small $\lambda$, the algorithm’s run time gets larger as $T\propto 1/\eta\propto L\propto(\lambda+\mathrm{other~{}terms})$.
Thus, although the objective function is highly non-linear, we can show that the Gradient Descent of our Algorithm 1 can converge to a good solution of the Attention Weights Pruning problem.

After solving the optimization problem, we obtain a pruning mask with real-valued entries. In practice, however, this pruning mask must be converted into a binary form, specifically $M\in\{0,1\}^{d\times d}$. We define the pruning ratio $\rho$ 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 $\rho$-th percentile and to one for those above. This ensures that only the specified proportion of weights are pruned.

###### Theorem 4.1 (Main result, formal version of Theorem 1.3 ) .

###### Proof.

###### Remark 4.2 .

## 5 Technique Overview

In Section 5.1, we introduce some useful tools from previous work.
In Section 5.2, we derive the close form of the gradient of Attention Weights Pruning.
In Section 5.3, we calculate the Lipschitz constant of that gradient.
In Section 5.4, we prove the PL inequality for our loss function.

### 5.1 Previous Tools on Convergence of GD

To analyze the convergence behavior of GD for our optimization problem (Definition 3.2), we first introduce the concept of $g$-proxy, $\xi$-optimal Polyak–Łojasiewicz(PL) inequality [, , ], under which GD will converge:

###### Definition 5.1 ( $g$ -proxy, $\xi$ -optimal PL inequality, Definition 1.2 in [ 29 ] ) .

We say that a function $f:\mathbb{R}^{p}\to\mathbb{R}$ satisfies a $g$-proxy, $\xi$-optimal Polyak–Łojasiewicz inequality with parameters $\alpha>0$ and $\mu>0$ (in short, $f$satisfies the $(g,\xi,\alpha,\mu)$-PL inequality) if there exists a function $g:\mathbb{R}^{p}\to\mathbb{R}$ and scalars $\xi\in\mathbb{R}$, $\mu>0$ such that for all $w\in\mathbb{R}^{p}$,

$$ $\displaystyle\|\nabla f(w)\|^{\alpha}\geq\frac{1}{2}\mu(g(w)-\xi).$ $$

PL inequality is a powerful tool for studying non-convex optimization, and it has been used in recent studies on provable guarantees for neural networks trained by gradient
descent [, , , , , , , , , ].
It provides a 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, [] derives the following GD convergence guarantees.

###### Theorem 5.2 (Theorem 3.1 in [ 29 ] ) .

Suppose $f(w)$ satisfies the $(g(\cdot),\xi,\alpha,\mu)$-proxy PL inequality for some function $g(\cdot):\mathbb{R}^{p}\to\mathbb{R}$. Assume that $f$ is non-negative and has $L$-Lipschitz gradients. Then for any $\epsilon>0$, provided $\eta<1/L$, GD with fixed step size $\eta$ and run for $T=2\eta^{-1}(\mu\epsilon/2)^{-2/\alpha}f(w^{(0)})$ iterations results in the following guarantee,

$$ $\displaystyle\min_{t<T}g(w^{(t)})\leq\xi+\epsilon.$ $$

The above theorem establishes that under the $(g,\xi,\alpha,\mu)$-PL inequality and Lipschitz continuity of the gradient, GD converges to a point where the proxy function $g(w)$ is within $\epsilon$ of $\xi$.
To apply this result to our specific problem, we need to verify these conditions for our loss function ${\cal L}(M)$.

###### Definition 5.1 ( g 𝑔 g italic_g -proxy, ξ 𝜉 \xi italic_ξ -optimal PL inequality, Definition 1.2 in [ 29 ] ) .

###### Theorem 5.2 (Theorem 3.1 in [ 29 ] ) .

### 5.2 Closed Form of Gradient

As a first step, we compute the close form of the gradient $\nabla_{M}{\cal L}(M)$. The pruning mask $M$ is inside a non-linear function Softmax, which complicates our calculation.
We defer the proof to Section B.

###### Theorem 5.3 (Closed form of gradient, informal version of Theorem C.5 ) .

Let ${\cal L}(M)$ be defined in Definition 3.2.
Let $p$ be defined in Definition C.1.
Let $X\in\mathbb{R}^{n\times d}$, $M\in[0,1]^{d\times d}$, $W\in\mathbb{R}^{d\times d}$.
Then, we have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}(M)}{\mathrm{d}M}=W\circ(X^{\top}pX)+ \lambda M.$ $$

Based on Theorem 5.3, we calculate the gradient of the pruning mask from Line 10 to Line 15 in our Algorithm 1.

###### Theorem 5.3 (Closed form of gradient, informal version of Theorem C.5 ) .

### 5.3 Lipschitz of Gradient

Having obtained the close form of gradient, we proceed to investigate its Lipschitz continuity. We aim to show that the gradient $\nabla_{M}{\cal L}(M)$ is Lipschitz continuous with respect to $M$.

###### Theorem 5.4 (Lipschitz of the gradient, informal version of Theorem E.8 ) .

Let $R$ be some fixed constant satisfies $R>1$.
Let $X\in\mathbb{R}^{n\times d},W\in\mathbb{R}^{d\times d}$. We have $\|X\|_{F}\leq R$ and $\|W\|_{F}\leq R$.
Let ${\cal L}(M)$ be defined in Definition 3.2.
For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|\nabla_{M}\mathcal{L}(M)-\nabla_{M}\mathcal{L}(\widetilde{M})\| _{F}\leq(\lambda+30dn^{7/2}R^{6})\cdot\|M-\widetilde{M}\|_{F}.$ $$

We defer the proof to Section E.
Establishing the Lipschitz continuity of 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.

###### Theorem 5.4 (Lipschitz of the gradient, informal version of Theorem E.8 ) .

### 5.4 PL Inequality of Gradient

Next, we need to verify that our loss function satisfies the PL inequality with appropriate parameters.
To complete the verification of the conditions required for convergence, we demonstrate that ${\cal L}(M)$ satisfies the PL inequality.
We show that $\nabla_{M}\mathcal{L}(M)$ satisfies the PL inequality in this lemma:

###### Lemma 5.5 (PL inequality, informal version of Lemma F.10 ) .

Let $M,X,W,\widetilde{D},\widetilde{A},~{}\lambda,~{}\mathcal{L},~{}\mathcal{L}_{
\mathrm{attn}}$ be defined in Definition 3.2.
Assume that $XX^{\top}\succeq\beta I$ and $\min_{i,j\in[n]}(\widetilde{D}^{-1}\widetilde{A})_{i,j}\geq\delta>0$.
Also,

- •
Let $\mu=2\min_{i,j\in[d]}\{|W_{i,j}|\}\cdot\beta\cdot\delta$.
- •
Let $\xi=12\sqrt{n}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot\lambda d/\mu$.

We have

$$ $\displaystyle\|\nabla_{M}\mathcal{L}(M)\|_{F}^{2}\geq\frac{1}{2}\mu(2\mathcal{ L}_{\mathrm{attn}}(M)+\frac{2\lambda^{2}}{\mu}\|M\|_{F}^{2}-\xi).$ $$

We defer the proof to Section F.
By confirming that ${\cal L}(M)$ satisfies the PL inequality and that its gradient is Lipschitz continuous, we then apply Theorem 5.2 to conclude that GD will converge to a 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 in our Theorem 4.1, $XX^{\top}\succeq\beta I$ and $\min_{i,j\in[n]}(\widetilde{D}^{-1}\widetilde{A})_{i,j}\geq\delta>0$.

###### Lemma 5.6 (Informal version of Lemma F.4 ) .

Let $B\in\mathbb{R}^{n\times n}$ and $X\in\mathbb{R}^{n\times d}$. Assume that $XX^{\top}\succeq\beta I$.
Then, we have

$$ $\displaystyle\|X^{\top}BX\|_{F}\geq\beta\|B\|_{F}.$ $$

###### Lemma 5.7 (Informal version of Lemma F.7 ) .

Let $B\in\mathbb{R}^{n\times n}$ and each row summation is zero, i.e., $B\cdot{\bf 1}_{n}={\bf 0}_{n}$. Let $\widetilde{B}\in[0,1]^{n\times n}$ and each row summation is 1, i.e., $\widetilde{B}\cdot{\bf 1}_{n}={\bf 1}_{n}$. Assume that $\min_{i,j\in[n]}\widetilde{B}_{i,j}\geq\delta>0$.
Then, we can show

$$ $\displaystyle\|B\circ\widetilde{B}-\operatorname{diag}((B\circ\widetilde{B}) \cdot{\bf 1}_{n})\widetilde{B}\|_{F}\geq\delta\cdot\|B\|_{F}.$ $$

Figure: (a) Varying $\lambda$.
Refer to caption: x2.png

###### Lemma 5.5 (PL inequality, informal version of Lemma F.10 ) .

###### Lemma 5.6 (Informal version of Lemma F.4 ) .

###### Lemma 5.7 (Informal version of Lemma F.7 ) .

## 6 Experiment

In this section, we discuss the experiments conducted to illustrate the effectiveness of our Algorithm 1. We first introduce our settings in Section 6.1. Then, we present our results in Section 6.2.

### 6.1 Settings

#### Method and evaluation.

We implement our method following the pseudocode in Algorithm 1, using 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 Definition 3.2 as our loss function, optimizing over the pruning mask $M$ using gradient descent based on the closed-form expression derived in Theorem 5.3. To accelerate convergence, we leverage momentum into the optimization process and fix the momentum parameter at $0.9$. After obtaining the optimal pruning mask, we convert $M$ to a binary pruning mask to prune $W$, maintaining sparsity at the desired pruning ratio $\rho$. We use the relative error as our evaluation metric, which is defined as

$$ $\displaystyle\frac{\|\widetilde{D}^{-1}\widetilde{A}-D^{-1}A\|_{F}^{2}}{\|D^{- 1}A\|_{F}^{2}},$ $$

where $\widetilde{D}$, $\widetilde{A}$, $D$, $A$ are defined in Definition 3.2.

#### Baselines.

We compare our method with two linear pruning approaches, namely Wanda [] and SparseGPT []. Wanda is a 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 a 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 $W_{Q}$ and $W_{K}$ separately (see Figure 1).

#### Data.

In order to assess the efficacy of different methods in approximating the attention matrix, we construct the data via a carefully defined generating process. Specifically, we create multiple independent random Gaussian matrices $G\in\mathbb{R}^{d\times d}$, where each entry of $G$ drawn from a normal distribution i.e., $G_{i,j}\sim\mathcal{N}(0,1)$ for $i,j\in[d]$. Then, we perform singular value decomposition (SVD) on matrix $G$, i.e., $U,S,V^{\top}=\mathrm{SVD}(G)$. We retain the first four singular values in $S$ and set others to zero, constraining the rank to four. Our $W_{Q}$ and $W_{K}$ are then constructed as $U\operatorname{diag}(S)V^{\top}$. The weight matrix $W$ used in our setting is formed by taking the product $W=W_{Q}W_{K}^{\top}$. For $X\in\mathbb{R}^{n\times d}$, we generate it as a full-rank Gaussian random matrix.

#### Setup.

In our experiments, the weight matrix dimension $d=64$ is kept as constant across all figures, and we simulate two datasets of size $k=16$ and $k=64$.
We set the input sequence length $n=128$ for experiments (a) and (c) in Figure 2. The pruning ratio $\rho=0.5$ is set for experiments (a) and (b) in Figure 2.
For our method, the regularization coefficient $\lambda:=\widetilde{\lambda}/n$ where we abuse the notation to denote $\widetilde{\lambda}$ as the same used in Definition 3.2 and $\lambda$ here is the parameter we really control in experiments. $\lambda$ 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 as $T=100$. The step size is set as $\eta=0.1/\lambda$ because Theorem 4.1 indicates that $\eta$ is inversely proportional to $\lambda$ with some constant, i.e., $\eta\propto 1/L\propto 1/(\lambda+\mathrm{other~{}terms})$.

### 6.2 Results

Overall, the results in Figure 2 show that our Algorithm 1 outperforms Wanda and SparseGPT with a large margin, which supports our theoretical analysis in Theorem 4.1. In the following, we will discuss each setting in detail.

#### Relation with regularization coefficient λ 𝜆 \lambda italic_λ .

The leftmost column of Figure 2 investigates the impact of the regularization coefficient $\lambda$ on relative error. As $\lambda$ increases from very small values, the relative error initially decreases sharply for our algorithm, reaching a minimum before gradually rising again, which forms a $U$ shape curve. This behavior indicates that there is an optimal $\lambda$ where our algorithm achieves its best performance around $2^{-4}$.
The U-shape curve phenomena are well-known in most hyper-parameter choosing, e.g., regularization coefficient.

#### Relation with input sequence length n 𝑛 n italic_n .

The center column of Figure 2 explores how the relative error changes with respect to the input sequence length $n$. As $n$ increases, the relative error for all three methods grows, though at different rates. Our method demonstrates a slower increase, maintaining a 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 to SparseGPT as $n$ is relatively small.

#### Relation with pruning ratio ρ 𝜌 \rho italic_ρ .

The rightmost column of Figure 2 illustrates the relationship between the relative error and the pruning ratio $\rho$ for the three methods under comparison: our algorithm, Wanda, and SparseGPT. As the pruning ratio $\rho$ increases, all methods exhibit a rise in relative error, indicating a degradation in approximation accuracy. However, our algorithm consistently outperforms both Wanda and SparseGPT across the range of $\rho$, with a lower relative error. SparseGPT and Wanda follow a similar trend, closely tracking each other.

## 7 Conclusion

This paper introduced a novel approach to LLM weight pruning that directly optimizes for approximating the attention matrix. We provided theoretical guarantees for the convergence of our Gradient Descent-based algorithm to a near-optimal pruning mask solution. Preliminary results demonstrated the method’s effectiveness in maintaining model performance while reducing computational costs. This work establishes a new theoretical foundation for pruning algorithm design in LLMs, potentially enabling more efficient inference on resource-constrained devices.

## Acknowledgement

Research is partially supported by the National Science Foundation (NSF) Grants 2023239-DMS, CCF-2046710, and Air Force Grant FA9550-18-1-0166.

## Appendix A Preliminary

In Section A.1, we introduce some notations we use in this paper.
In Section A.2, we provide some basic facts.

### A.1 Notations

For any positive integer $n$, we use $[n]$ to denote set $\{1,2,\cdots,n\}$.
For two vectors $x\in\mathbb{R}^{n}$ and $y\in\mathbb{R}^{n}$, we use $\langle x,y\rangle$ to denote the inner product between $x,y$, i.e., $\langle x,y\rangle=\sum_{i=1}^{n}x_{i}y_{i}$.
For each $a,b\in\mathbb{R}^{n}$, we use $a\circ b\in\mathbb{R}^{n}$ to denote the Hadamard product, i.e. the $i$-th entry of $(a\circ b)$ is $a_{i}b_{i}$ for all $i\in[n]$.
We use $e_{i}$ to denote a vector where only $i$-th coordinate is $1$, and other entries are 0 0.
We use ${\bf 1}_{n}$ to denote a length-$n$ vector where all the entries are ones.
We use $\|x\|_{p}$ to denote the $\ell_{p}$ norm of a vector $x\in\mathbb{R}^{n}$, i.e. $\|x\|_{1}:=\sum_{i=1}^{n}|x_{i}|$, $\|x\|_{2}:=(\sum_{i=1}^{n}x_{i}^{2})^{1/2}$, and $\|x\|_{\infty}:=\max_{i\in[n]}|x_{i}|$.
For $A\in\mathbb{R}^{m\times n}$, let $A_{i}\in\mathbb{R}^{n}$ denote the $i$-th row and $A_{*,j}\in\mathbb{R}^{m}$ denote the $j$-th column of $A$, where $i\in[m]$ and $j\in[n]$.
For a square matrix $A$, we use $\operatorname{tr}[A]$ to denote the trace of $A$, i.e., $\operatorname{tr}[A]=\sum_{i=1}^{n}A_{i,i}$.
For two matrices $X,Y\in\mathbb{R}^{m\times n}$, the standard inner product between matrices is defined by $\langle X,Y\rangle:=\operatorname{tr}[X^{\top}Y]$.
We use $\exp(A)$ to denote a matrix where $\exp(A)_{i,j}:=\exp(A_{i,j})$ for a matrix $A\in\mathbb{R}^{n\times d}$.
For $k>n$, for any matrix $A\in\mathbb{R}^{k\times n}$, we use $\|A\|$ to denote the spectral norm of $A$, i.e. $\|A\|:=\sup_{x\in\mathbb{R}^{n}}\|Ax\|_{2}/\|x\|_{2}$.
We use $\|A\|_{\infty}$ to denote the $\ell_{\infty}$ norm of a matrix $A\in\mathbb{R}^{n\times d}$, i.e. $\|A\|_{\infty}:=\max_{i\in[n],j\in[d]}|A_{i,j}|$.
We use $\|A\|_{F}$ to denote the Frobenius norm of a matrix $A\in\mathbb{R}^{n\times d}$, i.e. $\|A\|_{F}:=\sqrt{\sum_{i\in[n]}\sum_{j\in[d]}|A_{i,j}|^{2}}$.
For a symmetric matrix $A\in\mathbb{R}^{n\times n}$, we use $A\succeq 0$ (positive semidefinite (PSD)), if for all $x\in\mathbb{R}^{n}$, we have $x^{\top}Ax\geq 0$.
We use $\lambda_{\min}(A)$ and $\lambda_{\max}(A)$ to denote the minimum and the maximum eigenvalue of the square matrix $A$, respectively.
Let $A\in\mathbb{R}^{n\times d}$. We use $a:=\operatorname{vec}(A)$ to denote a length $nd$ vector. We stack rows of $A$ into a column vector, i.e. $\operatorname{vec}(A):=[a_{1}^{\top},a_{2}^{\top},\dots,a_{n}^{\top}]^{\top}$ where $a_{i}^{\top}$ is the $i$-th row of $A$, or simply $\operatorname{vec}(A)_{j+(i-1)d}:=A_{i,j}$ for any $i\in[n],j\in[d]$.

### A.2 Facts

###### Fact A.1 (Indexing) .

Suppose we have matrices $U\in\mathbb{R}^{n\times m},V\in\mathbb{R}^{m\times d}$.
We define

$$ $\displaystyle\underbrace{X}_{n\times d}:=\underbrace{U}_{n\times m}\underbrace {V}_{m\times d}.$ $$

Then, we have the following:

- •
Indexing for one row: $X_{i}=V^{\top}U_{i}\in\mathbb{R}^{d}$, i.e. $X_{i}^{\top}=U_{i}^{\top}V$, for $i\in[n]$.
- •
Indexing for one column: $X_{*,j}=UV_{*,j}\in\mathbb{R}^{n}$ for $j\in[d]$.

###### Fact A.2 .

We have

Part 1.Suppose we have vectors $u\in\mathbb{R}^{n},v\in\mathbb{R}^{n}$. For $i\in[n]$, we define

$$ $\displaystyle x_{i}:=u_{i}v_{i}.$ $$

Then we have the following:

- •
$\underbrace{x}_{n\times 1}=\underbrace{u\circ v}_{n\times 1}=\underbrace{
\operatorname{diag}(u)}_{n\times n}\underbrace{v}_{n\times 1}=\underbrace{
\operatorname{diag}(v)}_{n\times n}\underbrace{u}_{n\times 1}$

Part 2.Suppose we have matrix $W\in\mathbb{R}^{n\times n}$, vector $u\in\mathbb{R}^{n}$. For $i\in[n]$, we define

$$ $\displaystyle X_{*,j}=W_{*,j}u_{j}.$ $$

Then we have the following:

- •
$X=W\operatorname{diag}(u)$

###### Fact A.3 (Calculus) .

We have

Part 1. (Scalar calculus)
For any $t\in\mathbb{R}$, function $f:\mathbb{R}\to\mathbb{R}$, we have

- •
$\frac{\mathrm{d}f^{n}(t)}{\mathrm{d}t}=nf^{n-1}(t)\frac{\mathrm{d}f(t)}{
\mathrm{d}t}$.

Part 2. (Vector calculus)
For any $x,y\in\mathbb{R}^{n}$, $t\in\mathbb{R}$, we have

- •
$\frac{\mathrm{d}(x\circ y)}{\mathrm{d}t}=\frac{\mathrm{d}x}{\mathrm{d}t}\circ y
+\frac{\mathrm{d}y}{\mathrm{d}t}\circ x$. (Product rule of vector Hadamard product)
- •
$\frac{\mathrm{d}\langle x,y\rangle}{\mathrm{d}t}=\langle\frac{\mathrm{d}x}{
\mathrm{d}t},y\rangle+\langle x,\frac{\mathrm{d}y}{\mathrm{d}t}\rangle$. (Product rule of inner product)
- •
$\frac{\mathrm{d}x}{\mathrm{d}x_{i}}=e_{i}$.

Part 3. (Matrix calculus)
For any $X,Y\in\mathbb{R}^{n\times m}$, $Z\in\mathbb{R}^{m\times d}$, $t\in\mathbb{R}$ which is independent of $Z$, function $f:\mathbb{R}\to\mathbb{R}^{n\times d}$, functions $f_{1}(t),f_{2}(t),\dots,f_{n}(t):\mathbb{R}\to\mathbb{R}^{n\times d}$, we have

- •
$\frac{\mathrm{d}(X\circ Y)}{\mathrm{d}t}=\frac{\mathrm{d}X}{\mathrm{d}t}\circ Y
+\frac{\mathrm{d}Y}{\mathrm{d}t}\circ X$. (Product rule of matrix Hadamard product)
- •
$\frac{\mathrm{d}\exp(f(t))}{\mathrm{d}t}=\exp(f(t))\circ\frac{\mathrm{d}f(t)}{
\mathrm{d}t}$, where $\exp(\cdot)$ is applied entry-wise.
- •
$\frac{\mathrm{d}(XZ)}{\mathrm{d}t}=\frac{\mathrm{d}X}{\mathrm{d}t}Z$.
- •
$\frac{\mathrm{d}(ZX^{\top})}{\mathrm{d}t}=Z\frac{\mathrm{d}X^{\top}}{\mathrm{d
}t}$.
- •
$\frac{\mathrm{d}}{\mathrm{d}t}\sum_{i=1}^{n}f_{i}(t)=\sum_{i=1}^{n}\frac{
\mathrm{d}f_{i}(t)}{\mathrm{d}t}$.
- •
$\underbrace{\frac{\mathrm{d}X}{\mathrm{d}X_{i,j}}}_{n\times m}=\underbrace{e_{
i}}_{n\times 1}\underbrace{e_{j}^{\top}}_{1\times m}$.

###### Fact A.4 (Basic algebra) .

Let $u\in\mathbb{R}^{n}$, $v\in\mathbb{R}^{n}$, $w\in\mathbb{R}^{n}$, $X\in\mathbb{R}^{n\times d}$, $Y\in\mathbb{R}^{n\times d}$, and $Z\in\mathbb{R}^{n\times n}$. Then, we have

- •
$\langle u,v\rangle=\langle v,u\rangle=u^{\top}v=v^{\top}u$
- •
$u\circ v=v\circ u=\operatorname{diag}(u)v=\operatorname{diag}(v)u$
- •
$\langle u,v\rangle=\langle u\circ v,{\bf 1}_{n}\rangle$
- •
$\langle u\circ v,w\rangle=\langle u\circ w,v\rangle=\langle w\circ v,u\rangle$
- •
$u^{\top}(v\circ w)=u^{\top}\operatorname{diag}(v)w$
- •
$(X\circ Y)^{\top}=X^{\top}\circ Y^{\top}$
- •
$X\circ e_{i}e_{j}^{\top}=X_{i,j}e_{i}e_{j}^{\top}$
- •
$\operatorname{diag}(u)Z\operatorname{diag}(v)=(uv^{\top})\circ Z$
- •
$XY^{\top}=\sum_{i\in[d]}X_{*,i}Y^{\top}_{*,i}$
- •
$X_{i,j}Y_{i,j}=(X\circ Y)_{i,j}$
- •
$\sum_{j\in[n]}u\circ A_{*,j}=u\circ\sum_{j\in[n]}A_{*,j}$
- •
$\|X\|_{F}^{2}=\operatorname{tr}[XX^{\top}]$
- •
$\operatorname{tr}[XY^{\top}]=\operatorname{tr}[Y^{\top}X]$
- •
$\|\operatorname{diag}(u)\|_{F}=\|u\|_{2}$

###### Fact A.5 (Norm bounds) .

For $a\in\mathbb{R}$, $u\in\mathbb{R}^{d}$, $X,Y\in\mathbb{R}^{n\times d}$, $Z\in\mathbb{R}^{d\times m}$ we have

- •
$\|aX\|_{F}=|a|\|X\|_{F}$ (absolute homogeneity).
- •
$\|X+Y\|_{F}\leq\|X\|_{F}+\|Y\|_{F}$ (triangle inequality).
- •
$|\langle X,Y\rangle|\leq\|X\|_{F}\cdot\|Y\|_{F}$ (Cauchy–Schwarz inequality).
- •
$\|X^{\top}\|_{F}=\|X\|_{F}$.
- •
$\|Xu\|_{2}\leq\|X\|\cdot\|u\|_{2}$
- •
$\|X\circ Y\|_{F}\leq\|X\|_{F}\cdot\|Y\|_{F}$.
- •
For any $i\in[n]$, $j\in[d]$, we have $|X_{i,j}|\leq\|X\|_{F}$.
- •
$\|X\|\leq\|X\|_{F}\leq\sqrt{k}\|X\|$ where $k$ is the rank of $X$.
- •
$\|Y\cdot Z\|_{F}\leq\|Y\|_{F}\cdot\|Z\|_{F}$.

###### Fact A.6 .

For matrices $A,B\in\mathbb{R}^{m\times n}$, we have

$$ $\displaystyle\|A+B\|_{F}^{2}=\|A\|_{F}^{2}+\|B\|_{F}^{2}+2\langle A,B\rangle.$ $$

###### Proof.

We can show

$$ $\displaystyle\|A+B\|_{F}^{2}=$ $\displaystyle~{}\operatorname{tr}[(A+B)^{\top}(A+B)]$ $\displaystyle=$ $\displaystyle~{}\operatorname{tr}[A^{\top}A+A^{\top}B+B^{\top}A+B^{\top}B]$ $\displaystyle=$ $\displaystyle~{}\operatorname{tr}[A^{\top}A]+\operatorname{tr}[B^{\top}B]+2 \operatorname{tr}[A^{\top}B]$ $\displaystyle=$ $\displaystyle~{}\|A\|_{F}^{2}+\|B\|_{F}^{2}+2\operatorname{tr}[A^{\top}B]$ $\displaystyle=$ $\displaystyle~{}\|A\|_{F}^{2}+\|B\|_{F}^{2}+2\langle A,B\rangle$ $$

where the first step follows from $\operatorname{tr}[A^{\top}A]=\|A\|_{F}^{2}$ for matrix $A\in\mathbb{R}^{m\times n}$, the second step follows from the basic algebra, the third follows from $\operatorname{tr}[X^{\top}Y]=\operatorname{tr}[XY^{\top}]$ for matrices $X,Y\in\mathbb{R}^{m\times n}$, the fourth step follows from $\operatorname{tr}[A^{\top}A]=\|A\|_{F}^{2}$ for matrix $A\in\mathbb{R}^{m\times n}$, and the last step follows from definition of inner product of matrices.
∎

###### Lemma A.7 .

Let $M\in\mathbb{R}^{n\times n}$. Let $X\in\mathbb{R}^{n\times n}$ be independent of $M$. We have

$$ $\displaystyle\frac{\mathrm{d}(M\circ X)}{\mathrm{d}M_{i,j}}=X_{i,j}e_{i}e_{j}^ {\top}$ $$

###### Proof.

We can show

$$ $\displaystyle\frac{\mathrm{d}(M\circ X)}{\mathrm{d}M_{i,j}}=$ $\displaystyle~{}M\circ\frac{\mathrm{d}X}{\mathrm{d}M_{i,j}}+X\circ\frac{ \mathrm{d}M}{\mathrm{d}M_{i,j}}$ $\displaystyle=$ $\displaystyle~{}X\circ\frac{\mathrm{d}M}{\mathrm{d}M_{i,j}}$ $\displaystyle=$ $\displaystyle~{}X\circ(e_{i}e_{j}^{\top})$ $\displaystyle=$ $\displaystyle~{}X_{i,j}e_{i}e_{j}^{\top}$ $$

where the first step, the second step and the third step follow from Fact A.3,
the fourth step follows from Fact A.4.

∎

###### Fact A.1 (Indexing) .

###### Fact A.2 .

###### Fact A.3 (Calculus) .

###### Fact A.4 (Basic algebra) .

###### Fact A.5 (Norm bounds) .

###### Fact A.6 .

###### Proof.

###### Lemma A.7 .

###### Proof.

## Appendix B Gradient Calculation

### B.1 Definitions

In this section, we introduce some definitions we used to compute $\frac{\mathrm{d}\mathcal{L}(M)}{\mathrm{d}M}$. First, we introduce the exponential function.

###### Definition B.1 (Exponential function $u$ , $\widetilde{u}$ ) .

If the following condition hold

- •
Let $X\in\mathbb{R}^{n\times d}$.
- •
Let $W\in\mathbb{R}^{d\times d}$.
- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $i_{0}\in[n]$.

We define $u\in\mathbb{R}^{n\times n}$ as follows

$$ $\displaystyle u:=\exp(XWX^{\top}).$ $$

We define $\widetilde{u}(M)\in\mathbb{R}^{n\times n}$ as follows

$$ $\displaystyle\widetilde{u}(M):=\exp(X(M\circ W)X^{\top}).$ $$

We define $i_{0}$-th row of $\widetilde{u}(M)$ as follows

$$ $\displaystyle\widetilde{u}(M)_{i_{0}}:=\exp(X(M\circ W)X^{\top})_{i_{0}}.$ $$

Then, we introduce the sum function.

###### Definition B.2 (Sum function of softmax $\alpha$ , $\widetilde{\alpha}$ ) .

If the following condition hold

- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $M_{c}\in\{0,1\}^{n\times n}$ be the causal attention mask defined in Definition 3.1.
- •
Let $u$, $\widetilde{u}(M)$ be defined as Definition B.1.
- •
Let $i_{0}\in[n]$.

We define $\alpha\in\mathbb{R}^{n}$ as follows

$$ $\displaystyle\alpha:=(u\circ M_{c})\cdot{\bf 1}_{n}.$ $$

We define $\widetilde{\alpha}(M)\in\mathbb{R}^{n}$ as follows

$$ $\displaystyle\widetilde{\alpha}(M):=(\widetilde{u}(M)\circ M_{c})\cdot{\bf 1}_ {n}.$ $$

We define $i_{0}$-th entry of $\widetilde{\alpha}(M)$ as follows

$$ $\displaystyle\widetilde{\alpha}(M)_{i_{0}}:=\langle(\widetilde{u}(M)\circ M_{c })_{i_{0}},{\bf 1}_{n}\rangle$ $$

Then, we introduce the Softmax probability function.

###### Definition B.3 (Softmax probability function $f$ , $\widetilde{f}$ ) .

If the following conditions hold

- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $M_{c}\in\{0,1\}^{n\times n}$ be the causal attention mask defined in Definition 3.1.
- •
Let $u$, $\widetilde{u}(M)$ be defined as Definition B.1.
- •
Let $\alpha$, $\widetilde{\alpha}(M)$ be defined as Definition B.2.
- •
Let $i_{0},j_{0}\in[n]$.

We define $f\in\mathbb{R}^{n\times n}$ for each $j\in[n]$ as follows

$$ $\displaystyle f:=\operatorname{diag}(\alpha)^{-1}(u\circ M_{c}).$ $$

We define $\widetilde{f}(M)\in\mathbb{R}^{n\times n}$ for each $j\in[n]$ as follows

$$ $\displaystyle\widetilde{f}(M):=\operatorname{diag}(\widetilde{\alpha}(M))^{-1} (\widetilde{u}(M)\circ M_{c}).$ $$

We define $i_{0}$-th row of $\widetilde{f}(M)$ as follows

$$ $\displaystyle\widetilde{f}(M)_{i_{0}}:=\widetilde{\alpha}(M)^{-1}_{i_{0}}( \widetilde{u}(M)\circ M_{c})_{i_{0}}.$ $$

We define the entry in $i_{0}$-th row, $j_{0}$-th column of $\widetilde{f}(M)$ as follows

$$ $\displaystyle\widetilde{f}(M)_{i_{0},j_{0}}:=\widetilde{\alpha}(M)^{-1}_{i_{0} }(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}.$ $$

Then, we introduce the one unit loss function.

###### Definition B.4 (One unit loss function $c$ ) .

If the following conditions hold

- •
Let $f$, $\widetilde{f}$ be defined in Definition B.3.
- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $i_{0},j_{0}\in[n]$.

We define $c(M)\in\mathbb{R}^{n\times n}$ as follows

$$ $\displaystyle c(M):=\widetilde{f}(M)-f$ $$

We define $i_{0}$-th row of $c(M)$ as follows

$$ $\displaystyle c(M)_{i_{0}}:=\widetilde{f}(M)_{i_{0}}-f_{i_{0}}$ $$

We define $j_{0}$-th column of $c(M)$ as follows

$$ $\displaystyle c(M)_{*,j_{0}}:=\widetilde{f}(M)_{*,j_{0}}-f_{*,j_{0}}$ $$

We define the entry in $i_{0}$-th row, $j_{0}$-th column of $c(M)$ as follows

$$ $\displaystyle c(M)_{i_{0},j_{0}}:=\widetilde{f}(M)_{i_{0},j_{0}}-f_{i_{0},j_{0}}$ $$

Then, we introduce the reconstruction error.

###### Definition B.5 (Reconstruction Error ${\cal L}_{\mathrm{attn}}$ ) .

If the following conditions hold

- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $c(M)$ be defined in Definition B.4.

We define ${\cal L}_{\mathrm{attn}}(M)\in\mathbb{R}$ as follows

$$ $\displaystyle{\cal L}_{\mathrm{attn}}(M):=\frac{1}{2}\|c(M)\|_{F}^{2}=\frac{1} {2}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}}^{2}.$ $$

Then, we introduce the regularization term.

###### Definition B.6 (Regularization Term ${\cal L}_{\mathrm{reg}}$ ) .

If the following conditions hold

- •
$M\in[0,1]^{d\times d}.$

We define ${\cal L}_{\mathrm{reg}}(M)\in\mathbb{R}$ as follows

$$ $\displaystyle{\cal L}_{\mathrm{reg}}(M):=\frac{1}{2}\lambda\|M\|_{F}^{2}.$ $$

Finally, we introduce the overall loss function.

###### Definition B.7 (Overall loss function ${\cal L}$ ) .

If the following conditions hold

- •
Let $M\in[0,1]^{d\times d}$.
- •
Let ${\cal L}_{\mathrm{attn}}(M)$ be defined in Definition B.5.
- •
Let ${\cal L}_{\mathrm{reg}}(M)$ be defined in Definition B.6.
- •
Let $\lambda\in\mathbb{R}_{+}$ be the regularization parameter.

We define ${\cal L}(M)$ as follows

$$ $\displaystyle{\cal L}(M):={\cal L}_{\mathrm{attn}}(M)+{\cal L}_{\mathrm{reg}}(M)$ $$

###### Definition B.1 (Exponential function u 𝑢 u italic_u , u ~ ~ 𝑢 \widetilde{u} over~ start_ARG italic_u end_ARG ) .

###### Definition B.2 (Sum function of softmax α 𝛼 \alpha italic_α , α ~ ~ 𝛼 \widetilde{\alpha} over~ start_ARG italic_α end_ARG ) .

###### Definition B.3 (Softmax probability function f 𝑓 f italic_f , f ~ ~ 𝑓 \widetilde{f} over~ start_ARG italic_f end_ARG ) .

###### Definition B.4 (One unit loss function c 𝑐 c italic_c ) .

###### Definition B.5 (Reconstruction Error ℒ attn subscript ℒ attn {\cal L}_{\mathrm{attn}} caligraphic_L start_POSTSUBSCRIPT roman_attn end_POSTSUBSCRIPT ) .

###### Definition B.6 (Regularization Term ℒ reg subscript ℒ reg {\cal L}_{\mathrm{reg}} caligraphic_L start_POSTSUBSCRIPT roman_reg end_POSTSUBSCRIPT ) .

###### Definition B.7 (Overall loss function ℒ ℒ {\cal L} caligraphic_L ) .

### B.2 Gradient for Each Row of X ⁢ ( M ∘ W ) ⁢ X ⊤ 𝑋 𝑀 𝑊 superscript 𝑋 top X(M\circ W)X^{\top} italic_X ( italic_M ∘ italic_W ) italic_X start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT

We introduce the Lemma of gradient for each row of $X(M\circ W)X^{\top}$.

###### Lemma B.8 .

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, we have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}(X(M\circ W)X^{\top})_{i_{0}}}{ \mathrm{d}M_{i_{1},j_{1}}}}_{n\times 1}=\underbrace{W_{i_{1},j_{1}}}_{\mathrm{ scalar}}\underbrace{X_{i_{0},i_{1}}}_{\mathrm{scalar}}\underbrace{X_{*,j_{1}}} _{n\times 1}$ $$

###### Proof.

We can simplify the derivative expression

$$ $\displaystyle\frac{\mathrm{d}(X(M\circ W)X^{\top})_{i_{0}}}{\mathrm{d}M_{i_{1} ,j_{1}}}=$ $\displaystyle~{}\frac{\mathrm{d}X(X(M\circ W))_{i_{0}}}{\mathrm{d}M_{i_{1},j_{ 1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}X(M\circ W)^{\top}X_{i_{0}}}{\mathrm{d}M_{i_{1 },j_{1}}}$ $\displaystyle=$ $\displaystyle~{}X\frac{\mathrm{d}(M\circ W)^{\top}}{\mathrm{d}M_{i_{1},j_{1}}} X_{i_{0}}$ (1) $$

where the first and second step follows from Fact A.1,
the third step follows from Fact A.3.

We further compute Eq. (B.2):

$$ $\displaystyle\frac{\mathrm{d}(M\circ W)^{\top}}{\mathrm{d}M_{i_{1},j_{1}}}=$ $\displaystyle~{}\frac{\mathrm{d}M^{\top}\circ W^{\top}}{\mathrm{d}M_{i_{1},j_{ 1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}M^{\top}\circ W^{\top}}{\mathrm{d}(M^{\top})_{ j_{1},i_{1}}}$ $\displaystyle=$ $\displaystyle~{}(W^{\top})_{j_{1},i_{1}}e_{j_{1}}e_{i_{1}}^{\top}$ $\displaystyle=$ $\displaystyle~{}W_{i_{1},j_{1}}e_{j_{1}}e_{i_{1}}^{\top}$ (2) $$

where the first follows from Fact A.4,
the second step follows from for any matrix $X$, $X_{i,j}=(X^{\top})_{j,i}$,
the third step follows from Fact A.7,
and the fourth step follows from for any matrix $X$, $X_{i,j}=(X^{\top})_{j,i}$.

Finally, we have

$$ $\displaystyle\frac{\mathrm{d}(X(M\circ W)X^{\top})_{i_{0}}}{\mathrm{d}M_{i_{1} ,j_{1}}}=$ $\displaystyle~{}XW_{i_{1},j_{1}}e_{j_{1}}e_{i_{1}}^{\top}X_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}W_{i_{1},j_{1}}(Xe_{j_{1}})(e_{i_{1}}^{\top}X_{i_{0}})$ $\displaystyle=$ $\displaystyle~{}W_{i_{1},j_{1}}X_{*,j_{1}}X_{i_{0},i_{1}}$ $$

where the first step follows from Eq. (B.2) and Eq. (B.2),
and the second step and the third step follows from basic algebra.
∎

We introduce the Lemma of gradient for each row of $\widetilde{u}(M)$.

###### Lemma B.8 .

###### Proof.

### B.3 Gradient for Each Row of u ~ ⁢ ( M ) ~ 𝑢 𝑀 \widetilde{u}(M) over~ start_ARG italic_u end_ARG ( italic_M )

###### Lemma B.9 .

If the following condition hold:

- •
Let $\widetilde{u}(M)$ be defined in Definition B.1.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, we have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}\widetilde{u}(M)_{i_{0}}}{\mathrm{d}M _{i_{1},j_{1}}}}_{n\times 1}=\underbrace{\widetilde{u}(M)_{i_{0}}}_{n\times 1} \circ(\underbrace{W_{i_{1},j_{1}}}_{\mathrm{scalar}}\underbrace{X_{i_{0},i_{1} }}_{\mathrm{scalar}}\underbrace{X_{*,j_{1}}}_{n\times 1})$ $$

###### Proof.

We have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}\widetilde{u}(M)_{i_{0}}}{\mathrm{d}M _{i_{1},j_{1}}}}_{n\times 1}=$ $\displaystyle~{}\underbrace{\frac{\mathrm{d}\exp(X(M\circ W)X^{\top})_{i_{0}}} {\mathrm{d}M_{i_{1},j_{1}}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}\exp(\underbrace{X}_{n\times d}\underbrace{(M\circ W)}_{d \times d}\underbrace{X^{\top}}_{d\times 1})_{i_{0}}\circ\underbrace{\frac{ \mathrm{d}(X(M\circ W)X^{\top})_{i_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}\underbrace{\widetilde{u}(M)_{i_{0}}}_{n\times 1}\circ \underbrace{\frac{\mathrm{d}(X(M\circ W)X^{\top})_{i_{0}}}{\mathrm{d}M_{i_{1}, j_{1}}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}\underbrace{\widetilde{u}(M)_{i_{0}}}_{n\times 1}\circ( \underbrace{W_{i_{1},j_{1}}}_{\mathrm{scalar}}\underbrace{X_{i_{0},i_{1}}}_{ \mathrm{scalar}}\underbrace{X_{*,j_{1}}}_{n\times 1})$ $$

where the first step follows from Definition B.1,
the second step follows from Fact A.3,
the third step follows from Definition B.1,
and the fourth step follows from Lemma B.8.
∎

###### Lemma B.9 .

###### Proof.

### B.4 Gradient for Each Entry of α ~ ⁢ ( M ) ~ 𝛼 𝑀 \widetilde{\alpha}(M) over~ start_ARG italic_α end_ARG ( italic_M )

We introduce the Lemma of gradient for each entry of $\widetilde{\alpha}(M)$.

###### Lemma B.10 .

If the following conditions hold:

- •
Let $\widetilde{u}(M)$ be defined in Definition B.1.
- •
Let $\widetilde{\alpha}(M)$ be defined in Definition B.2.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, we have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}\widetilde{\alpha}(M)_{i_{0}}}{ \mathrm{d}M_{i_{1},j_{1}}}}_{\mathrm{scalar}}=\langle\widetilde{u}(M)_{i_{0}} \circ(M_{c})_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $$

###### Proof.

We have

$$ $\displaystyle\underbrace{\frac{\mathrm{d}\widetilde{\alpha}(M)_{i_{0}}}{ \mathrm{d}M_{i_{1},j_{1}}}}_{\mathrm{scalar}}=$ $\displaystyle~{}\underbrace{\frac{\mathrm{d}\langle(\widetilde{u}(M)\circ M_{c })_{i_{0}},{\bf 1}_{n}\rangle}{\mathrm{d}M_{i_{1},j_{1}}}}_{\mathrm{scalar}}$ $\displaystyle=$ $\displaystyle~{}\langle\frac{\mathrm{d}(\widetilde{u}(M)\circ M_{c})_{i_{0}}}{ \mathrm{d}M_{i_{1},j_{1}}},{\bf 1}_{n}\rangle$ $\displaystyle=$ $\displaystyle~{}\langle\frac{\mathrm{d}\widetilde{u}(M)_{i_{0}}}{\mathrm{d}M_{ i_{1},j_{1}}}\circ(M_{c})_{i_{0}},{\bf 1}_{n}\rangle$ $\displaystyle=$ $\displaystyle~{}\langle\widetilde{u}(M)_{i_{0}}\circ(W_{i_{1},j_{1}}X_{i_{0},i _{1}}X_{*,j_{1}})\circ(M_{c})_{i_{0}},{\bf 1}_{n}\rangle$ $\displaystyle=$ $\displaystyle~{}\langle\widetilde{u}(M)_{i_{0}}\circ(M_{c})_{i_{0}},W_{i_{1},j _{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $$

where the first step follows from Definition B.2,
the second step follows from product rule of inner product in Fact A.3,
the third step follows from product rule of Hadamard product in Fact A.3,
the fourth step follows from Lemma B.9,
and the last step follows from Fact A.4.
∎

###### Lemma B.10 .

###### Proof.

### B.5 Gradient for Each Entry of f ~ ⁢ ( M ) ~ 𝑓 𝑀 \widetilde{f}(M) over~ start_ARG italic_f end_ARG ( italic_M )

We introduce the Lemma of gradient for each entry of $\widetilde{f}(M)$.

###### Lemma B.11 .

If the following conditions hold:

- •
Let $\widetilde{u}(M)$ be defined in Definition B.1.
- •
Let $\widetilde{\alpha}(M)$ be defined in Definition B.2.
- •
Let $\widetilde{f}(M)$ be defined in Definition B.3.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, $j_{0}\in[n]$, we have

$$ $\displaystyle\frac{\mathrm{d}\widetilde{f}(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1 },j_{1}}}=\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{j_{0} ,j_{1}}-\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W_{i_{1} ,j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $$

###### Proof.

We have

$$ $\displaystyle\frac{\mathrm{d}\widetilde{f}(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1 },j_{1}}}=$ $\displaystyle~{}\frac{\mathrm{d}\widetilde{\alpha}(M)^{-1}_{i_{0}}(\widetilde{ u}(M)\circ M_{c})_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}\widetilde{\alpha}(M)^{-1}_{i_{0}}}{\mathrm{d} M_{i_{1},j_{1}}}(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}+\frac{\mathrm{d}( \widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}} \widetilde{\alpha}(M)^{-1}_{i_{0}}$ $$

where the first step follows from Definition B.3,
and the second step follows from Fact A.3.

In the following part, we compute the two terms separately.

For the first term above, we have

$$ $\displaystyle~{}\frac{\mathrm{d}\widetilde{\alpha}(M)^{-1}_{i_{0}}}{\mathrm{d} M_{i_{1},j_{1}}}(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}$ $\displaystyle=$ $\displaystyle~{}(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}(-1)\widetilde{ \alpha}(M)^{-2}_{i_{0}}\frac{\mathrm{d}\widetilde{\alpha}(M)_{i_{0}}}{\mathrm{ d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}-(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}\langle\widetilde{ u}(M)_{i_{0}}\circ(M_{c})_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}} \rangle/\widetilde{\alpha}(M)^{2}_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}-(\widetilde{\alpha}(M)_{i_{0}}^{-1}(M_{c})_{i_{0},j_{0}} \widetilde{u}(M)_{i_{0},j_{0}})\langle\widetilde{\alpha}(M)_{i_{0}}^{-1} \widetilde{u}(M)_{i_{0}}\circ(M_{c})_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{ *,j_{1}}\rangle$ $\displaystyle=$ $\displaystyle~{}-\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}} ,W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ (4) $$

where the first step follows from Fact A.3,
the second step follows from Lemma B.10,
the third step follows from basic algebra,
and the fourth step follows from Definition B.3.

For the second term above, we have

$$ $\displaystyle~{}\frac{\mathrm{d}(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}}{ \mathrm{d}M_{i_{1},j_{1}}}\widetilde{\alpha}(M)^{-1}_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}\widetilde{u}(M)_{i_{0},j_{0}}(M_{c})_{i_{0},j _{0}}}{\mathrm{d}M_{i_{1},j_{1}}}\widetilde{\alpha}(M)^{-1}_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}(M_{c})_{i_{0},j_{0}}(\frac{\mathrm{d}\widetilde{u}(M)_{i_{0}} }{\mathrm{d}M_{i_{1},j_{1}}})_{j_{0}}\widetilde{\alpha}(M)^{-1}_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}((M_{c})_{i_{0},j_{0}}\widetilde{u}(M)_{i_{0},j_{0}}\widetilde {\alpha}(M)_{i_{0}}^{-1})W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{j_{0},j_{1}}$ $\displaystyle=$ $\displaystyle~{}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_ {j_{0},j_{1}}$ (5) $$

where the first step and the second step follow from basic algebra,
the third step follows from Lemma B.9,
and the fourth step follows from Definition B.3.

So, we have

$$ $\displaystyle\frac{\mathrm{d}\widetilde{f}(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1 },j_{1}}}=$ $\displaystyle~{}\frac{\mathrm{d}\widetilde{\alpha}(M)^{-1}_{i_{0}}}{\mathrm{d} M_{i_{1},j_{1}}}(\widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}+\frac{\mathrm{d}( \widetilde{u}(M)\circ M_{c})_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}} \widetilde{\alpha}(M)^{-1}_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_ {j_{0},j_{1}}-\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W_ {i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $$

where the first step follows from Eq. (B.5),
and the second step follows from Eq. (B.5) and Eq. (B.5).
∎

###### Lemma B.11 .

###### Proof.

### B.6 Gradient for Each Entry of c(M)

We introduce the Lemma of gradient for each entry of $c(M)$.

###### Lemma B.12 .

If the following conditions hold:

- •
Let $\widetilde{f}(M)$, $f$ be defined in Definition B.3.
- •
Let $c(M)$ be defined in Definition B.4.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, $j_{0}\in[n]$, we have

$$ $\displaystyle\frac{\mathrm{d}c(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}= \widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{j_{0},j_{1}}- \widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W_{i_{1},j_{1}}X _{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $$

###### Proof.

We have

$$ $\displaystyle\frac{\mathrm{d}c(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}=$ $\displaystyle~{}\frac{\mathrm{d}(\widetilde{f}(M)_{i_{0},j_{0}}-f_{i_{0},j_{0} })}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}\widetilde{f}(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i _{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_ {j_{0},j_{1}}-\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W_ {i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $$

where the first step follows from Definition B.4,
the second step follows from Fact A.3,
and the third step follows from Lemma B.11.
∎

###### Lemma B.12 .

###### Proof.

### B.7 Gradient for ℒ attn ⁢ ( M ) subscript ℒ attn 𝑀 {\cal L}_{\mathrm{attn}}(M) caligraphic_L start_POSTSUBSCRIPT roman_attn end_POSTSUBSCRIPT ( italic_M )

We introduce the Lemma of gradient for ${\cal L}_{\mathrm{attn}}(M)$.

###### Lemma B.13 .

If the following conditions hold:

- •
Let $\widetilde{f}(M)$ be defined in Definition B.3.
- •
Let $c(M)$ be defined in Definition B.4.
- •
Let ${\cal L}_{\mathrm{attn}}(M)$ be defined in Definition B.5.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, $j_{0}\in[n]$, we have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}_{\mathrm{attn}}(M)}{\mathrm{d}M_{i_{1},j _{1}}}=\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}B_{1}(M)+B_{2}(M)$ $$

where we have definitions:

- •
$B_{1}(M):=c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_
{0},i_{1}}X_{j_{0},j_{1}}$
- •
$B_{2}(M):=-c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f
}(M)_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$

###### Proof.

We have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}_{\mathrm{attn}}(M)}{\mathrm{d}M_{i_{1},j _{1}}}=$ $\displaystyle~{}\frac{1}{2}\frac{\mathrm{d}\|c(M)\|_{F}^{2}}{\mathrm{d}M_{i_{1 },j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{2}\frac{\mathrm{d}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n }(c(M)_{i_{0},j_{0}})^{2}}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{2}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}\frac{\mathrm{d }(c(M)_{i_{0},j_{0}})^{2}}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}}\frac{ \mathrm{d}c(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}$ $$

where the first step follows from Definition B.5,
the second step follows from the definition of Frobenius norm of matrix,
the third step follows from Fact A.3,
and the fourth step follows from Fact A.3.

Following Lemma B.12, we have

$$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}}\frac{ \mathrm{d}c(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}}( \widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{j_{0},j_{1}}- \widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W_{i_{1},j_{1}}X _{i_{0},i_{1}}X_{*,j_{1}}\rangle)$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1} }X_{i_{0},i_{1}}X_{j_{0},j_{1}}-c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0 }}\langle\widetilde{f}(M)_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $\displaystyle:=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}B_{1}(M)+B_{2}(M)$ $$

where the second step follows from basic algebra.
∎

###### Lemma B.13 .

###### Proof.

### B.8 Gradient for ℒ reg ⁢ ( M ) subscript ℒ reg 𝑀 {\cal L}_{\mathrm{reg}}(M) caligraphic_L start_POSTSUBSCRIPT roman_reg end_POSTSUBSCRIPT ( italic_M )

We introduce the Lemma of gradient for ${\cal L}_{\mathrm{reg}}(M)$.

###### Lemma B.14 .

If the following condition hold:

- •
Let ${\cal L}_{\mathrm{reg}}(M)$ be defined in Definition B.6.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, we have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}_{\mathrm{reg}}(M)}{\mathrm{d}M_{i_{1},j_ {1}}}=B_{3}(M)$ $$

where we have definition:

- •
$B_{3}(M):=\lambda M_{i_{1},j_{1}}$

###### Proof.

We have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}_{\mathrm{reg}}(M)}{\mathrm{d}M_{i_{1},j_ {1}}}=$ $\displaystyle~{}\frac{1}{2}\lambda\frac{\mathrm{d}\|M\|_{F}^{2}}{\mathrm{d}M_{ i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{2}\lambda(\frac{\mathrm{d}}{\mathrm{d}M_{i_{1},j_{1}} }\sum_{i_{0}=1}^{d}\sum_{j_{0}=1}^{d}M_{i_{0},j_{0}}^{2})$ $\displaystyle=$ $\displaystyle~{}\lambda M_{i_{1},j_{1}}$ $\displaystyle:=$ $\displaystyle~{}B_{3}(M)$ $$

where the first step follows from Definition B.6,
the second step follows from the definition of Frobenius norm of matrix,
and the third step follows from Fact A.3.
∎

###### Lemma B.14 .

###### Proof.

### B.9 Gradient for ℒ ⁢ ( M ) ℒ 𝑀 {\cal L}(M) caligraphic_L ( italic_M )

We introduce the Lemma of gradient for ${\cal L}(M)$.

###### Lemma B.15 .

If the following conditions hold:

- •
Let $\widetilde{u}(M)$ be defined in Definition B.1.
- •
Let $\widetilde{\alpha}(M)$ be defined in Definition B.2.
- •
Let $\widetilde{f}(M)$ be defined in Definition B.3.
- •
Let ${\cal L}_{\mathrm{attn}}(M)$ be defined in Definition B.5.
- •
Let ${\cal L}_{\mathrm{reg}}(M)$ be defined in Definition B.6.
- •
Let ${\cal L}(M)$ be defined in Definition B.7.

Let $i_{1}\in[d]$, $j_{1}\in[d]$, $i_{0}\in[n]$, $j_{0}\in[n]$, we have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}(M)}{\mathrm{d}M_{i_{1},j_{1}}}=\sum_{i_{ 0}=1}^{n}\sum_{j_{0}=1}^{n}(B_{1}(M)+B_{2}(M))+B_{3}(M)$ $$

where we have definitions:

- •
$B_{1}(M):=c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_
{0},i_{1}}X_{j_{0},j_{1}}$
- •
$B_{2}(M):=-c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f
}(M)_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$
- •
$B_{3}(M):=\lambda M_{i_{1},j_{1}}$

###### Proof.

$$ $\displaystyle\frac{\mathrm{d}{\cal L}(M)}{\mathrm{d}M_{i_{1},j_{1}}}=$ $\displaystyle~{}\frac{\mathrm{d}{\cal L}_{\mathrm{attn}}(M)+{\cal L}_{\mathrm{ reg}}(M)}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\frac{\mathrm{d}{\cal L}_{\mathrm{attn}}(M)}{\mathrm{d}M_{i_{1 },j_{1}}}+\frac{\mathrm{d}{\cal L}_{\mathrm{reg}}(M)}{\mathrm{d}M_{i_{1},j_{1}}}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}(B_{1}(M)+B_{2}(M))+B_{3}(M)$ $$

where the first step follows from Definition B.7,
the second step follows from Fact A.3,
and the third step follows from Lemma B.13 and Lemma B.14.
∎

###### Lemma B.15 .

###### Proof.

## Appendix C Matrix Form

### C.1 Matrix Form of B ⁢ ( M ) 𝐵 𝑀 B(M) italic_B ( italic_M )

Given the matrix form, we define $p$ to simplify the calculation.

###### Definition C.1 .

If the following conditions hold

- •
Let $X\in\mathbb{R}^{n\times d}$.
- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $W\in\mathbb{R}^{d\times d}$.
- •
Let $c(M)$ be defined in Definition B.4.
- •
Let $\widetilde{f}(M)$ be defined in Definition B.3.

We define $p_{1}$ as follows

$$ $\displaystyle p_{1}:=c(M)\circ\widetilde{f}(M)$ $$

We define the $j_{0}$-th column of $p_{1}$ as follows

$$ $\displaystyle(p_{1})_{*,j_{0}}:=(c(M)\circ\widetilde{f}(M))_{*,j_{0}}$ $$

We define $p_{2}$ as follows

$$ $\displaystyle p_{2}:=\operatorname{diag}(p_{1}\cdot{\bf 1}_{n})\widetilde{f}(M)$ $$

We define the $i_{0}$-th row of $p_{2}$ as follows

$$ $\displaystyle(p_{2})_{i_{0}}:={\bf 1}_{n}^{\top}(p_{1})_{i_{0}}\widetilde{f}(M )_{i_{0}}=\widetilde{f}(M)_{i_{0}}c(M)_{i_{0}}^{\top}\widetilde{f}(M)_{i_{0}}$ $$

We define $p$ as follows

$$ $\displaystyle p:=p_{1}-p_{2}=c(M)\circ\widetilde{f}(M)-\operatorname{diag}((c( M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n})\widetilde{f}(M)$ $$

We introduce the matrix view of $B_{1}(M)$ and its summation.

###### Lemma C.2 (Matrix view of $B_{1}(M)$ ) .

If we have the below conditions,

- •
Let $B_{1}(M,i_{1},j_{1}):=c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1}
,j_{1}}X_{i_{0},i_{1}}X_{j_{0},j_{1}}$, which is defined in Lemma B.15
- •
We define $C_{1}(M)\in\mathbb{R}^{d\times d}$. For all $i_{1},j_{1}\in[d]$, let $C_{1}(i_{1},j_{1})$ denote the $(i_{1},j_{1})$-th entry of $C_{1}(M)$. We define $C_{1}(i_{1},j_{1})=B_{1}(M,i_{1},j_{1})$.

Then, we can show that

- •
Part 1. For $i_{0},j_{0}\in[n]$
$\displaystyle C_{1}(M)=\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j
_{0}}}_{\mathrm{scalar}}(W\circ(X_{i_{0}}X^{\top}_{j_{0}}))$
- •
Part 2.
$\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}C_{1}(M)=W\circ(X^{\top}p_{1}X)$

###### Proof.

Part 1. We have

$$ $\displaystyle C_{1}(i_{1},j_{1})=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1} }X_{i_{0},i_{1}}X_{j_{0},j_{1}}$ $\displaystyle=$ $\displaystyle~{}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}}_ {\mathrm{scalar}}(\underbrace{X_{i_{0}}}_{d\times 1})_{i_{1}}(\underbrace{W_{* ,j_{1}}}_{d\times 1})_{i_{1}}\underbrace{(X_{j_{0}})_{j_{1}}}_{\mathrm{scalar}}$ $\displaystyle=$ $\displaystyle~{}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}}_ {\mathrm{scalar}}(\underbrace{\operatorname{diag}(X_{i_{0}})W_{*,j_{1}}}_{d \times 1})_{i_{1}}\underbrace{(X_{j_{0}})_{j_{1}}}_{\mathrm{scalar}}$ $$

where the first step follows from the definition of $C_{1}$,
the second step follows from Fact A.1,
and the third step follows from Fact A.2.

Following from Fact A.1, we can get $j_{1}$-th column of $C_{1}$

$$ $\displaystyle C_{1}(*,j_{1})=$ $\displaystyle~{}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}}_ {\mathrm{scalar}}\underbrace{\operatorname{diag}(X_{i_{0}})}_{d\times d} \underbrace{W_{*,j_{1}}}_{d\times 1}\underbrace{(X_{j_{0}})_{j_{1}}}_{\mathrm{ scalar}}$ $\displaystyle=$ $\displaystyle~{}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}}_ {\mathrm{scalar}}\underbrace{\operatorname{diag}(X_{i_{0}})}_{d\times d}( \underbrace{W}_{d\times d}\underbrace{\operatorname{diag}(X_{j_{0}})}_{d\times d })_{*,j_{1}}$ $$

where the second step follows from Fact A.2.

Following from Fact A.1, we can get $C_{1}(M)$

$$ $\displaystyle C_{1}(M)=$ $\displaystyle~{}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}}_ {\mathrm{scalar}}\underbrace{\operatorname{diag}(X_{i_{0}})}_{d\times d} \underbrace{W}_{d\times d}\underbrace{\operatorname{diag}(X_{j_{0}})}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}}_ {\mathrm{scalar}}(W\circ(X_{i_{0}}X^{\top}_{j_{0}}))$ (6) $$

where the second step follows from Fact A.4.

Part 2. We further compute the summation of $C_{1}(M)$.

$$ $\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}C_{1}(M)=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}\underbrace{c(M)_{i_{0},j_ {0}}\widetilde{f}(M)_{i_{0},j_{0}}}_{\mathrm{scalar}}(W\circ X_{i_{0}}X^{\top} _{j_{0}})$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}(W\circ(c(M)_{i_{0},j_{0}} \widetilde{f}(M)_{i_{0},j_{0}}X_{i_{0}}X^{\top}_{j_{0}}))$ $\displaystyle=$ $\displaystyle~{}W\circ\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}} \widetilde{f}(M)_{i_{0},j_{0}}X_{i_{0}}X^{\top}_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}W\circ\sum_{j_{0}=1}^{n}\sum_{i_{0}=1}^{n}((p_{1})_{*,j_{0}})_ {i_{0}}X_{i_{0}}X^{\top}_{j_{0}}$ $$

where the first step follows from Eq. (C.1),
the second step follows from basic algebra,
the third step follows from Fact A.4,
and the fourth step follows from Definition C.1.

Then following from Fact A.2, we have

$$ $\displaystyle~{}W\circ\sum_{j_{0}=1}^{n}\sum_{i_{0}=1}^{n}((p_{1})_{*,j_{0}})_ {i_{0}}X_{i_{0}}X^{\top}_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}W\circ\sum_{j_{0}=1}^{n}X^{\top}(p_{1})_{*,j_{0}}X^{\top}_{j_{ 0}}$ $\displaystyle=$ $\displaystyle~{}W\circ(X^{\top}p_{1}X)$ $$

∎

We introduce the matrix view of $B_{2}(M)$ and its summation.

###### Lemma C.3 (Matrix view of $B_{2}(M)$ ) .

If we have the below conditions,

- •
Let $B_{2}(M,i_{1},j_{1}):=-c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\langle
\widetilde{f}(M)_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ be defined in Lemma B.15.
- •
We define $C_{2}(M)\in\mathbb{R}^{d\times d}$. For all $i_{1},j_{1}\in[d]$, let $C_{2}(i_{1},j_{1})$ denote the $(i_{1},j_{1})$-th entry of $C_{2}(M)$. We define $C_{2}(i_{1},j_{1})=B_{2}(M,i_{1},j_{1})$.

Then, we can show that

- •
Part 1. For $i_{0},j_{0}\in[n]$
$\displaystyle C_{2}(M)=$
$\displaystyle~{}-c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{
\operatorname{diag}(X_{i_{0}})}_{d\times d}\underbrace{W}_{d\times d}
\underbrace{\operatorname{diag}(X^{\top}\widetilde{f}(M)_{i_{0}})}_{d\times d}$
- •
Part 2.
$\displaystyle\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}C_{2}(M)=-W\circ(X^{\top}p_{2
}X)$

###### Proof.

Part 1.
We have

$$ $\displaystyle-C_{2}(i_{1},j_{1})=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\langle \widetilde{f}(M)_{i_{0}},W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \widetilde{f}(M)_{i_{0}}^{\top}}_{1\times n}\underbrace{X_{*,j_{1}}W_{i_{1},j_ {1}}X_{i_{0},i_{1}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \widetilde{f}(M)_{i_{0}}^{\top}}_{1\times n}\underbrace{X_{*,j_{1}}(X_{i_{0}}) _{i_{1}}(W_{*,j_{1}})_{i_{1}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \widetilde{f}(M)_{i_{0}}^{\top}}_{1\times n}\underbrace{X_{*,j_{1}}( \operatorname{diag}(X_{i_{0}})W_{*,j_{1}})_{i_{1}}}_{n\times 1}$ $$

where the first step follows from the definition of $C_{2}$,
the second step, the third step and the fourth step follows from Fact A.4.

Following from Fact A.1, we can get $j_{1}$-th column of $C_{2}$

$$ $\displaystyle-C_{2}(*,j_{1})=$ $\displaystyle~{}\underbrace{\operatorname{diag}(X_{i_{0}})}_{d\times d} \underbrace{W_{*,j_{1}}}_{d\times 1}\underbrace{c(M)_{i_{0},j_{0}}\widetilde{f }(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0}}^{\top}}_{1\times n}\underbrace{X_{* ,j_{1}}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \operatorname{diag}(X_{i_{0}})}_{d\times d}\underbrace{W_{*,j_{1}}}_{d\times 1 }\underbrace{\widetilde{f}(M)_{i_{0}}^{\top}}_{1\times n}\underbrace{X_{*,j_{1 }}}_{n\times 1}$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \operatorname{diag}(X_{i_{0}})}_{d\times d}\underbrace{W_{*,j_{1}}}_{d\times 1 }\underbrace{(X^{\top}\widetilde{f}(M)_{i_{0}})_{j_{1}}}_{\mathrm{scalar}}$ $\displaystyle=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \operatorname{diag}(X_{i_{0}})}_{d\times d}(\underbrace{W\operatorname{diag}(X ^{\top}\widetilde{f}(M)_{i_{0}})}_{d\times d})_{*,j_{1}}$ $$

where the second step and the fourth step follows from Fact A.4,
and the third step follows from Fact A.1.

Following from Fact A.1, we can get $C_{2}$.

$$ $\displaystyle-C_{2}(M)=$ $\displaystyle~{}c(M)_{i_{0},j_{0}}\widetilde{f}(M)_{i_{0},j_{0}}\underbrace{ \operatorname{diag}(X_{i_{0}})}_{d\times d}\underbrace{W}_{d\times d} \underbrace{\operatorname{diag}(X^{\top}\widetilde{f}(M)_{i_{0}})}_{d\times d}$ (7) $$

Part 2. We further compute the summation of $C_{2}$

$$ $\displaystyle-\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}C_{2}(M)=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}} \widetilde{f}(M)_{i_{0},j_{0}}\underbrace{\operatorname{diag}(X_{i_{0}})}_{d \times d}\underbrace{W}_{d\times d}\underbrace{\operatorname{diag}(X^{\top} \widetilde{f}(M)_{i_{0}})}_{d\times d}$ $\displaystyle=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}c(M)_{i_{0},j_{0}} \widetilde{f}(M)_{i_{0},j_{0}}((X_{i_{0}}\widetilde{f}(M)_{i_{0}}^{\top}X) \circ W)$ $\displaystyle=$ $\displaystyle~{}W\circ\sum_{i_{0}=1}^{n}(X_{i_{0}}\widetilde{f}(M)_{i_{0}}^{ \top}X)\sum_{j_{0}=1}^{n}((p_{1})_{i_{0}})_{j_{0}}$ $$

where the first step follows from Eq. (7),
the second step and the third step follows from Fact A.4.

Following from Fact A.2, we have

$$ $\displaystyle~{}W\circ\sum_{i_{0}=1}^{n}(X_{i_{0}}\widetilde{f}(M)_{i_{0}}^{ \top}X)\sum_{j_{0}=1}^{n}((p_{1})_{i_{0}})_{j_{0}}$ $\displaystyle=$ $\displaystyle~{}W\circ\sum_{i_{0}=1}^{n}(X_{i_{0}}\widetilde{f}(M)_{i_{0}}^{ \top}X){\bf 1}_{n}^{\top}(p_{1})_{i_{0}}$ $\displaystyle=$ $\displaystyle~{}W\circ(X^{\top}\operatorname{diag}(p_{1}\cdot{\bf 1}_{n}) \widetilde{f}(M)X)$ $\displaystyle=$ $\displaystyle~{}W\circ(X^{\top}p_{2}X)$ $$

where the third step follows from Definition C.1.
∎

We introduce the matrix view of $B_{3}(M)$.

###### Lemma C.4 (Matrix view of $B_{3}(M)$ ) .

If the following conditions hold

- •
Let $B_{3}(M,i_{1},j_{1}):=\lambda M_{i_{1},j_{1}}$ be defined in Lemma B.15.
- •
We define $C_{3}(M)\in\mathbb{R}^{d\times d}$. For all $i_{1},j_{1}\in[d]$, let $C_{3}(i_{1},j_{1})$ denote the $(i_{1},j_{1})$-th entry of $C_{3}(M)$. We define $C_{3}(i_{1},j_{1})=B_{3}(M,i_{1},j_{1})$.

We can show that

$$ $\displaystyle C_{3}(M)=\lambda M.$ $$

###### Proof.

The proof is straightforward. By the definition of $C_{3}(M)$, for all $i_{1},j_{1}\in[d]$, the $(i_{1},j_{1})$-th entry of $C_{3}(M)$ is given by $C_{3}(i_{1},j_{1})=B_{3}(M,i_{1},j_{1})=\lambda M_{i_{1},j_{1}}$. Thus, the entire matrix $C_{3}(M)$ has entries that correspond to those of $\lambda M$. Therefore, we can conclude that $C_{3}(M)=\lambda M$ as required.
∎

###### Definition C.1 .

###### Lemma C.2 (Matrix view of B 1 ⁢ ( M ) subscript 𝐵 1 𝑀 B_{1}(M) italic_B start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT ( italic_M ) ) .

###### Proof.

###### Lemma C.3 (Matrix view of B 2 ⁢ ( M ) subscript 𝐵 2 𝑀 B_{2}(M) italic_B start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT ( italic_M ) ) .

###### Proof.

###### Lemma C.4 (Matrix view of B 3 ⁢ ( M ) subscript 𝐵 3 𝑀 B_{3}(M) italic_B start_POSTSUBSCRIPT 3 end_POSTSUBSCRIPT ( italic_M ) ) .

###### Proof.

### C.2 Matrix Form of d d ⁢ M ⁢ ℒ ⁢ ( M ) d d 𝑀 ℒ 𝑀 \frac{\mathrm{d}}{\mathrm{d}M}{\cal L}(M) divide start_ARG roman_d end_ARG start_ARG roman_d italic_M end_ARG caligraphic_L ( italic_M )

We introduce the matrix form of overall loss function.

###### Theorem C.5 (Close form of gradient, formal version of Theorem 5.3 ) .

If the following conditions hold

- •
Let ${\cal L}(M)$ be defined in Definition B.7.
- •
Let $p$ be defined in Definition C.1.
- •
Let $X\in\mathbb{R}^{n\times d}$.
- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $W\in\mathbb{R}^{d\times d}$.

We can show that

$$ $\displaystyle\frac{\mathrm{d}{\cal L}(M)}{\mathrm{d}M}=W\circ(X^{\top}pX)+ \lambda M.$ $$

###### Proof.

We have

$$ $\displaystyle\frac{\mathrm{d}{\cal L}(M)}{\mathrm{d}M}=$ $\displaystyle~{}\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{n}(C_{1}(M)+C_{2}(M))+C_{3}(M)$ $\displaystyle=$ $\displaystyle~{}W\circ(X^{\top}p_{1}X)-W\circ(X^{\top}p_{2}X)+\lambda M$ $\displaystyle=$ $\displaystyle~{}W\circ(X^{\top}(p_{1}-p_{2})X)+\lambda M$ $\displaystyle=$ $\displaystyle~{}W\circ(X^{\top}pX)+\lambda M$ $$

where the first step follows from Lemma B.15,
the second step follows from Lemma C.2, Lemma C.3, and Lemma C.4,
the third step follows from basic algebra,
and the fourth step follows from Definition C.1.
∎

###### Theorem C.5 (Close form of gradient, formal version of Theorem 5.3 ) .

###### Proof.

## Appendix D Bounds for Basic Functions

### D.1 Basic Assumptions

Here we introduce our bounded parameters assumption.

###### Assumption D.1 (Bounded parameters) .

We assume the following conditions

- •
Let $R$ be some fixed constant satisfies $R>1$.
- •
Let $X\in\mathbb{R}^{n\times d},W\in\mathbb{R}^{d\times d}$. We have $\|X\|_{F}\leq R$ and $\|W\|_{F}\leq R$.

Here we present the lemma of bounds for $M$ and $M_{c}$.

###### Lemma D.2 (Bounds for $M$ and $M_{c}$ ) .

Let $M\in[0,1]^{d\times d}$ and $M_{c}\in\{0,1\}^{n\times n}$ be the causal attention mask defined in Definition 3.1. For $M$, we have

$$ $\displaystyle\|M\|_{F}\leq d$ $$

For $M_{c}$, we have

$$ $\displaystyle\|M_{c}\|_{F}\leq n$ $$

###### Proof.

This Lemma simply follows from the definition of Frobenius norm, given that the max value of each entry in $M$ and $M_{c}$ is $1$.
∎

###### Assumption D.1 (Bounded parameters) .

###### Lemma D.2 (Bounds for M 𝑀 M italic_M and M c subscript 𝑀 𝑐 M_{c} italic_M start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT ) .

###### Proof.

### D.2 Bounds for Basic Functions

We first introduce the lemma of bounds for basic function.

###### Lemma D.3 .

Under Assumption D.1, for all $i_{0}\in[n]$, $j_{0}\in[n]$, $i_{1}\in[d]$, $j_{1}\in[d]$, we have the following bounds

- •
Part 1.
$\displaystyle\|\widetilde{f}(M)\|_{F}\leq\sqrt{n}$
- •
Part 2.
$\displaystyle\|c(M)\|_{F}\leq 2\sqrt{n}$
- •
Part 3.
$\displaystyle\|(c(M)\circ\widetilde{f}(M))\|_{F}\leq 2\sqrt{n}$
- •
Part 4.
$\displaystyle|\widetilde{f}(M)_{i_{0},j_{0}}|\leq 1$
- •
Part 5.
$\displaystyle|W_{i_{1},j_{1}}|\leq R$
- •
Part 6.
$\displaystyle|X_{i_{0},i_{1}}|\leq R$
- •
Part 7.
$\displaystyle\|\widetilde{f}(M)_{i_{0}}\|_{2}\leq 1$
- •
Part 8.
$\displaystyle|\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{j
_{0},j_{1}}|\leq R^{3}$
- •
Part 9.
$\displaystyle|\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W_
{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle|\leq R^{3}$
- •
Part 10.
$\displaystyle\|\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n}
)\|_{F}\leq 2n$

###### Proof.

Proof of Part 1.
Each entry in $\widetilde{f}(M)$ present a probability, thus for $i_{0}\in[n]$, $j_{0}\in[n]$, we have

$$ $\displaystyle 0\leq\widetilde{f}(M)_{i_{0},j_{0}}\leq 1.$ $$

For any $i_{0}$-th row of $\widetilde{f}(M)$, following from the definition of Softmax function, we know

$$ $\displaystyle\sum_{j_{0}=1}^{n}\widetilde{f}(M)_{i_{0},j_{0}}=1.$ $$

So we have

$$ $\displaystyle\sum_{j_{0}=1}^{n}\widetilde{f}(M)_{i_{0},j_{0}}^{2}\leq 1$ $$

which follows from $\widetilde{f}(M)_{i_{0},j_{0}}\leq(\widetilde{f}(M)_{i_{0},j_{0}})^{2}$.
Then, we can show

$$ $\displaystyle\|\widetilde{f}(M)\|_{F}=\sqrt{\sum_{i_{0}=1}^{n}\sum_{j_{0}=1}^{ n}\widetilde{f}(M)_{i_{0},j_{0}}^{2}}\leq\sqrt{n}$ $$

Proof of Part 2.
Following from Part 1., we can show

$$ $\displaystyle\|\widetilde{f}(M)\|_{F}\leq\sqrt{n}$ $$

and

$$ $\displaystyle\|f\|_{F}\leq\sqrt{n}.$ $$

Then we have

$$ $\displaystyle\|c(M)\|_{F}=$ $\displaystyle~{}\|\widetilde{f}(M)-f\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|\widetilde{f}(M)\|_{F}+\|f\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}2\sqrt{n}$ $$

where the first step follows from Definition B.4,
the second step follows from triangle inequality.

Proof of Part 3.
We have $0\leq\widetilde{f}(M)_{i_{0},j_{0}}\leq 1$, so we have

$$ $\displaystyle\|(c(M)\circ\widetilde{f}(M))\|_{F}\leq$ $\displaystyle~{}\|c(M)\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}2\sqrt{n}$ $$

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 D.1 and Fact A.5.

Proof of Part 6.
The proof simply follows from Assumption D.1 and Fact A.5.

Proof of Part 7.
See Proof of Part 1..

Proof of Part 8.
The proof simply follows from Part 4., Part 5., Part 6. and Part 7..

Proof of Part 9.
The proof simply follows from Part 4., Part 5., Part 6. and Part 7..

Proof of Part 10.
We have

$$ $\displaystyle\|\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n} )\|_{F}=$ $\displaystyle~{}\|(c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|{\bf 1}_{n}\|_{2}\|(c(M)\circ\widetilde{f}(M))\|_{F}$ $\displaystyle=$ $\displaystyle~{}\sqrt{n}\cdot 2\sqrt{n}$ $\displaystyle=$ $\displaystyle~{}2n$ $$

where the first step follows from Fact A.4
the second step follows from Fact A.5,
the third step follows from Part 3.,
and the last step follows from simple algebra.
∎

###### Lemma D.3 .

###### Proof.

### D.3 Bounds for Gradient of f ~ ⁢ ( M ) ~ 𝑓 𝑀 \widetilde{f}(M) over~ start_ARG italic_f end_ARG ( italic_M )

We introduce the lemma of bounds for gradient of $\widetilde{f}(M)$.

###### Lemma D.4 .

If the following conditions hold

- •
Let $\widetilde{f}(M)$ be defined in Definition B.3.
- •
Assumption D.1 holds.

Then we have

$$ $\displaystyle\|\frac{\mathrm{d}\operatorname{vec}(\widetilde{f}(M))}{\mathrm{d }\operatorname{vec}(M)}\|_{F}\leq 2dnR^{3}$ $$

###### Proof.

We have

$$ $\displaystyle~{}|\frac{\mathrm{d}f(M)_{i_{0},j_{0}}}{\mathrm{d}M_{i_{1},j_{1}}}|$ $\displaystyle=$ $\displaystyle~{}|\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X _{j_{0},j_{1}}-\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}},W _{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle|$ $\displaystyle\leq$ $\displaystyle~{}|\widetilde{f}(M)_{i_{0},j_{0}}W_{i_{1},j_{1}}X_{i_{0},i_{1}}X _{j_{0},j_{1}}|+|\widetilde{f}(M)_{i_{0},j_{0}}\langle\widetilde{f}(M)_{i_{0}} ,W_{i_{1},j_{1}}X_{i_{0},i_{1}}X_{*,j_{1}}\rangle|$ $\displaystyle\leq$ $\displaystyle~{}2R^{3}$ $$

For $\frac{\mathrm{d}\operatorname{vec}(\widetilde{f}(M))}{\mathrm{d}\operatorname{
vec}(M)}$, we can show

$$ $\displaystyle\|\frac{\mathrm{d}\operatorname{vec}(\widetilde{f}(M))}{\mathrm{d }\operatorname{vec}(M)}\|_{F}=$ $\displaystyle~{}\sqrt{\sum_{i_{2}=1}^{n^{2}}\sum_{j_{2}=1}^{d^{2}}|\frac{ \mathrm{d}\operatorname{vec}(\widetilde{f}(M))_{i_{0}}}{\mathrm{d} \operatorname{vec}(M)_{j_{0}}}|}$ $\displaystyle\leq$ $\displaystyle~{}2ndR^{3}$ $$

∎

###### Lemma D.4 .

###### Proof.

## Appendix E Lipschitz of Gradient

### E.1 Useful Facts

Here we introduce the fact of mean value theorem for matrix function.

###### Fact E.1 (Mean value theorem for matrix function, Fact C.6 in [ 61 ] ) .

If the following conditions hold

- •
Let $X,Y\in C\subset\mathbb{R}^{d\times d}$ where $C$ is an open convex domain.
- •
Let $g(X):C\to\mathbb{R}^{n\times n}$ be a differentiable matrix function on $C$.
- •
Let $\|\frac{\mathrm{d}\operatorname{vec}(g(X))}{\mathrm{d}\operatorname{vec}(X)}\|
_{F}\leq R$ for all $x\in C$.

We have

$$ $\displaystyle\|g(Y)-g(X)\|_{F}\leq R\|Y-X\|_{F}.$ $$

###### Proof.

For the convenience of proof, we define $x$ and $y$ as follows:

- •
$x:=\operatorname{vec}(X)$ and $y:=\operatorname{vec}(Y)$.
- •
$h(x):=\operatorname{vec}(g(X))$ and $h(y):=\operatorname{vec}(g(Y))$.
- •
$h^{\prime}(a)$ denotes a matrix which its $(i,j)$-th term is $\frac{\mathrm{d}h(a)_{j}}{\mathrm{d}a_{i}}$.

Assume we have $1$-variable function $\gamma(c)=f(x+c(y-x))$, we can apply Mean Value Theorem:

$$ $\displaystyle f(y)-f(x)=\gamma(1)-\gamma(0)=\gamma^{\prime}(t)(1-0)=\nabla f(x +t(y-x))^{\top}(y-x)$ (8) $$

where $t\in[0,1]$.
Let $G(c):=(h(y)-h(x))^{\top}h(c)$, we have

$$ $\displaystyle\|g(Y)-g(X)\|_{F}^{2}=$ $\displaystyle~{}G(y)-G(x)$ $\displaystyle=$ $\displaystyle~{}\nabla G(x+t(y-x))^{\top}(y-x)$ $\displaystyle=$ $\displaystyle~{}(\underbrace{h^{\prime}(x+t(y-x))}_{d^{2}\times n^{2}}\cdot \underbrace{h(y)-h(x)}_{n^{2}\times 1})^{\top}\cdot(y-x)$ $\displaystyle\leq$ $\displaystyle~{}\|h^{\prime}(x+t(y-x))\|\cdot\|h(y)-h(x)\|_{2}\cdot\|y-x\|_{2}$ $$

where the second step follows from Eq. (8),
the third step follows from chain rule,
the fourth step follows from Cauchy-Schwartz inequality.

By definition of matrix Frobenius norm and vector $\ell_{2}$ norm, we have

$$ $\displaystyle\|g(Y)-g(X)\|_{F}=\|h(y)-h(x)\|_{2}$ $$

and

$$ $\displaystyle\|Y-X\|_{F}=\|y-x\|_{2}$ $$

so, we can show

$$ $\displaystyle\|g(Y)-g(X)\|_{F}\leq R\|Y-X\|_{F}$ $$

which follows from $\|\frac{\mathrm{d}\operatorname{vec}(g(X))}{\mathrm{d}\operatorname{vec}(X)}\|
_{F}\leq R$ for all $x\in C$.
∎

Here we introduce the fact of Lipschitz for product of functions.

###### Fact E.2 (Lipschitz for product of functions, Fact H.2 in [ 22 ] ) .

Under following conditions

- •
Let $\{f_{i}(x)\}_{i=1}^{n}$ be a sequence of function with same domain and range.
- •
For each $i\in[n]$, we have
–
$f_{i}(x)$ is bounded: $\forall x$, $\|f_{i}(x)\|_{F}\leq R_{i}$ with $R_{i}\geq 1$.
–
$f_{i}(x)$ is Lipschitz continuous: $\forall x,y$, $\|f_{i}(x)-f_{i}(y)\|_{F}\leq L_{i}\|x-y\|_{F}$.

Then we have

$$ $\displaystyle\|\prod_{i=1}^{n}f_{i}(x)-\prod_{i=1}^{n}f_{i}(y)\|_{F}\leq 2^{n- 1}\cdot\max_{i\in[n]}\{L_{i}\}\cdot(\prod_{i=1}^{n}R_{i})\cdot\|x-y\|_{F}$ $$

###### Fact E.1 (Mean value theorem for matrix function, Fact C.6 in [ 61 ] ) .

###### Proof.

###### Fact E.2 (Lipschitz for product of functions, Fact H.2 in [ 22 ] ) .

### E.2 Lipschitz of f ~ ⁢ ( M ) ~ 𝑓 𝑀 \widetilde{f}(M) over~ start_ARG italic_f end_ARG ( italic_M )

We introduce the lemma about Lipschitz of $\widetilde{f}(M)$.

###### Lemma E.3 (Lipschitz of $\widetilde{f}(M)$ ) .

Under the following conditions

- •
Assumption D.1 holds.
- •
Let $\widetilde{f}(M)$ be defined as Definition B.3.

For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|\widetilde{f}(M)-\widetilde{f}(\widetilde{M})\|_{F}\leq 2dnR^{3 }\|M-\widetilde{M}\|_{F}$ $$

###### Proof.

We have

$$ $\displaystyle\|\widetilde{f}(M)-\widetilde{f}(\widetilde{M})\|_{F}\leq$ $\displaystyle~{}\|\nabla\widetilde{f}(M)\|_{F}\cdot\|M-\widetilde{M}\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}2dnR^{3}\cdot\|M-\widetilde{M}\|_{F}$ $$

where the first step follows from Fact E.1,
the second step follows from Lemma D.4.
∎

###### Lemma E.3 (Lipschitz of f ~ ⁢ ( M ) ~ 𝑓 𝑀 \widetilde{f}(M) over~ start_ARG italic_f end_ARG ( italic_M ) ) .

###### Proof.

### E.3 Lipschitz of c ⁢ ( M ) 𝑐 𝑀 c(M) italic_c ( italic_M )

We introduce the lemma about Lipschitz of $c(M)$.

###### Lemma E.4 (Lipschitz of $c(M)$ ) .

Under the following conditions

- •
Assumption D.1 holds.
- •
Let $c(M)$ be defined as Definition B.4.

For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|c(M)-c(\widetilde{M})\|_{F}\leq 2dnR^{3}\|M-\widetilde{M}\|_{F}$ $$

###### Proof.

We have

$$ $\displaystyle\|c(M)-c(\widetilde{M})\|_{F}\leq$ $\displaystyle~{}\|\nabla c(M)\|_{F}\cdot\|M-\widetilde{M}\|_{F}$ $\displaystyle=$ $\displaystyle~{}\|\nabla\widetilde{f}(M)\|_{F}\cdot\|M-\widetilde{M}\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}2dnR^{3}\cdot\|M-\widetilde{M}\|_{F}$ $$

where the first step follows from Fact E.1,
the second step follows from Lemma B.12,
the third step follows from Lemma D.4.
∎

###### Lemma E.4 (Lipschitz of c ⁢ ( M ) 𝑐 𝑀 c(M) italic_c ( italic_M ) ) .

###### Proof.

### E.4 Lipschitz of f ~ ⁢ ( M ) ∘ c ⁢ ( M ) ~ 𝑓 𝑀 𝑐 𝑀 \widetilde{f}(M)\circ c(M) over~ start_ARG italic_f end_ARG ( italic_M ) ∘ italic_c ( italic_M )

We introduce the lemma about Lipschitz of $\widetilde{f}(M)\circ c(M)$.

###### Lemma E.5 (Lipschitz of $\widetilde{f}(M)\circ c(M)$ ) .

Under the following conditions

- •
Assumption D.1 holds.
- •
Let $c(M)$ be defined as Definition B.4.
- •
Let $\widetilde{f}(M)$ be defined as Definition B.3.

For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(\widetilde{M})\circ c( \widetilde{M})\|_{F}\leq 6dn^{3/2}R^{3}\|M-\widetilde{M}\|_{F}$ $$

###### Proof.

We have

$$ $\displaystyle\mathrm{LHS}\leq$ $\displaystyle~{}\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(M)\circ c( \widetilde{M})\|_{F}+\|\widetilde{f}(M)\circ c(\widetilde{M})-\widetilde{f}( \widetilde{M})\circ c(\widetilde{M})\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|\widetilde{f}(M)\|_{F}\cdot\|c(M)-c(\widetilde{M})\|_{F}+\|c (\widetilde{M})\|_{F}\cdot\|\widetilde{f}(M)-\widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\sqrt{n}\cdot\|c(M)-c(\widetilde{M})\|_{F}+2\sqrt{n}\cdot\| \widetilde{f}(M)-\widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\sqrt{n}\cdot 2dnR^{3}\|M-\widetilde{M}\|_{F}+2\sqrt{n}\cdot 2 dnR^{3}\|M-\widetilde{M}\|_{F}$ $$

where the first step follows from triangle inequality,
the second step follows from Fact A.5,
the third step follows from Lemma D.3,
the fourth step follows from Lemma E.4 and Lemma E.3.

So we have

$$ $\displaystyle\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(\widetilde{M})\circ c( \widetilde{M})\|_{F}\leq 6dn^{3/2}R^{3}\|M-\widetilde{M}\|_{F}$ $$

∎

###### Lemma E.5 (Lipschitz of f ~ ⁢ ( M ) ∘ c ⁢ ( M ) ~ 𝑓 𝑀 𝑐 𝑀 \widetilde{f}(M)\circ c(M) over~ start_ARG italic_f end_ARG ( italic_M ) ∘ italic_c ( italic_M ) ) .

###### Proof.

### E.5 Lipschitz of diag ⁡ ( ( f ~ ⁢ ( M ) ∘ c ⁢ ( M ) ) ⋅ 𝟏 n ) diag ⋅ ~ 𝑓 𝑀 𝑐 𝑀 subscript 1 𝑛 \operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n}) roman_diag ( ( over~ start_ARG italic_f end_ARG ( italic_M ) ∘ italic_c ( italic_M ) ) ⋅ bold_1 start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT )

We introduce the lemma about Lipschitz of $\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n})$.

###### Lemma E.6 (Lipschitz of $\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n})$ ) .

If the following conditions hold

- •
Assumption D.1 holds.
- •
Let $c(M)$ be defined as Definition B.4.
- •
Let $\widetilde{f}(M)$ be defined as Definition B.3.

For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n })-\operatorname{diag}((\widetilde{f}(\widetilde{M})\circ c(\widetilde{M})) \cdot{\bf 1}_{n})\|_{F}\leq 6dn^{2}R^{3}\|M-\widetilde{M}\|_{F}$ $$

###### Proof.

We have

$$ $\displaystyle\mathrm{LHS}=$ $\displaystyle~{}\|(\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n}-(\widetilde{f}( \widetilde{M})\circ c(\widetilde{M}))\cdot{\bf 1}_{n}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\|((\widetilde{f}(M)\circ c(M))-(\widetilde{f}(\widetilde{M}) \circ c(\widetilde{M})))\cdot{\bf 1}_{n}\|_{2}$ $\displaystyle\leq$ $\displaystyle~{}\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(\widetilde{M})\circ c (\widetilde{M})\|\cdot\|{\bf 1}_{n}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\sqrt{n}\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(\widetilde{ M})\circ c(\widetilde{M})\|$ (9) $$

where the first step follows from Fact A.4,
the second step follows from basic algebra,
the third step follows from Fact A.5,
and the fourth step follows from $\|{\bf 1}_{n}\|_{2}=\sqrt{n}$.

Then we have

$$ $\displaystyle\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(\widetilde{M})\circ c( \widetilde{M})\|\leq\|\widetilde{f}(M)\circ c(M)-\widetilde{f}(\widetilde{M}) \circ c(\widetilde{M})\|_{F}$ (10) $$

which follows from Fact A.5.

Following Eq. (E.5), Eq. (10) and Lemma E.5, we have

$$ $\displaystyle\mathrm{LHS}\leq$ $\displaystyle~{}\sqrt{n}\cdot 6dn^{3/2}R^{3}\|M-\widetilde{M}\|_{F}=6dn^{2}R^{ 3}\|M-\widetilde{M}\|_{F}$ $$

∎

###### Lemma E.6 (Lipschitz of diag ⁡ ( ( f ~ ⁢ ( M ) ∘ c ⁢ ( M ) ) ⋅ 𝟏 n ) diag ⋅ ~ 𝑓 𝑀 𝑐 𝑀 subscript 1 𝑛 \operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n}) roman_diag ( ( over~ start_ARG italic_f end_ARG ( italic_M ) ∘ italic_c ( italic_M ) ) ⋅ bold_1 start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ) ) .

###### Proof.

### E.6 Lipschitz of diag ⁡ ( ( f ~ ⁢ ( M ) ∘ c ⁢ ( M ) ) ⋅ 𝟏 n ) ⁢ f ~ ⁢ ( M ) diag ⋅ ~ 𝑓 𝑀 𝑐 𝑀 subscript 1 𝑛 ~ 𝑓 𝑀 \operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n})\widetilde{f}%
(M) roman_diag ( ( over~ start_ARG italic_f end_ARG ( italic_M ) ∘ italic_c ( italic_M ) ) ⋅ bold_1 start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ) over~ start_ARG italic_f end_ARG ( italic_M )

We introduce the lemma about Lipschitz of $\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n})\widetilde{f}
(M)$.

###### Lemma E.7 (Lipschitz of $\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n})\widetilde{f} (M)$ ) .

If the following conditions hold

- •
Assumption D.1 holds.
- •
Let $c(M)$ be defined as Definition B.4.
- •
Let $\widetilde{f}(M)$ be defined as Definition B.3.

For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n })\widetilde{f}(M)-\operatorname{diag}((\widetilde{f}(\widetilde{M})\circ c( \widetilde{M}))\cdot{\bf 1}_{n})\widetilde{f}(\widetilde{M})\|_{F}\leq 24dn^{7 /2}R^{3}\|M-\widetilde{M}\|_{F}$ $$

###### Proof.

Following Fact E.2, we have

$$ $\displaystyle~{}\|\operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1} _{n})\widetilde{f}(M)-\operatorname{diag}((\widetilde{f}(\widetilde{M})\circ c (\widetilde{M}))\cdot{\bf 1}_{n})\widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}2^{1}\cdot\max\{6dn^{2}R^{3},6dn^{3/2}R^{3}\}\cdot(\sqrt{n} \cdot 2n)\|M-\widetilde{M}\|_{F}$ $\displaystyle=$ $\displaystyle~{}24dn^{7/2}R^{3}\|M-\widetilde{M}\|_{F}$ $$

where we have the upper bound in Lemma D.3, the Lipschitz of $\operatorname{diag}((\widetilde{f}(M)\circ c(M))$ and $\widetilde{f}(M)$ in Lemma E.3 and Lemma E.6.
∎

###### Lemma E.7 (Lipschitz of diag ⁡ ( ( f ~ ⁢ ( M ) ∘ c ⁢ ( M ) ) ⋅ 𝟏 n ) ⁢ f ~ ⁢ ( M ) diag ⋅ ~ 𝑓 𝑀 𝑐 𝑀 subscript 1 𝑛 ~ 𝑓 𝑀 \operatorname{diag}((\widetilde{f}(M)\circ c(M))\cdot{\bf 1}_{n})\widetilde{f}%
(M) roman_diag ( ( over~ start_ARG italic_f end_ARG ( italic_M ) ∘ italic_c ( italic_M ) ) ⋅ bold_1 start_POSTSUBSCRIPT italic_n end_POSTSUBSCRIPT ) over~ start_ARG italic_f end_ARG ( italic_M ) ) .

###### Proof.

### E.7 Lipschitz of Gradient

We introduce the lemma about Lipschitz of the gradient.

###### Theorem E.8 (Lipschitz of the gradient, formal version of Theorem 5.4 ) .

We can show $\nabla_{M}\mathcal{L}(M)$ is $L$-Lipschitz.

If the following conditions hold

- •
Assumption D.1 holds.
- •
Let $c(M)$ be defined as Definition B.4.
- •
Let $\widetilde{f}(M)$ be defined as Definition B.3.

For $M,\widetilde{M}\in\mathbb{R}^{d\times d}$, we have

$$ $\displaystyle\|\nabla_{M}\mathcal{L}(M)-\nabla_{M}\mathcal{L}(\widetilde{M})\| _{F}\leq(\lambda+30dn^{7/2}R^{6})\cdot\|M-\widetilde{M}\|_{F}$ $$

###### Proof.

We have

$$ $\displaystyle~{}\|\nabla_{M}\mathcal{L}(M)-\nabla_{M}\mathcal{L}(\widetilde{M} )\|_{F}$ $\displaystyle=$ $\displaystyle~{}\|W\circ(X^{\top}(c(M)\circ\widetilde{f}(M)-\operatorname{diag }(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{n})\widetilde{f}(M)$ $\displaystyle~{}-c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})+ \operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})\cdot{\bf 1 }_{n})\widetilde{f}(\widetilde{M}))X)+\lambda M-\lambda\widetilde{M}\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|W\circ(X^{\top}(c(M)\circ\widetilde{f}(M)-\operatorname{diag }(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{n})\widetilde{f}(M)$ $\displaystyle~{}-c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})+ \operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})\cdot{\bf 1 }_{n})\widetilde{f}(\widetilde{M}))X)\|_{F}+\|\lambda(M-\widetilde{M})\|_{F}$ (11) $$

where the first step follows from Theorem C.5,
and the second step follows from triangle inequality.
Now we proof these two terms separately.

For the first term, we have

$$ $\displaystyle~{}\|W\circ(X^{\top}(c(M)\circ\widetilde{f}(M)-\operatorname{diag }(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{n})\widetilde{f}(M)$ $\displaystyle~{}-c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})+ \operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})\cdot{\bf 1 }_{n})\widetilde{f}(\widetilde{M}))X)\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|W\|_{F}\cdot\|X^{\top}(c(M)\circ\widetilde{f}(M)- \operatorname{diag}(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{n})\widetilde{f}(M)$ $\displaystyle~{}-c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})+ \operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})\cdot{\bf 1 }_{n})\widetilde{f}(\widetilde{M}))X\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|W\|_{F}\cdot\|X\|_{F}^{2}\cdot\|c(M)\circ\widetilde{f}(M)- \operatorname{diag}(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{n})\widetilde{f}(M)$ $\displaystyle~{}-c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})+ \operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}(\widetilde{M})\cdot{\bf 1 }_{n})\widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|W\|_{F}\cdot\|X\|_{F}^{2}\cdot(\|c(M)\circ\widetilde{f}(M)-c (\widetilde{M})\circ\widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle~{}+\|\operatorname{diag}(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{ n})\widetilde{f}(M)-\operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}( \widetilde{M})\cdot{\bf 1}_{n})\widetilde{f}(\widetilde{M})\|_{F})$ $\displaystyle=$ $\displaystyle~{}R^{3}\cdot(\|c(M)\circ\widetilde{f}(M)-c(\widetilde{M})\circ \widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle~{}+\|\operatorname{diag}(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{ n})\widetilde{f}(M)-\operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}( \widetilde{M})\cdot{\bf 1}_{n})\widetilde{f}(\widetilde{M})\|_{F})$ $$

where the first step and the second step follows from Fact A.5,
the third step follows from triangle inequality,
and the fourth step follows from Assumption D.1.

Then we have

$$ $\displaystyle~{}R^{3}\cdot(\|c(M)\circ\widetilde{f}(M)-c(\widetilde{M})\circ \widetilde{f}(\widetilde{M})\|_{F}$ $\displaystyle~{}+\|\operatorname{diag}(c(M)\circ\widetilde{f}(M)\cdot{\bf 1}_{ n})\widetilde{f}(M)-\operatorname{diag}(c(\widetilde{M})\circ\widetilde{f}( \widetilde{M})\cdot{\bf 1}_{n})\widetilde{f}(\widetilde{M})\|_{F})$ $\displaystyle\leq$ $\displaystyle~{}R^{3}\cdot(24dn^{7/2}R^{3}\|M-\widetilde{M}\|_{F}+6dn^{3/2}R^{ 3}\|M-\widetilde{M}\|_{F})$ $\displaystyle\leq$ $\displaystyle~{}R^{3}\cdot(30dn^{7/2}R^{3}\|M-\widetilde{M}\|_{F})$ $\displaystyle=$ $\displaystyle~{}30dn^{7/2}R^{6}\|M-\widetilde{M}\|_{F}$ (12) $$

where the first step follows from Lemma E.5 and Lemma E.7,
the second step follows from $n\geq 1$.

For the second term, we have

$$ $\displaystyle\|\lambda(M-\widetilde{M})\|_{F}=\lambda\|M-\widetilde{M}\|_{F}$ (13) $$

which follows from Fact A.5.

Finally, we have

$$ $\displaystyle\|\nabla_{M}\mathcal{L}(M)-\nabla_{M}\mathcal{L}(\widetilde{M})\| _{F}\leq(\lambda+30dn^{7/2}R^{6})\cdot\|M-\widetilde{M}\|_{F}$ $$

which follows from Eq. (E.7), Eq. (E.7), and Eq. (13).
∎

###### Theorem E.8 (Lipschitz of the gradient, formal version of Theorem 5.4 ) .

###### Proof.

## Appendix F Convergence of Gradient Descent

### F.1 Helpful Statements

Here, we present useful fact that we use to prove our convergence result.

###### Fact F.1 .

We can show that for $a,b\in\mathbb{R}$

- •
Part 1.
$\displaystyle\sqrt{a^{2}+b^{2}}\geq\frac{|a|+|b|}{\sqrt{2}}$
- •
Part 2. Suppose $|a|>|b|$
$\displaystyle\sqrt{|a|-|b|}\geq\sqrt{|a|}-\sqrt{|b|}$

###### Proof.

Proof of Part 1.
Square both side of the inequality in Part 1., we have

$$ $\displaystyle\mathrm{LHS}=a^{2}+b^{2}$ $$

and

$$ $\displaystyle\mathrm{RHS}=\frac{a^{2}+2|a|\cdot|b|+b^{2}}{2}.$ $$

So we just need to prove

$$ $\displaystyle\mathrm{LHS}-\mathrm{RHS}=$ $\displaystyle~{}a^{2}+b^{2}-\frac{a^{2}+2|a|\cdot|b|+b^{2}}{2}$ $\displaystyle=$ $\displaystyle~{}\frac{a^{2}+b^{2}-2|a|\cdot|b|}{2}$ $\displaystyle=$ $\displaystyle~{}\frac{(|a|-|b|)^{2}}{2}$ $\displaystyle\geq$ $\displaystyle 0$ $$

which is hold because for any $x\in\mathbb{R}$, $x^{2}\geq 0$.

Proof of Part 2.
Square both side of the inequality in Part 2., we have

$$ $\displaystyle\mathrm{LHS}=|a|-|b|$ $$

and

$$ $\displaystyle\mathrm{RHS}=|a|+|b|-2\sqrt{|a||b|}$ $$

So we just need to prove

$$ $\displaystyle\mathrm{LHS}-\mathrm{RHS}=$ $\displaystyle~{}|a|-|b|-|a|-|b|+2\sqrt{|a||b|}$ $\displaystyle=$ $\displaystyle~{}2\sqrt{|a||b|}-2|b|$ $\displaystyle=$ $\displaystyle~{}2\sqrt{|b|}(\sqrt{|a|}-\sqrt{|b|})$ $\displaystyle\geq$ $\displaystyle~{}0$ $$

which is hold because $|a|>|b|$ and $|b|\geq 0$.
∎

###### Fact F.1 .

###### Proof.

### F.2 Lower Bound on Frobenius Norm

We present the lemma for the lower bound on the Frobenius norm in this section.

###### Lemma F.2 .

If the following conditions hold

- •
Let $B\in\mathbb{R}^{d\times d}$.
- •
Let $M\in[0,1]^{d\times d}$.
- •
Let $\lambda\in[0,1]$ be some constant.
- •
Suppose that $\|B\|_{F}\leq R$.

Then, we can show

- •
Part 1.
$\displaystyle\|B+\lambda M\|_{F}^{2}\geq\|B\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}
-2R\lambda d$
- •
Part 2.
$\displaystyle\|B+\lambda M\|_{F}\geq\frac{1}{\sqrt{2}}(\|B\|_{F}+\lambda\|M\|_
{F})-\sqrt{2R\lambda d}$

###### Proof.

Proof of Part 1.
We can show that

$$ $\displaystyle\|B+\lambda M\|_{F}^{2}=$ $\displaystyle~{}\|B\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}+2\langle B,\lambda M\rangle$ $\displaystyle\geq$ $\displaystyle~{}\|B\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}-2\|B\|_{F}\cdot\| \lambda M\|_{F}$ $\displaystyle\geq$ $\displaystyle~{}\|B\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}-2R\lambda d$ (14) $$

where the first step follows from Fact A.6,
the second step follows from Fact A.5,
the third step follows from the upper bound of $\|B\|_{F}$ and $\|M\|_{F}$.

Proof of Part 2.
Taking the square root on both sides, we get

$$ $\displaystyle\|B+\lambda M\|_{F}\geq$ $\displaystyle~{}\sqrt{\|B\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}-2R\lambda d}$ $\displaystyle\geq$ $\displaystyle~{}\sqrt{\|B\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}}-\sqrt{2R\lambda d}$ $\displaystyle\geq$ $\displaystyle~{}\frac{1}{\sqrt{2}}(\|B\|_{F}+\lambda\|M\|_{F})-\sqrt{2R\lambda d}$ $$

where the first step follows from Eq. (F.2),
the second step follows from Part 2. of Fact F.1,
and the third step follows from Part 1. of Fact F.1.

∎

###### Lemma F.2 .

###### Proof.

### F.3 Sandwich Lower Bound on Frobenius Norm

Here, we introduce a sandwich trace fact.

###### Fact F.3 .

If $A\succeq\beta I$, then $\operatorname{tr}[B^{\top}AB]\geq\beta\operatorname{tr}[B^{\top}B]$.

###### Proof.

As $A\succeq\beta I$, we have
$A-\beta I\succeq 0$. Multiplying both sides by $B^{\top}$ on the left and $B$ on the right (noting that these operations preserve the positive semidefiniteness), we have

$$ $\displaystyle B^{\top}(A-\beta I)B\succeq 0.$ $$

Taking the trace and utilizing the property that the trace of a positive semidefinite matrix is non-negative, we have

$$ $\displaystyle\operatorname{tr}[B^{\top}AB-\beta B^{\top}B]\geq 0,$ $$

which simplifies to

$$ $\displaystyle\operatorname{tr}[B^{\top}AB]-\beta\operatorname{tr}[B^{\top}B] \geq 0.$ $$

This concludes the proof.
∎

We establish a sandwich lower bound on the Frobenius norm.

###### Lemma F.4 (Formal version of Lemma 5.6 ) .

If the following conditions hold

- •
Let $B\in\mathbb{R}^{n\times n}$ and $X\in\mathbb{R}^{n\times d}$.
- •
Assume that $XX^{\top}\succeq\beta I$.

Then, we have

$$ $\displaystyle\|X^{\top}BX\|_{F}\geq\beta\|B\|_{F}$ $$

###### Proof.

We can show that

$$ $\displaystyle\|X^{\top}BX\|_{F}^{2}=$ $\displaystyle~{}\operatorname{tr}[X^{\top}BXX^{\top}B^{\top}X]$ $\displaystyle\geq$ $\displaystyle~{}\beta\cdot\operatorname{tr}[X^{\top}BB^{\top}X]$ $\displaystyle=$ $\displaystyle~{}\beta\cdot\operatorname{tr}[B^{\top}XX^{\top}B]$ $\displaystyle\geq$ $\displaystyle~{}\beta^{2}\cdot\operatorname{tr}[B^{\top}B]$ $\displaystyle=$ $\displaystyle~{}\beta^{2}\cdot\|B\|_{F}^{2}$ $$

where the first step, the third step and the fifth step follows from Fact A.4,
the second step and the fourth step follows from Fact F.3 and $XX^{\top}\succeq\beta I$.

Taking the square root of both side, we finish the proof.
∎

###### Fact F.3 .

###### Proof.

###### Lemma F.4 (Formal version of Lemma 5.6 ) .

###### Proof.

### F.4 Lower Bound on Hadamard Product Between Two Matrices

We present the lemma for lower bound on Hadamard product between two matrices in this section.

###### Lemma F.5 .

If the following conditions hold

- •
Let $B,W\in\mathbb{R}^{d\times d}$.

Then, we have

$$ $\displaystyle\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|B\|_{F}\geq\|W\circ B\|_{F} \geq\min_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|B\|_{F}.$ $$

###### Proof.

The proof directly follows from the definition of the Frobenius norm.
∎

###### Lemma F.5 .

###### Proof.

### F.5 Final Bound

We introduce some useful lemmas that we use to prove the final bound.

###### Lemma F.6 .

If the following conditions hold

- •
Let $b\in\mathbb{R}^{n}$ and $\langle b,{\bf 1}_{n}\rangle=0$.
- •
Let $f\in[\delta,1]^{n}$ and $\langle f,{\bf 1}_{n}\rangle=1$.

Then we have

$$ $\displaystyle\|(b-\langle b,f\rangle{\bf 1}_{n})\circ f\|_{2}\geq\delta\|b\|_{ 2}.$ $$

###### Proof.

Note that $\langle b,{\bf 1}_{n}\rangle=0$ so that $b$ and ${\bf 1}_{n}$ are orthogonal with each other. Then, we have

$$ $\displaystyle\|(b-\langle b,f\rangle{\bf 1}_{n})\circ f\|_{2}\geq$ $\displaystyle~{}\delta\|b-\langle b,f\rangle{\bf 1}_{n}\|_{2}$ $\displaystyle=$ $\displaystyle~{}\delta\sqrt{\|b\|_{2}^{2}+\|\langle b,f\rangle{\bf 1}_{n}\|_{2 }^{2}}$ $\displaystyle\geq$ $\displaystyle~{}\delta\|b\|_{2},$ $$

where the second step is from Pythagorean theorem.
∎

We present our final bound for proving the PL inequality.

###### Lemma F.7 (Formal version of Lemma 5.7 ) .

If the following conditions hold

- •
Let $B\in\mathbb{R}^{n\times n}$ and each row summation is zero, i.e., $B\cdot{\bf 1}_{n}={\bf 0}_{n}$.
- •
Let $\widetilde{f}(M)\in[0,1]^{n\times n}$ and each row summation is 1, i.e., $\widetilde{f}(M)\cdot{\bf 1}_{n}={\bf 1}_{n}$.
- •
Assume that $\min_{i,j\in[n]}\widetilde{f}(M)_{i,j}\geq\delta>0$.

Then, we can show

$$ $\displaystyle\|B\circ\widetilde{f}(M)-\operatorname{diag}((B\circ\widetilde{f} (M))\cdot{\bf 1}_{n})\widetilde{f}(M)\|_{F}\geq\delta\cdot\|B\|_{F}$ $$

###### Proof.

For any $i\in[n]$, let $B_{i}\in\mathbb{R}^{n}$ be the $i$-th row of $B$, and we have $\langle B_{i},{\bf 1}_{n}\rangle=0$ by the first condition.

For any $i\in[n]$, let $\widetilde{f}(M)_{i}\in\mathbb{R}^{n}$ be the $i$-th row of $\widetilde{f}(M)$, and we have $\langle\widetilde{f}(M)_{i},{\bf 1}_{n}\rangle=1$ by the second condition and $\widetilde{f}(M)_{i,j}\in[\delta,1]$ by the third condition.

By Lemma F.6, for any $i\in[n]$, we have

$$ $\displaystyle\|(B_{i}-\langle B_{i},\widetilde{f}(M)_{i}\rangle{\bf 1}_{n}) \circ\widetilde{f}(M)_{i}\|_{2}\geq\delta\|B_{i}\|_{2}.$ $$

Then, we have

$$ $\displaystyle~{}\|B\circ\widetilde{f}(M)-\operatorname{diag}((B\circ\widetilde {f}(M))\cdot{\bf 1}_{n})\widetilde{f}(M)\|_{F}^{2}$ $\displaystyle=$ $\displaystyle~{}\sum_{i\in[n]}\|(B_{i}-\langle B_{i},\widetilde{f}(M)_{i} \rangle{\bf 1}_{n})\circ\widetilde{f}(M)_{i}\|_{2}^{2}$ $\displaystyle\geq$ $\displaystyle~{}\sum_{i\in[n]}\delta^{2}\|B_{i}\|_{2}^{2}$ $\displaystyle=$ $\displaystyle~{}\delta^{2}\|B\|_{F}^{2}.$ $$

∎

###### Lemma F.6 .

###### Proof.

###### Lemma F.7 (Formal version of Lemma 5.7 ) .

###### Proof.

### F.6 PL Inequality

Here we present the bound for one unit loss function.

###### Lemma F.8 .

If the following conditions hold

- •
Let $c(M)$ be defined in Definition B.4.

We have

$$ $\displaystyle\|c(M)\|_{F}\leq~{}2\sqrt{n}.$ $$

###### Proof.

We have

$$ $\displaystyle\|c(M)\|_{F}\leq$ $\displaystyle~{}\|\widetilde{f}(M)\|_{F}+\|f\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}2\sqrt{n},$ $$

where the first step follows from Definition B.4 and triangle inequality,
the second step follows $x_{1}^{2}+\dots+x_{n}^{2}\leq(x_{1}+\dots+x_{n})^{2}$ when $x_{i}\geq 0$ for any $i\in[n]$.
∎

We present the lemma for proving the PL inequality.

###### Lemma F.9 .

If the following conditions hold

- •
Let $\widetilde{f}(M)$ be defined in Definition B.3.
- •
Let $c(M)$ be defined in Definition B.4.

We have

$$ $\displaystyle\|\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n} )\widetilde{f}(M)\|_{F}\leq\sqrt{n}.$ $$

###### Proof.

We have

$$ $\displaystyle\|\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n} )\widetilde{f}(M)\|_{F}\leq$ $\displaystyle\max_{i\in{n}}\{|(c(M)_{i}\circ\widetilde{f}(M)_{i})\cdot{\bf 1}_ {n}|\}\cdot\|\widetilde{f}(M)\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\|\widetilde{f}(M)\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\sqrt{n},$ $$

where the first step is by Frobenius norm definition and the second step follows from $\langle\widetilde{f}(M)_{i},{\bf 1}_{n}\rangle=1$ and $c(M)_{i}\in[-1,1]^{n}$ for any $i\in[n]$.
∎

Finally, we can show the lemma for PL inequality.

###### Lemma F.10 (PL inequality, formal version of 5.5 ) .

If the following conditions hold,

- •
Let $M\in[0,1]^{d\times d}$ .
- •
Let $\lambda\in[0,1]$ be some constant.
- •
Assume that $XX^{\top}\succeq\beta I$.
- •
Assume that $\min_{i,j\in[n]}\widetilde{f}(M)_{i,j}\geq\delta>0$.
- •
Let $\mathcal{L}(M)$ be defined in Definition B.7.

Furthermore,

- •
Let $\alpha=2$.
- •
Let $\mu=2\min_{i,j\in[d]}\{|W_{i,j}|\}\cdot\beta\cdot\delta$.
- •
Let $\xi=12\sqrt{n}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot\lambda d/\mu$.

We have

$$ $\displaystyle\|\nabla_{M}\mathcal{L}(M)\|_{F}^{\alpha}\geq\frac{1}{2}\mu(\|c(M )\|_{F}^{2}+\frac{2\lambda^{2}}{\mu}\|M\|_{F}^{2}-\xi).$ $$

###### Proof.

We have $\widetilde{f}(M)\cdot{\bf 1}_{n}={\bf 1}_{n}$ and $f\cdot{\bf 1}_{n}={\bf 1}_{n}$ by Definition B.3.
Note that $c(M)=\widetilde{f}(M)-f$ by Definition B.4.
Thus, we have $c(M)\cdot{\bf 1}_{n}={\bf 0}_{n}$.

On the other hand, we have

$$ $\displaystyle~{}\|W\circ(X^{\top}(c(M)\circ\widetilde{f}(M)-\operatorname{diag }((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n})\widetilde{f}(M))X)\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X^{\top}(c(M)\circ \widetilde{f}(M)-\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{ n})\widetilde{f}(M))X\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot\|c(M) \circ\widetilde{f}(M)-\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1 }_{n})\widetilde{f}(M)\|_{F}$ $\displaystyle\leq$ $\displaystyle~{}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot(\|c(M) \circ\widetilde{f}(M)\|_{F}+\|\operatorname{diag}((c(M)\circ\widetilde{f}(M)) \cdot{\bf 1}_{n})\widetilde{f}(M)\|_{F})$ $\displaystyle\leq$ $\displaystyle~{}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot(\|c(M)\|_ {F}+\|\operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n}) \widetilde{f}(M)\|_{F})$ $\displaystyle\leq$ $\displaystyle~{}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot(2\sqrt{n} +\sqrt{n})$ $\displaystyle=$ $\displaystyle~{}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot 3\sqrt{n}$ $$

where the first and forth steps follow Lemma F.5,
the second step follows from Frobenius norm property, the third step follows from triangle inequality,
the fifth step follows from Lemma F.8 and Lemma F.9.

Let $\alpha=2$. We have the following

$$ $\displaystyle~{}\|\nabla_{M}\mathcal{L}(M)\|_{F}^{2}$ $\displaystyle=$ $\displaystyle~{}\|W\circ(X^{\top}(c(M)\circ\widetilde{f}(M)-\operatorname{diag }((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n})\widetilde{f}(M))X)+\lambda M\|_ {F}^{2}$ $\displaystyle\geq$ $\displaystyle~{}\|W\circ(X^{\top}(c(M)\circ\widetilde{f}(M)-\operatorname{diag }((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n})\widetilde{f}(M))X)\|_{F}^{2}+ \lambda^{2}\|M\|_{F}^{2}-\alpha_{1}$ $\displaystyle\geq$ $\displaystyle~{}\alpha_{2}\cdot\|X^{\top}(c(M)\circ\widetilde{f}(M)- \operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n})\widetilde{f}( M))X\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}-\alpha_{1}$ $\displaystyle\geq$ $\displaystyle~{}\alpha_{2}\cdot\alpha_{3}\cdot\|c(M)\circ\widetilde{f}(M)- \operatorname{diag}((c(M)\circ\widetilde{f}(M))\cdot{\bf 1}_{n})\widetilde{f}( M)\|_{F}^{2}+\lambda^{2}\|M\|_{F}^{2}-\alpha_{1}$ $\displaystyle\geq$ $\displaystyle~{}\alpha_{2}\cdot\alpha_{3}\cdot\alpha_{4}\cdot\|c(M)\|_{F}^{2}+ \lambda^{2}\|M\|_{F}^{2}-\alpha_{1}$ $\displaystyle=$ $\displaystyle~{}\frac{1}{2}\mu(\|c(M)\|_{F}^{2}+\frac{2\lambda^{2}}{\mu}\|M\|_ {F}^{2}-\xi),$ $$

where the second step follows from Lemma F.2 and $\alpha_{1}=6\sqrt{n}\max_{i,j\in[d]}\{|W_{i,j}|\}\cdot\|X\|_{F}^{2}\cdot\lambda
d$, the third step follows from Lemma F.5 and $\alpha_{2}=\min_{i,j\in[d]}\{|W_{i,j}|\}$, the fourth step follows from Lemma F.4 and $\alpha_{3}=\beta$, the fifth step follows from Lemma F.7 and $\alpha_{4}=\delta$, and the last step follows from $\mu=2\alpha_{2}\cdot\alpha_{3}\cdot\alpha_{4}$ and $\xi=2\alpha_{1}/\mu$.

∎

###### Lemma F.8 .

###### Proof.

###### Lemma F.9 .

###### Proof.

###### Lemma F.10 (PL inequality, formal version of 5.5 ) .

###### Proof.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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?
- 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?
- 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.
soundness: 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.
- 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.
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.
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.
- 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. 
- 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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. How to revert the pruned fused matrix back to the query and key projection matrices respectively?
2. How does the proposed attention pruning affect length extrapolation ability?
soundness: 2
strengths: 1. Considering the similarity between attention matrices during pruning of the projection weight is novel and interesting.
2. The assumptions in theoretical analyses are mild and sensible.
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.
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.
2. Lacks of empirical results on pruning real-world LLM. The low error rates on synthetic data may not translate to real world dowstream performances.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Conside the weakness
soundness: 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.
2.  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.
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.
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.

2. 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.

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 3
questions: How does the convergence rate in Theorem 4.1 compare to the rates one can achieve with naive baselines?
soundness: 3
strengths: *Original problem.* The paper highlights an interesting, niche problem in the pruning literature, which, to my knowledge, has not been analyzed theoretically. 

*Clear presentation.* The paper is well-written and the structure is easy to follow. 

*Extensive theoretical analysis.* The authors provide an extensive theoretical analysis on the convergence of their pruning algorithm.
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.
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. 

*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).

Meta-review: This paper introduces a method for pruning query and key weight matrices in attention modules. The paper proposes a theoretical justification for the proposed scheme, with convergence guarantees. The method was initially evaluated on synthetic data, and as requested by the reviewers, also on real-world language models. The proposed approach seems to perform relatively well, but does not provide a viable strategy for pruning in real-world applications. As mentioned by reviewer opkW, most of the computation and parameters are in the feedforward layers. Tu summarize, this paper proposes an interesting finding, for which the empirical implications are yet a bit unclear. While this is a borderline paper, I recommend this paper for acceptance.

Decision: Accept (Poster)