SP: Enhancing Structured Pruning via PCA Projection (2024)

Abstract: Structured pruning is a widely used technique for reducing the size of pre-trained language models (PLMs), but current methods often overlook the potential of compressing the hidden dimension (d) in PLMs, a dimension critical to model size and efficiency. This paper introduces a novel structured pruning approach, Structured Pruning with PCA Projection (SP), targeting the effective reduction of d by projecting features into a space defined by principal components before masking. Extensive experiments on benchmarks (GLUE and SQuAD) show that SP can reduce d by 70%, compress 94% of the BERTbase model, and maintain over 96% accuracy and outperform other methods that compress d by 6% in accuracy at the same compression ratio. SP has also proven effective with other models, including OPT and Llama. Our data and code are available at ours repo.

heading: 1 Introduction
text: Pre-trained language models (PLMs) like BERT (Devlin et al., 2018), GPT-2 (Radford et al.), ChatGPT (Ouyang et al., 2022), and Llama (Touvron et al., 2023), leveraging transformer architectures, have gained prominence in Natural Language Processing due to their exceptional performance. However, their substantial storage requirements and long inference times present practical challenges, prompting the need for efforts to optimize their size and computational speed.
Among various model compression techniques, distillation (Sanh et al., 2019) and, pruning (Louizos et al., 2018), particularly structural pruning (Xia et al., 2022), are prominent. Distillation is adaptable, allowing for the design of
∗Corresponding author. †{huyuxuan1999,zhang-jing}@ruc.edu.cn,
nlpzhezhao@tencent.com, zhaochen100@ruc.edu.cn, chenxiaodong6677@163.com, {licuiping,chong}@ruc.edu.cn
compact student models while transferring knowledge from larger teacher models. Nevertheless, it might not fully leverage the teacher model’s parameters, often requiring the student model to undergo extensive training on large datasets. In contrast, structured pruning directly trims the original model, thereby utilizing its parameters more efficiently and lessening the dependency on large datasets.
The transformer architecture in PLMs is defined by weight matrices controlled by hidden dimension d, attention head size dh in multi-head attention layers, and filter dimension df in feed-forward network layers. Focused efforts in structured pruning have aimed at downsizing dh and df (Michel et al., 2019; Kwon et al., 2022; Voita et al., 2019; Wang et al., 2020). Yet, strategies for pruning d are underexplored, with limited compression success. For instance, CoFi (Xia et al., 2022) achieves a notable 95% sparsity and maintains over 90% accuracy but only marginally reduces d from 768 to 760. We posit that the hidden dimension d holds significant potential for compression. Recent research demonstrates that high-dimensional features produced by models can be projected into a subspace, leveraging their inherent low-rank properties (Chen et al., 2021). Our observations further reveal that the sizes of these subspaces can differ among various layers. This variability underscores the necessity of compressing d within each individual layer, as the distribution of principal components’ weights changes across layers, a phenomenon depicted in Fig. 1 (a). Unfortunately, presently, distillation remains the most effective means for compressing d. This paper aims to explore a potent structured pruning method for efficiently reducing d in PLMs.
We analyze and conjecture that existing pruning methods rarely compress d effectively for two main reasons. First, conventional mask-based pruning methods directly reduce dimensions in the original feature space, which doesn’t effectively eliminate unimportant dimensions. We interpret the prun-
3150
ing of d from the perspective of features’ principal components. Drawing inspiration from PCA theory (Maćkiewicz and Ratajczak, 1993), effective compression should preserve the principal components of the features. Considering that most dimensions in the original feature space are linked to these principal components, as indicated by the evenly distributed values across each row that represents a principal component in Fig. 1(b), traditional maskbased pruning in this space struggles to preserve the principal components’ integrity. Secondly, the residual structure in each block, including both the multi-head attention (MHA) and the feed-forward network (FFN), requires the input and output dimensions of each block to be identical, hindering the pruning of different d across various layers.
To overcome these limitations, we propose an
enhanced Structured Pruning approach using PCA Projection (SP3) to reduce the hidden dimension. This method involves projecting the original features into a new space, where principal components act as the basis vectors. In this way, mask-based structured pruning becomes more effective in eliminating less important dimensions while retaining the principal components. For the projection, we select several tokens from the dataset designated for pruning and process them through the transformer model to acquire their feature matrix. Subsequently, we apply SVD to this matrix to obtain the projection matrix, which is then integrated before the mask matrix in the conventional pruning methodology. This process enables the projection of the feature matrix into a space defined by the principal components (as depicted in Fig.1 (c)). Additionally, to enable different d in different layers, we introduce additional linear transformations in the residuals, resulting in layer-adapted structured pruning (as shown in Fig.1 (d)).
The primary contributions of this work are:
• We take a fresh perspective on structured pruning by reducing the hidden dimension of the transformer model, which has been rarely explored before.
• Our SP3 method enhances compression efficacy by projecting features into space formed by principal components before masking, which helps preserve the principal components of features. Furthermore, it integrates linear transformations in residual connections, enabling variable hidden dimensions across different layers for a more customized structured pruning approach.
• Extensive testing on the GLUE and SQuAD benchmarks reveals that SP3 can effectively reduce the hidden dimension by 70%, compressing 94% of the BERTbase model while still maintaining an accuracy of over 96%. When compared to other methods that also compress the hidden dimension (d), SP3 demonstrates a 6% improvement in accuracy at the same compression ratio.

heading: 2 Background

