Numerical Pruning for Efficient Autoregressive Models

Introduction

Transformers have been dominant in generative models. This includes Large Language Models (LLMs) (Vaswani et al. 2017; Touvron et al. 2023b) for language generation, as well as recent autoregressive image generation models (Van Den Oord, Vinyals et al. 2017; Esser, Rombach, and Ommer 2021; Ramesh et al. 2021; Yu et al. 2022). Notably, models such as LlamaGen (Sun et al. 2024), which use image tokenizers to convert continuous images into discrete tokens, have demonstrated the ability to surpass diffusion models (Ho, Jain, and Abbeel 2020; Rombach et al. 2022) in image generation tasks. The “next-token prediction” paradigm demonstrates significant capabilities in addressing both language and image generation tasks, enabling solutions that mimic human-like conversational interactions (Achiam, Adler et al. 2023; Li et al. 2024a).

Recognizing the capabilities of large autoregressive models pioneering works (Frantar and Alistarh 2023; Sun et al. 2023; Ma, Fang, and Wang 2023; Ashkboos et al. 2024; Zhan et al. 2021; Zhao et al. 2024; Zhan et al. 2024b) have sought to compress these models to enhance their execution efficiency. Compared to irregular pruning methods, structural pruning offers a more efficient reduction in both computational and memory overhead (Jian et al. 2021; Gong et al. 2022, 2023). By maintaining a consistent and regular structure, it simplifies implementation, accelerates processing, and leads to more predictable resource savings (Kong et al. 2022, 2023). However, most of these efforts focus solely on language models and language-related research areas. Consequently, their methods are not readily applicable to image generation tasks because of the fundamental differences in data structure and computational requirements between language and image processing (Reed et al. 2016; Parmar et al. 2018; Lee et al. 2022; Shen et al. 2024a, c, b, 2023b, 2023a; Li et al. 2023b, 2024e, 2024d). Therefore, it is crucial to explore the transformer architecture itself, rather than focusing on specific application models. This motivates us to develop a general method for compressing autoregressive models applicable to multiple kinds of generative tasks.

Additionally, the recovery of pruned models are crucial. Full-parameter retraining of large autoregressive models after pruning is often computationally prohibitive, making calibrations with a few samples a preferred approach. Previous work (Frantar and Alistarh 2023) employs the Optimal Brain Surgeon (OBS) technique (Hassibi, Stork, and Wolff 1993; LeCun, Denker, and Solla 1989) for weight updates during pruning. However, its heavy reliance on the approximation information increases sensitivity to noise and reduces robustness across different datasets. SliceGPT (Ashkboos et al. 2024) relies on a large number of samples for pruning and calibration, leading to overfitting on calibration data and limiting the generalization to other different datasets.

In this work, we present a novel structural pruning approach that leverages our proposed numerical score, combined with compensation techniques for performance recovery. We first calculate the numerical score for each layer through solving the optimal pruning mask for the minimization of pruning errors using the Newton’s method. By ranking these numerical scores of all layers, we generate the globally pruning mask with the specified pruning ratio. Additionally, we introduce a compensation algorithm to recover pruned models by updating the remaining weights to account for the loss caused by the pruned weights. We empirically evaluate our method using the LLaMA model family including LLaMA, LLaMA-2, and LLaMA-3 as representative LLMs and LlamaGen for image generation tasks. Experimental results show that our method outperforms other state-of-the-art approaches in both language and image generation tasks, validating the effectiveness of our proposed numerical score and compensation algorithm. Moreover, our method reduces GPU memory usage and accelerates generation without requiring any additional GPU-specific modifications. Our main contributions are summarized as follows,

We propose a numerical score, derived from the numerical solution of the optimal mask for minimizing pruning errors with Newton’s method.

We propose a compensation algorithm for the reconstruction of the pruned model, further enhancing the task performance of the pruned model.

Experimental results show that our method not only achieves state-of-the-art performance but also reduces memory usage and accelerates generation on GPUs.

Related Work

