TRANSFORMER-QL: A STEP TOWARDS MAKING TRANSFORMER NETWORK QUADRATICALLY LARGE
Year: 2020

Abstract:
Transformer networks have shown outstanding performance on many natural language processing tasks. However the context length (the number of previous tokens on which the output states depend) of a Transformer network grows at best linearly with the memory and computational power used. This limitation prevents a transformer network to have very long context in a resource limited application. In this work, we propose a class of transformer networks, namely Transformer-QL (Quadratically Large), in which, the context length can grow at best quadratically with the memory and computational power used. We have empirically evaluated a Transformer-QL model in three long range language modeling datasets. The results show that Transformer-QL can provide significant improvements over other state of the art networks.

heading: 1 INTRODUCTION
text: Since its introduction in Vaswani et al. (2017), Transformer networks have overtaken its predecessor Recurrent Neural Networks (RNN) in almost every natural language processing task. However, one limitation of Transformer network is its high requirement of memory and computational power. In a vanilla Transformer network, the memory and computational requirement grows quadratically with the sequence length, and thus with context length.
In an effort to overcome the above limitation, Transformer-XL (Dai et al., 2019) and Compressive Transformer (Rae et al., 2020) have been recently proposed. However, in both the network, the context length can grow at best linearly with the memory and computation usage. An alternative strategy have been explored in Li et al. (2019); Ye et al. (2019); Child et al. (2019); Zaheer et al. (2020); Beltagy et al. (2020); Wang et al. (2020b); Kitaev et al. (2020); Katharopoulos et al. (2020); Choromanski et al. (2020); Wang et al. (2020a). All these works have proposed to replace the vanilla self-attention network by a different one with linear or log-linear memory and computation complexity leading to novel transformer architectures with overall linear or log-linear cost. Although, they provide an improvement over the quadratic cost of the vanilla transformer network, the achieved cost is still, at best, linear. Besides, since those techniques are based on either sparsification or compression of the self attention mechanism, they struggle to accumulate long distance information (Gupta & Berant, 2020).
Several works such as Burtsev & Sapunov (2020); Ainslie et al. (2020); Gupta & Berant (2020) have proposed to increase the context length by introducing a global attention which attains to every input token, thus capable of capturing long distance dependency. However, capturing long distance dependency using those approaches involves extreme compression of state space by the global attention mechanism. Moreover, even though, they perform well on several tasks, their performance on language modeling task have not been tested. Another line of work (Zhang et al., 2019; Pappagari et al., 2019) have suggested to use hierarchical arrangement of transformer network to capture document-wide dependency. However, applicability of those networks requires hierarchical structure in the data itself. Moreover, those techniques have been proposed for document compression rather than language modeling.
In this paper, we propose a class of transformer architectures, namely Transformer-QL (Quadratically Large), to alleviate the problem of capturing long distance dependency. Similar to multi-scale transformer networks (Donahue et al., 2019; Subramanian et al., 2020; Zhao et al., 2020;
Dai et al., 2020), Transformer-QL captures the contextual information in multiple temporal scales - finer scales to capture recent past information and coarser scales to capture distance past information. Additionally like Transformer-XL, it keeps the hidden states of a past segment in memory and use it to process future segments causing the context length to grow beyond the current segment. Overall, the context length in Transformer-QL can grow up to quadratically with the memory/computational usage. The contributions of the work are as follows:
• We have proposed a novel class of transformer architectures, namely Transformer-QL, in which, the context length can be made to grow linearly with memory and computation cost. Further, employing a linear cost self attention layer like Wang et al. (2020b); Katharopoulos et al. (2020), the context length of Transformer-QL can be made to grow quadratically in both memory and computational cost.
• We have empirically evaluated a Transformer-QL model on three long range language modeling datasets. The results show significant improvement in perplexity score over Transformer-XL and Compressive Transformer.
The organization of the paper is as follows. In section 2, the proposed Transformer-QL architecture along with its background has been introduced. Section 3 provides empirical evaluation of Transformer-QL. The section also studies the sensitivity of Transformer-QL to several hyperparameters. Finally, in Section 4, the conclusion has been drawn and future directions of the work have been suggested.

