D E L IGH T: D EEP AND L IGHT - WEIGHT T RANSFORMER
A BSTRACT
We introduce a deep and light-weight transformer, DeLighT , that delivers similar or better performance than standard transformer-based models with signiﬁcantly fewer parameters.
DeLighT more efﬁciently allocates parameters both (1) within each Transformer block using the DeLighT transformation, a deep and lightweight transformation and (2) across blocks using block-wise scaling, that allows for shallower and narrower DeLighT blocks near the input and wider and deeper
DeLighT blocks near the output.
Overall, DeLighT networks are 2.5 totimes deeper than standard transformer models and yet have fewer parameters and operations.
Experiments on benchmark machine translation and language modeling tasks show that DeLighT matches or improves the performance of baseline Transformers with 2 to 3 times fewer parameters on average
I NTRODUCTION
Attention-based transformer networks (Vaswani et al., 2017) are widely used for sequence modeling tasks, including language modeling and machine translation.
To improve performance, models are often scaled to be either wider, by increasing the dimension of hidden layers, or deeper, by stacking more transformer blocks.
For example, T5 (Raffel et al., 2019) uses a dimension of 65K and GPT- (Brown et al., 2020) uses 96 transformer blocks.
However, such scaling increases the number of network parameters signiﬁcantly (e.g., T5 and GPT-3 have 11 billion and 175 billion parameters, respectively), and complicates learning, i.e., these models either require very large training corpora (Raffel et al., 2019;
Devlin et al., 2019; Brown et al., 2020) or careful regularization (Hinton et al., 2012; Wan et al., 2013; Merity et al., 2018a).
In this paper, we introduce a new parameter-efﬁcient attention-based architecture that can be easily scaled to be both wide and deep.
Our De ep and Ligh t-weight T ransformer architecture, DeLighT , extends the transformer architecture of Vaswani et al.
and delivers similar or better performance with signiﬁcantly fewer parameters and operations.
At the heart of DeLighT is the DeLighT transformation that uses the group linear transformations (GLTs) of Mehta et al.
(2018) with an expand-reduce strategy for varying the width and depth of the DeLighT block efﬁciently.
Since GLTs are local by nature, the DeLighT transformation uses feature shufﬂing, which is analogous to channel shufﬂing in convolutional networks (Zhang et al., 2018), to share information between different groups.
Such wide and deep representations facilitate replacing the multi-head attention and feed-forward layers in transformers with single headed attention and light-weight feed-forward layers, reducing total network parameters and operations.
Importantly, unlike transformers, the DeLighT transformation decouples the depth and width from the input size, allowing us to allocate parameters more efﬁciently across blocks by using shallower and narrower DeLighT blocks near the input and deeper and wider DeLighT blocks near the output.
We demonstrate that DeLighT models achieve similar or better performance than transformer models with signiﬁcantly fewer parameters and operations, on two common sequence modeling tasks, (i) machine translation and (ii) language modeling.
On the low resource WMT’16 En-Ro machine translation dataset, DeLighT attains transformer performance using 2 . 8 × fewer parameters.
On the high resource WMT’14 En-Fr dataset, DeLighT delivers better performance (+0.4 BLEU score) with 1 . 8 × fewer parameters than baseline transformers.
Similarly, on language modeling,
DeLighT matches the performance of Transformer-XL (Dai et al., 2019) with 1 . 5 × fewer parameters
R ELATED W ORK
Improving transformers: Several methods have been introduced to improve the transformer architecture.
The ﬁrst line of research addresses the challenge of computing self attention on long input sequences (Child et al., 2019; Kitaev et al., 2020; Beltagy et al., 2020).
These methods can be combined with our architecture.
The second line of research focuses on explaining multi-head attention (Raganato and Tiedemann, 2018; Brunner et al., 2020).
They show that increasing the number of transformer heads can lead to redundant representations (Voita et al., 2019a; Michel et al., 2019) and using ﬁxed attention heads with predeﬁned patterns (Raganato et al., 2020) or synthetic attention matrices (Tay et al., 2020) improves performance.
The third line of research focuses on improving transformers by learning better representations (Wu et al., 2019; 2020; So et al., 2019).
These works aim to improve the expressiveness of transformers using different transformations – for example, using convolutions (Wu et al., 2019; Gehring et al., 2017), gated linear units (Dauphin et al., 2017), or multi-branch feature extractors (So et al., 2019; Wu et al., 2020).
Our work falls into this category.
Unlike previous works, we show that it is possible to efﬁciently allocate parameters both at the block-level using the DeLighT transformation and across blocks using block-wise scaling.
Model scaling: Model scaling is a standard method to improve the performance of sequence models (Vaswani et al., 2017; Raffel et al., 2019;
Lan et al., 2020; Devlin et al., 2019; Shoeybi et al., 2019;
Tan and Le, 2019; Brown et al., 2020).
Model dimensions are increased in width-wise scaling (Vaswani et al., 2017; Devlin et al., 2019) while more blocks (e.g., Transformer blocks) are stacked in depth-wise scaling (Shoeybi et al., 2019; Brown et al., 2020; Wang et al., 2019).
In both cases (and their combination), parameters inside each block of the network are the same, which may lead to a sub-optimal solution.
To further improve the performance of sequence models, this paper introduces block-wise scaling that allows for variably-sized blocks and efﬁcient allocation of parameters in the network.
Our results show that (1) shallower and narrower DeLighT blocks near the input and deeper and wider DeLighT blocks near the output deliver the best performance, and (2) models with block-wise scaling coupled with model scaling achieve better performance compared to model scaling alone.
We note that convolutional neural networks (CNNs) also learn shallower and narrower representations near the input and deeper and wider representations near the output.
Unlike CNNs (e.g., ResNet of He et al. 2016) that perform a ﬁxed number of operations at each convolutional layer, the proposed block-wise scaling uses a variable number of operations in each layer and block.
Improving sequence models: There is also signiﬁcant recent work on other related methods for improving sequence models, including (1) improving accuracy using better token-level representations – for example, using BPE (Sennrich et al., 2016), adaptive inputs (Baevski and Auli, 2019) and outputs (Grave et al., 2017a), and DeFINE (Mehta et al., 2020), and (2) improving efﬁciency – for example, using compression (Chen et al., 2018; Sun et al., 2020), pruning (Han et al., 2016; Voita et al., 2019b), and distillation (Hinton et al., 2015; Sanh et al., 2019).
The closest to our work is the DeFINE transformation, which also learns representations using an expand-reduce strategy.
The key difference between the DeFINE transformation (Figure 1c) and the DeLighT transformation (Figure 1d) is that the DeLighT transformation more efﬁciently allocates parameters within expansion and reduction layers.
Unlike DeFINE, which uses fewer groups in group linear transformations to learn wider representations, DeLighT transformation uses more groups to learn wider representations with fewer parameters.
The DeLighT transformation achieves comparable performance to the DeFINE transformation but with signiﬁcantly fewer parameters.
D E L IGH T: D EEP AND L IGHT - WEIGHT T RANSFORMER
A standard transformer block (Figure 1a) comprises of multi-head attention that uses a query-keyvalue decomposition to model relationships between sequence tokens, and a feed forward network (FFN) to learn wider representations.
Multi-head attention obtains query Q , key K , and value V by applying three projections to the input, each consisting of h linear layers (or heads) that map the d m -dimensional input into a d h -dimensional space, where d h = d m /h is the head dimension.
The
FFN consists of two linear layers, where the ﬁrst expands the dimensions from d m to d f and the [Figure]
h
[Figure]
d
[Figure]
d f
[Figure]
M
[Figure]
lti - h ea d A tt e n ti on
F ee d F o r w a r d N e t w o r k
[Figure]
f
[Figure]
y e r s ( d e p t [Figure]
[Figure]
second reduces the dimensions from d f to d m .
The depth of a transformer block is 4, consisting of (1) three parallel branches for queries, keys, and values, (2) a fusion layer that combines the output of multiple heads, and (3) two sequential linear layers in the FFN.
In general, transformer-based networks sequentially stacks transformer blocks to increase network capacity and depth.
This paper extends the transformer architecture and introduces a deep and light-weight transformer,
DeLighT .
Our model uses a deep and light-weight expand-reduce transformation, DeLighT transformation (Section 3.1), that enables learning wider representations efﬁciently.
It also enables replacing multi-head attention and feed forward network (FFN) layers with single-head attention and a light-weight FFN (Section 3.2).
DeLighT transformation decouples attention dimensions from the depth and width, allowing us to learn representations efﬁciently using block-wise scaling instead of uniform stacking of transformer blocks (Section 3.3).
3. D E L IGH T T RANSFORMATION
DeLighT transformation maps a d m dimensional input vector into a high dimensional space (expansion) and then reduces it down to a d o dimensional output vector (reduction) using N layers of the group transformations of Mehta et al.
(2018), as shown in Figure 1d.
During these expansion and reduction phases, DeLighT transformation uses group linear transformations (GLTs) because they learn local representations by deriving the output from a speciﬁc part of the input and are more efﬁcient than linear transformations.
To learn global representations, the DeLighT transformation shares information between different groups in the group linear transformation using feature shufﬂing, analogous to channel shufﬂing in convolutional networks (Zhang et al., 2018).
A standard approach to increase the expressivity and capacity of transformers is to increase the input dimensions, d m .
However, increasing d m linearly also increases the number of operations in multihead attention ( O ( n 2 d m ) , where n is the sequence length) in a standard transformer block (Figure 1a).
In contrast, to increase the expressivity and capacity of the DeLighT block, we increase the depth and width of its intermediate DeLighT transformations using expansion and reduction phases.
This enables us to use smaller dimensions for computing attention, requiring fewer operations.
Formally, the DeLighT transformation is controlled by ﬁve conﬁguration parameters: (1) number of GLT layers N , (2) width multiplier w m , (3) input dimension d m , (4) output dimension d o , and [Figure] (5) maximum groups g max in a GLT.
In the expansion phase, the DeLighT transformation projects the d m -dimensional input to a high-dimensional space, d max = w m d m , linearly using (cid:100) N 2 (cid:101) layers.
In the reduction phase, the DeLighT transformation projects the d max -dimensional vector to a d o -dimensional space using the remaining N − (cid:100) N 2 (cid:101) GLT layers.
Mathematically, we deﬁne the output Y at each GLT layer l as:
[Equation]
(1) where W l = (cid:110) W l 1 , · · · , W lg l (cid:111) and b l = (cid:110) b l 1 , · · · , b lg l (cid:111) are the learnable weights and biases of group linear transformation F with g l groups at the l -th layer.
Brieﬂy, the F function takes the input X (cid:0) or H (cid:0) X , Y l − 1 (cid:1)(cid:1) and splits into g l non-overlapping groups such that X = (cid:8) X 1 , · · · , X g l (cid:9) .
The function F then linearly transforms each X i with weights W li and bias b li to produce output
Y li =
X i W li + b li .
The outputs of each group Y li are then concatenated to produce the output Y l .
The function H ﬁrst shufﬂes the output of each group in Y l − 1 and then combines it with the input X using the input mixer connection of Mehta et al.
(2020) to avoid vanishing gradient problems.
Figure 2 visualizes the expansion phase in the DeLighT transformation with group linear transformation, feature shufﬂing, and the input mixer connection.
The number of groups at the l -th GLT in DeLighT transformation are computed as:
[Equation]
In our experiments, we use g max = (cid:100)
d m 32 (cid:101) so that each group has at least 32 input elements.
3. D E L IGH T BLOCK
Figure 1b shows how we integrate DeLighT transformation into the transformer block to improve its efﬁciency.
The d m -dimensional inputs are ﬁrst fed to the DeLighT transformation to produce d o -dimensional outputs, where d o < d m .
These d o -dimensional outputs are then fed into a single head attention, followed by a light-weight FFN to model their relationships.
DeLighT layer and single head attention: Let us assume we have a sequence of n input tokens, each of dimensionality d m .
These n , d m -dimensional inputs are ﬁrst fed to the DeLighT transformation to produce n , d o -dimensional outputs, where d o < d m .
These n , d o -dimensional outputs are then projected simultaneously using three linear layers to produce d o -dimensional queries Q , keys K , and values V .
We then model contextual relationships between these n tokens using scaled dot-product attention (Eq. 3).
To enable the use of residual connections (He et al., 2016), the d o -dimensional outputs of this attention operation are linearly projected into a d m -dimensional space.
[Equation]
We hypothesize that the ability of DeLighT to learn wider representations allows us to replace multi-head attention with single-head attention.
The computational costs for computing attention in
[Figure]
(see Eq. 4)
[Figure]
the standard transformer and the DeLighT block are O ( d m n 2 ) and O ( d o n 2 ) respectively, where d o < d m .
Therefore, the DeLighT block reduces the cost for computing attention by a factor of d m /d o .
In our experiments, we used d o =
d m / 2 , thus requiring 2 × fewer multiplication-addition operations as compared to the transformer architecture.
Light-weight FFN:
Similar to FFNs in transformers, this block also consists of two linear layers.
Since the DeLighT block has already incorporated wider representations using the DeLighT transformation, it allows us to invert the functionality of FFN layers in the transformer.
The ﬁrst layer reduces the dimensionality of the input from d m to d m /r while the second layer expands the dimensionality from d m /r to d m , where r is the reduction factor (see Figure 1b).
Our light-weight
FFN reduces the number of parameters and operations in the FFN by a factor of rd f /d m .
In the standard transformer, the FFN dimensions are expanded by a factor of 4 .
In our experiments, we used r = 4 .
Thus, the light-weight FFN reduces the number of parameters in the FFN by 16 × .
Block depth: The DeLighT block stacks (1) a DeLighT transformation with N GLTs, (2) three parallel linear layers for key, query, and value, (3) a projection layer, and (4) two linear layers of a light-weight FFN.
Thus, the depth of DeLighT block is N + 4 .
Compared to the standard transformer block (depth is 4), DeLighT block is deeper.
3. B LOCK - WISE SCALING
Standard methods for improving the performance of sequence models include increasing the model dimensions (width scaling), stacking more blocks (depth scaling), or both.
However, such scaling is not very effective on small datasets.
For example, when a Transformer-Base ( d m = 512 ) network is replaced with Transformer-Large ( d m = 1024 ) on the WMT’16 En-Ro corpus, the number of parameters increases by approximately 4 × while the performance does not change appreciably (BLEU: 34.28 vs. 34.35).
We hypothesize that this happens because scaling model width and depth allocates parameters uniformly across blocks, which may lead to learning redundant parameters.
To create deep and wide networks, we extend model scaling to the block level (see Figure 3).
Scaling the DeLighT block: The DeLighT block learns deep and wide representations using the
DeLighT transformation, whose depth and width are controlled by two conﬁguration parameters: the number of GLT layers N and the width multiplier w m , respectively (Figure 3a).
These conﬁguration parameters allow us to increase the number of learnable parameters inside the DeLighT block independently of the input d m and output d o dimensions.
Such calibration is not possible with the standard transformer block because their expressiveness and capacity are a function of the input (input dimension = number of heads × head dimension).
Here, we introduce block-wise scaling that creates a network with variably-sized DeLighT blocks, allocating shallower and narrower DeLighT blocks near the input and deeper and wider DeLighT blocks near the output.
To do so, we introduce two network-wide conﬁguration parameters: minimum N min and maximum
N max number of GLTs in a DeLighT transformation.
For the b -th DeLighT block, we compute the number of GLTs N b and the width multiplier w bm in a DeLighT transformation using linear scaling (Eq. 4).
With this scaling, each DeLighT block has a different depth and width (Figure 3a).
[Equation]
Here, B denotes the number of DeLighT blocks in the network.
We add superscript b to number of
GLT layers N and width multiplier w m to indicate that these parameters are for the b -th block.
Network depth: The depth of transformer block is ﬁxed, i.e., 4.
Therefore, previous works (Raffel et al., 2019; Brown et al., 2020; Wang et al., 2019) have associated the depth of transformer-based networks with the number of transformer blocks.
In DeLighT , we present a different perspective to learn deeper representations, wherein each block is variably-sized.
To compute the network depth, we use the standard deﬁnition across different domains, including computer vision (e.g., ResNet of He et al. 2016) and theoretical machine learning (Telgarsky, 2016).
These works measures network depth as the number of sequential learnable layers (e.g., convolution, linear, or group linear).
Similarly, the depth of DeLighT and transformer networks with B blocks is (cid:80) B− 1 b =
( N b + 4) and 4 B , respectively.
E XPERIMENTAL RESULTS
We evaluate the performance of DeLighT on two standard sequence modeling tasks: (1) machine translation (Section 4.1) and (2) language modeling (Section 4.2).
4. M ACHINE T RANSLATION
Datasets and evaluation: We benchmark DeLighT models on four datasets: (1) IWSLT’
German-English (De-En), (2) WMT’16 English-Romanian
(En-Ro), (3) WMT’14 English-German (WMT’14 En-De), and (4) WMT’14 English-French (WMT’14 En-Fr).
For the IWSLT’14 De-En dataset, we replicate the setup of Wu et al.
and Edunov et al.
(2018), which uses 160K/7K/7K sentence pairs for training, validation, and testing with a joint BPE vocabulary of about 10K tokens, respectively.
For the WMT’14 English-German (En-De) dataset, we follow the setup of Vaswani et al.
The dataset has 3.9M/39K/3K sentence pairs for training, validation, and testing respectively with a joint BPE vocabulary size of 44K. 2 For the WMT’14 English-French
(En-Fr) dataset, we replicate the setup of Gehring et al.
(2017), which uses 36M/27K/3K sentence pairs for training, validation, and testing respectively with a joint BPE vocabulary size of 44K. The performance is evaluated in terms of BLEU (Papineni et al., 2002) (higher is better) on the test set.
We follow Wu et al.
for beam search related hyper-parameters.
Architecture: We follow the symmetric encoder-decoder architecture of Vaswani et al.
(2017) with sinusoidal positional encodings.
Both the encoder and the decoder have B DeLighT blocks.
Decoder blocks are identical to the encoder blocks (Figure 1b), except that they have an additional source-target single-head attention unit before the light-weight FFN.
In the source-target single-head attention unit, keys and values are projections over the encoder output (full details in Appendix A).
In our experiments, we use w m = 2 , N min = 4 , and N max = 8 for WMT’16 En-Ro, WMT’14 En-De, and WMT’14 En-Fr; resulting in 222 layer deep DeLighT networks.
For IWSLT’14 De-En, we used w m = 1 , N min = 3 , and N max = 9 for IWSLT’14 De-En; resulting in 289 layer deep network.
For simplicity, we set B = N max .
We use a learnable look-up table that maps every token in the vocabulary to a 128-dimensional vector.
We implement our models using Fairseq (Ott et al., 2019) and use their provided scripts for data pre-processing, training, and evaluation.
Training:
For IWSLT’14 De-En models, we follow the setup of Wu et al.
and train all our models for 50K iterations with a batch size of 4K tokens on a single NVIDIA GTX 1080 GPU.
For
WMT’16 En-Ro, we follow the training setup of Ghazvininejad et al.
(2019) and train models for 100K iterations on 16 NVIDIA Tesla V100 GPUs with an effective batch size of 64K tokens.
For
WMT’14 En-De and WMT’14 En-Fr, we follow the training set-up of Wu et al.
and train our models on 16 V100 GPUs for 30K and 50K iterations, respectively.
We use Adam (Kingma and Ba, 2015) to minimize cross entropy loss with a label smoothing value of 0.1 during training.
For a fair comparison, we trained baseline transformer models using the same training set-up.
[Table]
[Table] [Figure] [Table]
[Figure] [Figure]
4.1. R ESULTS
Comparison with baseline transformers: Table 1 compares the performance of DeLighT with the baseline transformers of Vaswani et al.
(2017) on different corpora.
DeLighT delivers better performance with fewer parameters than transformers, across different corpora.
Speciﬁcally, on low-resource (WMT’16 En-Ro) and high resource (WMT’14 En-De & WMT’14 En-Fr) corpora,
DeLighT delivers similar or better performance with 2 .
8 × and 1 .
8 × fewer parameters, respectively.
When the number of parameters are increased DeLighT outperforms transformers.
For example, on
WMT’14 En-Fr dataset
DeLighT is 3 .
7 × deeper than transformers and improves its BLEU score by 1.3 points yet with 13 million fewer parameters and 3 billion fewer operations (see Table 2)
Particularly interesting are the performance comparisons of DeLighT with the baseline transformers of Vaswani et al.
(2017) and its neural search variant, i.e., Evolved Transformer of So et al.
(2019), at two different parametric settings on WMT’14 En-De corpora in Figure 4.
For small models (< 10 M parameters), DeLighT models delivers better performance and for attaining the same performance as these models, DeLighT models requires fewer parameters.
Comparison with state-of-the-art methods: Most state-of-the-art methods have evaluated the performance on WMT’14 En-De while some have also evaluated on IWSLT’14 De-En.
Table 3 compares the performance of DeLighT with state-of-the-art methods on these two corpora.
DeLighT delivers similar or better performance than existing methods.
It is important to note that existing methods have improved baseline transformers with different design choices – for example, the asymmetric encoder-decoder structure (Wang et al., 2019) and neural architecture search (So et al., 2019).
We believe that DeLighT , in the future, would also beneﬁt from such design choices.
Scaling up DeLighT models: Figure 5 shows the performance of DeLighT models improves with increase in network parameters; suggesting their ability to learn representations across different corpora, including low-resource.
[Figure] [Figure]
4. L ANGUAGE M ODELING
Datasets and evaluation: We evaluate on the WikiText-103 dataset (Merity et al., 2017) that has 103M/217K/245K tokens for training, validation, and testing.
It has a word-level vocabulary of about 260K tokens.
Following recent works (Baevski and Auli, 2019; Dai et al., 2019), we report performance in terms of perplexity (lower is better) on the test set.
Architecture: We use the transformer-based decoder architecture of Baevski and Auli (2019) with B DeLighT blocks.
We use w m = 2 , N min = 4 , and N max = 12 .
We scale d m using values { 384 , 512 , 784 , 1024 } for increasing network parameters.
For simplicity, we set B = N max .
Following standard practice, we use adaptive input (Baevski and Auli, 2019) as a look-up table and adaptive output (Grave et al., 2017a) as the classiﬁcation layer with one head (head dimension is 128) and two tails (tail dimensions are 64 and 32).
We also share weights between the input and the output layers.
Training: We follow the training setup of Baevski and Auli (2019), except that we train our models on 8 NVIDIA Tesla V100 GPUs for 100K iterations with a context length of 512 and an effective batch size of 64K tokens.
We use Adam during training and use a context length of 480 during test.
Results: Table 4b compares the performance of DeLighT with previous methods on WikiText-103.
Table 4a plots the variation of perplexity with number of parameters for DeLighT and TransformerXL
(Dai et al., 2019) – which outperforms other transformer-based implementations (e.g., Baevski and
Auli 2019).
Both tables show that DeLighT delivers better performance than state-of-the-art methods (including Transformer-XL) and it does this using a smaller context length and signiﬁcantly fewer parameters, suggesting that the DeLighT transformation helps learn strong contextual relationships.
A NALYSIS AND D ISCUSSIONS ON C OMPUTATIONAL E FFICIENCY
Training time and memory consumption: Table 5 compares the training time and memory consumption of DeLighT with baseline transformers.
For an apples-to-apples comparisons, we implemented the Transformer unit without NVIDIA’s dedicated CUDA kernel, and trained both transformer and DeLighT full-precision networks for 30K iterations on 16 NVIDIA V100 GPUs. The transformer and DeLighT models took about 37 and 23 hours for training and consumed about 12.5 GB and 14.5 GB of GPU memory, respectively (R1 vs. R2). When we enabled the dedicated CUDA kernel provided by APEX library 3 for multi-head attention in Transformers, the training time of the
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
expect
[Table]
transformer model reduced from 37 to 16 hours while we did not observe any signiﬁcant change in memory consumption.
Motivated by this observation, we implemented dedicated CUDA kernels for grouping and ungrouping functions in GLTs (see Appendix E).
With these changes, training time and
GPU memory consumption of DeLighT reduced by about 4 hours and 3 GB, respectively.
We emphasize that grouping, linear transformation, feature shufﬂing, and ungrouping, can be implemented efﬁciently using a single CUDA kernel.
In future, we expect a dedicated CUDA kernel for these operations would further reduce the memory consumption as well as training/inference time.
Regularization: Table 6 shows that DeLighT delivers similar performance to baseline transformers, but with fewer parameters and less regularization.
This suggests that learning representations with better transformation functions alleviates the need for dropout.
C ONCLUSION
This paper introduces a deep and light-weight transformer architecture, DeLighT , that efﬁciently allocates parameters both within the DeLighT block and across DeLighT blocks.
Compared to state-of-the-art transformer models, DeLighT models are (1) deep and light-weight and (2) deliver similar or better performance.
In the future, we plan to apply DeLighT to other tasks, including language model pre-training, question answering, and language generation.
Acknowledgements: This research was supported by ONR N00014-18-1-2826, DARPA N6600119-2-403, NSF (IIS-1616112, IIS1252835), and an Allen Distinguished Investigator Award.
Authors would also like to thank members of the UW-NLP and the H2Lab at The University of Washington for their valuable feedback and comments.
A D E L IGH T A RCHITECTURES FOR L ANGUAGE M ODELING AND M ACHINE T RANSLATION
DeLighT architectures for language modeling and machine translation are shown in Figure 6.
For language modeling, we follow the architecture in Baevski and Auli (2019) while for machine translation, we follow the architecture in Vaswani et al.
Language modeling: Figure 6a shows the architecture for language modeling.
The architecture stacks B DeLighT blocks, the conﬁguration of each block is determined using block-wise scaling.
Each block has three sub-layers.
The ﬁrst layer is a DeLighT transformation that learns representations in high-dimensional space.
The second layer is a single-head attention that encodes contextual relationships.
The third layer is a position-wise light-weight feed-forward network.
Similar to Vaswani et al.
(2017), we employ a residual connections (He et al., 2016).
Similar to previous works (Baevski and Auli, 2019; Dai et al., 2019), we use tied adaptive input (Baevski and Auli, 2019) and adaptive softmax (Grave et al., 2017a) to map tokens to vectors and vectors to tokens, respectively.
Machine translation: Figure 6b shows the architecture for machine translation.
The encoder stacks B DeLighT blocks, the conﬁguration of each block is determined using block-wise scaling.
Similar to language modeling, each encoder block has three sub-layers.
The ﬁrst layer is a DeLighT transformation that learns representations in high-dimensional space.
The second layer is a single-head attention that encodes contextual relationships.
The third layer is a position-wise light-weight feed-forward network.
Similar to Vaswani et al.
(2017), we employ a residual connections (He et al., 2016).
We use learnable look-up table to map tokens
[Figure]
B ×
I npu t a nd ou t pu t w e i gh t s a r e
[Figure]
B × I npu t a n d ou t pu t w e i gh t s a r e
ti e d
[Figure]
to vectors.
Similar to the encoder, the decoder also stacks B blocks.
Decoder blocks are identical to encoder blocks, except that they have an additional source-target single-head attention unit before the light-weight FFN.
Keys and values in source-target single-head attention unit are projections over the encoder output.
We use standard learnable look-up table to map tokens to vectors and linear classiﬁcation layer to map vectors to tokens.
B G ROUP LINEAR TRANSFORMATION WITH I NPUT - MIXER CONNECTION
Group linear transformation (GLT) F splits a d m -dimensional input X into g non-overlapping groups such that X =
Concat ( X 1 , · · · , X g ) , where X i is a d m g -dimensional vector.
X i ’s are then simultaneously transformed using g linear transforms
W i ∈ R dmg × dog to produce g outputs Y i
=
X i W i .
Y i ’s are then concatenated to produce the ﬁnal d o -dimensional output Y =
Concat ( Y 1 , · · · , Y g ) .
Figure 7a shows an example of GLT in the expansion phase of DeLighT transformation.
For illustrative purposes, we have used the same dimensions in this example.
Recall that as we go deeper in the expansion phase, the number of groups increases.
In this example, the ﬁrst layer has one group, the second layer has two groups and the third layer has four groups.
GLTs learns group-speciﬁc representations and are local.
To allow
GLT to learn global representations, we use feature shufﬂe.
An example of GLT with feature shufﬂe is shown in Figure 7b.
Furthermore, training deep neural networks by merely stacking linear or group linear (with or without feature shufﬂe) is challenging because of vanishing gradient problem.
Residual connections introduced by He et al.
(2016) mitigates this problem and helps train deep neural networks.
However, such connections cannot be employed when input and output dimensions are not the same (e.g., during the expansion and reduction phases in DeLighT transformation).
To stabilize the training and learn deeper representations, we use input-mixer connection of Mehta et al.
Figure 7c shows an example of GLT with feature shufﬂe and input mixer connection.
[Figure]
G r oup s
[Figure]
G r oup s
[Figure]
G r o u p s =
G r o u p s =
G r o u p s =
[Figure]
C M ULTIPLICATION -A DDITION O PERATIONS IN D E L IGH T
The DeLighT block is built using linear transformations, GLTs, and scaled dot-product attention.
Total number of multiplication-addition operations (MACs) in a network is an accumulation of these individual operations.
Let n denotes the number of source tokens, m denotes the number of target tokens, d m denotes the input dimension, d o denotes the output dimension, and g denotes the number of groups in GLT.
The procedure for counting MACs for each of these operations is described below.
Group linear transformation (GLT): GLT F has g learnable matrices
W i ∈ R dmg × dog .
Therefore, GLT learns d m d o g parameters and performs d m d o g
MACs to transform d m -dimensional input to d o -dimensional output.
Following a standard practice, e.g., ResNet of He et al.
(2016), we count addition and multiplication as one operation instead of two because these operations can be fused in recent hardwares.
Importantly, when g = 1 , the GLT is the same as linear transformation.
Self-attention in DeLighT :
The scaled dot-product self-attention in DeLighT is deﬁned as:
[Equation]
where Q ∈ R n × d o , K ∈ R n × d o , V ∈ R n × d o denotes query, key, and value, respectively.
The attention operation involves two dot-products.
The ﬁrst dot product between Q and K while the second dot product is between the output of ﬁrst dot product and V .
Both dot products require d o n 2 MACs.
Therefore, total number of MACs in computing scaled dot-product self-attention are 2 d o n 2 .
In case of a source-target attention (as in machine translation), K ’s and V ’s are from the source (encoder) and Q ’s are incrementally decoded (one token at a time).
Therefore, the number of MACs required to decode m target tokens given n source tokens are m (cid:80) k
=1 2 knd o .
D A BLATIONS ON THE W IKI T EXT -103 DATASET
Table 7 studies the impact of DeLighT block parameters on the WikiText-103 dataset, namely (1) minimum number of GLTs N min , (2) maximum number of GLTs N max , (3) width multiplier w m , and (4) model dimension d m (see Figure 1b).
Figure 8, Figure 9, and Figure 10 shows the impact of the DeLighT transformation, feature shufﬂing, and the light-weight FFN.
Table 8 shows the effect of position of DeLighT transformation in the
DeLighT block while Figure 12 shows the effect of scaling DeLighT networks.
We choose the WikiTextdataset for ablations because it has very large vocabulary compared to other datasets (267K vs. 30-40K), allowing us to test the ability under large vocabulary sizes.
The performance is reported in terms of perplexity (lower is better) on the validation set.
In our ablation studies, we used the same settings for training as in Section 4.2 except that we train only for 50K iterations.
DeLighT block: Overall, Table 7 shows that scaling depth and width using DeLighT transformation and block-wise scaling improves performance.
We make following observations:
a) Block-wise scaling (R4, R5) delivers better performance compared to uniform scaling (R1-R3).
For instance, DeLighT with N min = 4 and N max = 8 (R4) is 1 .
25 × shallower than DeLighT with N min = 8 and N max = 8 (R2), but delivers better performance with a similar number of parameters and operations.
Scaling w m improves performance (R2 vs. R3), however, the improvement is signiﬁcantly lower than for the model with block-wise scaling (R3 vs. R5).
This suggests that non-uniform distribution of parameters across blocks allows the network to learn better representations.
b) Different ratios between N max and N min yields different results.
We observe signiﬁcant performance improvements when the ratio is greater than or equal to two.
For example, when we scale N max N min from 2 to(R6 vs. R8), the perplexity improves by ∼ 5 points with only a moderate increase in network parameters.
On the other hand, when the N max N min is close to 1 (R6 vs. R7), performance does not change appreciably.
This is likely because the allocation of parameters across blocks is close to uniform (Eq. 4).
This is consistent with our previous observation.
c) Learning shallower and narrower representations near the input and deeper and wider representations near the output achieves better performance.
For example, when we scaled N max from 8 to 12 for N min =(R6, R8), DeLighT delivered better performance with a similar number of parameters compared to a model with N min = 6 (R7, R9).
This is likely because the ratio of N max and N min is higher when N min = 4 , which helps allocate parameters per block more effectively.
d) Deeper and wider representations near the input and shallower and narrower representations near the output hurts performance (R13 vs. R16).
e) Scaling width using w m and d
m improves performance (R10-R15), however, their impact is different.
For example, when we scale w m and d m by two, the rate of increase in number of parameters and operations is more rapid with d m compared to w m .
DeLighT ’s ability to learn wider representations in different ways may be useful in selecting application speciﬁc models
Impact of DeLighT transformation: We replace DeLighT transformation in the DeLighT block (Figure 1b) with (1) the DeFINE transformation and (2) a stack of linear layers.
Figure 8 shows that DeLighT transformation delivers similar performance with signiﬁcantly fewer parameters compared to the DeFINE unit
[Table]
and linear layers.
In these experiments, the settings are the same as R13-R15 (Table 7), except, N max = 8 , because models with a stack of linear layers learn too many parameters.
Feature shufﬂing: Figure 9 shows that feature shufﬂing improves the performance of DeLighT by 1perplexity points.
Here, we use the same settings as in R13-R
(Table 7).
Light-weight FFN:
Figure 10 shows the impact of varying the reduction factor r in the light-weight FFN.
We use the same settings as in R13 (Table 7).
We did not observe any signiﬁcant drop in performance until r = 4 .
Beyond r = 4 , we see a drop in performance (perplexity increases by ∼ 2 points).
In such cases, the inner dimensions of the light-weight FFN are very small and hurt performance.
Notably, the light-weight FFN with r = 2 2 delivered the same performance as r = 2 − 2 , but with 1 .
28 × fewer network parameters.
At r = 2 − 2 , the light-weight FFN is the same as the FFN in Vaswani et al.
This suggests that the ability of
[Figure]
[Figure]
Reduction factor ( r )
Perplexity 26.
20.8 Parameters (in million)
[Figure]
B b l o c k s
[Equation]
(see Eq. 4)
[Equation]
Model dimension ( d m )
Perplexity
N min =4, N max =4, Mean=
N min =8, N max =8, Mean=
N min =4, N max =8, Mean=5.
N min =4, N max =12, Mean=7.
(b)
DeLighT transformation to learn representations in high-dimensional spaces efﬁciently allows us to reduce the computational burden on the FFN.
We also tested removing the light-weight FFN and while it reduced parameters by ∼ 0.5-1 M, performance dropped by about 2-3 perplexity points across different parametric settings.
Uniform vs. block-wise scaling: Figure 11 compares the performance of DeLighT with uniform and blockwise scaling.
For a given model dimension d m , DeLighT models with block-wise scaling delivers better performance.
Position of DeLighT transformation: We studied three conﬁgurations for the DeLighT transformation on the WikiText-103 validation set (Table 8): (1) DeLighT transformation followed by single-headed attention and light-weight FFN, (2) single-headed attention followed by DeLighT transformation, and (3) single-headed attention followed by DeLighT transformation and light-weight FFN.
For similar number of parameters, we found that (2) and (3) drops the performance of (1) signiﬁcantly across different parametric settings.
This suggests that deeper and wider representations helps learn better contextual representations; allowing us to replace multi-headed attention with single-headed attention.
Scaling up DeLighT : Figure 12 shows the results of DeLighT models obtained after varying conﬁguration parameters of DeLighT transformations ( N min ={4, 6}, N max ={8, 12}, w m ={2, 3, 4}, and d m ={256, 384, 512}).
We can see that scaling one conﬁguration parameter (e.g., d m ) while keeping other conﬁguration parameters constant (e.g., N min , N max , and w m ) consistently improves performance.
Conﬁguration
Parameters
Perplexity
DeLighT transformation + Single-head attention + Light-weight FFN 31 M 34.20 Single-head attention + DeLighT transformation 30 M 39.02 Single-head attention + DeLighT transformation + Light-weight FFN 31 M 39.
DeLighT transformation + Single-head attention + Light-weight FFN 99 M 23.16 Single-head attention + DeLighT transformation 96 M 28.33 Single-head attention + DeLighT transformation + Light-weight FFN 99 M 27.
[Figure]
[Figure]
exity            w m =w m =w m =
[Figure]
512 Model dimension [Figure]
Perpl exity            w m =w m =w m =
[Figure]
[Figure]
[Figure] [Figure]
w m = w m =w m
[Figure]
[Figure]
[Figure]
Perplexity            w m =w m =w m =
[Figure]
keeping other conﬁgu- ration parameters constant (e.g., N min , N max , and w m ) consistently improves performance.
The numbers on top of each bar represents network parameters (in million).
Lower value of perplexity means better performance.
This work investigates relationships between N min , N max , w m , and d m , manually.
We believe that a more principled approach, such as compound scaling of Tan and Le (2019), that establishes relationships between these parameters would produce more efﬁcient and accurate models.
E S OURCE CODE FOR GROUP LINEAR TRANSFORMATION
The source code for implementing group linear transformation (GLT) in PyTorch is shown in Listing 1.
The source code for efﬁciently implementing the grouping function in GLT is shown in Listing 2.
Since the ungrouping kernel is similar to grouping kernel, we have not shown it here.
The reshape and transpose operations in naive PyTorch implementation for grouping and ungrouping are replaced with a dedicated CUDA kernels, resulting in reduced memory footprint and faster training.
Listing 1: "Naive implementation of GLT
Pytorch" import torch def
’’’bsz = x
## GROUPING FUNCTION:
Converts [B x N] tensor to [g x B x N/g] ## # [B x N] -->
[B x g x N/g] x = x.contiguous().view(bsz, n_groups, -1)
# [B x g x N/g] -->
[g x B x N/g] x = x.transpose(0, 1) # transpose so group
## TRANSFORMATION FUNCTION:
Transforms from N/g-dimensional space to M/g-dimensional space ## #
[g x B x N/g] x [g x N/g x M/g] --> [g x B x M/g] x = torch.bmm(x, weights) # multiply with Weights # add bias if bias is not None: x = torch.add(x, bias) ## REGROUPING FUNCTION:
Converts [g x B x M/g] tensor to [B x M] ## # [g x B x M/g] -->
[B x g x M/g] x = x.transpose(0, 1) # transpose so that batch is first # [B x g x M/g] --
>
[B x M] x = x.contiguous().view(bsz, -1) return x
Listing 2: "Grouping kernel in CUDA" /* Grouping Kernel: Transforms input
N]
B x N/g] */ template<typename scalar t> __global__ void grouping_kernel_forward(const scalar_t* input, const int groups, const int total_elements, const int input_features, const int group_features, const int batch_size, scalar_t* output){ const int index = IMUL(blockIdx.x, blockDim.x) + threadIdx.x; if (index >= total_elements){ return; } const int b_idx = index / group_features; const int g_f_idx =
(index % group_features); int in_offset, out_offset; #pragma unroll for(int g=0; g < groups; g++){ in_offset = (b_idx * input_features) + (g * group_features) + g_f_idx; out_offset =
((g * batch_size + b_idx) * group_features) + g_f_idx; output[out_offset] = input[in_offset]; } }