The large number of parameters in LLMs motivates the need for pruning (Gong et al. 2020; Wu et al. 2022; Zhan et al. 2024a; Li et al. 2022; Zhang et al. 2022; Zhan et al. 2024c; Shen et al. 2024d) to improve efficiency. The work (Frantar and Alistarh 2023) introduces the Optimal Brain Surgeon (OBS) method (Hassibi, Stork, and Wolff 1993; LeCun, Denker, and Solla 1989) to compress the LLMs, which removes weights with minimal impact on the loss function. It then updates the remaining weights by utilizing the inverse of the Hessian matrix to mitigate errors caused by the pruning process. Unfortunately, this kind of pruning method is still irregular, meaning it does not lead to significant reductions in memory and computational requirements. Subsequent works, such as LLM-Pruner (Ma, Fang, and Wang 2023), SliceGPT (Ashkboos et al. 2024), and FLAP (An et al. 2023), propose structural pruning methods that effectively reduce memory usage and accelerate inference on GPUs. These methods offer significant advantages over irregular pruning by directly enhancing the utility and efficiency of the models. While autoregressive models excel in sequential data processing, such as text, the distinct nature of image data, where spatial relationships and pixel-level details are critical, demands different approaches. As a result, adapting these models to image generation introduces complexities that limit their scalability and effectiveness.

2 Autoregressive Models in Image Generation

Autoregressive models, initially renowned for their success with LLMs, have recently gained popularity in the image generation research area. Pioneering works (Van Den Oord, Vinyals et al. 2017; Esser, Rombach, and Ommer 2021) introduced image tokenizers that convert continuous images into discrete tokens. These tokenizers, which have been demonstrated to be effective by the following works (Ramesh et al. 2021; Yu et al. 2021, 2022), enable autoregressive models to generate image tokens using the next-token prediction approach. Recent work (Sun et al. 2024) delivers a series of image generation models with a new constructed image tokenizer. This research demonstrates the effectiveness of LLM frameworks in image generation tasks, validating their potential beyond traditional language applications. Additionally, the work (Li et al. 2024a) delves deeper into the continuous-valued domains of autoregressive models and removes the image tokenizers for image generation tasks. This work achieves stronger results while leveraging the speed advantage of sequence modeling, which further enhances the utilization and demonstrates the potential of autoregressive models in image generation tasks.

Methodology

2 Numerical Score

Note that, when we apply the mask column by column, the mask $M$ is a $D$-dimensional vector. Specifically, if $M_{j}=0$ for $j\in[D]$, we prune the entire row for $W$, i.e., $W_{j}=0$, and if $M_{j}=1$ we keep the original $W_{j}$.

To compute the numerical score, we explore the bound of the error (i.e., difference) between the original weights and pruned weights. For the bound of the error, we first formulate the error for $i\in[D^{\prime}]$ as

Simplify further, for $i\in[D^{\prime}]$, Eq. (1) can be transformed into the following,

In the above equation, the $\|{\bf 1}_{D}-M\|_{2}$ denotes the number of zero entries in $M$, which is corresponding to the simply $\rho\cdot D$. Furthermore, assuming $\|X\|\leq R$, we demonstrate that the following Lemma 1 holds,

We show that for $i\in[D^{\prime}]$ we have

In order to minimize the error, we regulate both $\rho$ and $\|W_{*,i}\|$. Then, we generalize the mask $M$ from binary value to real value for the calculation of the numerical score. Meanwhile, we set one threshold which converts the real-valued mask back into a binary mask. For mask $M\in^{D}$ and pruning ratio $\rho\in$, the calculation of the numerical score is formulated as follows,

To better solve Eq. (2), we define the numerical score $z\in^{D}$ and $r:=(1-\rho)D\in[0,D]$. The equality constraint in Eq. (2) is then equivalent to $\langle{\bf 1}_{D},z\rangle-r=0$.

Then, Eq. (2) becomes the minimization problem with the equality constraint. To efficiently solve such problem, we adopt the Newton’ method (Bubeck et al. 2015). By turning the equality constraint into a penalty term for regularization, we further generate the following equivalent problem,

