# delight_deep_and_lightweight_transformer__92464ad5.pdf

Published as a conference paper at ICLR 2021

DELIGHT: DEEP AND LIGHT-WEIGHT TRANSFORMER

Sachin Mehta1, Marjan Ghazvininejad2, Srinivasan Iyer2, Luke Zettlemoyer1, 2, and Hannaneh Hajishirzi1,3

1University of Washington 2Facebook AI Research 3Allen Institute for AI

We introduce a deep and light-weight transformer, De Ligh T, that delivers similar or better performance than standard transformer-based models with significantly fewer parameters. De Ligh T more efficiently allocates parameters both (1) within each Transformer block using the De Ligh T transformation, a deep and lightweight transformation and (2) across blocks using block-wise scaling, that allows for shallower and narrower De Ligh T blocks near the input and wider and deeper De Ligh T blocks near the output. Overall, De Ligh T networks are 2.5 to 4 times deeper than standard transformer models and yet have fewer parameters and operations. Experiments on benchmark machine translation and language modeling tasks show that De Ligh T matches or improves the performance of baseline Transformers with 2 to 3 times fewer parameters on average.

1 INTRODUCTION

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-3 (Brown et al., 2020) uses 96 transformer blocks. However, such scaling increases the number of network parameters significantly (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-efficient attention-based architecture that can be easily scaled to be both wide and deep.

Our Deep and Light-weight Transformer architecture, De Ligh T, extends the transformer architecture of Vaswani et al. (2017) and delivers similar or better performance with significantly fewer parameters and operations. At the heart of De Ligh T is the De Ligh T 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 De Ligh T block efficiently. Since GLTs are local by nature, the De Ligh T transformation uses feature shuffling, which is analogous to channel shuffling 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 De Ligh T transformation decouples the depth and width from the input size, allowing us to allocate parameters more efficiently across blocks by using shallower and narrower De Ligh T blocks near the input and deeper and wider De Ligh T blocks near the output.

We demonstrate that De Ligh T models achieve similar or better performance than transformer models with significantly 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, De Ligh T attains transformer performance using 2.8 fewer parameters. On the high resource WMT 14 En-Fr dataset, De Ligh T delivers better performance (+0.4 BLEU score) with 1.8 fewer parameters than baseline transformers. Similarly, on language modeling, De Ligh T matches the performance of Transformer-XL (Dai et al., 2019) with 1.5 fewer parameters

Published as a conference paper at ICLR 2021

on the Wiki Text-103 dataset. Our source code is open-source and is available at: https://github.com/ sacmehta/delight

2 RELATED WORK

Improving transformers: Several methods have been introduced to improve the transformer architecture. The first 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 fixed attention heads with predefined 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 efficiently allocate parameters both at the block-level using the De Ligh T 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 efficient allocation of parameters in the network. Our results show that (1) shallower and narrower De Ligh T blocks near the input and deeper and wider De Ligh T 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., Res Net of He et al. 2016) that perform a fixed 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 significant 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 De FINE (Mehta et al., 2020), and (2) improving efficiency 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 De FINE transformation, which also learns representations using an expand-reduce strategy. The key difference between the De FINE transformation (Figure 1c) and the De Ligh T transformation (Figure 1d) is that the De Ligh T transformation more efficiently allocates parameters within expansion and reduction layers. Unlike De FINE, which uses fewer groups in group linear transformations to learn wider representations, De Ligh T transformation uses more groups to learn wider representations with fewer parameters. The De Ligh T transformation achieves comparable performance to the De FINE transformation but with significantly fewer parameters.

3 DELIGHT: DEEP AND LIGHT-WEIGHT TRANSFORMER

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 dm-dimensional input into a dh-dimensional space, where dh = dm/h is the head dimension. The FFN consists of two linear layers, where the first expands the dimensions from dm to df and the

Published as a conference paper at ICLR 2021

Multi-head Attention Feed Forward Network (FFN)

Attention ops:

FFN params:

(a) Transformer block

De Ligh T transformation with Single-head Attention Light-weight FFN

Attention ops:

FFN params:

Depth = 4 + Nb

(b) De Ligh T block

Input (dm-dimensional)

Output (do-dimensional)

Expansion Reduction

No. of layers (depth) = N

(c) De FINE transformation

Input (dm-dimensional)

Output (do-dimensional)

Expansion Reduction

No. of layers (depth) = N

(d) De Ligh T transformation

Figure 1: (a, b) Block-wise comparison between the standard transformer block of Vaswani et al. (2017) and the De Ligh T block. In the De Ligh T transformation, the number of operations in computing attention are reduced by half while the number of parameters (and operations) in the FFN are reduced by 16 . Transformations with learnable parameters ( Linear and De Ligh T ) are shown in color. The shape of linear transformations indicate their operation (expansion, reduction, etc.). (c, d) compares the De FINE transformation (Mehta et al., 2020) with the De Ligh T transformation. Compared to the De FINE transformation, the De Ligh T transformation uses group linear transformations (GLTs) with more groups to learn wider representations with fewer parameters. Different colors are used to show groups in GLTs. For simplicity, feature shuffling is not shown in (d).

