Transformer-QL: A Step Towards Making Transformer Network Quadratically Large

1 INTRODUCTION . 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 . 2 METHOD . 2.1 TERMINOLOGY AND NOTATIONS . 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 . 2.2 BACKGROUND . 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 . 2.3 THE MODEL . 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 .

Summary: 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.