To explain how we solve this, we define the loss function for $i\in[D^{\prime}]$ as follows,

Meanwhile, for regularization term, we define as follows,

Combining Lemma 12 and Lemma 13 at Appendix D.3, we compute the gradient of Eq. (4) and Eq. (5) as follows,

Combining Lemma 14 and Lemma 15 at Appendix D.3, we compute the Hessian of Eq. (4) and Eq. (5) as follows,

Subsequently, using Algorithm 1, we efficiently compute the optimal numerical $z$ in $O(TD^{3})$, where $T$ represents the number of iterations for Newton’s Method, typically around 50 in practice. Besides, we derive the following Theorem 2.

Let $r\in[0,D]$ denote the number of ones (it can be a fractional number).

Let $\lambda>0$ denote a regularization co-efficients.

There exists an algorithm (Algorithm 1) that can get the optimal $z$ in $O(TD^{3})$ for Eq. (3).

3 Global Pruning

where $D_{h}$ denotes the dimension of each head, $h\in[H]$ is the head index.

When the threshold $\eta$ is determined, we prune the heads in the Attention module and the channels in the MLP module across all layers based on the strategy that removes heads or channels with numerical scores below the threshold.

4 Compensation for Pruning

With the above discussion, we obtain the pruning mask with Newton’s method. To further improve the model performance, we modify the remaining weights in the model to compensate the loss of the pruned weights.

Note that to align the internal computations in the attention and MLP modules, we prune the rows of the output layers in the modules and the columns in other layers of the modules. If the columns of a layer with $W$ is pruned in $XW$, the corresponding columns of the output also become zero and we are not able to compensate its loss, since modifying other unpruned columns can not change the zero output for the pruned columns. Thus, we only update the weights of the output layers with row pruning in the Attention and MLP modules. We modify the remaining rows based on pruned rows in $W$. For layers with column pruning, we do not modify their unpruned weights.

It can be further transformed to the following,

Eq. (3.4) can be solved analytically with the following Theorem 3. The detailed proof is shown in Appendix B.

The optimal solution for Eq. (3.4) can be derived as the following,

The optimal loss of Problem (3.4) corresponding to the optimal weight perturbation can be expressed as

The sum in Eq. (12) is computed over $D^{\prime}$ (the number of columns in $W$), i.e., $i\in[D^{\prime}]$.

If the rank of $2X^{\top}X$ is not full so that the inversion $(2X^{\top}X)^{-1}$ is unavailable, we apply the dampening method to compute $(2X^{\top}X+\gamma\cdot I)^{-1}$ instead of $(2X^{\top}X)^{-1}$, with $\gamma$ as the dampening ratio.

5 Complexity Analysis

For the computation of numerical score, according to the Lemma 1 and Theorem 2, the complexity is $O(TD^{3})$ where $T$ represents the number of iterations for Newton’s Method, typically around 50 in practice. Additionally, for the compensation method, as demonstrated in Eq. (11), the complexity is $O(D^{3})$ as we need to compute the inverse of a matrix. The matrix multiplication with $M_{p}$ or $M_{p}^{T}$ just selects the columns or rows of a matrix, without the need of actual multiplication. The complexity for numerical score calculation and compensation is the same with state-of-the-art methods, such as SparseGPT (Frantar and Alistarh 2023). In practice, the compensation is finished with just a few data samples on only the output projection layers of the Attention module and the down projection layers of the MLP module, which is more efficient compared with other recovery methods such as LLM-Pruner (Ma, Fang, and Wang 2023) to finetune the whole model on whole dataset, or SliceGPT (Ashkboos et al. 2024) to adopt a large amount of samples for calibration.

Experimental Results