second reduces the dimensions from df to dm. 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, De Ligh T. Our model uses a deep and light-weight expand-reduce transformation, De Ligh T transformation (Section 3.1), that enables learning wider representations efficiently. 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). De Ligh T transformation decouples attention dimensions from the depth and width, allowing us to learn representations efficiently using block-wise scaling instead of uniform stacking of transformer blocks (Section 3.3).

3.1 DELIGHT TRANSFORMATION

De Ligh T transformation maps a dm dimensional input vector into a high dimensional space (expansion) and then reduces it down to a do 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, De Ligh T transformation uses group linear transformations (GLTs) because they learn local representations by deriving the output from a specific part of the input and are more efficient than linear transformations. To learn global representations, the De Ligh T transformation shares information between different groups in the group linear transformation using feature shuffling, analogous to channel shuffling in convolutional networks (Zhang et al., 2018).

A standard approach to increase the expressivity and capacity of transformers is to increase the input dimensions, dm. However, increasing dm linearly also increases the number of operations in multihead attention (O(n2dm), where n is the sequence length) in a standard transformer block (Figure 1a). In contrast, to increase the expressivity and capacity of the De Ligh T block, we increase the depth and width of its intermediate De Ligh T transformations using expansion and reduction phases. This enables us to use smaller dimensions for computing attention, requiring fewer operations.

Formally, the De Ligh T transformation is controlled by five configuration parameters: (1) number of GLT layers N, (2) width multiplier wm, (3) input dimension dm, (4) output dimension do, and

Published as a conference paper at ICLR 2021

GLT (Groups = 1)

Input mixer

GLT (Groups = 2)

GLT (Groups = 4)

Input mixer

Figure 2: Example illustrating the expansion phase in the De Ligh T transformation that uses GLTs, feature shuffling, and an input mixer connection, to learn deeper and wider representations efficiently. For illustrative purposes, we have used the same input and output dimensions.

(5) maximum groups gmax in a GLT. In the expansion phase, the De Ligh T transformation projects the dm-dimensional input to a high-dimensional space, dmax = wmdm, linearly using N

2 layers. In the reduction phase, the De Ligh T transformation projects the dmax-dimensional vector to a do-dimensional space using the remaining N N

2 GLT layers. Mathematically, we define the output Y at each GLT layer l as:

Yl = F X, Wl, bl, gl , l = 1 F H X, Yl 1 , Wl, bl, gl , Otherwise (1)

where Wl = n Wl 1, , Wl gl o and bl = n bl 1, , bl gl o are the learnable weights and biases of

group linear transformation F with gl groups at the l-th layer. Briefly, the F function takes the input X or H X, Yl 1 and splits into gl non-overlapping groups such that X = X1, , Xgl . The function F then linearly transforms each Xi with weights Wl i and bias bl i to produce output Yl i = Xi Wl i + bl i. The outputs of each group Yl i are then concatenated to produce the output Yl. The function H first shuffles the output of each group in Yl 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 De Ligh T transformation with group linear transformation, feature shuffling, and the input mixer connection.

The number of groups at the l-th GLT in De Ligh T transformation are computed as:

gl = min(2l 1, gmax), 1 l N/2 g N l, Otherwise (2)

In our experiments, we use gmax = dm

32 so that each group has at least 32 input elements.

3.2 DELIGHT BLOCK

Figure 1b shows how we integrate De Ligh T transformation into the transformer block to improve its efficiency. The dm-dimensional inputs are first fed to the De Ligh T transformation to produce do-dimensional outputs, where do < dm. These do-dimensional outputs are then fed into a single head attention, followed by a light-weight FFN to model their relationships.

De Ligh T layer and single head attention: Let us assume we have a sequence of n input tokens, each of dimensionality dm. These n, dm-dimensional inputs are first fed to the De Ligh T transformation to produce n, do-dimensional outputs, where do < dm. These n, do-dimensional outputs are then projected simultaneously using three linear layers to produce do-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 do-dimensional outputs of this attention operation are linearly projected into a dm-dimensional space.

Attention(K, Q, V) = softmax QKT

We hypothesize that the ability of De Ligh T to learn wider representations allows us to replace multi-head attention with single-head attention. The computational costs for computing attention in

Published as a conference paper at ICLR 2021

Block-wise Uniform

N B 1 = Nmax

(see Eq. 4)

(a) Uniform vs. block-wise

B0 B1 B2 B3 B4 B5 B6 B7 Decoder blocks

No. of parameters

(in thousand)

Block-wise Uniform

B0 B1 B2 B3 B4 B5 B6 B7 Decoder blocks

No. of operations

(in million)

Block-wise Uniform

(b) Distribution of parameters and operations within each block

Figure 3: Block-wise scaling efficiently allocates parameters and operations across blocks, leading to shallower and narrower De Ligh T blocks near the input and deeper and wider De Ligh T blocks near the output. In (b), De Ligh T networks with both uniform (N=Nmin=Nmax=8) and block-wise (Nmin=4, Nmax=8) scaling have about 16.7 M parameters and perform 3.5 B operations (computed for a sequence length of n = 30), however, the De Ligh T network with block-wise scaling delivered 2 points better perplexity.