heading: 2.1 Transformer Architecture
text: A typical transformer comprises a stack of transformer layers, each containing an MHA block and an FFN block. We take BERT (Devlin et al., 2018) as an example to describe this architecture.
Multi Head Attention (MHA). The MHA block takes x ∈ Rd×N as input and consists of H attention heads that facilitate interactions between tokens, with a LayerNorm (Ba et al., 2016) step.
xM = LN(x+MHA(x)), MHA(x) = H∑
i=1
Att(i)(x),
Att(i)(x) = W (i)⊤ O W (i) V x · Softmax(
(W (i) K x) ⊤(W (i)Q x)√ dh ),
(1)
where W (i)Q ,W (i) K ∈ RdQK×d, W (i) V ,W (i) O ∈ RdVO×d denote the query, key, value, and output matrices, respectively, d denotes the hidden dimension, dh = d/H denotes the head size, N denotes the sequence length, and dQK and dVO denote the intermediate dimensions.
Feed Forward Network (FFN). The FFN block takes xM as input and generates xF as output:
xF = LN(xM + FFN(xM)), FFN(xM) = W ⊤ D gelu(WUxM),
(2)
where gelu is the activation function, and WU ,WD ∈ Rdf×d are parameters of an FFN block, df indicates the intermediate dimension.

heading: 2.2 Mask-based Structured Pruning
text: Mask-based structured pruning approaches eliminate parameters by integrating additional mask variables z into the original model’s parameters, as follows:
xM = diag(zout)LN(diag(zin)(x+ zMHAMHA(x))), MHA(x) = H∑
i=1
z (i) headAtt (i)(x),
xF = diag(zout)LN(diag(zin)(xM + zFFNFFN(xM))), FFN(xM) = W ⊤ D diag(zf)gelu(WUxM),
(3)
where zMHA, zFFN ∈ {0, 1}, zhead ∈ {0, 1}H , zf ∈ {0, 1}df , zin, zout ∈ {0, 1}d are used to remove entire MHA block, FFN block, attention head, intermediate dimensions of FFN block, and hidden dimensions, respectively. Then, the pruning process can be modeled as an optimization problem with sparsity constraints as follows,
minL(θ) s.t. sparsity(z) ≥ t, (4)
where θ denotes the parameters of the model, and t denotes the pre-defined target sparsity.
3 SP3
Workflow of SP3. The SP3 we propose unfolds in three phases, as depicted in Fig. 2. Initially, PCA projection is conducted by selecting a subset of the training data, processing it through the transformer model to yield the feature matrix X for each block per layer, and then determining each X’s principal components U via SVD decomposition. Following
this, a projection matrix, initialized with these principal components, is integrated into the model’s computational framework, coupled with a mask for structured pruning, detailed in Section 3.1. Subsequently, the model undergoes pruning based on the masks tailored to a specific task and pruning objective, as outlined in Section 3.2. The final stage involves merging these masks and projection matrices with the original model parameters, thereby creating the compressed model, which is further elaborated in Section 3.3.