We conduct the experiments on LLaMA model families including LLaMA-1 (Touvron et al. 2023a), LLaMA-2 (Touvron et al. 2023b), and LLaMA-3 (Meta 2024) for the language generation tasks. For evaluations, we compare the perplexity of the models on the WikiText2 (Merity et al. 2016), PTB (Marcus, Santorini, and Marcinkiewicz 1993), and C4 (Raffel et al. 2020) datasets with the 2048 sequence length. We also follow LLM-Pruner to evaluate the zero-shot accuracy on common sense reasoning zero-shot classification datasets including BoolQ (Clark et al. 2019a), PIQA (Bisk et al. 2020), HellaSwag (Zellers et al. 2019), WinoGrande (Sakaguchi et al. 2021), ARC-easy (Clark et al. 2018), ARC-challenge (Clark et al. 2018), and OpenbookQA (Mihaylov et al. 2018). For experiments, we adopt 128 samples from training dataset of WikiText2 to compute the numerical score and compensate the pruned models. For fairness, we also adopt 128 samples for other methods.

As for the image generation tasks, we adopt the LlamaGen (Sun et al. 2024) model family with LlamaGen-XXL and LlamaGen-3B to verify the effectiveness of our method on image generation tasks. We adopt the Fréchet inception distance (FID) (Heusel et al. 2017), Inception Score (IS) (Salimans et al. 2016), sFID (Nash et al. 2021), and Precision/Recall (Kynkäänniemi et al. 2019) as the evaluation metrics on ImageNet dataset (Deng et al. 2009). For all evaluations, we utilized ADM’s TensorFlow scripts (Dhariwal and Nichol 2021) to ensure fair and consistent comparisons. Given that LLM-Pruner requires a backward process and SliceGPT has slow pruning, we further implement FLAP for comparative analysis in image generation tasks. In pratical, we generate 128 images for each class of ImageNet with LlamaGen models for the computation of numerical score and compensation. Same strategy for FLAP for fairness.

2 Results of LLMs

For the LLaMA models, we present the results with different pruning ratios varying from 10% to 70% in Table 1. Based on the perplexity results evaluated with 2048 sequence length on three datasets, our method consistently outperforms other methods across all pruning ratios, demonstrating the effectiveness of our proposed approach. Full results with more sparse ratios are included in Table 5 of Appendix A.1. Results show that for the larger model LLaMA-65B with pruning ratio of 70%, both LLM-Pruner and FLAP fail to produce an effective pruned model with their respective methods. In contrast, our method successfully maintains the most of the model’s capabilities.

We further evaluate the zero-shot capabilities of the pruned model across seven downstream tasks. The results of LLaMA-7B model are shown in Table 2. Full results, including additional pruning ratios and the LLaMA-13B model, are detailed in Table 9 in Appendix A.5. Our method demonstrates superior performance compared to the other three methods on those common sense reasoning zero-shot classification datasets. Besides, we show the results with LLaMA and LLaMA-2 models of our method on MMLU (Hendrycks et al. 2021) and GSM8K (Cobbe et al. 2021) datasets in Table 8 of Appendix A.4, which demonstrates that our method retains both generative and mathematical capabilities.

We show the results for LLaMA-2 and LLaMA-3 models with 2048 sequence length on WikiText2 dataset in Figure 2. The detailed perplexity results for both model families on three datasets are shown in Table 6 and Table 7 of Appendix A.2 and A.3. The blue line representing our method’s results consistently appears at the lowest position on the graphs, indicating its superior performance compared to the other methods with all model families.

3 Results of Image Generation

We implement the FLAP pruning method on LlamaGen model and compare this method on image generation task. We show the sparse results with LlamaGen-XXL (1.4B) and LlamaGen-3B models on ImageNet with 384$\times$384 resolution in Table 3. We observe that for the smaller model LlamaGen-XXL (1.4B), our method shows a distinct advantage at higher pruning ratios. For the larger model LlamaGen-3B, our method consistently outperforms across all pruning ratios, effectively preserving most of the original model’s capabilities. We further visualize the images generated by 10% sparsity models in Figure 3 with additional visualizations provided in Figure 6 of Appendix A.6. We observe that our method generates better image results compared to FLAP method in most cases.

4 Ablation Study