the standard transformer and the De Ligh T block are O(dmn2) and O(don2) respectively, where do < dm. Therefore, the De Ligh T block reduces the cost for computing attention by a factor of dm/do. In our experiments, we used do = dm/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 De Ligh T block has already incorporated wider representations using the De Ligh T transformation, it allows us to invert the functionality of FFN layers in the transformer. The first layer reduces the dimensionality of the input from dm to dm/r while the second layer expands the dimensionality from dm/r to dm, 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 rdf/dm. In the standard transformer, the FFN dimensions are expanded by a factor of 4.1 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 De Ligh T block stacks (1) a De Ligh T 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 De Ligh T block is N + 4. Compared to the standard transformer block (depth is 4), De Ligh T block is deeper.

3.3 BLOCK-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 (dm = 512) network is replaced with Transformer-Large (dm = 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 De Ligh T block: The De Ligh T block learns deep and wide representations using the De Ligh T transformation, whose depth and width are controlled by two configuration parameters: the number of GLT layers N and the width multiplier wm, respectively (Figure 3a). These configuration parameters allow us to increase the number of learnable parameters inside the De Ligh T block independently of the input dm and output do 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 De Ligh T blocks, allocating shallower and narrower De Ligh T blocks near the input and deeper and wider De Ligh T blocks near the output.

To do so, we introduce two network-wide configuration parameters: minimum Nmin and maximum Nmax number of GLTs in a De Ligh T transformation. For the b-th De Ligh T block, we compute

1Transformer-base uses dm=512 and df=2048 while Transformer-large uses dm=1024 and df=4096.

Published as a conference paper at ICLR 2021

the number of GLTs N b and the width multiplier wb m in a De Ligh T transformation using linear scaling (Eq. 4). With this scaling, each De Ligh T block has a different depth and width (Figure 3a).

N b = Nmin + (Nmax Nmin) b

B 1 , wb m = wm + (Nmax Nmin) b

Nmin(B 1) , 0 b B 1 (4)

Here, B denotes the number of De Ligh T blocks in the network. We add superscript b to number of GLT layers N and width multiplier wm to indicate that these parameters are for the b-th block.

Network depth: The depth of transformer block is fixed, 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 De Ligh T, we present a different perspective to learn deeper representations, wherein each block is variably-sized. To compute the network depth, we use the standard definition across different domains, including computer vision (e.g., Res Net 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 De Ligh T and transformer networks with B blocks is PB 1 b=0 (N b + 4) and 4B, respectively.

4 EXPERIMENTAL RESULTS

We evaluate the performance of De Ligh T on two standard sequence modeling tasks: (1) machine translation (Section 4.1) and (2) language modeling (Section 4.2).

4.1 MACHINE TRANSLATION

Datasets and evaluation: We benchmark De Ligh T models on four datasets: (1) IWSLT 14 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. (2019) 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. (2017). 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. (2019) 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 De Ligh T 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 wm = 2, Nmin = 4, and Nmax = 8 for WMT 16 En-Ro, WMT 14 En-De, and WMT 14 En-Fr; resulting in 222 layer deep De Ligh T networks. For IWSLT 14 De-En, we used wm = 1, Nmin = 3, and Nmax = 9 for IWSLT 14 De-En; resulting in 289 layer deep network. For simplicity, we set B = Nmax. 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. (2019) 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. (2019) 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.

2We use training and validation data that is compatible with the Tensor2Tensor library (Vaswani et al., 2018) in order to have fair comparisons with recent works (e.g., Evolved Transformer).

Published as a conference paper at ICLR 2021

IWSLT 14 De-En WMT 16 En-Ro

Model # Params Ratio BLEU BLEU # Params Ratio BLEU BLEU

Transformer (Vaswani et al., 2017) 34.4 62 M 34.3

Transformer (Our impl.) 42 M 1.0 34.3 62 M 1.0 34.3 De Ligh T 14 M 0.3 33.8 -0.5 22 M 0.35 34.3 0.0 De Ligh T 30 M 0.7 35.3 +1.0 53 M 0.85 34.7 +0.4

(a) Results on small corpora

WMT 14 En-De WMT 14 En-Fr

Model # Params Ratio BLEU BLEU # Params Ratio BLEU BLEU

Transformer (Vaswani et al., 2017) 62 M 27.3 62 M 38.1

Transformer (Our impl.) 67 M 1.0 27.7 67 M 1.0 39.2 De Ligh T 37 M 0.55 27.6 -0.1 37 M 0.55 39.6 +0.4 De Ligh T 54 M 0.80 28.0 +0.3 54 M 0.80 40.5 +1.3

(b) Results on large corpora

Table 1: Comparison with baseline transformers on machine translation corpora. De Ligh T models require significantly fewer parameters to achieve similar performance. Here, and indicate the best reported transformer baselines from Wu et al. (2019) and Ghazvininejad et al. (2019), respectively.

Depth # Params # MACs BLEU

Transformer 60 67 M 11.1 B 39.2 De Ligh T 222 37 M 5.6 B 39.6 De Ligh T 222 54 M 8.1 B 40.5

Table 2: De Ligh T networks are deep, lightweight and efficient as compared to transformers. BLEU score is reported on the WMT 14 En-Fr dataset. To compute network depth, we count the number of sequential layers in the network (Section 3.3). We used 20 source and 20 target tokens for computing multiplication-addition operations (MACs). See Appendex C for details.

0 10 20 30 40 50 60 Parameters (in million)

BLEU (WMT'14 En-De)

1.8x fewer parameters

De Ligh T Transformer Evolved Trans.

Figure 4: Comparison of De Ligh T with Transformers and Evolved Transformers at two different settings, on the WMT 14 En-De corpus: (1) the number of parameters is the same and (2) the performance is the same.

4.1.1 RESULTS

Comparison with baseline transformers: Table 1 compares the performance of De Ligh T with the baseline transformers of Vaswani et al. (2017) on different corpora. De Ligh T delivers better performance with fewer parameters than transformers, across different corpora. Specifically, on low-resource (WMT 16 En-Ro) and high resource (WMT 14 En-De & WMT 14 En-Fr) corpora, De Ligh T delivers similar or better performance with 2.8 and 1.8 fewer parameters, respectively. When the number of parameters are increased, De Ligh T outperforms transformers. For example, on WMT 14 En-Fr dataset, De Ligh T 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 De Ligh T 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), De Ligh T models delivers better performance and for attaining the same performance as these models, De Ligh T 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 De Ligh T with state-of-the-art methods on these two corpora. De Ligh T 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 De Ligh T, in the future, would also benefit from such design choices.

Scaling up De Ligh T models: Figure 5 shows the performance of De Ligh T models improves with increase in network parameters; suggesting their ability to learn representations across different corpora, including low-resource.

Published as a conference paper at ICLR 2021

Model # Params BLEU

Transformers (Vaswani et al., 2017) 42 M 34.3 Variational Attention (Deng et al., 2018) 33.1 Dynamic convolutions (Vaswani et al., 2017) 43 M 35.2 Lite Transformer (Wu et al., 2020) 33.6 De Ligh T (Ours) 30 M 35.3

(a) IWSLT 14 De-En

Model # Params BLEU

Transformer (Vaswani et al., 2017) 62 M 27.3 DLCL (Wang et al., 2019) 62 M 27.3 Evolved Transformer (So et al., 2019) 46 M 27.7 Lite Transformer (Wu et al., 2020) 26.5 De Ligh T (Ours) 37 M 27.6

(b) WMT 14 En-De

Table 3: Comparison with state-of-the-art methods on machine translation corpora. De Ligh T delivers similar or better performance than state-of-the-art models with fewer parameters. Here, indicates that the network uses neural architecture search (NAS) and indicates that full network parameters are not reported.

0 5 10 15 20 25 30 Parameters (in million)

(a) IWSLT 14 De-En

0 10 20 30 40 50 60 Parameters (in million)

(b) WMT 16 En-Ro

0 10 20 30 40 50 60 Parameters (in million)

(c) WMT 14 En-De

0 10 20 30 40 50 60 Parameters (in million)

(d) WMT 14 En-Fr

Figure 5: Scaling up De Ligh T models. The performance of De Ligh T improves with an increase in the number of network parameters, across different corpora, including low-resource (WMT 16 En-Ro).

4.2 LANGUAGE MODELING

Datasets and evaluation: We evaluate on the Wiki Text-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 De Ligh T blocks. We use wm=2, Nmin=4, and Nmax=12. We scale dm using values {384, 512, 784, 1024} for increasing network parameters. For simplicity, we set B = Nmax. 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 classification 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 De Ligh T with previous methods on Wiki Text-103. Table 4a plots the variation of perplexity with number of parameters for De Ligh T and Transformer XL (Dai et al., 2019) which outperforms other transformer-based implementations (e.g., Baevski and Auli 2019). Both tables show that De Ligh T delivers better performance than state-of-the-art methods (including Transformer-XL) and it does this using a smaller context length and significantly fewer parameters, suggesting that the De Ligh T transformation helps learn strong contextual relationships.

5 ANALYSIS AND DISCUSSIONS ON COMPUTATIONAL EFFICIENCY

Training time and memory consumption: Table 5 compares the training time and memory consumption of De Ligh T 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 De Ligh T full-precision networks for 30K iterations on 16 NVIDIA V100 GPUs. The transformer and De Ligh T 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 library3 for multi-head attention in Transformers, the training time of the

3https://github.com/NVIDIA/apex

Published as a conference paper at ICLR 2021

20 40 60 80 100 120 140 Parameters (in million)

De Ligh T (Ours) Transformer-XL

(a) De Ligh T vs. Transformer-XL

Method Network Context # Params Perplexity Depth Length (in million) (Test)

LSTM (Grave et al., 2017b) 48.70 LSTM + Neural Cache (Grave et al., 2017b) 40.80 QRNN (Merity et al., 2018b) 151 M 33.00

Transformer-XL (Dai et al., 2019) 64 640 151 M 24.03 Transformer-XL (Our impl.) 64 640 151 M 24.34 Transformer-XL (Our impl.) 64 480 151 M 24.91 De Ligh T (Ours) 158 480 99 M 24.14

(b) Comparison with existing methods

Table 4: Results on the Wiki Text-103 dataset. Compared to Transformer-XL, De Ligh T delivers similar or better performance (lower perplexity) with fewer parameters. For Transformer-XL, we reproduce results using the official source code. For evaluating Transformer-XL with a context length of 480, we set the mem_len hyper-parameter to 480 in the official evaluation scripts.

Row # Model # Params BLEU Training Memory (in million) (WMT 14 En-Fr) time (in GB)

R1 Transformer (unoptimized) 67 M 39.2 37 hours 12.5 GB R2 De Ligh T (unoptimized) 54 M 40.5 23 hours 14.5 GB R3 Transformer (w/ Apex optimized) 67 M 39.2 16 hours 11.9 GB R4 De Ligh T (w/ optimized grouping) 54 M 40.5 19 hours 11.5 GB

Table 5: Comparison with baseline transformers in terms of training speed and memory consumption. In R4, we implemented CUDA kernels for grouping and ungrouping functions only (see Appendix E). We expect De Ligh T to be more efficient with a single and dedicated CUDA kernel for grouping, transformation, feature shuffling, and ungrouping. Memory consumption is measured on a single NVIDIA GP100 GPU (16 GB memory) with a maximum of 4096 tokens per batch and without any gradient accumulation.

Model Dropout BLEU

Transformer (62 M) 0.10 27.3 Transformer (62 M) 0.30 27.7 De Ligh T (37 M) 0.05 27.6

Table 6: De Ligh T requires less regularization as compared to baseline transformers (Dataset: WMT 14 En-De).

transformer model reduced from 37 to 16 hours while we did not observe any significant 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 De Ligh T reduced by about 4 hours and 3 GB, respectively. We emphasize that grouping, linear transformation, feature shuffling, and ungrouping, can be implemented efficiently 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 De Ligh T 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.

6 CONCLUSION

This paper introduces a deep and light-weight transformer architecture, De Ligh T, that efficiently allocates parameters both within the De Ligh T block and across De Ligh T blocks. Compared to state-of-the-art transformer models, De Ligh T models are (1) deep and light-weight and (2) deliver similar or better performance. In the future, we plan to apply De Ligh T 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.

Published as a conference paper at ICLR 2021

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pages 5998 6008, 2017.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. ar Xiv preprint ar Xiv:1910.10683, 2019.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam Mc Candlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. ar Xiv preprint ar Xiv:2005.14165, 2020.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), 2019.

Geoffrey E Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhutdinov. Improving neural networks by preventing co-adaptation of feature detectors. ar Xiv preprint ar Xiv:1207.0580, 2012.

Li Wan, Matthew Zeiler, Sixin Zhang, Yann Le Cun, and Rob Fergus. Regularization of neural networks using dropconnect. In International conference on machine learning, pages 1058 1066, 2013.

Stephen Merity, Nitish Shirish Keskar, and Richard Socher. Regularizing and optimizing LSTM language models. In International Conference on Learning Representations, 2018a. URL https://openreview.net/ forum?id=Syy GPP0TZ.

Sachin Mehta, Rik Koncel-Kedziorski, Mohammad Rastegari, and Hannaneh Hajishirzi. Pyramidal recurrent unit for language modeling. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, 2018.

Xiangyu Zhang, Xinyu Zhou, Mengxiao Lin, and Jian Sun. Shufflenet: An extremely efficient convolutional neural network for mobile devices. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 6848 6856, 2018.

Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Association for Computational Linguistics, 2019.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. ar Xiv preprint ar Xiv:1904.10509, 2019.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.

Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. ar Xiv:2004.05150, 2020.

Alessandro Raganato and Jörg Tiedemann. An analysis of encoder representations in transformer-based machine translation. In Proceedings of the 2018 EMNLP Workshop Blackbox NLP: Analyzing and Interpreting Neural Networks for NLP, November 2018.

Gino Brunner, Yang Liu, Damian Pascual, Oliver Richter, Massimiliano Ciaramita, and Roger Wattenhofer. On identifiability in transformers. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=BJg1f6EFDB.

Elena Voita, Rico Sennrich, and Ivan Titov. The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), 2019a.

Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? In Advances in Neural Information Processing Systems, pages 14014 14024, 2019.