heading: 2 METHOD

heading: 2.1 TERMINOLOGY AND NOTATIONS
text: In a transformer network, the input sequence are partitioned into smaller segments of fixed length. Each segment is processed independently of other segments. We refer to the number of tokens in each segment as segment length. In a transformer network with recurrent memory like TransformerXL, the hidden states of the recent past segments are preserved in a fixed length memory. We refer to the number of tokens in each layer of the memory unit as memory length. For an output state (i.e. the output states of the last layer), we use the term context length to refer to the number of past tokens on which the output state depends. In transformer network, different output states might have different context length. We respectively refer the minimum and maximum of the context lengths of all the output states in a network as minimum context length and maximum context length of the network. We refer the sum of segment length and the memory length using the term window length.
We denote the segment length, memory length, window length and model dimension by ns, nm, nw and dm respectively. Thus, we have nw = ns + nm. We also use the notations slt and m l t to denote the output and memory of l-th layer at time step t for l = 1, 2, · · · , L where L is the total number of Layers. The output and memory of embedding layer at time step t have been denoted by s0t and m0t respectively. The number of heads in the self attention layers has been denoted by H .

heading: 2.2 BACKGROUND
text: Transformer A transformer network consists of stacked collection of multiple transformer layers. Each transformer layer contains a multi-head self-attention layer followed by a position-wise feed forward layer. Though the memory and computational cost of position-wise feed forward layer is linear in the length of input sequence, the multi-head self attention layer has a quadratic cost. The transformer network tackles the quadratic memory and computational cost by dividing the input sequence into smaller segments and applying the transformer network on each segment independently. However, such method limits the context lengths to the segment length. Dai et al. (2019) has named this problem as context fragmentation problem.
Transformer-XL Dai et al. (2019) has proposed Transformer-XL to solve the context fragmentation problem. In Transformer-XL, instead of discarding the hidden states after the computation of a segment, they are saved in memory (please refer to Figure 3). During the computation of the following segments, the self attention is applied over the hidden states of both the current segment and the
memory, thus has an increased context length without quadratic increase in the memory and computational cost. In fact, the memory/computational cost of the self-attention layer of Transformer-XL grows only quadratically only with the segment size ns and linearly with the memory length nm. On the other hand, the context lengths get increased by a length of nm per layer. By keeping ns small and making nm large enough, the memory and computational cost of Transformer-XL can be made close to linear with respect to the context lengths. Rae et al. (2020) has proposed to improved the memory/computational cost of Transformer-XL further by keeping the part of the memory states in a compressed form. However, even with this improvement, the memory and computational cost can be at best linear in the context length.