heading: 3.1 PCA Projection
text: Current structured pruning methods typically compress PLMs by directly applying a mask to their hidden dimensions. However, this approach overlooks the relationship between each feature dimension and its corresponding principal component. Principal component analysis of these features reveals that a few principal components hold the majority of the weight in every layer (as indicated in Fig. 1(a)). Additionally, these principal components are connected to a large number of dimensions, as evidenced by the uniformly distributed values across each row representing a principal component (see Fig. 1(b)). Consequently, directly pruning dimensions can adversely affect several principal components, thereby complicating the compression of hidden dimensions.
To address this, we incorporate a projection matrix prior to each masking matrix in the computational process of PLMs, referred to as PCA Projection. This matrix is designed to project features into a space determined by their principal components. Subsequently, these principal components can be pruned directly using masks. It’s crucial to ensure that these added projection matrices do not alter the output of the PLM and can be seamlessly integrated with the model’s parameters. In the following section, we detail the methodology for injecting these projection matrices to meet both of these conditions.
Hidden Dimension PCA Projection. To decrease the hidden dimension (d) between layers, our initial step involves selecting calibration data from the training dataset. We then derive the corresponding feature matrix X ∈ Rd×T from either the MHA block or the FFN block, along with their respective residuals. Here, T represents the number of sampled tokens. We choose T to be marginally greater than d to guarantee that the number of tokens ex-
ceeds the number of feature dimension. This feature matrix X acts as the input for the LayerNorm layer. Subsequently, we apply SVD to extract the principal components of the feature. Each column of the matrix U represents a principal component, obtained from the SVD of the matrix X , as described by U,Σ, V ⊤ = SVD(X − X), where X represents the matrix obtained by broadcasting the average values of each column.
Next, we introduce the matrix U both before and after each LayerNorm layer as follows. Initially, we recognize:
LN(X) = diag(γ)norm(X −X) = diag(γ)norm(UUT (X −X)),
(5)
where LN signifies LayerNorm module (the bias term is omitted for simplicity), norm represents the normalization function defined as norm(x) = x/∥x∥∗d , and UUT = I because columns of U are orthogonal. Given the following formulas,
norm(Ux) = Ux/∥Ux∥ ∗ d = Ux/ √ (x⊤U⊤Ux) ∗ d
= Ux/ √
(x⊤x) ∗ d = Unorm(x),
(6)
X −X = (I − 1 d 11⊤)X
= RX, (7)
where we define R = (I− 1 d 11⊤) using I to represent the identity matrix and 1 to denote a vector filled with ones. Then, we have,
LN(X) = diag(γ)Unorm(U⊤RX). (8)
Based on this, we assign projection matrices Pin = U⊤R and Pout = diag(γ)U , respectively.
The computation process after introducing the hidden dimension PCA Projection is as follows,
xM = P M outdiag(z M out)LN(diag(z M in )P M in (x+ zMHAMHA(x))), xF = P F outdiag(z F out)LN(diag(z F in)P F in(xM + zFFNFFN(xM))),
(9)
where the superscripts M and F denote the projection and mask matrices within the MHA and FFN blocks, respectively.
MHA Intermediate Dimension PCA Projection. While our focus is primarily on compressing the hidden dimension, our method is also applicable to compressing the intermediate dimension of the MHA block. We begin by collecting the feature
matrix X ∈ Rd×N from the output preceding the MHA block. We aim to project the features W (i)Q X , W
(i) K X and W (i) V X using their principal components
U (i) Q , U (i) K , and U (i) V respectively. They should satisfy the following equations:
W (i) V X = U (i)⊤ V U (i) V W (i) V X,
X⊤W (i)⊤Q W (i) K X = X ⊤W (i)⊤Q U (i)⊤ Q U (i) K W (i) K X,
(10)
where U (i)V in the equation can be solved directly by SVD, i.e., U (i)V = U , where U,Σ, V ⊤ = SVD(WV X). The solution of U (i)Q , U (i) K in Eq. 10 is a more complex, and we follow the method in DRONE (Chen et al., 2021) to compute U (i)Q , U (i) K :
U (i) Q = Σ
(i) 1 2 Z U (i)⊤ Z Σ (i)−1 1 U (i)⊤ 1 ,
U (i) K = Σ
(i) 1 2
Z V (i) Z Σ (i)−1 2 U (i)⊤ 2 ,
(11)
where
U (i) 1 ,Σ (i) 1 , V (i)⊤ 1 = SVD(W (i) Q X), U (i) 2 ,Σ (i) 2 , V (i)⊤ 2 = SVD(W (i) K X),
Z(i) = Σ (i)⊤ 1 U (i)⊤ 1 U (i) 2 Σ (i) 2 ,
U (i) Z ,Σ (i) Z , V (i)⊤ Z = SVD(Z (i)).
(12)
We then assign P (i)Q = U (i) Q , P (i) K = U (i) K , P (i) V = U (i)⊤ V , and P (i)O = U (i) V as our target projection matrices. For the proof of Eq. 11, refer to Appendix A. The computation process incorporating MHA intermediate dimension PCA Projection is as follows,
Att(i)(x) = W (i)⊤ O P (i)⊤ O diag(z (i) O ) ⊤diag(z(i)V )P (i) V W (i) V x·
Softmax( (W
(i) K x) ⊤P (i)⊤K diag(z (i) K ) ⊤diag(z(i)Q )P (i) Q (W (i) Q x)√
dh ).
(13)
Model Structure. We integrate the aforementioned PCA projections across various dimensions and present the final structure of our model in Fig. 3. Unlike mask-based pruning methods such as CoFi, which employ a shared mask in each block, our approach utilizes distinct masks for different blocks and incorporates specific projection matrices for each block (as contrasted in Eq. 3 and Eq. 9). This strategy enables us to independently compress the hidden dimension of each block. Additionally, we insert projection matrices in the MHA block to compress its intermediate dimension (as detailed in Eq. 13). The transformer’s residual structure necessitates the same hidden dimension across different blocks. By introducing
a linear transformation in the residuals, we effectively circumvent this constraint (as illustrated in Figure 1(d)). Readers interested in the algorithmic details can find the pseudo-code for the PCA Projection in the Appendix B.
Complexity. The computational complexity of PCA Projection primarily hinges on three factors: the number of layers n, the hidden dimension d, and the number of sampled tokens T , resulting in an approximate complexity of O(n ∗ d2 ∗ T ). In practice, for a model with 100M parameters, the PCA Projection process takes around 5 minutes. For a larger model with 7B parameters, we estimate that the PCA Projection would require roughly 1 hour, significantly less than the time required for model pruning.

heading: 3.2 Model Pruning
text: Objective Function. Pruning is governed by the masks incorporated into the model, and the model’s sparsity is determined based on these masks. Throughout the training process, all mask variables are treated as continuous real numbers. Upon completion of the training, those mask variables falling below a certain threshold — set in accordance with the desired sparsity level — are converted to 0, thereby establishing the final pruned structure. We use pruning loss (Wang et al., 2020) that forces the expected sparsity of the model to be close to the desired sparsity:
Lpruning = λ1 · (ŝ− t) + λ2 · (ŝ− t)2, (14)
where ŝ is the expected sparsity, t is the target sparsity, and λ1, λ2 are two Lagrange multipliers. Please refer to Appendix C for more detail of ŝ. Additionally, we incorporate the same task-specific loss, Ltask, as used in CoFi (Xia et al., 2022), applied to the task-oriented dataset during the pruning process to maintain model performance on specified tasks.
Multi-stage Pruning. In SP3, we implement masks at three hierarchical levels: dimension-level (zin, zout, zf , zQ, zK, zV, zO), head-level (zhead), and layer-level (zMHA, zFFN). Our findings reveal that during training, the model prioritizes headlevel and layer-level masks for compression, while dimension-level masks are sometimes less emphasized. Consequently, we have divided the training process into two distinct stages. In the first stage, we exclusively employ dimension-level masks, and
then, in the subsequent stage, we add head-level and layer-level masks. Empirically, the initial stage spans 5 epochs, followed by the second stage, which lasts for 15 epochs.

heading: 3.3 Model Fusing
text: After training, we first prune projection matrices, attention heads, MHA blocks, and FFN blocks based on different masks. After that, we merge the projection matrices with weight matrices. Suppose function Pr denotes pruning all zero rows and columns of the matrix.
MHA Layer Fusing. For the MHA layer, we have
Ŵ (i) Q = Pr(diag(zQ)PQW (i) Q P M outdiag(z M out)), Ŵ (i) K = Pr(diag(zK)PKW (i) K P M outdiag(z M out)), Ŵ (i) V = Pr(diag(zV)PV W (i) V P M outdiag(z M out)),
Ŵ (i)⊤ O = Pr(diag(z F in)P F inW (i)⊤ O POdiag(zO)),
(15)
where Ŵ (i)Q , Ŵ (i) K , Ŵ (i) V , and Ŵ (i) O denotes the compressed weight matrices of the MHA block.
FFN Layer Fusing. For the FFN layer, we have
ŴU = Pr(diag(zf )WUP M outdiag(z M out)), Ŵ⊤D = Pr(diag(z F in)P F inW ⊤ D diag(zf )),
(16)
where ŴU and ŴD denotes the compressed weight matrices of the FFN block.
Residual Modification. Since we adaptively learn different hidden dimensions in different layers, we need to add weight matrices to the residual. For example, for i-th block of Transformer, we have
x(i) =LN(W (i) R x+Block (i)(x(i−1)))), (17)
where WR is a matrix added to the residual, which is determined by multiplying the pruned projection matrices before and after the block (refer to Figure 1 (d)), i.e.,
W (i) R = Pr(diag(z (i) in )P (i) in P (i−1) out diag(z (i−1) out )). (18)