To demonstrate the effectiveness of our method for short sequence lengths, we present the results generated with a sequence length of 128 in Table 4 using the LLaMA-7B model and the WikiText2 dataset. Comprehensive results, including additional pruning ratios and datasets, are provided in Table 10 of Appendix A.7. As observed, our method consistently performs the best across all pruning ratios.

To verify the efficiency of the compensation process for our method, we conducted experiments using different numbers of samples. The results of these experiments are shown in Figure 4. The results demonstrate that the performance difference between compensation with 128 samples versus 512 or even 1024 samples is minimal across all pruning ratios. This indicates that 128 samples are sufficient for our compensation method, highlighting its efficiency.

We show the memory reduction and generation acceleration in Figure 5. The results are obtained using an NVIDIA A100 GPU with a sentence consisting of 64 tokens as the model input. The results show that as the pruning ratio increases, there is a corresponding decrease in GPU memory usage and an increase in generation speed, which validates the effectiveness of our method.

Conclusion and Limitation

In this paper, we propose the numerical score which is calculated through Newton’s Method for the minimization of pruning errors. We further sort numerical scores across all model layers for global pruning. Additionally, we introduce a compensation algorithm to reconstruct weights in pruned models. Experimental results show that our method achieves the state-of-the-art performance, which demonstrates the effectiveness of our method. Meanwhile, our method reduces memory usage and accelerates generation on GPUs without requiring additional implementations. One limitation of our method is its reduced effectiveness with smaller LlamaGen models for image generation tasks, primarily due to the usage of the discrete image tokenizer, which tends to lose important details as the sparsity increases.

References

Appendix A More Experimental Results

We show the full results of LLaMA family models on three different datasets in Table 5 varying sparsity ratio from 10% to 70% with all kinds of model size. Our method consistently achieves better performance than all other three state-of-the-art methods.

A.2 LLaMA-2 Results

We show the detailed perplexity results of LLaMA-2 family models on three different datasets in Table 6. Our methods achieves better performance and shows better generalization on different datasets than all the other methods.

A.3 LLaMA-3 Results

We further show the detailed perplexity results of LLaMA-3 family models on three different datasets in Table 7. LLM-Pruner codebase does not support LLaMA-3 models. We achieve consistent better performance than all the other models on three datasets with all kinds of model sizes.

A.4 Results on Generation and Math Datasets

We provide the results of our method on generation and math datasets in Table 8. The results demonstrate that our pruning method effectively preserves both the model’s generation capabilities and its mathematical performance.

A.5 Full Results for Common Sense Reasoning Datasets

We provide the task performance on common sense reasoning dataset with LLaMA-7B and LLaMA-13B models with 10% and 20% sparsity in Table 9. The results show that our method can perform better than other three methods.

A.6 More Results of Image Generation Tasks

We visualize more image generation results in Figure 6.

A.7 Full Results with 128 Sequence Length

We provide the full perplexity results with 128 input sequence length on three different datasets in Table 10. The results show the effectiveness of our method for short input sequence.

Appendix B Derivation of Theorem 3

The Lagrange function for Problem (3.4) is

The gradients with reference to $\delta W$ should be 0, i.e.,

By applying Equation (15) in Eq. B, we have the following,

Note that $e_{i}^{\top}e_{i}=1$ and $e_{i}^{\top}e_{j}=0$, if $i\neq j$. Besides, we can switch the position of two terms in the trace function, such as $X(2X^{\top}X)^{-1}M_{p}\lambda_{i}$ and $\lambda_{i}^{\top}M_{p}^{\top}(2XX^{\top})^{-1}X^{\top}$. Further, after switching the two terms in the trace function, if the output is a scale, we can omit the trace function.

The gradients with reference to $\lambda$ should be 0, i.e.,

We can obtain the optimal $\lambda$ as below,

The optimal $\delta W$ can be derived as below,

The minimal loss/error corresponding to the optimal $\delta W$ can be obtained by

As demonstrated in Eq. (B), since $e_{i}$ is a one-hot vector, $Ae_{i}e_{i}^{\top}$ only has non-zero values in the $i$-th column with all zeros for all other columns. Thus, in the sum of Eq. (B), each term indexed by $i$ just computes the $i$-th column of the output. Furthermore, the computation of the $i$-th column does not affect the $j$-th column, $\forall j\neq i$. For each column, we have the following,

