E XPLORING EXTREME PARAMETER COMPRESSION FOR PRE - TRAINED LANGUAGE MODELS
A BSTRACT
Recent work explored the potential of large-scale Transformer-based pre-trained models, especially Pre-trained Language Models (PLMs) in natural language processing.
This raises many concerns from various perspectives, e g., ﬁnancial costs and carbon emissions.
Compressing PLMs like BERT with negligible performance loss for faster inference and cheaper deployment has attracted much attention.
In this work, we aim to explore larger compression ratios for PLMs, among which tensor decomposition is a potential but under-investigated one.
Two decomposition and reconstruction protocols are further proposed to improve the effectiveness and efﬁciency during compression.
Our compressed BERT 1 with 1 / 7 parameters in
Transformer layers performs on-par with, sometimes slightly better than the original BERT in GLUE benchmark.
A tiny version achieves 96.7% performance of
BERT-base with 1 / 48 encoder parameters (i e., less than 2M parameters excluding the embedding layer) and 2 .
7 × faster on inference.
To show that the proposed method is orthogonal to existing compression methods like knowledge distillation, we also explore the beneﬁt of the proposed method on a distilled BERT.
I NTRODUCTION
Pre-trained Language Models such as BERT (Devlin et al., 2018) and ALBERT (Lan et al., 2019) have signiﬁcantly improved various NLP tasks with signiﬁcant improvement.
Much recent work (Brown et al., 2020; Narayanan et al., 2021; Fedus et al., 2021) explores the potential of super large-scale
PLMs.
However, such large-scale PLMs are both economically and ecologically unfriendly (Bender et al., 2021; Patterson et al., 2021).
Furthermore, deployment of large-scale PLMs is also challenging since (1) a model cannot be fully deployed or stored in a single GPU server, model parallelism would consume extra time for network communication among many servers; (2) edge devices may not have enough space for storing models; (3) the long inference time cannot support real-time feedback.
Scaling down a model with negligible performance drop would facilitate the real-world applications of PLMs in a smaller size, faster inference time, and less network communication cost.
For example, recent work explores quantization (Zhang et al., 2020; Bai et al., 2020), weights pruning (Hou et al., 2020), and knowledge distillation (Jiao et al., 2020; Sanh et al., 2020) for BERT (one of the most popular PLMs).
We argue that existing methods cannot largely compress large-scale PLMs as stated in Sec. 2.
In this paper, we aim to explore extreme parameter compression (i.e., bigger compression ratios) although they are by deﬁnition challenging.
The parameter redundancy in PLMs was demonstrated by (Kovaleva et al., 2019; Michel et al., 2019; Voita et al., 2019; Cordonnier et al., 2021), for which we divide into two groups: intramatrix redundancy and inter-matrix redundancy .
The former happens in different heads that are calculated separately, e.g., attentions among heads act on a similar subspace and are therefore lowrank (Cordonnier et al., 2021) – we relate this phenomenon to the so-called ‘decomposability’ deﬁned in this paper.
Like self-attention layers, decomposability also holds in FFN layers – each FFN layer could be decomposed to many independent sub-FFNs (as explained in Appendix B).
One example of inter-matrix redundancy happens across different layers , e.g., attention maps among layers might be similar (Clark et al., 2019; Vig, 2019; Rogers et al., 2020).
Exploration of main weight matrices in Transformer layers ﬁnds that these weight matrices are possible to be approximated in a low-rank manner – evidencing the possible intra-matrix redundancy and inter-matrix redundancy .
We comprehensively analyze and compare different decomposition methods for parameter compression including matrix decomposition (denoted as II), tensor train decomposition (Oseledets, 2011) (denoted as III) and Tucker decomposition (De Lathauwer et al., 2000) (denoted as IV).
The fundamental difference between them is as below.
II conducts matrix factorization (e.g., SVD) for each weight matrix thanks to intra-matrix redundancy .
Regarding intermatrix redundancy , III shares the head and tail matrices while keeping the core matrix individual; IV introduces ‘matrix bank’ to make parameter scale being nearly constant w.r.t.
the number of layers.
It is concluded that Tucker decomposition (IV) is more parameter-efﬁcient than others in terms of compression ratios.
ALBERT (Lan et al., 2019) and III can be considered as special cases of IV.
The practical challenges of matrix/tensor decomposition for compression are twofold.
First, the decomposition may result in a discrepancy between the raw weights and approximated weights, and exact decomposition is impossible with large compression ratios.
Instead, Knowledge Distillation (KD) is used on the compressed model to simulate the predictions of the raw model in a loss-aware manner.
Second, reconstruction may lead to additional computation costs.
An efﬁcient reconstruction protocol is implemented by reordering multiplication operations that also preserve the same results.
The contributions of this work are (1) we propose a formal framework with standardized terminology to comprehensively discuss matrix/tensor decomposition methods to compress Transformer-based language models; (2) we adopt tensor decomposition for compressing PLMs which is also faster, while existing work (Ma et al., 2019; Liu et al., 2021) did not show the potential for speedup in PLMs; (3) our compressed BERT with 1 / 7 parameters in Transformer layers performs on-par with the original BERT in GLUE benchmark.
Also, a tiny version achieves 96.7% performance of BERT-base with only 1 / 48 parameters in Transformer layers and 2 .
7 × faster on inference.
We directly use the proposed methods on TinyBERT (Jiao et al., 2020) that is purely based on KD, since our work is complementary to existing compression methods like KD.
R ELATED W ORK
Compressing PLMs Although various work was proposed to design a new efﬁcient Transformer (Tay et al., 2020), e.g., (Ma et al., 2019; Choromanski et al., 2020; Wang et al., 2020;
Kitaev et al., 2020; Zaheer et al., 2020; Cao et al., 2020), in this paper, we are focusing on the compression of Transformer-based pre-trained language models.
The difference is that the latter expects to reuse well-trained models, e.g., BERT and GPT (and even GPT3 (Radford et al., 2019) and
PanGu- α (Zeng et al., 2021)) with manageable computing resources, which typically does not change the original Transformer architecture.
Taking BERT, one of the most commonly-used pre-trained language models, as an example.
Existing work explored quantization (Zhang et al., 2020; Bai et al., 2020), weights pruning (Lagunas et al., 2021), knowledge distillation (Jiao et al., 2020; Sanh et al., 2020), progressive module replacing (Xu et al., 2020), neural architecture search (Xu et al., 2021;
Yin et al., 2021) and matrix decomposition (Noach & Goldberg, 2020).
We argue that existing compression methods (see Tab. 1) may be inadequate for extreme parameter compression, which is under-investigated.
The reasons are manifold, ﬁrst, the knowledge distillationbased method generally learns a new student model from scratch, which cannot inherit too much knowledge from the teacher model before distillation.
Second, some methods have upper bounds of compression ratio.
For example, layer-sharing ALBERT (Lan et al., 2019) shares parameters in L layers with maximum L times compression.
Quantization replaces existing 32-bit parameters with binary parameters with a maximum of 32 times reduction.
Moreover, quantization needs further hardware support, which is usually ad hoc to speciﬁc platforms.
Weight pruning arguably cannot achieve a big compression ratio (McCarley et al., 2019; Michel et al., 2019).
[Table]
Matrix/tensor decomposition for compression Tensor/matrix decomposition aims to approximate a given tensor using a set of smaller tensors/matrices.
It has been investigated to compress and speed up CNNs, RNNs, and Transformers for many years (Lebedev et al., 2014;
Yang et al., 2017; Ye et al., 2018; Denton et al., 2014; Winata et al., 2019).
Matrix decomposition (e.g., ALBERT (Lan et al., 2019) in the embedding layer and (Noach & Goldberg, 2020)) could decrease parameter scale with a linear factor depending on the selected rank.
More advanced tensor decomposition approaches can be implemented by tensor network, which has recently been used to compress general neural networks (Gao et al., 2020; Novikov et al., 2015), compress embedding layer (Khrulkov et al., 2019;
Hrinchuk et al., 2020; Panahi et al., 2019).
Recently, Ma et al.
(2019) redesigned a new ‘Self-Attention Network’ (SAN) in Transformer architecture inspired by block-term tensor decomposition.
The compression ratio is limited since the majority of parameters in Transformer comes from another module called ‘Feed-Forward Network’ (FFN) instead of SAN; moreover, the model does not have the potential for speedup regardless of compression ratios since the time complexity is closed to vanilla Transformer.
Noach & Goldberg (2020) reparameterized each weight matrix using matrix decomposition and further distill the compressed models, which nearly achieves a compression ratio of 1.7.
Liu et al.
(2021) adopted matrix product operators to reparameterize each group of weight matrices in embedding, SAN, and FFN, and only a small part of tensors of MPO (called ‘auxiliary tensors’) are ﬁne-tuned in downstream tasks.
The compression ratio of the total parameter is negligible and the inference might be slow.
Those works inspire us to explore in depth extreme parameter compression for large-scale PLMs.
We argue that compression ratio of existing work using matrix/tensor decomposition (Ma et al., 2019;
Liu et al., 2021; Noach & Goldberg, 2020) for PLMs is relatively-small; most of them do not have speedup effect, limiting their applications in large-scale PLMs.
The potential to compress PLMs with matrix/tensor decomposition is under-investigated.
In this work, we adopt tensor decomposition, to cubically compress the parameters of PLMs.
M OTIVATIONS FOR PARAMETER COMPRESSION
Pre-trained language models are typically a stack of multiple Transformer (Vaswani et al., 2017) layers that consist of a Self-Attention Network (SAN) module and a Feed-Forward Network (FFN) module, see App.
A for Transformer.
Sec.
3.1 will introduce an important property called ‘decomposability’ for SAN and FFN, which indicates each sub-component in SAN or FFN is independently calculated without interactions between sub-components and they may therefore be redundant.
3. D ECOMPOSABILITY IN T RANSFORMER
A computing module f is decomposable if its sub-components { g 1 , g 2 , ⋯ g H } could be independently calculated without interactions: f ( x )
=
δ ( g 1 ( x ) , g 2 ( x ) , ⋯ , g H ( x )) .
Usually, δ is a simple operation that has negligible computing cost compared to { g h } .
Especially, backpropagation between sub-components is independent if δ is concatenation or addition.
Sub-components in f could be calculated in parallel without interactions.
We will examine if SAN and FFN are decomposable .
Decomposability in SAN
Following (Hou et al., 2020), SAN could be decomposed as a sum of the output of every head.
For the query/key/value/output transformations parameterized by
    
 3&$FRPSRQHQWV     F DS W X U HG  Y D U L DQ F H DYHUDJH (
[Figure]
W Q / W K / W V / W O , we divide them as N h heads:
[Equation]
[Equation] A single head is calculated as
[Equation]
Then SAN ( X ) =
∑ N h h = 1 Att h ( X ) , indicating SAN is decomposable.
Cordonnier et al.
(2021) argued that the attentions among heads learn redundant key/query projections due to independent calculation.
Decomposability in FFN The two weight matrices in FFN, denoted as W In and W Out .
[Equation]
As Geva et al.
(2020) points out, FFN also operates as a key-value mechanism similar to SAN: one can consider the input of FFN as a query vector, and two linear layers of FFN as keys and values, respectively.
There may exist some similar key-value pairs that may introduce redundancy.
Remarks on decomposability Decomposability does not necessarily make sub-components being complementary, especially without orthogonality-like constraints between them.
Sub-components may learn similar patterns, leading to some redundancy (Michel et al., 2019).
For example, Hou et al.
(2020) demonstrated that BERT could be pruned in both width and height without performance drop.
Moreover, BERT is over-parameterized in the sense that the number of training examples in downstream tasks is much fewer (393k training examples for the biggest task in GLUE, i.e., MNLI) comparing to the enormous number of parameters (110M parameters for BERT-base) – this was thought problematic according to the traditional machine learning theory (Vapnik, 2013).
We believe that the pre-trained language models in downstream tasks could be compressed.
3. E XPLORATION OF AN PRE - TRAINED T RANSFORMER
Beyond the decomposability of Transformer, we conduct exploration on a widely-used Transformerbased PLM, i.e., BERT (Devlin et al., 2018).
Technically, we calculate the captured variance ratio by
Principal Component Analysis (PCA), as an indicator to measure the parameter redundancy.
The main weight matrices of Transformer layers are { W Q , W K , W V , W O ,
W
In , W Out } 2 .
Sec.
2 and Appendix B show that FFN could be separately calculated in a multi-head fashion .
We could sequentially split both { W In } and { W Out } into four groups like { W Inh } h = 4 and { W Outh } h = reconcile
Here we could get 12 D × D weight blocks for each Transformer layer, 4 for SAN, and 8 for FFN
In Fig.
1 we could ﬁnd the intra-matrix and inter-matrix redundancy : Figure 1a shows that half dimensions could capture more than 90% variance of all weight matrices, this conﬁrms our statement in Sec. 3.1.
Furthermore, we also study the redundancy between two matrices by conducting
PCA on the concatenated matrix between two arbitrarily paired weight matrices.
See Fig.
1b, half dimensions could capture nearly 80% variance, which suggests some possibility to compress intermatrix redundancy .
This inter-matrix redundancy may be twofold: (1) subFFNs are decomposable; (2) calculations (e.g., attention maps) in different layers may be similar.
Regarding the latter, some existing works like RNNs (Goyal & Bengio, 2020), Neural Ordinary Differential Equations (ODE) (Chen et al., 2018), and cross-layer sharing ALBERT (Lan et al., 2019) show that it could work even with the parameter equivalence hypothesis among layers.
A GENERAL FRAMEWORK FOR PARAMETER COMPRESSION
The observation that the main weight blocks in BERT could be approximated in a low-rank manner (thanks to the intra-matrix and inter-matrix redundancy) inspires us to use decomposition.
Here we introduce and compare some standard decomposition methods (see App. D) in compressing PLMs.
4. E XPLORING PARAMETER COMPRESSION
In principle, SANs and FFNs could be separately compressed; in this work, we additionally explore stacking SAN weights and FFN weights together as a uniﬁed protocol since each weight block has an identical shape (i.e., D × D ).
The main weight matrices in the j -th Transformer layer are
[Equation]
Weights of a L -layers Transformer are stacked as a 3-rd order tensor in R 12 L × D × D .
The original non-decomposed weights is called I: W I = { W ( j ) }
Lj = 1 ∈ R 12 LD 2 .
Each weight matrix in W I is W
I
i =
W i ∈ R D × D .
Here, we explore standard decomposition methods including matrix decomposition, tensor train decomposition (Oseledets, 2011) and Tucker decomposition (De Lathauwer et al., 2000).
II: matrix decomposition Motivated by intra-matrix redundancy , one can adopt Matrix decomposition to factorize/approximate a matrix into some smaller ones.
A typical example is singular value decomposition (SVD), called ‘II- α ’, for each D × D matrix W i ∈
W I ,
[Equation]
[Equation] One can also drop the diagonal Σ i by decomposing it into two parts that are multiplied to U i and V i, namely W
i
≈ U
i V i , denoted as ‘II- β ’.
U i
∈
R D ×
d and V i ∈ R
d × D and usually d < D .
Since the compression ratio is D 2 2 Dd with reducing the rank from D to d , the preserved rank of the approximated matrix linearly decreases with the compressing rates.
[Table]
raw (Devlin et al., 2018)
[Table]
24LDd w/
[Table]
(Ben Noach & Goldberg,
[Table]
[Table] III: tensor train decomposition Inspired by the inter-matrix redundancy , one could expect to share weights among matrices.
The biggest terms in Eq.
4 are U i and V i while { Σ i } is relatively small since Σ i ∈ R d × d and d is relatively small compared to D. We could share { U i } and { V i } among matrices to save parameters.
This results in
[Equation]
Here, { Σ i } are not necessarily diagonal.
This results in a tensor-train (TT) decomposition (Oseledets, 2011) 4 .
One can also consider higher-order TT decomposition (i.e., a longer chain for tensor multiplications) which could be more parameter-efﬁcient; this often needs to reshape the raw tensor into a higher-order tensor with heuristics.
However, it is more time-consuming and costs more GPU memory during training, which we leave as future work.
IV: Tucker decomposition In Eq.
5, the biggest term is the { Σ i } ∈ R 12 L × d 2 , especially the number of layers may be large in practice (e.g., L = 24 for BERT-large).
To this end, we propose a ﬁxed-sized matrix bank such that a weight matrix is considered as a linear combination of these matrices inside the bank, making the parameter scale become nearly a constant with respect to the number of layers.
Namely,
[Equation]
[Equation] where C ∈ R l × d 2 is a matrix bank with a size of l , each matrix is assigned with a weight vector
P i ∈ R 1 × l .
ALBERT (Lan et al., 2019) could be considered as a special case of IV, see App.
G.
4. C OMPARISON BETWEEN I,II, III, AND IV
The comparison in parameter scale between these decomposition methods is in Tab. 2.
Since D > d and L > l , we could generally conclude that the parameter scales decrease from I, II, III to IV.
We can observe that marginal parameter cost to add a new layer in IV is nearly 12 l , which is negligible compared to the other parameters.
During the inference phase, the terms that do not involve batch size b or sequence length n could be calculated in an ofﬂine way only once before starting inference, which costs more storage but gets slightly acceleration – since the main purpose of this work is to compress models, we ignore it in this work but encourage doing it in speed-sensitive scenarios.
E XTREMELY COMPRESSING BERT USING T ENSOR D ECOMPOSITION 5. D ECOMPOSITION PROTOCOL
IV reduces space complexity from O ( 12 LD 2 ) to O ( ld 2 + 12 Ll + 2 Dd ) where d < D and l < 12 L .
l determines to which degree we want to share Transformer parameters among all modules, a ﬂexible factor to smoothly transform vanilla BERT to layer-shared BERT (or called ‘ALBERT’ (Lan et al., 2019)).
d determines the expressive power (rank) of each linear transformation (originally D × D ).
The decomposition protocol does not change the raw architecture of BERT, alternatively, it introduces a new reparameterization of the existing weights.
However, the approximated weights W IV usually
[Table]
are not exactly equal to the raw weights W I .
Moreover, the tiny decomposition discrepancy of weight matrices in low-layer may lead to an accumulated difference in the ﬁnal output due to the multiple-layer neural network architecture 5 .
In this work, we propose to use knowledge distillation to simulate the ﬁnal output of raw models.
[Equation]
f W I is the raw BERT model and f W IV is the compressed one.
We argue that approximation in prediction (like knowledge distillation in Eq. 7) is more important than approximation in weights.
Such a loss-aware strategy in compression could be found in quantization (Hou et al., 2018).
5. R ECONSTRUCTION PROTOCOL
A slice of D × D parameter block is represented as matrix product, W IV i ≈ U
( P i C ) V ∈ R D × D .
For an input X ∈ R b × n × D where b is the batch size and n is the sequence length, an output of linear transformation between X and a D × D parameter block will be Y =
XW i ;; =
X U ( P i C ) V .
Since matrix multiplication is associative 6 , different multiplication order will not affect the ﬁnal result but their computational complexity may be different 7 .
One can see the computational complexity for multiplication order in Tab. 3.
In practice, the batch size b will be set as big as possible to increase data throughput and make training more stable, we could conclude that IV-3 is more efﬁcient than
IV-2.
IV-3 is more efﬁcient than IV-1 when D > 2 d ; in practice, D is typically much bigger than d and D > 2 d .
In conclusion, setting IV-3 is most efﬁcient in this scenario.
E XPERIMENTS 6. S ETTINGS
Decomposition For BERT-base ( L = 12 , D = 768 , ),
W I ∈ R 144 D 2 is decomposed into a core tensor and three factor matrices (see Fig. 2), its reconstruction could be seen in Sec.
Knowledge distillation As (Jiao et al., 2020; Zhang et al., 2020; Bai et al., 2020) did, we use two-stage knowledge distillation for the compressed model.
At General Distillation (GD) stage, we adopt Knowledge Distillation (KD) for the compressed model to simulate the last-layer hidden states and last-layer attention maps of the general teacher model (BERT-base).
At the second stage, we adopt Task-speciﬁc Distillation (TD) to simulate the logits of a task-speciﬁc BERT model (e.g., ﬁne-tuned on MNLI task).
In GD, compressed models are trained with two epochs.
In TD, we also augment training data by randomly replacing a random word with a similar word according to either word vector similarity using Glove (Pennington et al., 2014) or the predicted logistics of BERT when masking the target word, see more details in (Jiao et al., 2020).
GLUE evaluation GLUE (Wang et al., 2018)
(see App.
I for more details) includes datasets for single document classiﬁcation and sentence pair classiﬁcation.
Fine-tuning and evaluation on GLUE follows the settings from Huggingface (Wolf et al., 2019).
The best-performed model is selected according to the dev set, where we select the learning rate in [ 1 e -5, 2 e -5] and batch size in [ 16 , 32 ] .
[Table]
[Table]
6. R ESULTS
As shown in Tab.
[Table]
our decomposed BERT with layer rank 144 and dimension rank 384, called ‘
BERT-III-384’, outperforms the BERT-base, with only 1 / 7 parameters in Transformer layers and slightly bigger throughout.
BERT-IV-72-384 performs on-par with raw BERT, which is slightly worse than BERT-III-384 due to the smaller size.
Observe that a bigger rank (both for layer mode and dimension mode) usually consistently leads to better performance.
BERT-III-64 achieves 96.7% performance (82.7 vs. 80.0) with only 1 / 48 parameters of Transformer layers and 2 .
7 × speedup.
Tab. 4 shows BERT-IVs are smaller than existing parameter sharing method (Lan et al., 2019) and decomposition method (Noach & Goldberg, 2020; Mao et al., 2020).
BERT-III -384/ BERT-IV -72-384 achieve comparable or even slightly better results than (Noach & Goldberg, 2020).
BERT-IV outperforms (Mao et al., 2020) with a large margin.
BERT-IV outperforms ALBERT – the latter needs training from scratch (1M training steps) while the former does not (less than 0.2M steps).
Observe that BERT-III-384 outperforms BERT-IV-72-384 since the former has more parameters and therefore is more expressive.
Note that BERT-III-384 and BERT-IV-d-384 have nearly identical
RPS and inference latency.
Note, we take BERT-IV-36-128 as an example to compare with matrix decomposition (a.k.a, II, which is implemented by Noach & Goldberg (2020) with a rank of 245, denoted as BERT-IV-245), BERT-IV-36-128 is faster (see RPS in Table 4 and inference time in 8), smaller, and better-performed (see Table 10 for full performance comparison) than BERT-IV-245, evidencing the advantage of BERT-IV over matrix decomposition for compression.
To contextualize BERT-III/BERT-IV with other compression methods like knowledge distillation,
Tab. 4 shows that BERT-III/BERT-IV achieves comparable performance with knowledge distillation methods (Sun et al., 2019; Xu et al., 2020; Jiao et al., 2020) while with fewer parameters.
Our results
[Table]
[Table] is also comparable to quantization methods Zhang et al.
; Bai et al.
(2020) that use 2-bit or 3-bit weights, and pruning Lagunas et al.
(2021) (see Tab. 11 in App.
L).
To show the proposed method is orthogonal to existing compression methods like pure knowledge distillation , we further explore the proposed method in TinyBERT (Jiao et al., 2020), called ‘TinyBERT-IV’.
Tab. 5 shows performance loss to compress TinyBERT (degrading from 83.0 to 81.
in TinyBERT-IV-72-384) is bigger than compressing raw BERT (degrading from 82.7 to 82.6 in
BERT-IV -72-384).
This is probably due to smaller redundancy in TinyBERT compared to BERT.
A NALYSIS
Ablation on the necessity of knowledge distillation Tab. 6 shows that both GD and TD are essential for an effective decomposed BERT.
In particular, the overall performance decreases from 82.6 to 80.0 by removing TD.
Note that the model will collapse if we directly take decomposed
BERT for ﬁne-tuning without knowledge distillation.
Decomposition on FFNs or SANs For FFNs and SANs, we use the half size of matrix bank (i.e., 24 for SANs and 48 for FFNs) and half dimension rank (i.e., 384) respectively.
The two settings are called ‘BERT-IV-FFN-48-384’ and ‘BERT-IV-SAN-24-384’.
Tab.
7 shows that solely compressing
SANs or FFNs could nearly achieve on par with the raw model since smaller compression ratios are achieved.
In detail, FFNs are slightly easier to be compressed even with a big compression ratio comparing to SANs.
It is intriguing to notice that BERT-IV-72-384 outperforms BERT-IV-SAN-24384, although the former additionally compresses FFNs and has much fewer parameters.
The reason may be that the size of the matrix bank in BERT-IV-72-384 (i.e., 72), which is shared between FFNs and SANs, is bigger than its counterpart in BERT-IV-SAN-24-384 (i.e., 24).
This can shed some light on the beneﬁt to stacking FFNs and SANs together, see more discussions in App F.
C ONCLUSION AND F UTURE W ORK
To largely compress PLMs, we also comprehensively compare many matrix/tensor decomposition methods and conclude tensor decomposition has the potential for extreme parameter compression.
We therefore efﬁciently implemented tensor decomposition inspired compression for BERT, thanks to an efﬁcient reconstruction protocol.
To compensate for the decomposition discrepancy, knowledge distillation is used to simulate the output of raw BERT, which is demonstrated to be essential in ablation study.
Our method with 1 / 7 parameters could achieve comparable with the original model, with slight speedup.
A tiny version achieves more than 96.7% performance of BERT-base with 1 /parameters in Transformer layers and 2 .
7 × faster on inference.
In the future, we expect compression of PLMs to shift purely encoder-based language models (e.g., BERT) to decoder language models, since the latter has been designed as big as we could afford, e.g. GPT3 (Radford et al., 2019).
The potential of the proposed methods for compressing larger model is discussed in App.
N. Furthermore, hybrid methods by mixing knowledge distillation, quantization, parameter sharing, weight pruning, and matrix/tensor decomposition together are potential in practice.
A B ACKGROUND OF T RANSFORMER AND BERT A. T RANSFORMER
A Transformer layer (see Fig. 3) consists of a self-attention (SAN) module and a feed-forward network (FFN) module.
An input X for SAN will be linearly transformed into query, key, value, and output space { Q, K, V } as below 8 :
[Equation]
The self-attention mechanism (a.k.a Scaled Dot-Product Attention) is calculated as
[Equation]
For a multi-head version of the self-attention mechanism, it linearly projects Q, K, V with h times using individual linear projections to smaller dimensions (e.g. d k = d model h ), instead of performing a single attention function with d model -dimensional keys, values and queries.
Finally, the output of
SAN is
[Equation]
where
Q =
[ Q 1 ; ⋯ Q h ] , K
=
[ K 1 ; ⋯ K h ] , and V =
[ V 1 ; ⋯ V h ] .
The individual attention heads are independently calculated, Cordonnier et al.
(2021) claims that there is some redundancy in multi-head attention.
Since the output of SAN is a linear transformation (using W O ) of V , which is a weighted sum of
V .
A stack of many purely SAN layers is not expressive (Dong et al., 2021), since it is equivalent to a single linear transformation.
To this end, a feed-forward network with non-linear activation is alternately used with each SAN layer,
[Equation]
Since some neurons after the activation function (e.g., δ is ReLU or GELU (Hendrycks & Gimpel, 2016)) become inactivated (zero), d in is usually bigger than d model to avoid the low-rank bottleneck, typically, d in = 4 × d model =
d out .
Other tricks, such as layer normalization, residual connection, dropout, and weight decay are also adopted to relieve the optimization and overﬁtting problems when it goes deeper.
A
. BERT
BERT is a Transformer architecture-based pre-trained language model trained on plain corpora by using a masked language model pre-training objective to the capacity and scalability of the
Transformer (Devlin et al., 2018).
BERT signiﬁcantly improved the SOTA performance of many downstream tasks, including classiﬁcation benchmark GLUE (Wang et al., 2018).
Note that the parameters of SAN and FAN linearly increase with the number of layers while the embedding layers keeps constant with respect to the layer number.
B ‘M ULT - HEAD ’ F EED F ORWARD N EURAL N ETWORK
[Figure]
them
In the multi-head attention mechanism, individual attention heads are separately calculated, however the calculation is low-rank due to the redundancy among heads (Cordonnier et al., 2021).
In this paper, we argue that such redundancy may also appear in the feed-forward neural networks.
In the feed-forward neural networks, element-wise activation functions are usually adopted, e.g.,
GELU (Hendrycks & Gimpel, 2016) and Relu (Agarap, 2018); that is, each activation can be independently calculated.
By partitioning the original W in into H column groups and W out into H row groups, one can revise a feed-forward neural layer (i.e. FFN ( X ) =
δ ( X W in ) W out ) as a sum of H independent ‘thin’ sub-FFNs with a dimension of D H = 4 D / H as below:
[Equation]
Where W in h ∈ R D
× D h and W out h ∈ R D h × D , W in =
[ W in 1 ⊕ ⋯ ⊕ W in H ] ∈ R D × 4 D and W out =
[ W out 1 ⊕ ⋯ ⊕ W out H ] ∈ R 4 D × D .
In this paper, we set H = 4 , since W in h , W out h ∈ R D
× D and each transformation in sub-FFNs layer are full-rank transformations.
See Fig 4 for graphical illustration.
One can refer to block partitioned matrix multiplication to understand the equivalence between Fig.
4a and 4b.
C CONCATENATION OVER ROWS OR COLUMNS
PCA on columnly-stacked and rowly-stacked matrices are shown in Fig.
5a and Fig.
5b.
Observe that there is no too much difference between them.
[Figure]
D M ATRIX AND TENSOR DECOMPOSITION
Tensor is a generalized ‘matrix’ with typically a dimension of 3 or more; sometimes, one can also call a matrix as a bi-dimensional ‘tensor’ and a vector as a one-dimensional ‘tensor’.
Formally, an n -dimensional tensor is an array with n indexes as
[Equation]
An entry in X is accessed by selecting n ordered index ( i 1 , ⋯ , i n ) , and { 0 ≤ i k < I n , i k ∈ N } .
In this section, we mainly discuss order-3 tensors for simplicity while higher-order tensors also hold.
Typical tensor decomposition includes Canonical Polyadic (CP) decomposition (Carroll & Chang, 1970; Harshman, 1972) and Tucker decomposition (Tucker, 1963).
CP decomposition approximates a high-dimension tensor with as a sum of many rank-one tensors.
In the case of three-dimensional tensor decomposition,
[Equation]
Where A r ∈ R
I 1 , B r ∈ R I 2 , C r ∈ R I 3 and A ∈ R R × I 1 , B ∈ R R
×
I 2 and C ∈ R R × I 3 .
⊗ is the tensor product 9 .
Tucker decomposition decomposes a tensor into a set of factor matrices and one small low-rank core tensor,
[Equation]
15) where A ∈ R R 1 ×
I 1 , B ∈ R R 2 ×
I 2 , C ∈ R R 3 × I 3 and G ∈ R R 1 × R 2 × R 3 .
× 1 , × 2 , and × 3 are mode- k products 10 .
{ R 1 , R 2 , R 3 } is sometimes called Tucker ranks .
An entry with index ( i 1 , i 2 , i 3 ) is calculated as
[Equation]
[Equation]
⎤⎥⎥⎥⎥⎥⎥⎥⎥⎥⎥⎥⎦
Tucker decomposition will degrade to CP decomposition,where the core tensor G is constrained to be super-diagonal 11 and R = R
= R 2 = R 3 .
A k -slice in the ﬁrst mode, i.e., a matrix with a size of I 2 ×
I 3 , would be ˆ X k =
R
∑ i 1 =
R
∑ i 1 =
R 3 ∑ i 3 =
G i 1 ,i 2 ,i 3 A i 1 ,k
( B
[Equation]
C i 3 ) , ˆ X k ∈ R I 2 × I(17)
All slices in a speciﬁc mode share factor matrices (i.e., B , C ) in other modes.
Therefore, there exists inter-correlation between these slices.
These shared factor matrices not only make the compression ratio of tensor decomposition being much bigger than multiple independent matrix decomposition (Noach & Goldberg, 2020), but also can be beneﬁcial especially when there is redundancy among these parameter matrices and factor matrices can be utilized common features.
E D IFFERENCE WITH EXISTING WORK E. D IFFERENCE WITH T ENSOR DECOMPOSITION USED FOR COMPRESSING CNN S , RNN S AND E MBEDDINGS
Compressing pre-training models is a new scenario.
We believe that exploring tensor decomposition in pre-trained language models is non-trivial.
In the pre-trained language models, we could test tensor decomposition in very deep and wide networks like GPT 3 (96 layers and a hidden dimension of 12288), while this work is the ﬁrst step.
Existing work Ma et al.
Liu et al.
; Gao et al.
Khrulkov et al.
; Hrinchuk et al.
Panahi et al.
(2019) which use tensor network for compressing neural networks do not have the potential for acceleration.
The bottleneck in speed limits the application of tensor decomposition in real-world applications: compressing models but consuming longer inference time seems be useful in very rare scenarios.
We argue that it is nontrivial to compress neural networks using tensor decomposition with acceleration effects, as this work did.
Work mechanisms for compression are totally different, previous works compress each weight matrix ( W Q , W K , W V , W O , W in , W out in each layer) individually using matrix/tensor decomposition.
They are making use of local redundancy inside each matrix.
While in big models (PLMs), we believe that making use of cross-matrix/cross-layer redundancy is also, sometimes more, beneﬁcial.
We believe that using tensor decomposition for cross-matrix/cross-layer redundancy is a signiﬁcant difference.
E. D IFFERENCE WITH C ORDONNIER ET AL . (2021)
(Cordonnier et al., 2021) is quite impressive and inspiring.
It does inspire this paper, however, we want to highlight the difference with (Cordonnier et al., 2021).
Motivation The motivation is different, (Cordonnier et al., 2021) found the redundancy in different heads.
We make use of redundancy of both intra-matrix redundancy and inter-matrix redundancy (including cross-layer redundancy).
Method The architecture in (Cordonnier et al., 2021) is slightly different from Transformer (or
BERT) since it redesigns the self-attention network with collaborated attention.
While our architecture is nearly the same as the raw model, we simulate each matrix multiplication of BERT with a product of many smaller matrices.
Goal Our work aims to extremely compress neural networks, which cannot be achieved by (Cordonnier et al., 2021).
(Cordonnier et al., 2021) is to explore the possibility to share key/query projection in SANs.
Note that SAN only has 1/3 parameters in Transformer, this proportion even becomes smaller when considering the embedding layer.
By compressing only SANs with 1/3 parameters, its overall compression ratio is limited.
Potential for efﬁciency (Cordonnier et al., 2021) is faster than the raw model only if D k is small.
The smallest setting with D k = 64 has 20% FLOPs reduction.
This shows its potential for acceleration is limited.
A similar setting with D k = 64 , we have nearly 80% reduction in terms of
FLOPs.
Potential for compression ratios The smallest model for BERT-base (110M) in (Cordonnier et al., 2021) has 96.6M when D k = 128 ; its best compression ratio is 1.14.
While in our model, the smallest model has a much bigger compression ratio, while being faster and performing better.
Potential for effectiveness The model in Cordonnier et al.
(2021) ( 96.6M parameters) achieves 93.5% (77.6/83.0) performance with BERT base, while our smallest model with 25M parameters (plus parameters in the embedding layer) achieves 97.7% performance (80.8/82.7) of BERT base.
A slight difference is that we test our model on the test dataset through GLUE online benchmark while
Cordonnier et al.
(2021) test their model on ofﬂine dev dataset through average results for three runs, so we use the relative performance w.r.t.
the original BERT-base.
F M ORE DISCUSSIONS TO COMPRESS WEIGHTS MATRICES IN SAN S AND FFN S TOGETHER F. S HARING LEADS TO A BIGGER MATRIX BANK
If we separately compress SANs and FFNs, we would have two matrix banks: one for SANs and one
FFNs: P FFN C FFN and P SAN C SAN .
Each weight matrix in FFN(or SAN) is speciﬁc to a matrix bank C FFN for FFN (or C SAN for SAN) and its weight vector over the bank P FFNi (or P SANi ).
Note that the two matrix banks have the most parameters since it is three-order tensor while others ( U , V , P ) are matrices.
Note that matrices in two matrix banks have the same shape, one could merge (share) the two matrix banks (a m -size d × d matrix bank and a n -size m -size d × d matrix bank) to get a single bigger ( ( m + n ) -size) matrix bank , this could boost the expressive power for both FFNs and SANs due to the bigger matrix bank (denoted as [ C FFN ; C SAN ] ).
The unmerged one is a special case of the merged one.
Let us deﬁne a new P ′ , each element in which is deﬁned as below:
[Equation]
[Table]
P ′
i is the new weight vector over the shared ( m + n ) -size matrix bank.
Without losing any generality, we could relieve the zero-constrains in P ′ to get a general form that each element in P ′ is not forced to be zero.
This could make FFNs and SANs share more matrix bases and therefore get better capacity.
The beneﬁt could be empirically evidenced in Table 7: solely compressing SANs (without compressing FFNs) underperforms both compressing FFNs and SANs; although the former has much more parameters.
Since in the shared setting, the shared matrix bank could compensate the counterpart of
SANs.
Another na¨ıve motivation is to design a uniﬁed protocol for both SANs and FFNs.
F. D ISCUSSIONS ON THE SCALE DIFFERENCE BETWEEN WEIGHTS MATRICES IN SAN S AND FFN S
We want to clarify that we do not intend to approximate the raw weights (and their norm/scale).
Previously, we tried to add a regularizer in the training loss function to force the reconstructed weights using decomposition to be as close as the raw weights.
                    
W Q W K
W V W o
W In W Out
[Equation]
This does not improve the performance, but worsen the performance.
The reason may be as below.
Since we cannot perfectly approximate the raw weights with a decent compression ratio, there always exists some difference between the raw weights and the reconstructed weights using decomposition.
However, even a slight difference might lead to a big difference in output.
Plus, we also observe that the ﬁnally learned W IV has a big difference with W I .
So we give up our efforts to approximate the raw weights.
For example, we found that even randomly initializing the compressed BERT nearly achieves identical performance, compared to decomposition from a given model.
This also shows that approximating raw weights is not beneﬁcial.
Instead, we use knowledge distillation to simulate the input-output mapping from the original model and tolerate the difference between raw weights and reconstructed weights.
We believe that the former makes more sense than the latter.
Empirical results show the former performs well.
Norm difference in decomposition Decomposition is to approximate a three-order tensor with three factor tensors and a core tensor.
[Equation]
Each matrix has its speciﬁc factor vector P i , which does not have any constraints.
The ﬂexibility of norms in P i could compensate the norm difference for the original matrices.
In case the readers may be curious, we show the norms in Figure 6.
The norms in different matrices do have some differences.
We also compare the norms between the raw weight matrices and the reconstructed weights matrices in Figure 7.
It shows that the norms of reconstructed weights also have a big difference compared to that of corresponding original weight matrices.
We argue that the knowledge distillation does not aim to approximate the raw weight matrices but only approximate its input-output mapping functions.
[Figure]
       
W Q W K W V
W O W
In W Out
W Q ,9 W K ,9 W V ,
W O ,9 W
In ,9 W Out ,
G L INKING IV TO ALBERT
For IV, we have W IV =
U ( P C ) V .
Considering a speciﬁc case when U =
V =
I ( D ) , l = 12 (i.e.,
C ∈ R 12 × D 2 ) and P =
⎡⎢⎢⎢
⎢⎢⎢⎢⎢⎢⎢⎣
I (
[Equation]
I
[Equation]
[Equation]
⋯ I ( 12 )
[Equation]
R 12 L × 12 .
Where I ( i ) is an i × i identity matrix.
Namely
[Equation]
Then
[Equation]
Eq.
19 results in a layer-shared BERT like ALBERT (Lan et al., 2019).
(Lan et al., 2019) additionally compresses the embedding layer using matrix factorization.
H H OW IV GENERALIZES III
III (tensor train decomposition) is a special case of IV (tucker decomposition) in this paper .
Note that, we are not stating the correspondence between tensor train decomposition and tucker decomposition in general case, but a three-order tensor train decomposition is a special case of tucker decomposition.
Let us recall their deﬁnition:
[Equation]
Or in another format:
[Equation]
× 2 and × 3 are the mode- 2 and mode- 3 multiplication.
The only difference Σ ∈ R 12 L × d × d vs.
( P C ) ∈
R 12 L × d × d .
In the latter, P ∈ R 12 L × l and C ∈ R l × d × d .
In a special case of tucker
[Table]
[Figure]
decomposition when 1) 12 L = l and 2) P =
I 12 L , it is equivalent to the 3-order tensor train decomposition (III).
I
L is a (square) identity matrix with a size of 12 L × 12 L , in which the diagonal element equals one and 0 elsewhere.
In our experiments
[Figure]
there existed two settings which set 12 L = l = 144 that satisﬁed 1).
In such settings, the rigorous difference with BERT-III and BERT-IV is that we do not force P to be equal to an identity matrix, instead, P is ﬂexible to be trained.
We argue that adding such a ﬂexible P could enable crossing layer sharing and enhance the capacity of III.
Interestingly, once we ﬁnish training BERT-IV-144-384 and BERT-IV-144-64, before starting inference, one could merge P and and C as a single matrix: R 12 L × 12 L × R 12 L × d 2 → R 12 L × d 2 , which does not affect the model but slightly improve inference speed.
By doing so, it is rigorously equivalent to III.
Thus, we rename
BERT-IV-144-384/BERT-IV-144-64 as BERT-III-384 and BERT-III-384.
I GLUE BENCHMARK
The data statistics in GLUE is shown in Tab. 8.
J I NFERENCE TIME
As shown in Fig.
8, the raw BERT base is the slowest one.
BERT-III/BERT-IV with dimension rank d = 384 is slight faster than raw BERT.
The inference time consistently decreases when d becomes smaller.
Especially, BERT-III-64 is 2.6 × faster than BERT, which is consistent to RPS shown in Tab.
Note that BERT-III -384 and BERT-IV -72-384 are faster than BERT in terms of inference time and
RPS (see RPS in Tab. 4), although they have the identical FLOPS.
A similar observation could be found in ALBERT (Lan et al., 2019)
– layer-shared BERT has the same FLOPS with the raw BERT while the former is faster than the latter.
The better inference time and RPS with the same FLOPS may be due to that the relatively smaller model consumes less memory and bandwidth inside GPU, see https://openreview.net/forum?id=H1eA7AEtvS .
This speedup effect in this matter depends on speciﬁc platforms and hardware.
BERT-II implemented by Noach & Goldberg (2020) has a rank for hidden dimension of 245, there we called ‘BERT-II-245’.
It inference time is close to our BERT-IV-36-256, this is because the inference time is mainly related to the dimension rank d .
K O N THE COMPRESSION OF THE EMBEDDING LAYER
In this work, we did not compress the embedding layer.
The reasons to not compress the embedding layer are manyfold.
First, compressing the embedding layer will deﬁnitely increase training/inference time, since a single embedding lookup is fast enough.
Secondly, the embedding layer does not increase the total parameters when BERT has more transformer layers in the encoder.
Since the parameters in the embedding layer are constant with respect to network depth.
Table 9 shows the parameter numbers when compressing main weight matrices with a half dimension rank (when the hidden dimension is 768, its half is 384 as the rank d, meanwhile keeping layer rank unchanged).
Here, we also consider embedding layer for a fair comparison as suggested.
models
[Table]
BERT-base-uncased
[Table]
BERT-large-uncased
[Table]
Note that the parameters of embedding becomes more negligible when PLMs have more layers and bigger hidden dimension, in which case the compression ratio will approximate an ideal upper bound ( 8 = 2 × 2 2 , which is linear to the deduction times of layer rank and quadratic to the deduction times of dimension rank ; in practice, we could use bigger deduction in both ranks, as we did in the experiments).
Plus, the shape of an embedding layer ( V D ) is related to the size of vocabulary, which is heterogeneous to other weight matrices in Self-attention network and Feed-forward network ( D 2 or 4 D 2 ).
Therefore it is incompatible with the current compression protocol.
To additionally compress the embedding layer, we might have to design a totally different compression protocol for the embedding layer, which makes this work more complicated.
Finally, the embedding layer is relatively easy to compress, see tensorized embedding (Hrinchuk et al., 2020) and word2ket (Panahi et al., 2019) which compress embedding with maximum 307 and 93,675 times respectively with a slight performance drop.
We believe that it is trivial to additionally compress the embedding layer.
Thus, we leave compressing the embedding layer as future work.
[Table]
[Table]
L C OMPARISON TO METHODS THAT USES NON - STANDARD METRICS IN GLUE
(Noach & Goldberg, 2020) uses non-standard setting for GLUE tasks.
Metrics in (Noach & Goldberg, 2020) are Accuracy (MNLI (average of MNLI match and MNLI mis-match), QNLI, RTE, SST2), Avg of Accuracy and F1 (MRPC, QQP), Matthew’s correlation (CoLA), Avg of Pearson and
Spearman correlations (STS-B).
Tab.
10 tailored our result in their evaluation protocol.
In this work, we do not include CoLA dataset since training samples in CoLA are too few to have stable evaluations.
M A NALYSIS OF THE LEARNED MODEL
Observing that in IV, each matrix is parameterized as W IV i
=
U ( P i C ) V .
With shared U and V , each matrix has a speciﬁc parameter P
i (called ‘factor vectors’ later), we analyze { P i } to shed some light on understanding what BERT-IV learns.
In BERT-IV-72-384, we calculate the cosine distances between any two factor vectors.
d ( P i , P j ) = 1 − cos ( P i , P j ) , lead to a 144 × 144 matrix.
We could observe that there is a cross-layer pattern that { W In } (the weight matrices in the layer of FFNs) among layers are relatively close.
This pattern becomes weaker in deep (top) layers, this might suggest that top layers may learn more diverse feature projections.
The reason for this pattern needs further investigation.
We also visualize the in-matrix pattern in each layer (e.g. in the 1st, 4-th, 8-th, and 12-th layer) in
Figure 10.
It shows that a pair between W Inh and W Outh is relatively close, e.g., W
In 1 and W Out,
W In 2 and W Out, etc.
N P OTENTIAL OF COMPRESSION IN LAGER MODELS
In this section, we compare various compression ratios of II, III, and IV, see the parameters and their compression ratios (this additionally considers parameters in the embedding layer) in Table 12.
It shows that III and IV could achieve much bigger compression ratios than II when using the same rank for hidden dimension (i.e., 2048).
We claim that the proposed model has better potential to compress bigger models.
Q K V O W
In 1 W
In 2 W
In 3 W
In 3 W Out 1 W Out 2 W Out 3 W Out
Q K V
O W
In 1 W
In
W
In 3 W
In 3 W Out
W Out 2 W Out 3 W Out0.
[Figure]
Q
[Figure]
Q K V
O W
In 1 W
In
W
In 3 W
In 3 W Out
W Out 2 W Out 3 W Out0.
[Figure]
Q K V O W
In 1 W
In 2 W
In 3 W
In 3 W Out 1 W Out 2 W Out 3 W Out
Q K V
O W
In 1 W
In
W
In 3 W
In 3 W Out
W Out 2 W Out 3 W Out0.
[Figure]
Q K V O W
In 1 W
In 2 W
In 3 W
In 3 W Out 1 W Out 2 W Out 3 W Out
Q K V
O W
In 1 W
In
W
In 3 W
In 3 W Out
W Out 2 W Out 3 W Out0.
[Figure]
[Table]
in each layer of a trained
BERTIV-72-
[Table]
The order is listed as [ W Q , W K , W V , W O , W
In 1 , W
In 2 , W
In 3 , W
In 4 , W Out, W Out, W Out, W Out] .
[Table]