heading: 4 Experiments

heading: 4.1 Setup
text: Datasets. We evaluate our approach on GLUE (Wang et al., 2018) tasks and SQuAD (Rajpurkar et al., 2016) v1.1. GLUE tasks include SST-2 (Socher et al., 2013), MNLI (Kim et al., 2019), QQP (Wang et al., 2017), QNLI, MRPC (Dolan and Brockett, 2005), STS-B, and RTE. Each comparison model is evaluated on the test set of these
datasets. For the GLUE benchmark, accuracy metrics are used for MNLI, QQP, QNLI, SST2, MRPC, and RTE tasks, while Spearman’s correlation is reported for STS-B. In the case of the SQuAD benchmark, we present the F1 score. For comprehensive details about these datasets, please refer to Appendix D.1.
Baselines. We evaluate the proposed SP3 against established distillation methods like TinyBERT (Jiao et al., 2020), WID (Wu et al., 2023b), and the pruning approach CoFi (Xia et al., 2022).
Settings. Our primary experiments focus on compressing the BERTbase model. SP3 and CoFi perform model pruning using the previously mentioned training datasets. In contrast, WID and TinyBERT additionally use more general pre-training data. For a fairer comparison, we have excluded the data augmentation component from TinyBERT in our experiments. To generate calibration data for the SVD decomposition in SP3, we randomly choose 512 samples from each dataset and concatenate them, and then from them, we randomly select 4,096 tokens as our calibration dataset. For detailed information on our experimental methods, please see the Appendix D.2.

heading: 4.2 Main Results
text: As shown in Table 1, we compress the BERTbase model and compare the performance of SP3 with other methods under similar sparsity. First, compared to the original model, SP3 retains 96% of the model performance while removing 94% of the model parameters. Second, compared to other baselines, SP3 has extra degrees of freedom. Compared to the pruning method CoFi, SP3 can additionally compress the hidden dimensions and intermediate dimensions of the MHA block. Unlike distillation methods, SP3 can use different hidden dimensions at different layers with additional head-level and layer-level pruning. Overall, SP3 obtains the best performance on all datasets, indicating that utilizing these extra degrees of freedom can benefit the performance during compression. The comparisons with additional baselines on BERTbase and experiments on TinyBERT with augmentation and OPT-125m are presented in Appendix E.
We delve into the impact of specifically compressing the hidden size d in our SP3 . This entails retaining only the projection matrices P ·in, P · out and masks z·int, z · out in SP 3 (refer to Eq. 9), while main-
taining the masks zint, zout in the CoFi approach (see Eq. 3). As indicated in Table 2, SP3 demonstrates significant superiority over CoFi, which relies solely on masks for compressing the hidden dimensions. These experimental findings underscore the efficacy of PCA Projection in the compression of hidden dimensions, affirming its effectiveness in enhancing model performance.
We also apply our method to Large Language Models (LLMs), with experimental settings and results detailed in Appendix F. Table 8 illustrates the effective compression achieved by our method on TinyLlaMa. Additionally, thanks to the preRMSNorm architecture of TinyLlaMa, we propose Group PCA Projection. This approach unifies the computation of main components of features across several consecutive blocks, effectively reducing parameters introduced at the residuals and improving performance compared to distinct main components in each block, as demonstrated in Table 9.

heading: 4.3 Ablation Study on PCA Projection
text: PCA Projection employs principal components as the projection matrix. To assess the importance of these principal components, we conduct an experiment where we initialize the projection matrices
as identity matrices. To prevent other masks from interfering with the role of PCA Projection, we use only dimension-level masks in our experiments. The results, as depicted in Table 3, indicate that excluding PCA Projection results in a significant performance decline on the RTE and MRPC datasets, with a slight decrease observed on the SST-2 and STS-B datasets. This is because, for larger datasets, the model can directly optimize suitable compactor parameters from the data. The results underscore the significant role of PCA Projection in enhancing the performance of the compressed model.
Furthermore, PCA Projection relies on calibration data derived from the training dataset. We examine how the number of tokens in this data affects model performance. As shown in Table 4, the accuracy of the model varies with different sample token sizes. Considering that a higher number of samples adds complexity to the SVD computation, choosing 4,096 tokens strikes an optimal balance between accuracy and computational efficiency. We also conduct experiments using different random seeds to sample calibration data and show its robustness. Appendix G presents the details as well as the studies about different masks.

heading: 4.4 Structures of Pruned Model
text: We study the pruned structures produced by SP3. Take the MRPC dataset as an example. Figure 4 shows the structural information of the pruned model. More results are shown in Appendix H.
From Fig. 4 (b) and (c), as well as figures for other datasets, it’s evident that the model structure varies across different datasets. However, a consistent observation across these structures is that layers nearer the output are more compressed than those closer to the inputs. Additionally, the intermediate dimensions of the FFN block are notably more compressed across all datasets compared to the intermediate dimensions of the MHA block. This distinction is highlighted when comparing the green bars to the blue and red bars in Fig. 4 (b). The compression patterns we observe are consistent with those reported in prior pruning studies (Xia et al., 2022). In addition, Fig. 4 (a) and similar figures for other datasets indicate a trend where the model’s hidden dimension reduces as the number of layers increases. This trend suggests that the model incrementally compresses features into more compact dimensions during its inference process.