heading: 2.3 THE MODEL
text: Overview In this paper, we explore to increase the context length by compressing the hidden states hierarchically. The high level view of our architecture is shown in Figure 1. As shown in the figure, the model processes the input sequence in several scales of temporal granularity. Each scale consists of several transformer layers with recurrent memory. The output of the last layer of one scale is compressed causing the temporal granularity as well as the segment length to reduce. As the segment length reduces, we also simultaneously increase the memory length to keep the total length (i.e. segment length + memory length) of the layer constant. Then the new segment and memory is fed as the input to the first layer of the next scale. The resulting architecture is similar to the multiscale transformer architectures (Donahue et al., 2019; Subramanian et al., 2020; Zhao et al., 2020). Additionally, Transformer-QL keeps recurrent memory to store hidden states of previous segments. Therefore, in Transformer-QL, the layers belonging to a finer scale process contextual information in fine-grained manner, but have a smaller context length. On the other hand, a layer belonging to a coarser scale process information in coarse-grained manner, but have a longer context length (please refer to Figure 5 for a detailed illustration of the context lengths of Transformer-QL layers). To get the final output of the network, we causally combine the (possibly over-sampled) outputs of the last layer of each scale and pass those through several transformer layers (following Subramanian et al. (2020); Zhao et al. (2020)) to learn deep representation of the output.
The Compression Function For compression, we use one of average pooling and max pooling with pool size and stride both equal to c where c is the rate by which we compress the states while transitioning from one scale to the next. Let slt and m l t be the output and memory states of l-th layer with length nls and n l m respectively. We apply the compression function on the concatenation of m l t and slt to get the output s l+1 t of length n l+1 s = (n l s + n l m)/c (for simplicity assume that n l s + n l m is divisible by c). If nl+1s > n l s, we take the last n l s elements of s l+1 t to form the output of the
compression layer. Finally, we keep a recurrent memory ml+1t of length n l s + n l m − nl+1s making nls + n l m = n l+1 s + n l+1 m to hold.
The Memory Updates In Transformer-XL with segment length ns and memory length nm, the segment of length ns is get shifted into the memory. In other words, the memory for the next time step is computed as mlt+1 = concat(m l t, s l t)[−nm :] for all layer l. However, in Transformer-QL, the granularity of layers belonging to different scales are different. More precisely, a segment of length n0s belonging to scale 1 is compressed into a segment of length n 0 s/c
i−1 at a layer belonging to scale i. Thus, in Transformer-QL, we update the memory of a layer l belonging to scale i as mlt+1 = concat(m l t, s l t[: n i h])[−nim :] where nih = n0s/ci−1 and nim are the shift length and the memory length at scale i for i = 0, 1, · · · respectively. The complete algorithm of Transformer-QL is shown in Figure 2.
Droppath Regularization Since the output of the last layer of every scale is summed in the accumulation layer, the path through a higher scale forms to a deeper network while the path through a lower scale forms to a shallower network. Consequently, layers in the higher scales might remain under-fitted due to lack of gradient flow through the deeper network while the layers in the lower scales get over-fitted. To alleviate this problem, we have introduced droppath regularization. In the accumulation layer, let each output be computed as so = 1l ∑l i=1 s
i where si represents the (possibly over-sampled) output of scale i and l is the total number of scales. In droppath regularization with droppath probability p, we drop the output of all the scales below j with a probability p/(l − 1) for j = 2, 3, · · · , l from the accumulated output. More precisely, we generate a random number u from uniform probability distribution and compute the output as so = 1l−j+1 ∑l i=j s i if
u ∈ [ (j−2)p l−1 , (j−1)p l−1 ] . For u ≥ p, no droppath is applied.

heading: 2.4 THE COMPLEXITY
text: The memory/computational complexity of a Transformer-XL network (Dai et al., 2019) with segment length ns, memory length nm and L layer is Θ((α(nm, ns) + ns)L) where α(·, ·) is the complexity of self-attention layer. The context length nc of the network is Θ(nmL). Since α(nm, ns) = Ω(nm + ns) (Li et al., 2019; Ye et al., 2019; Child et al., 2019; Zaheer et al., 2020; Beltagy et al., 2020; Wang et al., 2020b; Kitaev et al., 2020; Katharopoulos et al., 2020; Choromanski et al., 2020), the memory and computational complexity of Transformer-XL in terms of context length is Ω(nc). Similarly, the memory and computational complexity of Compressive Transformer (Rae et al., 2020) in term of context length is Ω(nc/c) where c is the compression rate. Therefore, the memory/computational complexity of both Transformer-XL network and Compressive Transformer network in term of the context length is at least linear. Consequently, increasing the context length in both the networks requires at least linear increase in the amount of both memory and computational requirements.
On the other hand, a Transformer-QL network with L Transformer-XL layers and i compression layers, the context length nc becomes Θ(ci(ns +nm)) = O(clogcns(ns +nm)) = O(ns(ns +nm)) where ns = n0s and nm = n 0 m are the segment and memory length in scale 1 of the network. Note that, since at most i = logcns compression layer can be used in Transformer-QL, we have ci = O(clogcns) = O(ns). If we set nm = O(ns), we have nc = O((ns)2). However, the time and memory complexity of a Transformer-QL network is Θ(α(ns, nm)L + (ns + nm)i)) = Θ(α(ns, nm)(L + i)). Since α(ns, nm) = Ω(ns + nm) and we set nm = O(ns), the memory/computational complexity of Transformer-QL becomes Ω(ns(L + i)). Therefore, the memory/computational complexity of Transformer-QL in terms of context length is Ω( √ nc(L + i)) =
Ω (√ nc(L+ logcns) ) . Thus, the complexity of Transformer-QL can be at best sub-linear. Moreover, if we set the compression rate c to ns, the memory and computational complexity can be at best Θ( √ nc) or, in other words, the context length can be at best quadratic in the memory and computational cost. In Appendix B, we provide an algorithm to compute a tight estimation of the context length of a Transformer-QL network. In the appendix, we have also provided a detailed illustration of the dependency structure of the hidden states of a Transformer-QL network on the past tokens.

