DictFormer: Tiny Transformer with Shared Dictionary

1 INTRODUCTION . 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 difficult 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 ; Mehta et al. , 2021 ) or careful regularization ( Merity et al. , 2017 ; Mehta et al. , 2021 ) are required . Furthermore , the trained model is over-parameterized ( Reid 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-efficient 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 efficient pre-trained transformer encoders like BERT on self-supervised learning task . However , directly sharing all weights in an encoder-decoder transformer on a sequence-to-sequence task like machine translation will dramatically decrease transformer ’ s 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 significantly 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 coefficients instead of weights sharing . Specifically , a shared dictionary among all encoder/decoder blocks can significantly reduce the parameter redundancy and Under review as a conference paper at ICLR 2022 therefore compresses the model size . Few unshared linear projection with coefficients 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 coefficients are not trivial since indices used for looking up dictionary are not differentiable . We convert coefficients and indices into a sparse matrix and train it with l1 norm relaxation and convert this sparse matrix into dense coefficients and indices during inference . Extensive experiments demonstrate that our DictFormer provides significant 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 , or achieves 2.2 higher BLEU score under similar model size ; on WMT 2014 GermanEnglish , DictFormer brings about 4.9× model compression and 1.9× computation reduction with even 0.7 higher BLEU score ; on WMT 2014 English-French , DictFormer obtains consistent performance improvements : 4.9× model compression and 1.9× less computation with a similar or better BLEU score , or 2.1 BLEU score improvement with similar model size . 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 , or achieves ∼ 1.7 lower perplexity than transformer on WikiText-103 benchmark . 2 RELATED WORK AND MOTIVATION . Related Lite-weight Transformers . Several methods have been proposed to design lite-weight transformers . The first 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 can not reduce model size . The second line of research is model compression , e.g. , quantization ( Prato et al. , 2020 ) , pruning ( Behnke & Heafield , 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 efficient 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 can not 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 significantly 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 difficult due to the large and dynamic search space that is dependent on the specific 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 can not reduce Mult-Adds numbers and training/inference time . Figure 8 ( a ) in Appendix shows the comparisons of Transformers with weights sharing and our dictionary sharing . Weights sharing techniques can not solve the deployment challenges of transformers on resource-limited devices for real-time NLP applications . Under review as a conference paper at ICLR 2022 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 ) significantly reduce the model size over vanilla transformer and even one recent lightweight transformer DeLighT ( Mehta et al. , 2021 ) , they can not reduce the # Mult-Adds and even suffer from larger # Mult-Adds , e.g. , Universal Transformer . Particularly , Tied Transformer ( Xia et al. , 2019 ) and Subformer ( Reid et al. , 2021 ) compress transformer model by 1.7× ∼ 2.6× but can not reduce Mult-Adds . Universal Transformer ( Dehghani et al. , 2019 ) achieves ∼ 1 BLEU score improvement with 1.4× less parameters , but the Mult-Adds numbers are increased by ∼ 4.3× , thereby significantly 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 lite-weight transformer architecture to meet the defined mobile settings , e.g. , 10 million parameters , but their tiny architectures have to 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 . 3 DICTFORMER . 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 flexible 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 coefficients , shown in Figure 2 . For example , N -layer weights attention WAi and FFN WFi , where i ∈ [ 0 , N − 1 ] , in Transformer ( Figure 2 ( a ) ) are represented by smaller dictionaries DA , DF that are shared by N blocks , and N -layer unshared coefficients CAi and C F i 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 ( d2N ) to O ( d ( m + tN ) ) . This is a O ( dN/ ( m + tN ) ) × model size reduction since dictionary size m < d , coefficient size t < < d , where the embedding size is d. Dictformer also reduces # Mult-Adds from O ( d2Nn ) to O ( dNn ( m+ t ) ) , where n is input sequence length . The details on how to define and calculate # Params and # Mult-Adds are shown in Appendix A.1 . Shared-dictionary Attention . Given a N -layer transformer model , we define that Qi , Ki , and Vi 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 . Attention ( Qi , Ki , Vi ) = softmax ( Qi ·KTi√ d ) · Vi ( 1 ) Our DictFormer utilizes equation 2 instead of MultiHead ( Qi , Ki , Vi ) = MHi · WOi used in Transformer ( Vaswani et al. , 2017 ) to compute multi-head values , where MHi is derived according Under review as a conference paper at ICLR 2022 FFN Attention Embed Attention FFN Sh ar ed E nc od er T x Embed Attention FFN Sh ar ed D ec od er T xAttention Shared Encoder Shared Encoder Embed ... t x t x

Summary: The authors proposed an efficient transformer layer based on a dictionary of shared parameters instead of standard self-attention.  The goal is to reduce redundant parameters in transformer models.  The main contributions are: a lite transformer model, modification of the self-attention parameters, and evaluation on language dowstream tasks. The proposed transformer model outperforms related work on the machine translation and language modelling tasks.