D ICT F ORMER : T INY T RANSFORMER WITH S HARED D ICTIONARY
A BSTRACT
We introduce DictFormer with efﬁcient shared dictionary to provide a compact, fast, and accurate transformer model.
DictFormer signiﬁcantly reduces the redundancy in the transformer’s parameters by replacing the prior transformer’s parameters with compact, shared dictionary, few unshared coefﬁcients and indices.
Also,
DictFormer enables faster computations since expensive weights multiplications are converted into cheap shared look-ups on dictionary and few linear projections.
Training dictionary and coefﬁcients are not trivial since indices used for looking up dictionary are not differentiable.
We adopt a sparse-constraint training with l 1 norm relaxation to learn coefﬁcients and indices in DictFormer.
DictFormer is ﬂexible to support different model sizes by dynamically changing dictionary size.
Compared to existing lightweight Transformers, DictFormer consistently reduces model size over Transformer on multiple tasks, e.g., machine translation, abstractive summarization, and language modeling.
Extensive experiments show that
DictFormer reduces 3 .
6 × to 8 .
9 × model size with similar accuracy over multiple tasks, compared to Transformer.
I NTRODUCTION
Transformer (Vaswani et al., 2017) has been widely used in natural language processing (NLP) for its superior capability in capturing long-distance dependencies.
However, its good performance comes with the giant model size.
For example, T5 (Raffel et al., 2019) with a hidden dimension of 65K and GPT-3 (Brown et al., 2020) with 96 transformer blocks have 11 billion and 175 billion parameters, respectively.
These large Transformers suffer from multiple severe issues, such as complicated learning and difﬁcult deployment on mobile/IoT devices.
First, during the training of a big transformer model, large training corpora (Raffel et al., 2019; Brown et al., 2020) or careful regularization (Merity et al., 2017; Mehta et al., 2021) are required.
Furthermore, the trained model is over-parameterized (Reid et al., 2021; Zhao et al., 2021).
The large model size with 11 billion and 175 billion parameters is beyond the capabilities of many edge devices including mobile devices and
IoTs.
Therefore, there is an urgent need to design parameter-efﬁcient and fast transformer model that eliminates redundant parameters and enables real-time NLP applications on the edge.
Weights sharing is a choice proved by ALBERT (Lan et al., 2020) for designing compact and efﬁcient pre-trained transformer encoders, like BERT, on self-supervised learning task.
However, directly sharing all weights in an encoder-decoder transformer for a sequence-to-sequence task like machine translation will dramatically decrease performance and accuracy (Reid et al., 2021).
Although a recent framework Universal Transformer (UT) (Dehghani et al., 2019) shows that a vanilla transformer’s accuracy can be improved by using recursive weight sharing, UT’s high accuracy comes at the cost of deeper blocks or wider hidden dimensions, which signiﬁcantly enlarges the computational cost and does not necessarily reduce the model size.
This paper introduces a new compact, fast, and accurate transformer architecture, DictFormer, that can be easily trained and deployed on edge devices.
DictFormer depends on dictionary sharing and unshared linear projection coefﬁcients instead of weights sharing.
Speciﬁcally, a shared dictionary among all encoder/decoder blocks can signiﬁcantly reduce the parameter redundancy and therefore compresses the model size.
Few unshared linear projection with coefﬁcients on the shared dictionary enable each encoder/decoder block to have distinct layer-wise feature representations, thus improving the representation abilities compared to prior weights sharing.
Also, DictFormer provides a method to dynamically control the layer-wise representation abilities by using the groupwise shared dictionary for layers with large-dimension features, e.g., Feed-forward Network (FFN).
Last but not least, we show that training dictionary and coefﬁcients are not trivial since indices used for looking up dictionary are not differentiable.
We convert coefﬁcients and indices into a sparse matrix and train it with l 1 norm relaxation and convert this sparse matrix into dense coefﬁcients and indices during inference.
Extensive experiments demonstrate that our DictFormer provides signiﬁcant improvements over existing transformers on three sequence-to-sequence tasks, (i) machine translation, (ii) abstractive summarization, and (iii) language modeling.
For machine translation, on IWSLT 2014 GermanEnglish, DictFormer attains transformer performance with 8 .
9 × fewer parameters and 2 × fewer
Multi-Adds; on WMT 2014 German-English, DictFormer brings about 4 .
9 × model compression and 1 . 9 × computation reduction; on WMT 2014 English-French, DictFormer obtains consistent performance improvements: 4 . 9 × model compression and 1 .
9 × less computation with a similar score.
For abstractive summarization, DictFormer reduces the model size by more than 4 .
7 × on
CNN-DailyMail dataset.
For language modeling, DictFormer matches the performance of transformer with 3 . 6 × to 5 .
7 × fewer parameters than transformer on WikiText-103 benchmark.
R ELATED W ORK AND M OTIVATION
Related lightweight Transformers.
Several methods have been proposed to design lightweight transformers.
The ﬁrst line of research is to reduce the transformer computation complexities by redesigning self-attention mechanism including (Katharopoulos et al., 2020;
Zhou et al., 2021; Raganato et al., 2020; You et al., 2020; Correia et al., 2019; Kaiser et al., 2018).
These methods cannot reduce model size.
The second line of research is model compression, e.g., quantization (Prato et al., 2020; Lou et al., 2020b), pruning (Behnke & Heaﬁeld, 2020), low-rank factorization (Ma et al., 2019), and knowledge distillation (Wang et al., 2020).
These two directions of research can be combined with our Dictformer.
The third line of research is efﬁcient architecture design (So et al., 2019;
Wu et al., 2020; Mehta et al., 2021) by improving the expressiveness of transformers.
The forth line of research is weights sharing (Xia et al., 2019;
Ma et al., 2019; Dehghani et al., 2019; Reid et al., 2021; Takase & Kiyono, 2021) by reusing parameters across transformer blocks.
Weights sharing cannot reduce computations.
Our Dictformer falls into the third and forth category.
We show that our Dictformer with Dictionary sharing can reduce both model size and computations.
Transformers with Weights Sharing.
Weight sharing is surprisingly effective to compress model size for discriminate NLP models based on Transformer encoders, e.g., BERT.
For example, prior work ALBERT (Lan et al., 2020) shows that even sharing all parameters across layers does not introduce any accuracy reduction.
However, for generative sequence-to-sequence models based on transformer’s encoders and decoders, sharing all parameters will signiﬁcantly decrease accuracy on multiple standard machine translation or language modelling tasks (Reid et al., 2021; Takase &
Kiyono, 2021).
To match vanilla Transformer’s accuracy, multiple works (Reid et al., 2021; Takase & Kiyono, 2021; Xia et al., 2019; Ma et al., 2019) only share weights across partial layers instead of all layers.
However, partial weights sharing remarkably brings down the model size compression effect of weights sharing.
Also, how to decide which layers should be shared in partial weights sharing is difﬁcult due to the large and dynamic search space that is dependent on the speciﬁc tasks.
Transformer with all-parameters sharing such as Universal Transformer (Dehghani et al., 2019) matches or improves transformer’s performance at the cost of a wider or deeper transformer architecture.
A wider transformer with a larger embedding dimension enlarges the model size and brings larger computations (Mult-Adds).
A deeper transformer with more encoder/decoder blocks does not only increase model size, but also introduces more computations.
Importantly, weights sharing techniques cannot reduce Mult-Adds numbers and training/inference time.
Figure 6(a) in Appendix shows the comparisons of Transformers with weights sharing and our dictionary sharing.
Weights sharing techniques cannot solve the deployment challenges of transformers on resource-limited devices for real-time NLP applications.
[Figure]
Motivation of DictFormer with Dictionary Sharing.
Figure 1 shows that although transformers with weights sharing including Universal Transformer (Dehghani et al., 2019), Tied Transformer (Xia et al., 2019), and Subformer (Reid et al., 2021) signiﬁcantly reduce the model size over vanilla transformer and lightweight transformer DeLighT (Mehta et al., 2021), they cannot reduce the #Mult-Adds and even suffer from larger #MultAdds.
Particularly, Tied Transformer (Xia et al., 2019) and Subformer (Reid et al., 2021) compress transformer model by 1 .
7 × ∼ 2 .
6 ×
but cannot reduce #Mult-Adds.
Universal Transformer (Dehghani et al., 2019) achieves ∼ 1 BLEU score improvement with 1 . 4 × less parameters, but the #Mult-Adds is increased by ∼ 4 .
3 × , thereby signiﬁcantly prolonging the NLP inference latency and restraining the deployment of real-time NLP applications on edge devices.
To enable the deployment of transformer models on mobile devices, recent work Evolved Transformer (So et al., 2019) and Lite Transformer (Wu et al., 2020) try to design new lightweight transformer architecture to meet the deﬁned mobile settings, e.g., 10 million parameters, but their tiny architectures suffer from a large accuracy decrease.
As Figure 1 shows, Evolved Transformer and Lite Transformer lose 2.9 and 2.4 BLEU score compared to base transformer, respectively.
D ICT F ORMER
Overview.
We propose DictFormer with dictionary sharing to enable a fast, compact, and accurate transformer.
When matching transformer’s accuracy, DictFormer can reduce more than 3 .
6 × parameters and ∼ 3 × Mult-Adds shown in Figure 1, which outperforms prior transformers with a higher
BLEU score.
Also, DictFormer is ﬂexible to compress model size given an accuracy threshold.
For instance, when DictFormer matches the accuracy of Lite Transformer, it could further compress the model of lite transformer by ∼ 2 . 2 × .
Given a N -layer transformer model, we can easily transform it into DictFormer by converting all the weights in N blocks into one shared dictionary and few unshared look-up coefﬁcients, shown in Figure 2.
For example, N -layer weights attention W Ai and
FFN W Fi , where i ∈ [0 , N − 1] , in Transformer (Figure 2(a)) are represented by smaller dictionaries D A , D F that are shared by N blocks, and N -layer unshared coefﬁcients C Ai and C Fi for
DictFormer (Figure 2(b)).
Meanwhile, the attention and FFN operations in vanilla Transformer are replaced with shared-dictionary attention and group-wise shared-dictionary FFN whose details are introduced in our following contents.
Dictformer reduces Transformer’s #Params from O ( d 2 N ) to
O ( d ( m + tN )) .
This is a O ( dN/ ( m + tN ))
× model size reduction since dictionary size m < d , coefﬁcient size t << d , where the embedding size is d .
Dictformer also reduces #Mult-Adds from
O ( d 2 Nn ) to O ( dNn ( m + t )) , where n is input sequence length.
The details on how to deﬁne and calculate #Params and #Mult-Adds are shown in Appendix A.1.
[Table] [Table]
[Table]
Encoder Block
N x
FFN
Attention Embed
N x Attention
Block #Params:  (12d 2 N) #Mult
[Figure]
Adds:  (12d 2 Nn)
Dict-FFN
Dict-Attention Embed
Encoder Block
N x
Dict-FFN
Dict-Attention Embed
N x Dict-Attention
Decoder Block d m D A t dC iA d m
D F
C
i F t 4d
[Figure]
Decoder
[Figure]
[Figure] [Figure]
Embed
Encoder
[Figure]
d
[Figure]
dC
[Figure]
Shared-dictionary Attention.
Given a N -layer transformer model, we deﬁne that Q i , K i , and V i are the i -th layer query, key, and values.
Same to Transformer (Vaswani et al., 2017), we also use equation 1 to calculate the attention scores once we have query, key, and values.
[Equation]
Our DictFormer utilizes equation 2 instead of MultiHead ( Q i , K i , V i ) =
MH i · W Oi used in
Transformer (Vaswani et al., 2017) to compute multi-head values, where MH i is derived according to equation 3, and D A , C Oi , and I
Oi are used to linearly project MH i , instead of W Oi .
The attention value head ji for each head j in layer i is computed by equation 4 instead of head ji = Attention ( Q i ·
W Q j i , K i · W K j i , V i · W V j i ) used in existing transformers.
[Equation]
The reason why we use the shared attention dictionary D A , indices I i and coefﬁcients C i to replace the larger and unshared weights W Q j i , W K j i , W V
j
i and W O j i used in previous transformers is that the linear projections of lookup of D A with coefﬁcients C Xi can have similar representation abilities as W X j i , where X represents inputs type, e.g.,
X = Q means input is query.
To be speciﬁc
, D A , I i , and C
i have size of d × m A , t A × d , and t A × d , respectively.
m A and t A are the dictionary size and coefﬁcient size in attention.
As equation 5 shows, the linear projections in transformer’s attention, e.g., W Q j i · Q j , are replaced by our lightweight shared dictionary projection function ( SD ), e.g.,
SD ( Q j , D A , C Q
j i , I Q j i ) , derived by equation 6 and equation 7.
SD ( Q j , D A , C Q
j i , I Q j i ) replaces
W Q j i · Q j since W Q j i can be replaced by looking up D A and few linear projections with coefﬁcients
C Q j i without accuracy decrease.
The following paragraphs and Figure 3(a) introduce the details of look-up and scaling using equation 6 and equation 7.
[Equation]
In equation 6, the lookup of D A by indices I Q
j
i is deﬁned by D A
[: , I Q
j i [ t, i d ]] that can fetch the
I Q j i [ t, i d ]] -column vector from D A ; Unshared linear projection that is used to enlarge the representation ability of shared dictionary is depicted by (cid:93) W Q j i = (cid:80)
[Equation]
i [ t, i d ] (cid:12) D
A [: , I Q
j i [ t, i d ]] , where (cid:12) represents scaling a fetched vector from a dictionary with a scalar in coefﬁcients.
Therefore, linearly projecting Q j with W i according to W Q j i · Q j is replaced by (cid:93) W Q j i · Q j in equation 6 when we ﬁnd proper D A , I Q
j i , C Q
j i to meet that (cid:93) W Q j i and W iQ j have similar representation abilities, e.g., they have matched accuracy.
Directly generating (cid:93) W iQ j and multiplying it with Q j potentially increases the computations.
To tackle this problem, we compute the multiplications between Q i and dictionary D A as O i ﬁrst according to equation 7, and reuse O i for the following look-ups and scaling with C Q j i .
[Equation]
i d
[Equation] [Equation]
We use Figure 3(a) to show an example of equation 5 about how to lookup dictionary with ind
Shared Dictionary D A d d W iQ 1 13+ + 0.2 0.
N-layer Weight W Q t A =3 d
Index I iQ [1, 13, 35] d t A =
Coefficients C iQ [0.2, 0.3, 0.5] m A  4d
Group-2 Dictionary D F1
+ + 0.3 0.
d
Index I iF[1, 7, 27] d
Coeffi.
[0.3, 0.1, 0.6] d [0.1, 0.6, 0.3]
Coeffi.
d 4d
N-layer Weight W F 2 m F
W iF
C iF 2 g=
C iF 2 g= t Ft Ft F
[Figure]
dices and scaling it with coefﬁcients.
This process can create a N -layer weight W Q .
For example, given the shared dictionary D A , to generate the ﬁrst column of weight W Qi , the ﬁrst column of index matrix I Qi is taken out, e.g., [1 , 13 , 35] , and is used as indices to fetch the the corresponding columns from D A ,
e.g., D A
[:][1] , D A
[:][13] , D A [:][35] .
Then the ﬁrst column of coefﬁcients C Qi , e.g., [0 . 2 , 0 . 3 , 0 . 5] , are multiplied with D A [:][1] ,
D A
[:][13] , D A [:][35] respectively, and the sum of multiplication works as the ﬁrst column of W Qi .
In this way, weights W in attention with 4 d 2 N parameters are compressed into m A , I Q , and C Q with size dm A + 8 t
A dN .
For example, DictFormer can reduce 8.7 × model size when d = 512 , m A = 256 , t A = 24 , N = 6 .
More calculation details can be found at appendix A.1.
Group-wise Shared-dictionary FFN.
The FFN in prior transformer includes two-layer computations: (i) F 1 = max (0 , X i · W F 1 i + b 1 ) , (ii) F 2 = F 1 · W F 2 i + b 2 .
Instead of the regular linear projections X i · W F 1 i + b 1 and F 1 · W F 2 i + b 2 , DictFormer uses a new lightweight projection called group-wise shared dictionary projection ( GSD ), i.e., GSD ( X i , D, C F 1 i , I F 1 i , G ) and n
O i t A d
Index I i [1, 13, 35] d t A
Coefficients C i [0.2, 0.3, 0.5] 1 130.2 0.
+ + d n
Attention Results S
d m A [0.2, 0.3, 0.5]
Sparse Coefficients Z n
Shared O i .
m
A m A 
GSD ( F 1 , D, C F 2 i , I F 2 i , G ) , in equation 8 and equation 9 to compute the FFN.
Here GSD projection is same as SD in shared-dictionary attention when G = 1 .
In SD projection, it works well to replace a N weights of d × d with a d × m dictionary where each column of dictionary is multiplied with a scalar in coefﬁcients.
However, when the column of dictionary is large, e.g., 4 × d in the second layer of FFN, SD projection is difﬁcult to achieve an accurate model since an vector of dictionary with 4 × d , e.g. 2048, multiplied by a same scalar is not ﬂexible enough.
To increase the ﬂexibility of SD projection and improve performance, GSD ﬁrstly divides each column of dictionary into G groups equally and assign different scalars to multiply the numbers in each group.
[Equation]
[Equation]
Equation 10 and equation 11 show the computations of GSD .
Dictionary D F 2 [:, :] is divided equally into G groups and the g -th group is deﬁned as D F 2 [ Index g , :] , where Index g =
[( g − 1) 4 dG , g 4 dG ) and g ∈ [1 , G ] .
Also, the multiplication result between g -th group dictionary and input is deﬁned as O g [: , b ] shown in equation 11, where b ∈ [1 , m F 2 ] .
Different groups of dictionary share the same look-up indices
I F 2 i .
The t F 2 queried vectors in each group will be scaled by its corresponding coefﬁcients C F g 2 i and the scaled results are then accumulated to derive the g -th GSD result shown in equation 10.
In addition, G group-wise GSD are further summed to compute the ﬁnal GSD result.
[Equation]
Figure 3(b) shows an example of group-2 shared dictionary for the the second layer of FFN.
The dictionary D F 2 is equally split into two parts that share the same indices I F 2 i but each group has unique coefﬁcients, e.g., C F g =
i and C F g =22 i .
The ﬁrst group and the second group use the corresponding coefﬁcients, e.g., [0 . 3 , 0 . 1 , 0 . 6] and [0 . 1 , 0 . 6 , 0 . 3] , to scale the queried dictionary vectors, respectively.
The accumulation of scaled vectors can be used to represent the N -layer weight W F 2 .
Training DictFormer via Constraints and Relaxation.
DictFormer replaces transformer weights with linear projections of a dictionary, thereby having three-step computations including small projections, lookup and scale.
DictFormer computes a small multiplication between input and dictionary, generating intermediate variable O , looking up O and scaling it with coefﬁcients.
To train the
DictFormer, one should jointly optimize dictionary, index I and coefﬁcients C .
Directly training the Dictformer leads to a combinatorial optimization problem since index I is non-continuous.
Although one could use AutoML, such as evolutionary method and Reinforcement learning to jointly learn dictionary, I and C , these methods suffer from a large training time with low performance.
To workaround the training of index I , we reformulate the method to a regular linear projection with sparse constraints, which can efﬁciently train the DictFormer.
We convert the index
I i and coefﬁcients C i into sparse coefﬁcients Z , so that the training of index
I i and coefﬁcients
C i is replaced with training sparse coefﬁcients Z .
During the deployment phase, the sparse coefﬁcients Z are reversely transformed into index
I i and coefﬁcients C
i shown in Figure 4.
In particular, the shape of coefﬁcients C and index I is t A × d × N .
The shape of sparse Z is m A × d × N .
There are two steps to derive Z : Initializing all elements as 0 and copying C elements to Z according to the index values in I i .
For example, since the ﬁrst column of I 1 and Cis [1 , 13 , 35] and [0 . 2 , 0 . 3 , 0 . 5] , all entries of the ﬁrst column in Z are zeros except the 1st entry, 13th entry, and the 35th entry are 0 .
2 , 0 . 3 , and 0 .
Therefore, the lookup and scale of O i can be converted into the matrix multiplication between O i and Z i shown in equation 12.
The new coefﬁcient C has a sparsity constraint that requires that the non-zero elements ( l 0 norm ) of each column is t A << m A shown in equation 13.
Now we can train Z with || Z Q j i
[: , i d ] || l 0 =
t A , instead of training I and C .
[Equation]
The l 0 norm sparsity constraint in equation 13 is non-differentiable and we cannot directly train Z .
To achieve the training with sparsity, we are inspired by existing sparse training methods Tsuruoka et al.
; Li et al.
(2016) to ﬁrstly loose this constraint to a l 1 norm constraint shown in equation 14 to penalize the non-zero parameters and clip the unimportant values that are near to zeros using equation 16.
To be speciﬁc, the gradients of coefﬁcients are calculated by equation 15, where parameter λ is used to control the trade-off between the gradient of loss L , δLδZ , and l 1 norm sparsity constraint.
[Equation]
Equation 16 is a threshold function during training.
During the backpropagation, the derivative of threshold function is 1 if | x | > value ( ρ ) , 0 otherwise.
During the forward phase, Equation 16 can be used to globally clip the near-zero values to zero given a ratio ρ , where value ( ρ ) derives the value at ratio ρ in the ascending order.
One also can force that at most t of the elements in each column of
Z are non-zero while the rest are set to 0 by letting value ( ρ ) is the t -th largest value of each column.
[Equation]
E XPERIMENTAL M ETHODOLOGY
Machine Translation Dataset.
Three machine translation benchmarks are tested: IWSLT’
German-English (De-En), WMT’14 English to German (En-De), and WMT’14 English to France (En-Fr).
For IWSLT’14 De-En, we adopt the same settings in Wu et al.
(2020) with 160K training sentence pairs and 10K joint byte pair encoding (BPE) vocabulary in lower case.
For WMT’
En-De, our models are trained with 4.5M sentence pairs, validated on newstest2013, and tested on newstest2014.
Also, a 32K joint source and target BPE is used in the vocabulary.
For WMT’14 EnFr, we follow the setup in Wu et al.
(2020) by training models on 36M training sentence pairs from
WMT’14, validating on newstest2012 and 2013, and testing on newstest2014.
The vocabulary has a size of 40K and is based on a joint source and target BPE factorization.
The evaluation setting used in Vaswani et al.
(2017) is utilized with a beam size of 4 and a length penalty of 0.6.
We replicate the same BLEU calculations in Wu et al.
(2020) with case-sensitive tokenization.
The last 10 model checkpoints are averaged for testing and the lowest-perplexity model is picked for validation.
Abstractive Summarization Dataset .
We evaluate DictFormer on CNN-DailyMail dataset (Chen et al., 2016) that has 280K news articles with multi-sentence summaries.
We follow the settings metric.
Language Modeling Dataset .
We also evaluate DictFormer on WIKITEXT-103 (Merity et al., 2016) that has a 260K BPE vocabulary and contains 103M/217K/245K tokens for training, validation, and testing.
Language modeling performance about perplexity (ppl) is reported.
Architecture and Evaluation.
For machine translation tasks, we follow the sequence-to-sequence transformer (Vaswani et al., 2017) with encoder and decoder to develop our model.
For IWSLT dataset, the general settings are same as (So et al., 2019).
For WMT dataset, the setting is based on (Wu et al., 2020).
Also, the same model for machine translation is used for summarization task.
For language modelling, we use the transformer-based decoder architecture of (Baevski & Auli, 2019) but with 512 model dimensions and 12 blocks that is same to (Wu et al., 2020).
Fairseq’s transformer implementation (Ott et al., 2019) is used as the backbone for the baseline model.
In our DictFormer architecture, encoder and decoder have three dictionaries, respectively, i.e. , encoder attention dictionary D Aenc , encoder F 1 dictionary D F 1 enc , encoder F 2 dictionary D F 2 enc , and decoder attention dictionary D Adec , decoder F 1 dictionary D F 1 dec , decoder F 2 dictionary D F 2 dec .
Each of them is shared by all blocks in encoder or decoder.
For example, D Aenc is shared by all attention blocks in encoder.
We study the effects of DictFormer with various dictionary size from 80 toand different coefﬁcient size from 12 to 84 shown in Figure 8.
One can pick up a dictionary size and coefﬁcient size to control the trade-off between accuracy and model size according to the requirements.
For example, one could get the results in our tables when dictionary size is set as m A = d 2 .
5 , m F 1 = d 2 . 5 , m F 2 = d 2 , and t A =
d 10 , t F 1 = t F 2 = d 15 .
In this paper
[Table]
#Params means parameter numbers of transformer without including embedding layers; #TParams shows parameter numbers of model and embedding layers.
The Mult-Adds calculation is tested when a model is used to translating two sentences with the length of 30 in default (same to previous work Wu et al.
[Table] [Table]
R ESULTS
Machine Translation.
We ﬁrst report the machine translation results and compare them with prior works.
Our baseline transformer model settings are in line with Lite Trans.
(Wu et al., 2020) which provides the best results under mobile settings.
Our DictFormer generally outperforms the state-ofthe-art transformers on IWSLT’14 De-En, WMT’14 En-De, and WMT’14 En-Fr.
Table 1 represents the quantitative results of our DictFormer.
Compared to Transformer (Vaswani et al., 2017), DictFormer compresses 8 .
9 × model size, reduces 1 .
9 × #Mult-Adds, and improves 2.2 BLEU score on
IWSLT’14 De-En, which achieves better trade-off between performance and model size than previous works.
Similar to IWSLT’14 De-En, our experiments on larger corpora including WMT’
En-De and En-Fr also show DictFormer obtains more compact and accurate model than existing
Transformers.
For instance, only our DictFormer can compress more than 6 × model size over
Transformer without decreasing any model performance under mobile settings.
[Table]
Abstractive Summarization and Language Modeling.
We then report the abstractive summarization and language modeling results and compare them with prior works.
Abstractive summarization model follows the transformer model
[Table]
but language modeling model follows the adaptive inputs model (Baevski & Auli, 2019).
Our DictFormer consistently outperforms prior transformers on both summarization and language modeling tasks.
Table 2 represents the quantitative results of our
DictFormer.
For abstractive summarization on CNN-DailyMail, DictFormer achieves similar F1Rouge score with Transformer but requires 4 .
7 × less model size and ∼ 4 × less computations.
For language modeling on WIKITEXT103, compared to adaptive inputs model (Baevski & Auli, 2019),
DictFormer compresses 3 . 8 × model size with less computations and matched testing perplexity.
Sensitive Study and Ablation Study.
We study the DictFormer’s performance under different model size and hyper-parameters in Appendix.
We also use table 3 in Appendix to show the ablation study of DictFormer’s techniques.
The performance of DictFormer improves with an increase in the number of model parameters, across different corpora.
DictFormer achieves similar BLEU score with Transformer using fewer model parameters.
Although DictFormer’s performance improves with an increase in the size of dictionary and coefﬁcients, the model size is also enlarged.
So one should pick up a proper dictionary or coefﬁcient size according to accuracy requirements.
C ONCLUSION
We propose DictFormer with dictionary sharing to enable a fast, compact, and accurate transformer.
DictFormer signiﬁcantly reduces both model size and computation parameters of prior transformers with similar performance.
Also, DictFormer can consistently improve the performance on multiple tasks, such as machine translation, abstractive summarization, and language modeling with similar model size.
DictFormer code is available at https://github.com/SamNLP/DictFormer .
R
A A PPENDIX A. D ICT F ORMER BENEFITS WITH DICTIONARY SHARING
The block comparison of Transformer
DictFormer
Embed Attention FFN S ha r ed E n c ode r
T x
Embed Attention FFN S ha r ed D e c ode r
T x Attention
SharedEncoder SharedEncoder Embed ... t x t x N/t
SharedDecoder SharedDecoder Embed ... t x t x N/t
Embed Dict.
Atten.
N x Sole Coeff.
Dict.
FFN Sole Coeff.
Embed Dict.
Atten.
Query & Scale Dict.
FFN Query & Scale N x (a) Transformers with T-layer weight sharing (b) N/t t-layer partial weight sharing (c) N-layer dictionary sharing and sole coefficients
Embed Attention FFN S ha r ed E n c ode r
T x
Embed Attention FFN S ha r ed D e c ode r
T x Attention
SharedEncoder SharedEncoder Embed ... t x t x N/t
SharedDecoder SharedDecoder Embed ... t x t x N/t
Embed Shared Dictionary
N x Unique Query & Scale (a) Transformers with T-layer weight sharing (b) N/t t-layer partial weight sharing (c) N-layer dictionary sharing and sole coefficients Embed Shared Dictionary N x Unique Query & Scale
[Figure]
[Table]
(b) N/t t-layer partial weight sharing (c) N-layer dictionary sharing and sole coefficients head A tt en t i on
FF N
W i Q j
W i K j
W i V j Attention
Concat
Add
[Table]
[Figure]
[Table]
D F d d d d d d
[Table]
F
[Table] [Table]
Paras #: 4Nd
Ops: 4Nnd
S D P A tt en t i on
G
[Table]
FF
[Table]
D Attention Concat
Add
D F
Add
D FFN
GSDP (Eq. 13)
SDP (Eq. 7)
SDP (Eq. 7) SDP (Eq. 7)
Paras #: 8ds+md Ops:
Nnd(m+s)
Paras #: 3dp+32ds
Ops:
Nnd(3p+32s)
d m d d p d (a) Transformer block
M u l t ihead
A tt en t i on
FF N
W i Q j W i K j
W i V j
Attention
Concat
Add
W i O
W i F W i F
Add (b) Our DictFormer block d d d d d d d F = 4d d
Paras #: 8Nd
Ops: 8Nnd
Paras #: 4Nd
Ops: 4Nnd 2 +dn
S D
A tt en t i on
G S D
FF N
D
Attention
Concat
Add
D
Add
D
Lookup GSD (Eq. 9)
Lookup SD (Eq. 5)
SD (Eq. 5) SD (Eq. 5)
Paras #: 8Ndt A +m A d
Ops:
Nnd(m A +t A )
Paras #: 3dm F +32Ndt F
Ops:
Nnd(3m F +32t F + dnd
m A d d d
F
F
A m F 4d m F
Figure 5 further depicts the comparison of transformer encoder block and DictFormer encoder block.
The i -th Transformer block including attention and FFN has four weights W Ai and W Fi .
W Ai has four weights including query W Q j i , key W K j i , value W V j i , and output W O j i , thereby having 4 Ndparameters and 4 Nnd 2 Mult-Adds given a sequence of size n .
W Fi has two parts: W F 1 i and W F 2 i .
N -layer FFN blocks have 8 Nd 2 parameters and 8 Nnd 2 Mult-Adds operations.
In contrast, DictFormer signiﬁcantly reduces the parameters and Mult-Adds of Transformer.
Speciﬁcally, each DictFormer block replaces the attention and FFN with shared dictionary (SD) attention and group-wise shared dictionary (GSD) FFN.
Linear projection is performed by looking up dictionary D A using
Eq.5.
SD Attention has 8 Ndt A
+ m A d parameters and Nnd ( m A + t A ) operations.
GSD FFN has 3 dm F + 32 Ndt F parameters and Nnd (3 mF + 32 t F ) operations, where t , m are the size of coefﬁcient and dictionary.
To simplify the calculations of #params and #Mult-Adds, we can let t
A = t F , m A = m F .
Dictformer reduces Transformer’s #Params from O ( d 2 N ) to O
( d ( m + tN )) .
This is a O ( dN/ ( m + tN ))
× model size reduction since dictionary size m < d , coefﬁcient size t << d , where the embedding size is d .
Dictformer also reduces #Mult-Adds from O ( d 2 Nn )
to
O ( dNn ( m + t )) .
A. T HE COMPARISON BETWEEN WEIGHT SHARING AND DICTIONARY SHARING
Weight sharing is surprisingly effective to compress model size for discriminate NLP models based on Transformer encoders, e.g., BERT.
For example, prior work ALBERT (Lan et al., 2020) shows that even sharing all parameters across layers does not introduce any accuracy reduction.
However, for generative sequence-to-sequence models based on transformer’s encoders and decoders, sharing all parameters will signiﬁcantly decrease accuracy on multiple standard machine translation or language modelling tasks (Reid et al., 2021; Takase & Kiyono, 2021).
To match vanilla Transformer’s accuracy, multiple works (Reid et al., 2021; Takase & Kiyono, 2021; Xia et al., 2019; Ma et al., 2019) only share weights across partial layers instead of all layers.
However, partial weights sharing remarkably brings down the model size compression effect of weights sharing.
Also, how to decide which layers should be shared in partial weights sharing is difﬁcult due to the large and dynamic search space that is dependent on the speciﬁc tasks.
[Table]
S ha r ed E n c ode r
T x
Embed Attention FFN S ha r ed D e c ode r
T x Attention
SharedEncoder SharedEncoder Embed ... t x t x N/t
SharedDecoder SharedDecoder Embed ... t x t x N/t
Embed Dict.
Atten.
N x Sole Coeff.
Dict.
FFN Sole Coeff.
Embed Dict.
Atten.
Query & Scale Dict.
FFN Query & Scale N x
[Table]
a) Transformers
[Table]
T-layer weight
[Table]
) N/t t-layer partial weight sharing (c) N
[Table]
layer dictionary sharing and
[Table]
coefficients
[Table]
S ha r ed E n c ode r
T x
[Table]
FFN S ha r ed D e c ode r
T x Attention
SharedEncoder SharedEncoder
[Table]
t x t x N/t
SharedDecoder SharedDecoder Embed ... t x t x N/t
Embed Shared Dictionary
N x Unique Query & Scale
Embed Shared Dictionary
N x Unique Query & Scale
[Table]
[Table] [Table]
Transformer with all-parameters sharing such as Universal Transformer (Dehghani et al., 2019) matches or improves transformer’s performance at the cost of a wider or deeper transformer architecture.
A wider transformer with a larger embedding dimension enlarges the model size and brings larger computations (Mult-Adds).
A deeper transformer with more encoder/decoder blocks does not only increase model size but also introduces more computations.
Importantly, weights sharing techniques cannot reduce Mult-Adds numbers and training/inference time.
Figure 6(a) shows the comparisons of Transformers with weights sharing and our DictFormer with dictionary sharing.
Weights sharing techniques cannot solve the deployment challenges of transformers on resourcelimited devices and obtain real-time NLP applications.
To reduce both model size and computations of Transformers, we introduce DictFormer with dictionary sharing instead of previous weights sharing.
In particular, DictFormer shares dictionary across all layers so there is no need to decide which layers should be shared.
Also, DictFormer with few unshared look-up coefﬁcients does not require a wider embedding size or deeper encoder/decoder to improve accuracy.
Rather than compressing model size by parameters sharing, DictFormer with dictionary sharing can also enable computations sharing to reduce running latency.
Therefore, our DictFormer provides a compact, fast, and accurate transformer model for sequence-to-sequence NLP tasks.
A. T RAINING SETTINGS
To fairly compare with our baselines, all of our training settings are in line with (Wu et al., 2020).
For machine translation tasks, a dropout of 0.3 is used and the dropout ratio is linearly scaled down when we shrink the dimension of the embeddings for WMT datasets.
The learning rate linearly warms up from 10 − 7 to 10 − 3 followed by a cosine annealing with a single cycle and Adam optimizer same as (So et al., 2019) and (Wu et al., 2020).
Inverse linear square root learning rate scheduling is used for IWSLT De-En.
The summarization training settings are the same as machine translation.
To train models on language modeling task, we follow the training setting in (Baevski & Auli, 2019), but decrease the dropout rate by half in FFN layers.
The training experiments of WMT, summarization, and language modeling are conducted on 8 NVIDIA Tesla V100 GPUs.
IWSLT
De-En is trained on two GPUs.
Machine translations tasks are trained for 50K steps, but language modelling tasks are trained for 286K steps.
We further describe deployment settings in Appendix.
[Figure] [Figure]
[Figure] [Figure]
Deploying deep learning models including transformer on powerful cloud-based servers as a service suffers from data privacy issues when users’ data is private or conﬁdential (Lou & Jiang, 2019;
Lou et al., 2019).
In addition, the inference cost of privacy-sensitive applications that performed over encrypted data using fully homomorphic encryption is prohibitively expensive (Feng et al., 2021; Lou et al., 2021; 2020a).
(i) Mobile settings.
Scaling up transformer models leads to higher machine translation performance, but those large architectures are not suitable for real-world applications, especially when we should deploy them on edge devices where computation and memory size are highly constrained.
Multiple prior works like (Wu et al., 2020) formalize the deployment of lightweight transformers on mobile devices by deﬁning the mobile settings regarding the amount of computation and parameter numbers.
Particularly, given the ARM Cortex-A72 mobile CPU with computation performance of 48G FLOPS, the computation constraint for machine translation should be under 500M Mult-Adds and the parameter number should be around 10M with a sequence oftokens that are the general length for machine translation.
(ii) System-independent settings.
Other than deployed on mobile devices, DictFormer also supports system-independent deployment.
In this setting, we directly choose prior transformer architectures like DelighT (Mehta et al., 2021) and
Subformer (Reid et al., 2021) as baselines and compare our DictFormer with them under the the same performance constraint.
A. S ENSITIVITY S TUDY
We study the DictFormer’s performance under different model size as Figure 7 shows.
This is important since system-independent settings require scalable DictFormer.
The performance of DictFormer improves with an increase in the number of model parameters, across different corpora.
DictFormer achieves similar BLEU score with Transformer using 3 × to 10 × less model parameters.
In Figure 8, we also report the hyper-parameter’s effects on DictFormer’s performance with
WMT’14 En-Fr dataset with embedding size d = 512 .
Figure 8(a) shows the dictionary size study when the coefﬁcient size is ﬁxed to 60.
All dictionaries share the same size, and all coefﬁcients also share same size.
The Figure 8(b) shows the coefﬁcient size study when the dictionary size is ﬁxed 240.
Although DictFormer’s performance improves with an increase in the size of dictionary and coefﬁcients, the model size is also enlarged.
So one should pick up a proper dictionary or coefﬁcient size according to accuracy requirements.
A. A BLATION S TUDY
We also use table 3 to show the ablation study of DictFormer’s techniques.
Transformer Vaswani et al.
(2017) adopts the base architecture with hidden size 512.
Transformer-Flatten ﬂattens the feed-forward hidden size of transformer from 2048 to 512, thus reducing the parameters frommillions to 25 millions without hurting the accuracy. Only-Dictionary is based on TransformerFlatten and uses unshared dictionary architecture.
This method achieves 33.2 BLEU with 9.7 million model parameters.
The #TParams is 14.3 million and 27.4 million when adding embedding parameters on IWSLT’14 De-En and WMT’14 En-Fr, respectively. When sharing dictionaries in attention, Shared-Attention achieves similar BLEU score with 4.2 million model parameters. Mean-
[Figure]
[Figure]
P a r a m e t e r s ( i n m illi o n )
[Figure]
BLEU          12 24 36 48 60
B L E U
P a r a m e t e r s ( i n m illi o n )
Coefficient size (t)
Model Size
BLEU while, Shared-FFN only shares FFN using group-wise shared dictionaries (G is group number).
DictFormer that shares both FFN and attention with G=2, denoted by Shared-both , requires 2.6M model parameters and matches the accuracy of unshared DictFormer.
Improved-Embed represents that the embedding layers of Shared-both are further compressed by using existing method Mehta et al.
(2021) so that the #TParams is reduced to 5.1M from 7.1M.
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]