heading: 3 EMPIRICAL EVALUATION
text: In this section, we empirically evaluate the efficacy of Transformer-QL for long range language modeling task. Towards that goal, we compare the results of Transformer-QL with that of TransformerXL (Dai et al., 2019) and Compressive Transformer (Rae et al., 2020). Then we evaluate the sensitivity of Transformer-QL to several hyper-parameters.

heading: 3.1 COMPARISON WITH STATE OF THE ART METHODS
text: State of the Art Methods We compare Transformer-QL network with the following two networks:
Transformer-XL (Dai et al., 2019) Transformer-XL is similar to vanilla Transformer with two modifications. It uses recurrent memory to store and access the hidden states of the past time steps. The recurrent memory enables to increase the minimum context length up to nmL where nm is the memory length and L is the number of layers. It also uses relative positional embedding of token instead of absolute positional embedding.
Compressive Transformer (Rae et al., 2020) Like Transformer-XL, Compressive Transformer also uses recurrent memory. However, Compressive Transformer keeps part of the recurrent memory in compressed format, thus has an increased context length over Transformer-XL.
Datasets We compare Transformer-QL against the above two networks on three long range language modeling datasets: SimpleBooks-2 (Nguyen, 2019), SimpleBooks-92 (Nguyen, 2019) and WikiText-103 (Merity et al., 2017). SimpleBooks-2 and SimpleBooks-92 are created from Gutenberg book corpus (www.gutenberg.org) while WikiText-103 are created from Wikipedia articles. All
the three datasets preserve paragraph and section structures of their sources making those suitable for long range language modeling task. The statistics of the three datasets are shown in Table 4 of Appendix C.
Experimental Details For the experiments of Transformer-XL and Compressive Transformer, we have used an 8-layer network. And for the experiments of Transformer-QL, we have used a network with 3 layers in scale 1, 3 layers in scale 2 and 2 layers in the output block. Thus, the TransformerQL has a total of eight layers as in Transformer-XL and Compressive Transformer. We set the compression rate of Compressive Transformer to 2. In Transformer-QL, we have used max-pooling layer with pool size 2 and stride 2 as the compression layer. Thus, both the Transformer-QL and Compressive Transformer have a compression rate of 2. For the experiments on the SimpleBooks92 and WikiText-103, we have set the model dimension to 1536 and used an initial learning rate of 1 × 10−4. On the other hand, for the experiments on SimpleBooks-2, we have set the model dimension to 256 and learning rate to 2.5 × 10−4. All the models have been trained using Adam optimizer. We set the droppath probability of Transformer-QL to 0.3. The details of other hyperparameters can be found in Appendix E .
Results The results of the comparison are shown in Table 1. The results are grouped by the window length nw of the test model as the lower bound of memory and computation requirement directly depends on it. In all the datasets and settings, Transformer-XL performs worst among all three. Worst performance of Transformer-XL is not surprising as it has smallest average context length (shown in the fifth column) for a given nw. However, Compressive Transformer has a slightly larger average context length than Transformer-QL. Yet, Transformer-QL has performed similarly or significantly better than Compressive Transformer in all the setting which indicates that Transformer-QL can exploits the contextual information more effectively than Compressive Transformer.