Published as a conference paper at ICLR 2021

Alessandro Raganato, Yves Scherrer, and Jörg Tiedemann. Fixed encoder self-attention patterns in transformerbased machine translation. ar Xiv preprint ar Xiv:2002.10260, 2020.

Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models. ar Xiv preprint ar Xiv:2005.00743, 2020.

Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and Michael Auli. Pay less attention with lightweight and dynamic convolutions. In International Conference on Learning Representations, 2019.

Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long-short range attention. In International Conference on Learning Representations, 2020.

David So, Quoc Le, and Chen Liang. The evolved transformer. In Proceedings of the 36th International Conference on Machine Learning, pages 5877 5886, 2019.

Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N Dauphin. Convolutional sequence to sequence learning. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 1243 1252. JMLR. org, 2017.

Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pages 933 941. JMLR. org, 2017.

Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations, 2020.

Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick Le Gresley, Jared Casper, and Bryan Catanzaro. Megatron-lm: Training multi-billion parameter language models using gpu model parallelism. ar Xiv preprint ar Xiv:1909.08053, 2019.

Mingxing Tan and Quoc V. Le. Efficientnet: Rethinking model scaling for convolutional neural networks. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, 2019.

Qiang Wang, Bei Li, Tong Xiao, Jingbo Zhu, Changliang Li, Derek F. Wong, and Lidia S. Chao. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770 778, 2016.