heading: 5 Related Work
text: Distillation. Knowledge distillation (Hinton et al., 2015; Hou et al., 2020; Jiao et al., 2020; Sanh et al., 2019) is a model compression approach that trans-
fers knowledge from a larger teacher model to a smaller student model. In contrast to the distillation methods, our approach uses less computational cost while achieving the same performance.
Pruning. Existing pruning methods can be broadly divided into two categories: unstructured and structured. Unstructured pruning (Gale et al., 2019; Frankle and Carbin, 2018; Kurtic et al., 2022; Louizos et al., 2018; Sanh et al., 2020) aims to remove unimportant scalar values from the model’s parameters. In contrast, structured pruning (Lin et al., 2020; Lagunas et al., 2021; Wang et al., 2020; Xia et al., 2022) are proposed to remove weight blocks in PLMs, including the entire layer (Fan et al., 2019; Sajjad et al., 2023, 2020), attention heads (Michel et al., 2019; Voita et al., 2019), and filters (McCarley et al., 2019; Prasanna et al., 2020). Unlike the previous structured pruning, we focus on exploring how to compress the hidden dimensions of the model efficiently.
Low-Rank Factorization (LRF). LRF methods compress the PLMs by decomposing the weight matrices (Hua et al., 2022; Xiao et al., 2023) based on low-rank property. In contrast, we propose PCA Projection based on the low-rank property of the features. For more discussions refer to Appendix I.
Re-parameterization (Re-p). Re-p methods have been proposed to improve existing model compression methods (Ding et al., 2021; Lin et al., 2021; Wu et al., 2023b). Its core idea is to represent the compressed model based on the original model. Therefore, these methods can utilize the parameters of the model more efficiently, which leads to better performance. SP3 can also be regarded as a Re-p method, and the major difference between SP3 and the existing methods is that we better initialize the parameters based on the PCA theory.

heading: 6 Conclusion
text: This study introduces SP3, an enhanced structured pruning approach for compressing PLMs. SP3 employs PCA Projection, facilitating easier optimization of the compression model, thereby boosting its
performance. Concurrently, SP3 uses non-shared masks in the hidden dimension. This allows the model to selectively determine its hidden dimensions size at different layers in line with the desired sparsity. When applied to BERTbase and evaluated on the GLUE and SQuAD benchmarks, SP3 notably achieves a 94% sparsity with only a minor 4% reduction in accuracy.
Limitations
Our proposed SP3 introduces extra weight matrices in the residual parts when merging the inserted compactor matrices with the weight matrices. When the model is compressed to 5M, these extra parameters account for more than half of the model. This predominance hinders further compression. In future research, we aim to explore strategies to eliminate these extraneous parameters. Meanwhile, the efficient acceleration of the model mainly relies on coarse-grained pruning (head-level, layer-level), while our approach uses more fine-grained pruning, and how to balance the fine-grained and coarsegrained pruning to achieve higher acceleration also needs further exploration.
When applying to large language models, integrating compactors during training in our methodology may lead to increased memory consumption, necessitating efforts to reduce memory overhead. Additionally, experimental results indicate that our tuning-based pruning technique may face challenges with overfitting, particularly when tuning using a limited SFT dataset. Therefore, investigating a tuning-free pruning technique is warranted.
Ethical Considerations
Intellectual Property. The datasets we use include oasst_top1_2023-08-25, MNLI, QNLI, QQP, SST-2, MRPC, STS-B, RTE, SQuAD, ARC-e, ARC-c, BoolQ, HellaSwag, openbookQA, PIQA, and WinoGrande. The models we use include BERTbase, OPT125m, and TinyLlama. These are publicly accessible and well-established resources aimed at facilitating diverse AI and NLP research endeavors.
Data annotation. We utilize the annotations provided by existing datasets, thereby eliminating the need for manual annotation in our study.
Intended Use. SP3 is a model pruning technique
designed to compress the parameters of a transformer model while preserving performance, effectively reducing the computational resources needed for deploying the model.
Misuse risks. SP3 is a pruning method for transformer models and incorrect use of SP3 might degrade the performance of some applications.
Misuse Control. We intend to make our approach available to the open-source community, enabling users to gain a deeper understanding of our methodology and mitigate the risk of misuse.

heading: Acknowledgments
text: This work is supported by the National Key Research & Develop Plan (2023YFF0725100) and the National Natural Science Foundation of China (62322214, U23A20299, 62076245, 62072460, 62172424, 62276270). This work is supported by Public Computing Cloud, Renmin University of China.

heading: A MHA Intermediate Dimension PCA Projection
text: We can verify the effectiveness of the Eq. 11 in the following way,
∥X⊤W (i)⊤Q W (i) K X −X⊤W (i)⊤ Q U (i)⊤ Q,:kU (i) K,:kW (i) K X∥,
(19)
Eq. 19 measures the difference between the inner product computed using the first k principal components of U (i)Q , U (i) K and the original result.
Next, we prove that U (i)Q and U (i) K satisfy the
following equation,
Mk = U (i)⊤ Q,:kU (i) K,:k (20)
where,