heading: 3.2 EFFECT OF MODEL DIMENSION
text: In this section, we investigate the effect of model dimension on the performance of TransformerQL. Towards that goal, we have performed experiments on WikiText-103 dataset with varying model dimension. For each experiment, we compared the test perplexity of Transformer-QL with that of Transformer-XL. The results are shown in Table 2. The improvement in test perplexity of Transformer-QL over Transformer-XL has been computed by subtracting the test perplexity of Transformer-QL from that of Transformer-XL. The relative improvement is computed by
Relative improvement = Improvement× 100
Test perplexity of Transformer-XL
As shown in the table, Transformer-QL performs relatively worse for small model dimension like 512 and relative improvement increases as the model dimension increases. We speculate that the relatively worse performance of Transformer-QL for smaller model dimension is caused by the difficulty in compressing hidden states during switching from one scale to the next. To alleviate the problem, Donahue et al. (2019) have proposed to increase the model dimension as the model transits from a lower scale to a higher one. On the other hand, Dai et al. (2020) have suggested a novel query-only-pooling to solve the problem. We take it as a future work to try those approaches in Transformer-QL.

heading: 3.3 EFFECT OF CONTEXT LENGTH
text: In this section, we study the relative improvement in perplexity scores obtained by Transformer-QL over Transformer-XL for varying context length. The results are shown in Table 3. From the table, it can be noticed that relative improvement obtained by Transformer-QL is more when the context length of the Transformer-XL networks are smaller in the first place. For example, for test ns = 08 and nm = 08, the relative improvement is as high as 8.80%. On the other hand, for the test ns = 02 and nm = 30, the relative improvement is only 2.76%. This can be explained by the fact that for the segment and memory length 02 and 30, the average context length of Transformer-XL is already large enough (241) to provide good result. By extending the average context length from 241 to 332, Transformer-QL provides only a small improvement following the law of diminishing return (Hestness et al., 2017).

heading: 4 CONCLUSION AND FUTURE WORK
text: In the work, we have proposed a class of transformer networks namely Transformer-QL in which the context length can grow quadratically in memory and computational usage. Our empirical evaluation shows that Transformer-QL can perform significantly better than other long range language modeling networks like Transformer-XL and Multi-scale Transformer by exploiting longer context length. Further more, it can perform significantly better than Compressive Transformer by exploiting the contextual information more effectively.
In our empirical evaluation, we have evaluated a Transformer-QL network with only one compression layer. In future, we want to evaluate a network with more then one compression layers. Also, we have empirically found that the performance of Transformer-QL network can be worse than that of Transformer-XL when the model dimension is small. As our future work, we want explore different methods for removing this limitation.

heading: A TRANSFORMER-XL ALGORITHM
text: The algorithm for Transformer-XL is shown in Figure 3.

heading: B CONTEXT LENGTH OF TRANSFORMER-QL
text: A tight estimate of the minimum context length of a Transformer-QL network can be computed using algorithm of Figure 4. For simplicity, we have assumed that all the division operations result into integer output. We have also assumed that there is at least one layer in every scale. The maximum context length can be obtained by adding ns to the minimum context length.
Additionally, in Figure 5, we have shown the detailed computation of minimum/maximum context length with an example. In the figure, the notation Slt1:t2 used to denote a hidden states of l-th layer and the state depends on the t1-th to t2-th tokens of the input sequence. In the example of the figure, each output state depends on at least 44 previous tokens. In other words, minimum context length of the network is 44. On the other hand, in a Transformer-XL network of same segment length, memory length and number of layers, the minimum context length would have been 4× 4 = 16.

heading: C STATISTICS OF DATASETS
text: The statistics of the datasets are shown in Table 4.

heading: D COMPARISON WITH MULTI-SCALE TRANSFORMER
text: In this section, we empirically compare Transformer-QL with Multi-scale Transformer (Subramanian et al., 2020). Our implementation of Multi-scale Transformer is same as Transformer-QL without any recurrent memory. The resultant Multi-scale Transformer is similar to the button-up model of Subramanian et al. (2020). We have used hyperparameter settings same as TransformerQL to train the Multi-scale Transformer. The result is shown in Table 5. From the table, we can see that Multi-scale Transformer has been widely bitten by Transformer-QL even when the Multi-scale Transformer has been trained and tested with a larger window length.