Rico Sennrich, Barry Haddow, and Alexandra Birch. Neural machine translation of rare words with subword units. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), August 2016.

Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2019.

Édouard Grave, Armand Joulin, Moustapha Cissé, David Grangier, and Hervé Jégou. Efficient softmax approximation for GPUs. In International Conference on Machine Learning, 2017a.

Sachin Mehta, Rik Koncel-Kedziorski, Mohammad Rastegari, and Hannaneh Hajishirzi. De FINE: Deep Factorized Input Token Embeddings for Neural Sequence Modeling. In International Conference on Learning Representations, 2020.

Patrick Chen, Si Si, Yang Li, Ciprian Chelba, and Cho-Jui Hsieh. Groupreduce: Block-wise low-rank approximation for neural language model shrinking. In Advances in Neural Information Processing Systems, 2018.

Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: a compact task-agnostic bert for resource-limited devices. In Association for Computational Linguistics (ACL), 2020.

Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. In International Conference for Representation Learning, 2016.

Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019b.

Published as a conference paper at ICLR 2021

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. In NIPS Deep Learning and Representation Learning Workshop, 2015.

Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. In 5th Workshop on Energy Efficient Machine Learning and Cognitive Computing - Neur IPS, 2019.

Matus Telgarsky. Benefits of depth in neural networks. COLT, 2016.

