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: This work proposes a modification of the original Transformer architecture by replacing attention layers and layers in its Feed-Forward Networks across all of its blocks with learned shared dictionaries. The proposed model, called DictFormer, has a smaller number of parameters and uses a smaller amount of computational operations when compared to the original Transformer and some of its variations. When evaluated against these models on popular machine translation, summarization, and language modeling benchmarks, DictFormer achieves comparable or better performance.