heading: Mk = argmin
text: M ∥X⊤W (i)⊤Q W (i) K X −X⊤W (i)⊤ Q MW (i) K X∥,
s.t. rank(M) = k (21)
Suppose
XQ = W (i) Q X, XK = W (i) K X, U1,Σ1, V ⊤ 1 = SVD(XQ), U2,Σ2, V ⊤ 2 = SVD(XK).
(22)
Then, we have
Mk = argmin M ∥X⊤QXK −X⊤QMXK∥
= argmin M ∥V ⊤1 X⊤QXKV2 − V ⊤1 X⊤QMXKV2∥ = argmin M ∥Σ⊤1 U⊤1 U2Σ2 − Σ⊤1 U⊤1 MU2Σ2∥ (23)
Define
Z = Σ⊤1 U ⊤ 1 U2Σ2, UZ ,ΣZ , V ⊤ Z = SVD(Z),
(24)
then, based on the PCA theory, we know that Zk = UZ,:kΣZ,:kV ⊤ Z,:k is the solution of
argmin Z∗
∥Z − Z∗∥, s.t. rank(Z∗) = k. (25)
Therefore, we have
Σ⊤QU ⊤ 1 MkU2Σ2 = UZ,:kΣZ,:kV ⊤ Z,:k
⇒ Mk = U1Σ−1⊤1 UZ,:kΣZ,:kV ⊤Z,:kΣ−12 U⊤2 . (26)
Refer to Eq. 11,
U (i) Q = Σ
(i) 1 2 Z U (i)⊤ Z Σ (i)−1 1 U (i)⊤ 1 ,
U (i) K = Σ
(i) 1 2
Z V (i) Z Σ (i)−1 2 U (i)⊤ 2 ,
(27)
we also have
U (i) Q,:k = Σ
(i) 1 2 Z,:kU (i)⊤ Z,:k Σ (i)−1 1 U (i)⊤ 1 ,
U (i) K,:k = Σ
(i) 1 2
Z,:kV (i) Z,:kΣ (i)−1 2 U (i)⊤ 2 .
(28)
Therefore,
Mk = U (i)⊤ Q,:kU (i) K,:k (29)

heading: B Pseudo Code of PCA Projection
text: The pseudo-code of PCA Projection is shown in Algorithm 1.

heading: C Sparsity
text: The expected sparsity ŝ is computed as follows:
ŝ = 1
M (
L∑
i
H∑
j
d∑
k
dh∑
l
z (i) MHA · z (i,j) head · z (i,k) out,FFN,i−1 · z (i,l) Q +
L∑
i
H∑
j
d∑
k
dh∑
l
z (i) MHA · z (i,j) head · z (i,k) out,FFN,i−1 · z (i,l) K +
L∑
i
H∑
j
d∑
k
dh∑
l
z (i) MHA · z (i,j) head · z (i,k) out,FFN,i−1 · z (i,l) V +
L∑
i
H∑
j
d∑
k
dh∑
l
z (i) MHA · z (i,j) head · z (i,k) in,MHA,i · z (i,l) O +
L∑
i
d∑
k
df∑
l
z (i) FFN · z (i,k) out,MHA,i · z (i,l) f +
L∑
i
d∑
k
df∑
l
z (i) FFN · z (i,k) in,FFN,i · z (i,l) f +
L∑
i
d∑
k
z (i,k) out,FFN,i−1 · z (i,k) in,MHA,i+
L∑
i
d∑
k
z (i,k) out,MHA,i · z (i,k) in,FFN,i),
(30)
where M denotes the total number of parameters of the model.

heading: D Experiment Setting Details
text: D.1 Datasets
GLUE (Wang et al., 2018) benchmark consists of various tasks related to sentence similarity calculation, classification, textual entailment, and natural language inference. It includes 10 tasks, namely AX, COLA, QQP, MNLI, MRPC, QNLI, QQP, RTE, SST-2, STS-B, and WNLI. The number of training examples for each task is as follows: 1.1k, 10.7k, 432k, 5.8k, 105k, 364k, 3k, 70k, 67k, and 852, respectively. SQuAD 1.1 (Rajpurkar et al., 2016) dataset involves question answering tasks, containing 88K training examples.
D.2 Experiment Setup
Our SP3 is developed using PyTorch (Paszke et al., 2019) and executed on a server equipped with four NVIDIA 3090 GPU cards for all experiments. We source the BERT model from the HuggingFace (Wolf et al., 2019) Transformers library. To prepare for compression, we first fine-tune BERT to create a task-specific model. This fine-tuning occurs on the training datasets from the GLUE benchmark and
SQuAD, spanning 3 epochs, with a batch size of 32 and learning rates of 1e-5 and 2e-5, respectively. We maintain the default hyperparameters provided by HuggingFace.
The pruning of BERT follows the same configuration as the model fine-tuning process. We first fine-tune SP3 with the task-specific objective for 2 epochs, then dimension-level pruning begins at the 3rd epoch, followed by the head-level and layerlevel pruning at the 8th epoch.