Sergey Edunov, Myle Ott, Michael Auli, David Grangier, and Marc Aurelio Ranzato. Classical structured prediction losses for sequence to sequence learning. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), 2018.

Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, and Jakob Uszkoreit. Tensor2tensor for neural machine translation. Co RR, abs/1803.07416, 2018. URL http://arxiv.org/abs/1803. 07416.

Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. In Proceedings of the 40th annual meeting on association for computational linguistics, pages 311 318. Association for Computational Linguistics, 2002.

Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. Fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.

Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pages 6114 6123, 2019.

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.

Yuntian Deng, Yoon Kim, Justin Chiu, Demi Guo, and Alexander Rush. Latent alignment and variational attention. In Advances in Neural Information Processing Systems, pages 9712 9724, 2018.

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2017.

Edouard Grave, Armand Joulin, and Nicolas Usunier. Improving neural language models with a continuous cache. In International Conference on Learning Representations, 2017b.

Stephen Merity, Nitish Shirish Keskar, and Richard Socher. An analysis of neural language modeling at multiple scales. ar Xiv preprint ar Xiv:1803.08240, 2018b.

Published as a conference paper at ICLR 2021

A DELI G HT ARCHITECTURES FOR LANGUAGE MODELING AND MACHINE TRANSLATION

De Ligh T 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. (2017).

Language modeling: Figure 6a shows the architecture for language modeling. The architecture stacks B De Ligh T blocks, the configuration of each block is determined using block-wise scaling. Each block has three sub-layers. The first layer is a De Ligh T 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 De Ligh T blocks, the configuration of each block is determined using block-wise scaling. Similar to language modeling, each encoder block has three sub-layers. The first layer is a De Ligh T 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

Inputs (shifted right)

Adaptive Inputs

Positional Encoding

Masked Singlehead Attention

Light-weight FFN

Adaptive Softmax

Input and output weights are tied

(a) Language Modeling

Look-up Table

Positional Encoding

Single-head Attention

Light-weight FFN

Outputs (shifted right)

Look-up Table

Positional Encoding

Masked Singlehead Attention

Single-head Attention

Light-weight FFN

Input and output weights are tied

(b) Machine translation

Figure 6: Sequence modeling with De Ligh T. Here, green color hexagon represents the De Ligh T transformation.

Published as a conference paper at ICLR 2021

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 classification layer to map vectors to tokens.

B GROUP LINEAR TRANSFORMATION WITH INPUT-MIXER CONNECTION

Group linear transformation (GLT) F splits a dm-dimensional input X into g non-overlapping groups such that X = Concat(X1, , Xg), where Xi is a dm

g -dimensional vector. Xi s are then simultaneously transformed

using g linear transforms Wi R dm

g to produce g outputs Yi = Xi Wi. Yi s are then concatenated to produce the final do-dimensional output Y = Concat(Y1, , Yg).

Figure 7a shows an example of GLT in the expansion phase of De Ligh T 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 first layer has one group, the second layer has two groups and the third layer has four groups. GLTs learns group-specific representations and are local. To allow GLT to learn global representations, we use feature shuffle. An example of GLT with feature shuffle is shown in Figure 7b. Furthermore, training deep neural networks by merely stacking linear or group linear (with or without feature shuffle) 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 De Ligh T transformation). To stabilize the training and learn deeper representations, we use input-mixer connection of Mehta et al. (2020). Figure 7c shows an example of GLT with feature shuffle and input mixer connection.

Groups = 1 Groups = 2 Groups = 4