Thus we can obtain the following optimal solution,

Appendix C Preliminary

In this section, we present preliminary concepts, several basic facts, and definitions for our paper.

C.2 Facts

$\|a\circ b\|_{2}\leq\|a\|_{2}\cdot\|b\|_{2}$.

$a\circ b=b\circ a=\operatorname{diag}(a)b=\operatorname{diag}(b)a$.

$\operatorname{diag}(a)C\operatorname{diag}(b)=(ab^{\top})\circ C$.

$AB^{\top}=\sum_{i\in[d]}A_{*,i}B_{*,i}^{\top}$.

C.3 Internal Computation Alignment

where $Q=X\cdot W_{q},K=X\cdot W_{k},V=X\cdot W_{v}$. Multiplying $W_{o}$, we have the final output:

Appendix D Theory

Optimization is a cornerstone of computer science and mathematics, encompassing various techniques to find optimal solutions. Linear Programming (LP) and Semi-Definite Programming (SDP) form a fundamental basis in this field, addressing problems with linear/quadratic objective functions and constraints, offering powerful tools for diverse applications (Anstreicher 2000; d’Aspremont et al. 2006; Amini and Wainwright 2009; Diakonikolas et al. 2019; Dong, Hopkins, and Li 2019; Cohen, Lee, and Song 2021; Jambulapati, Li, and Tian 2020; Jiang et al. 2020a; Gu and Song 2022; Song, Ye, and Zhang 2023). The concept of dynamic maintenance has gained prominence, focusing on efficient solution updates as input data changes (Cohen, Lee, and Song 2021; Lee, Song, and Zhang 2019; Brand 2020; Jiang et al. 2020b; Brand et al. 2020a; Jiang et al. 2020c; Song and Yu 2021; Dong, Lee, and Ye 2023; van den Brand 2020; Jiang et al. 2020a; Huang et al. 2021a; Gu and Song 2022). In machine learning, optimization techniques have significantly impacted Support Vector Machines (SVMs) (Chang and Lin 2001; Joachims 2006; Gu, Song, and Zhang 2023; Gao et al. 2023a; Tarzanagh et al. 2023; Brand, Song, and Zhou 2023; Li, Song, and Zhou 2023; Gao, Mahadevan, and Song 2023) and Empirical Risk Minimization (ERM) (Nesterov 1983; Vapnik 1991; Polyak and Juditsky 1992; Bartlett, Bousquet, and Mendelson 2005; Bottou and Bousquet 2007; Nemirovski et al. 2009; Moulines and Bach 2011; Feldman et al. 2012; Nesterov 2013; Johnson and Zhang 2013; Vapnik 2013; Shalev-Shwartz and Zhang 2013; Défossez and Bach 2014; Frostig et al. 2015; Zhang and Xiao 2017; Jin et al. 2018; Lee, Song, and Zhang 2019; Lianke et al. 2023; Bian, Song, and Yin 2023), enhancing their efficiency and performance. Recent work continues to expand the application of optimization across various domains, including graph algorithms, numerical methods, and advanced machine learning architectures (Li and Liang 2018; Du et al. 2018; Arora et al. 2019a, b; Song and Yang 2019; Cai et al. 2019; Zhang, Martens, and Grosse 2019; Cao and Gu 2019; Zou and Gu 2019; Oymak and Soltanolkotabi 2020; Ji and Telgarsky 2019; Lee et al. 2020; Huang et al. 2021b; Zhang et al. 2020b; Brand et al. 2020b; Zhang et al. 2020a; Song, Zhang, and Zhang 2021; Alman et al. 2022; Zhang 2022; Gao, Mahadevan, and Song 2023; Li, Song, and Zhou 2023; Qin, Song, and Yang 2023; Chu, Song, and Yang 2024; Shrivastava, Song, and Xu 2023; Qin, Song, and Sun 2023; Deng, Mahadevan, and Song 2023; Liang et al. 2024a; Li et al. 2024c), demonstrating its ongoing significance in advancing computational capabilities.