heading: E More Experimental Results
text: E.1 Comparing SP3 with More Baselines The comparison results between SP3 and more baselines are shown in Table 5.
E.2 Comparing SP3 with TinyBERT We skip the data augmentation step in TinyBERT, following the approach of CoFi (Xia et al., 2022), as this data augmentation significantly enlarges the training dataset, leading to a substantial increase in training time. Here, we just present the results of TinyBERT with data augmentation (DA) for reference. The comparison results are shown in Table 6. Even with data augmentation, our proposed SP3 demonstrates comparable performance.
E.3 Experiment Results on OPT125m The results of the experiment on OPT125m are shown in Table 7. For OPT125m, we use the same experimental setup as for BERTbase.
F SP3 for Large Language Model
SP3 can also be used to prune large language model (LLM), e.g., Llama (Touvron et al., 2023)). To mitigate computational overhead of LLM, our approach restricts using SFT data exclusively for model pruning. Concurrently, to reduce pruning complexities,
we only prune the hidden dimensions and the filters within the FFN block, while ignoring the pruning of attention heads and layers, because coarse-grained pruning requires more training data to recover the performance of the model. We also ignore the pruning of the intermediate dimensions of the MHA block, as it may conflict with the RoPE positional encoding (Su et al., 2024) which is commonly used for LLM. The structure used to prune the LLM is shown in Fig. 5.
F.1 Group PCA Projection As shown in section 3.3, to compress the hidden dimensions of the model, we need to add additional matrices to each residual. These additional added matrices increase the total number of parameters in the model by 15%. Fortunately, we find that if we use the same principal components to initialize the compactors of successive blocks, we can avoid introducing additional matrices at the residuals. Taking Llama as an example, for the i-th block of the Llama model, we have
x(i) = Block(i)(RMSNorm(x(i−1))) + x(i−1). (31)
After computing the principal component matrix U , we have
x(i) = Block(i)(RMSNorm(UU⊤x(i−1))) + x(i−1), (32)
= Block(i)(U · RMSNorm(U⊤x(i−1))) + x(i−1). (33)
Suppose x̂(i−1) = U⊤x(i−1), x̂(i) = U ′⊤x(i). Then, we have
x̂(i) = U ′⊤Block(i)(U · RMSNorm(x̂(i−1))) + U ′⊤Ux̂(i−1).
(34)
It can be seen that if we make U ′ equal to U , we don’t need to add extra matrices at the residuals, because at this point we have U
′⊤U = I . We refer to this unified computation of the main components of features in several consecutive blocks as Group PCA Projection. For LLM, the matrix introduced when the group size is 8 only increases the number of model parameters by 1%.
It is also worth noting that Group PCA Projection requires models to use RMSNorm with a prenorm architecture, but fortunately, based on existing work (Jiang et al., 2023), all models that use LayerNorm can be directly converted to use RMSNorm. Moreover, nearly all LLMs use the pre-norm architecture.
F.2 Experiments
Setup. Following existing works on pruning LLM (Xia et al., 2023; Ma et al., 2023). we consider the commonsense reasoning and question-answering tasks including Hellaswag (Zellers et al., 2019), OpenBookQA (Mihaylov et al., 2018), WinoGrande (Sakaguchi et al., 2021), ARC-Easy and ARC-Challenge (Clark et al., 2018), BoolQ (Clark et al., 2019), and PIQA (Bisk et al., 2020). We evaluate our models using the Language Model Evaluation Harness (Gao et al., 2023) framework. We use the SFT dataset oasst_top1_2023-08-25 to prune the TinyLlaMa-1.1B-Chat-v0.3 (Peiyuan Zhang and Lu, 2023). Without additional declarations, for TinyLlama we use Group PCA Projection with group size 8.
Results. The experimental results on TinyLlaMa are shown in Table 8. We finetune the model on the dataset for 2 epochs with a learning rate of 1e-5 and a batch size of 4 to obtain the pruned model. Then, we test the performance of the pruned model under a zero-shot setting. Our method updates both the model parameters and the mask. Since for limited SFT data, training too many epochs may lead to model overfitting, while too few updates can lead to non-convergence of the mask. Therefore we use smaller batch sizes so that convergence of the mask is guaranteed in case the model does not overfit the data.
Impact of Group Size. We also explore the effect of using different group sizes in Group PCA Projection on model performance. The experimental results are shown in Table 9. To prevent the interference of additional factors, we only consider the compression of the hidden dimension in this experiment. From the experimental results, it can be seen that using a smaller group size does not improve the performance of the model. On the contrary, a group size of 8 can be utilized without introducing too many additional parameters and performs better than a smaller group size. We speculate that larger group sizes with fewer additional parameters at residuals act as a form of regularization, helping to prevent overfitting of the calibration data by the results of principal component analysis.
Combining with Existing Pruning Methods. Our method can also be combined with other pruning methods, e.g., LLM-Pruner (Ma et al., 2023). For these methods, we first initialize the compactors
using PCA projection, after which the compactors are directly fused to the model. Additionally, to maintain compatibility with existing pruning methods and avoid altering the model’s structure, we utilize the same principal component matrix across all layers. In this way, SP3 can be viewed as a re-representation of the model making it easier to be pruned. In this experiment, we combine SP3 with LLM-Pruner. The experimental results are shown in Table 10. It can be seen that combining our method can improve the performance of the LLM-pruner.

heading: G More Ablation Experiments
text: G.1 Impact of Multi-level Masks
SP3 uses different levels of masks to compress the model. To explore the impact of different levels of masks on model performance, we conduct the following experiments: (1) Use all levels of masks. (2) Ignore head-level masks. (3) Ignore layer-level masks. (4) Ignore head-level and layer-level masks.
The experimental results are presented in Table 11. Observations indicate superior model performance on the SST-2 and RTE datasets when all mask levels are utilized. For the STS-B datasets, the removal of the head-level mask results in the most precise models. Notably, the optimal performance on the MRPC dataset is achieved by a model that excludes both the head-level and layer-level masks. Given the data volume in each dataset, we hypothesize that minor alterations in the head-level and layer-level masks can significantly influence model outputs compared to the dimension-level mask. This implies that the head-level and layerlevel masks might be more challenging to optimize. Consequently, removing either the head-level or layer-level mask in smaller datasets can stabilize the optimization process, leading to a more precise model. As dataset sizes increase, the need for model compression flexibility becomes evident, with multi-level masking yielding superior outcomes.
G.2 Impact of Different Random Seeds
SP3 needs to use the calibration data to initialize all the compactors. To verify the robustness of SP3 under different calibration data, we conduct experiments using different random seeds to sample the calibration data. We use SP3 to prune BERTbase and OPT125m on datasets QNLI, SST-2, MRPC, STS-B, and RTE using a range of random seeds
(0, 1, 2, 3, 4, 42), where BERTbase are pruned to 5M and OPT125m are pruned to 7M. The results of the experiment are shown in Table 12. Employing these six random seeds has produced various calibration sets, and our experimental findings indicate that these different calibration sets have a minimal impact on the model’s performance after compression due to their small standard deviation.