Groups = 1 Groups = 2 Groups = 4

(b) GLT with feature shuffle

Groups = 1 Groups = 2

Input mixer

Input mixer

(c) GLT with feature shuffle and input mixture connection

Figure 7: This figure visualizes different variants of group linear transformations that are used in the De Ligh T transformation.

C MULTIPLICATION-ADDITION OPERATIONS IN DELIGHT

The De Ligh T 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, dm denotes the input dimension, do 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 Wi R dm

fore, GLT learns dmdo

g parameters and performs dmdo

g MACs to transform dm-dimensional input to

Published as a conference paper at ICLR 2021

do-dimensional output. Following a standard practice, e.g., Res Net 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 De Ligh T: The scaled dot-product self-attention in De Ligh T is defined as:

Attention(K, Q, V) = softmax QKT

where Q Rn do, K Rn do, V Rn do denotes query, key, and value, respectively.

The attention operation involves two dot-products. The first dot product between Q and K while the second dot product is between the output of first dot product and V. Both dot products require don2 MACs. Therefore,

total number of MACs in computing scaled dot-product self-attention are 2don2 .

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

k=1 2kndo .

D ABLATIONS ON THE WIKITEXT-103 DATASET

Table 7 studies the impact of De Ligh T block parameters on the Wiki Text-103 dataset, namely (1) minimum number of GLTs Nmin, (2) maximum number of GLTs Nmax, (3) width multiplier wm, and (4) model dimension dm (see Figure 1b). Figure 8, Figure 9, and Figure 10 shows the impact of the De Ligh T transformation, feature shuffling, and the light-weight FFN. Table 8 shows the effect of position of De Ligh T transformation in the De Ligh T block while Figure 12 shows the effect of scaling De Ligh T networks. We choose the Wiki Text-103 dataset 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.

De Ligh T block: Overall, Table 7 shows that scaling depth and width using De Ligh T 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, De Ligh T with Nmin = 4 and Nmax = 8 (R4) is 1.25 shallower than De Ligh T with Nmin = 8 and Nmax = 8 (R2), but delivers better performance with a similar number of parameters and operations. Scaling wm improves performance (R2 vs. R3), however, the improvement is significantly 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 Nmax and Nmin yields different results. We observe significant performance improvements when the ratio is greater than or equal to two. For example, when we scale Nmax

Nmin from 2 to 3 (R6 vs. R8), the perplexity improves by 5 points with only a moderate increase in network parameters. On the other hand, when the Nmax

Nmin 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 Nmax from 8 to 12 for Nmin = 4 (R6, R8), De Ligh T delivered better performance with a similar number of parameters compared to a model with Nmin = 6 (R7, R9). This is likely because the ratio of Nmax and Nmin is higher when Nmin = 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 wm and dm improves performance (R10-R15), however, their impact is different. For example, when we scale wm and dm by two, the rate of increase in number of parameters and operations is more rapid with dm compared to wm. De Ligh T s ability to learn wider representations in different ways may be useful in selecting application specific models.

Impact of De Ligh T transformation: We replace De Ligh T transformation in the De Ligh T block (Figure 1b) with (1) the De FINE transformation and (2) a stack of linear layers. Figure 8 shows that De Ligh T transformation delivers similar performance with significantly fewer parameters compared to the De FINE unit

Published as a conference paper at ICLR 2021

Row # Nmin Nmax wm dm Depth Parameters MACs Perplexity

Uniform vs. block-wise scaling

R1 4 4 2 256 43 14.1 M 2.96 B 56.19 R2 8 8 2 256 115 16.6 M 3.49 B 48.58 R3 8 8 4 256 115 22.1 M 4.64 B 45.10