heading: E HYPERPARAMETER SETTING
text: We used the following values for hyperparameter for the experiments on SimpleBooks-2 datasets:
Hyperparameter Transformer-XL Compressive Transformer-QL Multi-scale Transformer Transformer d model 256 256 256 256 d embed 256 256 256 256 div val 1 1 1 1 untie r False False False False proj same dim True True True True n head 4 4 4 4 d head 64 64 64 64 d inner 1024 1024 1024 1024 train batch size 128×8ns 128×8 ns 128×8 ns 128×8 ns train ns 08 04 08 08 train nm 08 06 08 - train ncm - 06 - - pre lnorm True True True True warmup steps 0 0 0 0 train steps 60, 000 60, 000 60, 000 60, 000 learning rate 0.00025 0.00025 0.00025 0.00025 min lr ratio 0.004 0.004 0.004 0.004 clip 0.25 0.25 0.25 0.25 dropout 0.1 0.1 0.1 0.1 dropatt 0.1 0.1 0.1 0.1 droppath - - 0.3 0.3 init std 0.02 0.02 0.02 0.02 proj init std 0.01 0.01 0.01 0.01 recons loss weight - 0.01 - -
For the SimpleBooks-92 datasets and model dimension 1536, the following values of hyperparameters are used:
Hyperparameter Transformer-XL Compressive Transformer-QL Multi-scale Transformer Transformer d model 1536 1536 1536 1536 d embed 1536 1536 1536 1536 div val 4 4 4 4 untie r False False False False proj same dim True True True True n head 16 16 16 16 d head 96 96 96 96 d inner 6144 6144 6144 6144 train batch size 512×8ns 512×8 ns 512×8 ns 512×8 ns train ns 08 04 08 08 train nm 08 06 08 08 train ncm - 06 - - pre lnorm True True True True warmup steps 0 0 0 0 train steps 250, 000 250, 000 250, 000 250, 000 learning rate 0.0001 0.0001 0.0001 0.0001 min lr ratio 0.004 0.004 0.004 0.004 clip 0.1 0.1 0.1 0.1 dropout 0.15 0.15 0.15 0.15 dropatt 0.15 0.15 0.15 0.15 droppath - 0.3 0.3 0.3 init std 0.02 0.02 0.02 0.02 proj init std 0.01 0.01 0.01 0.01 recons loss weight - 0.01 - -
For the WikiText-103 datasets and model dimension 1536, the following values of hyperparameters are used:
Hyperparameter Transformer-XL Compressive Transformer-QL Multi-scale Transformer Transformer d model 1536 1536 1536 1536 d embed 1536 1536 1536 1536 div val 4 4 4 4 untie r False False False False proj same dim True True True True n head 16 16 16 16 d head 96 96 96 96 d inner 6144 6144 6144 6144 train batch size 512×16ns 512×16 ns 512×16 ns 512×16 ns train ns 16 08 16 16 train nm 16 12 16 - train ncm - 12 - - pre lnorm True True True True warmup steps 0 0 0 0 train steps 350, 000 350, 000 350, 000 350, 000 learning rate 0.0001 0.0001 0.0001 0.0001 min lr ratio 0.004 0.004 0.004 0.004 clip 0.1 0.1 0.1 0.1 dropout 0.15 0.15 0.15 0.15 dropatt 0.15 0.15 0.15 0.15 droppath - - 0.3 0.3 init std 0.02 0.02 0.02 0.02 proj init std 0.01 0.01 0.01 0.01 compression rate - 2 2 2 recons loss weight - 0.01 - -
For training models of model dimensions 512 and 1024, we have used initial learning rate of 0.0005 and 0.00025 respectively keeping the rest of the hyper-parameters same.

Review summaries:
The work introduces the Transformer-QL, a transformer-based model that aims to capture long distance dependencies in the input. The network processes the information defining multiple temporal scales, with finer scales for nearby elements, and coarser scales for distant information. It also includes the recurrent memory extension from Transformer-XL from Dai et al. The model is tested in a long range language modeling task.