heading: H Structures of Pruned Model
text: The structure of pruned models on RTE, SST-2, STS-B, MNLI, QNLI, QQP and SQuAD are shown in Fig. 6 and Fig. 7. We show the model structure in terms of dimension-level and head-level sparsity. Rather than directly displaying layer-level sparsity, we indirectly illustrate it through histograms depicting dimension-level and head-level sparsity. If a value at a certain position is 0, it suggests layerlevel pruning has taken place at that position.
I Difference Between SP3 and Low-rank Factorization Methods
Existing low-rank factorization methods (Garg et al., 2019; Liu and Ng, 2022; Ma et al., 2019; Riera et al., 2022; Yin et al., 2022; Zhou et al., 2019; Zhang and Wang, 2022) primarily focus on decomposing the weight matrix using PCA. In contrast, our approach applies PCA to the features generated during the model inference process and re-represents these features based on the principal components. This feature representation enhances compression efficacy and offers a distinct perspective from traditional weight matrix decomposition using PCA.
SliceGPT (Ashkboos et al., 2024) stands out as the most pertinent work to our research, as it focuses on model compression via principal component analysis of the feature matrix. However, notable distinctions exist between our approach and SliceGPT. Firstly, SliceGPT overlooks the variability in principal component distributions across layers, opting for a uniform compression rate of hidden dimension d across all layers, while our approach adaptively determines the compression rate for different layers through learnable masks. Secondly, it neglects strategies for eliminating parameters introduced by residuals, which affects the compression rate of the compressed model. To address this problem, we instead propose the Group PCA Projection approach to reduce the number of parameters added at the residuals. Lastly, SliceGPT
exclusively addresses decoder-only model architectures, omitting consideration for other model configurations. In contrast, we consider models with a decoder architecture, such as Llama, and OPT, as well as models with an encoder architecture, such as BERT.
Algorithm 1: PCA Projection Input: T : sampled tokens size, D: calibration data, L(1∼N): model Output: Projection Matrix Pin, Pout, PQ, PK , PV , PO XQ, XK, XV, XM, XF ← [], [], [], [], [] for j ← 1 to |D| do
x(0) ← D[j] for i← 1 to N do
[WQ,WK ,WV ,WO]← L(i) for h← 1 to dh do
X (h,i) Q ← Concate(X (h,i) Q ,W (h) Q x (i−1)) X (h,i) K ← Concate(X (h,i) K ,W (h) Q x (i−1)) X (h,i) V ← Concate(X (h,i) V ,W (h) Q x
(i−1)) end xM = MHA (i)(x(i−1)) + x(i−1) xLNM = LN(xM) xF = FFN (i)(xM) + xM xLNF = LN(xF) x(i) = xLNF X
(i) M ← Concate(X (i) M , xM)
X (i) F ← Concate(X (i) F , xF)
end end for i← 1 to N do
[WQ,WK ,WV ,WO]← L(i) for h← 1 to dh do
X (h,i) V ← Sample(X (h,i) V , T ) X (h,i) Q ← Sample(X (h,i) Q , T ) X (h,i) K ← Sample(X (h,i) K , T ) UV ,ΣV , V ⊤ V ← SVD(X(h,i)V ) U1,Σ1, V ⊤ 1 ← SVD(X(h,i)Q ) U2,Σ2, V ⊤ 2 ← SVD(X(h,i)K ) Z = Σ⊤1 U ⊤ 1 U2Σ2 UZ ,ΣZ , V ⊤ Z = SVD(Z) UQ = Σ 1 2 ZU ⊤ Z Σ −1 1 U ⊤ 1 UK = Σ 1 2 ZVZΣ −1 2 U ⊤ 2 P (h,i) V , P (h,i) O ← U⊤V , UV P (h,i) Q , P (h,i) K ← UQ, UK
end [γM , γF ]← L(i) R← I − 1
d 11⊤
X (i) M ← Sample(X (i) M , T ) X (i) F ← Sample(X (i) F , T ) UM,ΣM, V ⊤ M ← SVD(X(i)M ) UF,ΣF, V ⊤ F ← SVD(X(i)F ) P (M,i) in , P (M,i) out ← U⊤MR, diag(γM)UM P (F,i) in , P (F,i) out ← U⊤F R, diag(γF)UF
end

heading: Future Work
text: In future research, we aim to explore strategies to eliminate these extraneous parameters. Meanwhile, the efficient acceleration of the model mainly relies on coarse-grained pruning (head-level, layer-level), while our approach uses more fine-grained pruning, and how to balance the fine-grained and coarsegrained pruning to achieve higher acceleration also needs further exploration. When applying to large language models, integrating compactors during training in our methodology may lead to increased memory consumption, necessitating efforts to reduce memory overhead. Additionally, experimental results indicate that our tuning-based pruning technique may face challenges with overfitting, particularly when tuning using a limited SFT dataset. Therefore, investigating a tuning-free pruning technique is warranted. Ethical Considerations
Intellectual Property. The datasets we use include oasst_top1_2023-08-25, MNLI, QNLI, QQP, SST-2, MRPC, STS-B, RTE, SQuAD, ARC-e, ARC-c, BoolQ, HellaSwag, openbookQA, PIQA, and WinoGrande. The models we use include BERTbase, OPT125m, and TinyLlama. These are publicly accessible and well-established resources aimed at facilitating diverse AI and NLP research endeavors. Data annotation. We utilize the annotations provided by existing datasets, thereby eliminating the need for manual annotation in our study. Intended Use. SP3 is a model pruning technique
designed to compress the parameters of a transformer model while preserving performance, effectively reducing the computational resources needed for deploying the model. Misuse risks. SP3 is a pruning method for transformer models and incorrect use of SP3 might degrade the performance of some applications. Misuse Control. We intend to make our approach available to the open-source community, enabling users to gain a deeper understanding of our methodology and mitigate the risk of misuse