R4 4 8 2 256 92 16.7 M 3.51 B 46.30 R5 4 12 2 256 158 21.0 M 4.41 B 41.18 Varying depth (Nmin and Nmax (Eq. 4)

R6 4 8 2 256 92 16.7 M 3.51 B 46.30 R7 6 8 2 256 102 16.5 M 3.46 B 46.68 R8 4 12 2 256 158 21.0 M 4.41 B 41.18 R9 6 12 2 256 172 20.0 M 4.20 B 42.26 Varying De Ligh T transformation s width wm (Eq. 4)

R10 4 12 2 256 158 21.0 M 4.41 B 41.18 R11 4 12 3 256 158 23.8 M 4.99 B 39.92 R12 4 12 4 256 158 27.1 M 5.69 B 39.10 Varying model width dm R13 4 12 2 256 158 21.0 M 4.41 B 41.18 R14 4 12 2 384 158 29.9 M 6.28 B 35.14 R15 4 12 2 512 158 43.8 M 9.20 B 30.81 Deeper and wider near the Input

R16 12 4 2 256 158 21.0 M 4.41 B 43.10

Table 7: Ablations on different aspects of the De Ligh T block, including uniform vs. block-wise scaling, depth scaling, and width scaling. Rows partially highlighted in color have the same configuration (repeated for illustrating results). Our experimental setup is similar to Section 4, except that we train our models for 50K iterations. Multiplication and addition operations (MACs) are computed for 20 time steps.

and linear layers. In these experiments, the settings are the same as R13-R15 (Table 7), except, Nmax = 8, because models with a stack of linear layers learn too many parameters.

Feature shuffling: Figure 9 shows that feature shuffling improves the performance of De Ligh T by 1-2 perplexity points. Here, we use the same settings as in R13-R15 (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 significant 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 = 22 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. (2017). This suggests that the ability of

20 40 60 Parameters (in million)

De FINE De Ligh T Linear

Figure 8: Impact of different transformations. De Ligh T transformations are more parametric efficient than De FINE and linear transformations. Lower perplexity value means better performance.

20 25 30 35 40 Parameters (in million)

w/ shuffle w/o shuffle

Figure 9: Impact of feature shuffling. Feature shuffling allows us to learn representations from global information and improves performance. Lower perplexity value means better performance.

Published as a conference paper at ICLR 2021

2 2 20 21 22 23 Reduction factor (r)

20.8 Parameters (in million)

Figure 10: Impact of reduction factor r in light-weight FFN. The ability of De Ligh T transformation to learn representations in high-dimensional spaces efficiently allows us to reduce the computational burden on the FFN. Lower perplexity value means better performance.

Block-wise Uniform

N B 1 = Nmax

(see Eq. 4)

128 256 384 Model dimension (dm)

Nmin=4, Nmax=4, Mean=4 Nmin=8, Nmax=8, Mean=8 Nmin=4, Nmax=8, Mean=5.6 Nmin=4, Nmax=12, Mean=7.6

Figure 11: Uniform vs. block-wise scaling. (a) contrasts the uniform and block-wise scaling methods. (b) compares the results of De Ligh T with uniform and block-wise scaling methods on the Wiki Text-103 dataset. De Ligh T networks with block-wise scaling delivers better performance across different settings. Lower perplexity value means better performance.

De Ligh T transformation to learn representations in high-dimensional spaces efficiently 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 De Ligh T with uniform and blockwise scaling. For a given model dimension dm, De Ligh T models with block-wise scaling delivers better performance.

Position of De Ligh T transformation: We studied three configurations for the De Ligh T transformation on the Wiki Text-103 validation set (Table 8): (1) De Ligh T transformation followed by single-headed attention and light-weight FFN, (2) single-headed attention followed by De Ligh T transformation, and (3) single-headed attention followed by De Ligh T transformation and light-weight FFN. For similar number of parameters, we found that (2) and (3) drops the performance of (1) significantly 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 De Ligh T: Figure 12 shows the results of De Ligh T models obtained after varying configuration parameters of De Ligh T transformations (Nmin={4, 6}, Nmax={8, 12}, wm={2, 3, 4}, and dm={256, 384, 512}). We can see that scaling one configuration parameter (e.g., dm) while keeping other configuration parameters constant (e.g., Nmin, Nmax, and wm) consistently improves performance.

Published as a conference paper at ICLR 2021

Configuration Parameters Perplexity

De Ligh T transformation + Single-head attention + Light-weight FFN 31 M 34.20 Single-head attention + De Ligh T transformation 30 M 39.02 Single-head attention + De Ligh T transformation + Light-weight FFN 31 M 39.43

De Ligh T transformation + Single-head attention + Light-weight FFN 99 M 23.16 Single-head attention + De Ligh T transformation 96 M 28.33 Single-head attention + De Ligh T transformation + Light-weight FFN 99 M 27.94

Table 8: Effect of the position of De Ligh T transformation. Lower value of perplexity means better performance.

128 256 384 512 Model dimension (dm)

wm = 2 wm = 3 wm = 4

(a) Nmin=4, Nmax=8

128 256 384 512 Model dimension (dm)

wm = 2 wm = 3 wm = 4

(b) Nmin=6, Nmax=8

128 256 384 512 Model dimension (dm)

wm = 2 wm = 3 wm = 4

(c) Nmin=4, Nmax=12

128 256 384 512 Model dimension (dm)

wm = 2 wm = 3 wm = 4

(d) Nmin=6, Nmax=12

Figure 12: Scaling up De Ligh T. Scaling one configuration parameter (e.g., dm) while keeping other configuration parameters constant (e.g., Nmin, Nmax, and wm) 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 Nmin, Nmax, wm, and dm, 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 efficient and accurate models.

E SOURCE CODE FOR GROUP LINEAR TRANSFORMATION

The source code for implementing group linear transformation (GLT) in Py Torch is shown in Listing 1. The source code for efficiently 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 Py Torch implementation for grouping and ungrouping are replaced with a dedicated CUDA kernels, resulting in reduced memory footprint and faster training.

Published as a conference paper at ICLR 2021

Listing 1: "Naive implementation of GLT in Pytorch" import torch def glt_function(x, n_groups, weights, bias=None):

:param x: Input tensor of size [B x N], where B is batch size and N

is input dimension :param n_groups: number of groups in GLT :param weights: glt weights [g x N/g x M/g] :param bias: GLT bias (optional) of size [g x 1 x M/g] :return: output tensor of size [B x M] bsz = x.size(0)

## 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 that group is first

## 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 from [B x N] to [g x 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(block Idx.x, block Dim.x) + thread Idx.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]; } }