Attention mechanisms have been extensively studied and developed over the years, becoming a fundamental component in various neural network architectures. (Deng, Li, and Song 2023; Li et al. 2023a; Gu et al. 2024a) investigate the softmax regression problem, while (Song, Ye, and Zhang 2023) propose and analyze the attention kernel regression problem. The rescaled hyperbolic functions regression is examined by (Gao, Song, and Yin 2023). Following this, (Song, Wang, and Yin 2023; Li et al. 2023c) delve into two-layer attention regression problems. (Gu et al. 2024b) demonstrate that attention layers in transformers learn two-dimensional cosine functions. Additionally, (Deng et al. 2023) explore data recovery using attention weights, and (Kacham, Mirrokni, and Zhong 2023; Song, Xu, and Yin 2023) investigate the replacement of the softmax unit with a polynomial unit. Moreover, some works theoretically explore variations or combinations of the attention mechanism with other techniques, such as quantum attention (Gao et al. 2023b, 2024), tensor attention (Alman and Song 2024; Liang et al. 2024e), and differentially private attention (Liang et al. 2024d; Gao et al. 2023c; Gu et al. 2024c) and other applications such as (Clark et al. 2019b; Tenney, Das, and Pavlick 2019; Hewitt and Liang 2019; Vig and Belinkov 2019; Belinkov 2022; Brand, Song, and Zhou 2023; Chen et al. 2024c, b, a; Liang et al. 2024b; Li et al. 2024b; Liang et al. 2024c, f, e).

D.2 Error Bound for Masked Weight

In this section, we show how to derive the error bound for masked weight. We assume $B=1$ for simplicity of proofs.

Let $M\in\{0,1\}^{D}$ and $\rho\in$ be ratio of number of zeros to number of entries in $M$.

We can show that for $i\in[D^{\prime}]$ we have

where the first step follows from basic algebra, the second step and the third step follow from Fact 6, the last step follows from we assume $\|X\|\leq R$ and the definition of $\rho$. ∎

D.3 Find Optimal Mask

In this section, we show how to use Newton’s method to find the numerical score defined in Section 3.2.

Let $r\in[0,D]$ denote the number of ones (it can be a fractional number).

Let $\lambda>0$ denote a regularization co-efficients.

There exist an algorithm (Algorithm 1) that can get the optimal $z$ such that

We can use Newton’s method to solve this problem. To use Newton’s method we need to compute the gradient and Hessian. After calculation in Lemma 12, 13, 14, and 15, we use Algorithm 1 to get optimal $z$. ∎

Then, we calculate the gradient for each loss.

Let $L(z)_{i}:=0.5\|XW_{*,i}-X(z\circ W_{*,i})\|_{2}^{2}$ for $i\in[D^{\prime}]$.

Part 1. For each $i\in[D^{\prime}],j\in[D]$

Part 2. For each $i\in[D^{\prime}],j\in[D]$

where the first step follows from Definition 11, the second step follows from Fact 8, and the last step follows from simple algebra.

where the first step follows from definition of $L(z)_{i}$, the second step follows from Fact 8, the third step follows from simple algebra, and the last step follows from Part 1.

where the first step follows from Part 2, the second step follows from property of transpose, the third step follows from simple algebra, the fourth step follows from Fact 7, and the last step follows from Fact 7.

where the first step follows from definition of $L(z)$, the second step follows from basic calculus, the third step follows from Part 3, and the last step follows from Fact 7. ∎

Proof of part 2. Using Part 1, we can show

Hessian Calculation

Let $L(z)_{i}:=0.5\|XW_{*,i}-X(z\circ W_{*,i})\|_{2}^{2}$ for $i\in[D^{\prime}]$.

where the first step follows from basic calculus, the second step follows from Lemma 12, and the last step follows from Fact 8. ∎

where the first step follows from basic calculus, the second step follows from Lemma 13, the third step follows from basic calculus, and the last step follows from basic calculus. ∎