Sub-Linear Memory: How to Make Performers SLiM

1 Introduction . The Transformer architecture [ 38 ] has changed the landscape of deep learning for sequential data . A computational advantage of Transformers over conventional methods such as recurrent neural networks ( RNNs ) [ 17 , 9 ] is parallelization over the sequence dimension , meaning that the training speed can be increased by simply using more compute resources . However , this parallel-friendly structure of self-attention comes at a cost of quadratic Θ ( L2 ) time and memory complexity , where L is the length of the Transformer ’ s input sequence . A recent line of work aimed to address this restriction , using either structured sparsity [ 8 ] , truncated back-propagation [ 12 ] , clustering [ 20 , 31 ] or linear attention methods [ 18 , 10 , 11 , 33 , 23 ] . For a detailed overview of efficient Transformers , see [ 37 ] . We refer to the family of linear attention architectures as Performers ( also known as Linear Transformers ) , following [ 11 ] , since their generic kernel formulation covers all the aforementioned linear attention methods . Performers reduce time and memory complexity to linear O ( L ) and can provably approximate conventional quadratic Transformers [ 11 ] , demonstrating strong performance in a systematic comparison of efficient Transformers [ 36 ] . This recent trend of feeding longer sequences into Transformers , coupled with the use of deeper models , introduces new challenges for researchers and practitioners . Whereas conventional Transformer setups benefit from large-batch optimization [ 42 ] , long sequence modelling necessitates smaller batch sizes in order to fit the model into memory . For example , recently proposed efficient Transformers 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . operating on long sequences used moderately small batch sizes of 1-8 instances [ 20 , 18 , 11 ] . Aiming to use larger batch sizes , practitioners introduced various tricks – e.g . [ 28 ] introduced gradient accumulation ( included in the popular Fairseq library , [ 27 ] ) , which splits the batch into smaller chunks which are evaluated sequentially , then the resulting batch gradient is accumulated . Gradient accumulation allows to decrease memory usage at the cost of longer time , but it can only be applied when the batch size is bigger than 1 . In this paper , we are discussing a situation when even a batch size of 1 is prohibitive , while longer processing times are affordable e.g . when finetuning a pretrained Transformer on low-memory devices ( e.g . smartphones , embedded devices or microcontrollers ) on the client-generated data without additional server computations . Heuristics , such as chunking the input into independent subsegments or truncated back-propagation [ 12 ] , limit gradient propagation across the whole input , and , consequently , impair long-context learning . We propose a solution based on the analysis of Performers . We discover a remarkable property : even for batch size of 1 , a user can decrease memory consumption at the cost of smaller parallel bandwidth of the model . Notably , no approximations are introduced , so the obtained gradient is correct and backward-compatible . Our proposed long-sequence training algorithm can be used for training or fine-tuning on a low-memory device , thus contributing towards decentralized and democratized deep learning . The algorithm has the following advantages : 1 . The integer parameter C , 1 ≤ C ≤ L , controls a tradeoff between the memory , scaling as O ( C ) in addition to a negligible input sequence storage , and parallel running time , scaling as O ( ( L/C ) logC ) . When C = 1 , the algorithm consumes as much memory as if a single token were fed into Performer , plus a small fully characterized addition . 2 . For any C , the algorithm requires as many floating point operations ( FLOPs ) as two standard forward and one backward pass plus a small addition . We evaluate the proposed tradeoff empirically , and confirm backward-compatibility for the synthetic Copying Task and language modelling on Penn Treebank [ 25 ] and Enwik8 [ 24 ] datasets.1 1Code : https : //github.com/google-research/google-research/tree/master/performer/ models/slim_performer . 2 Background : linear self-attention and Performer . We commence by defining linear self-attention [ 18 , 11 , 33 , 23 ] . Consider a sequence of length L and three matrices : queries Q ∈ RL×d , keys K ∈ RL×d and values V ∈ RL×d . Then linear self-attention is defined as a functional producing Y = Att ( Q , K , V ) ∈ RL×d , ∀l ∈ { 1 , . . . , L } : Yl = ∑l l′=1 Vl′ · ( g ( Kl′ ) > g ( Ql ) ) ∑l l′=1 g ( Kl′ ) > g ( Ql ) = ( ∑l l′=1 Vl′ × g ( Kl′ ) > ) × g ( Ql ) ( ∑l l′=1 g ( Kl′ ) ) > g ( Ql ) , ( 1 ) where by Zl ∈ Rd2× ... we denote slice Zl , : , ... , : of Z ∈ Rd1×d2× .... g : Rd → RM+ s a positive mapping for a stable division in ( 1 ) . M is typically much smaller than L , e.g . M = d. For needs of autoregressive generative modelling , when each element depends only on previous elements of the sequence [ 38 ] , Yl only depends on inputs at indices { 1 , . . . , l } . Self-attention of type ( 1 ) was proposed for processing sequences of the long length L , since the original self-attention from [ 38 ] scales as O ( L2 ) . The second transition in ( 1 ) , which is due to associativity of matrix multiplication , suggests an algorithm to compute linear self-attention efficiently in subquadratic time . For tensors Z ( 1 ) , . . . , Z ( n ) of the same shape , let Z = ( Z ( i ) ) ni=1 be a tensor such that for all 1 ≤ i ≤ n , Zi , : , ... , : = Z ( i ) . By R ∈ RL×d×M , S ∈ RL×M denote a tensor and a matrix such that R = PS ( ( Vl × g ( Kl ) > ) Ll=1 ) , S = PS ( ( g ( Kl ) ) Ll=1 ) , ( 2 ) where PS ( Z ) = ( ∑i i′=1 Zi′ ) n i=1 is a prefix sum along the first dimension of Z . Next , compute ∀1 ≤ l ≤ L : Yl = ( Rl × g ( Ql ) ) / ( S > l g ( Ql ) ) . ( 3 ) Depending on the prefix-sum algorithm used in ( 2 ) , we can obtain different complexity estimates for linear self-attention . [ 18 ] propose to iterate through l = 1 , . . . , L maintaining only current Rl , Sl , and compute and store the result Yl . This way , tensors R , PS ( R ) ∈ RL×d×M are not stored in memory , resulting in O ( L ) time complexity and O ( L ( d+M ) + dM ) memory complexity . [ 18 ] also propose a similar computation of gradients through ( 2-3 ) ; see Appendix C for more discussion . Alternatively , [ 11 ] employ a parallel prefix-sum algorithm [ 21 , 39 ] , which , for a tensor Z ∈ RL× ... , finds PS ( Z ) inO ( logL ) parallel time andO ( L ) memory . Applying this algorithm for PS ( R ) , PS ( S ) and then computing ( 3 ) results in only O ( logL ) parallel time complexity and O ( LdM ) memory consumption . A similar approach was proposed in the context of parallel RNNs [ 26 ] . Linear self-attention is a part of Performers [ 11 ] – efficient Transformers . Below we outline Performer architecture [ 11 ] for language modelling , while our analysis can be applicable in broader setups . Let p ∈ ΣL be an input sequence of length L , where Σ is a finite alphabet . By emb ( pl , l ) ∈ Rdmodel , 1 ≤ l ≤ L , denote a linear combination of the pl token ’ s learned embedding and positional embedding of l ’ s position ( sinusoids with different frequencies , as in [ 38 ] ) . Then Performer is defined as a parametrized mapping from X ( 0 ) = ( emb ( pl , l ) ) Ll=1 ∈ RL×dmodel into X ( out ) ∈ RL×|Σ| through a sequence of hidden representations X ( 1 ) , . . . , X ( s ) ∈ RL×dmodel . For each 1 ≤ r ≤ s , X ( r ) is obtained from X ( r−1 ) by applying a Performer layer consisting of two parts . The first part is a multi-head self-attention which comprises of k linear self-attentions applied in parallel ( dmodel = kd ) plus a skip connection . The second part is a feedforward network applied independently to each element of the sequence . See a formal definition of the architecture in Appendix A . For each l , 1 ≤ l ≤ L− 1 , X ( out ) l are predicted logits of the next token pl+1 . Let Ll ( X ( out ) l ) denote the cross-entropy loss with respect to pl+1 , or zero when l = L. The loss is defined as L = ( L− 1 ) −1 · ( L1 ( X ( out ) 1 ) + · · ·+ LL ( X ( out ) L ) ) . ( 4 ) 3 Low-memory back-propagation algorithm . 3.1 Compact notation for Performer . We rewrite transformations X ( 0 ) → X ( 1 ) → · · · → X ( s ) in the following form . For each 1 ≤ r ≤ s , T ( r−1 ) , Γ ( r−1 ) =F ( r ) ( X ( r−1 ) ; θ ) , U ( r−1 ) =PS ( T ( r−1 ) ) , X ( r ) =G ( r ) ( U ( r−1 ) , Γ ( r−1 ) ; θ ) . ( 5 ) Algorithm 1 Low-memory forwardbackward pass . See Algorithm 2 for updateProc . Compared to notation from the text , redundant indices are dropped and tensor names are reused here and in Algorithm 2 . Input : p ∈ ΣL , θ ∈ Rnparam , C ∈ N. Output : loss L , gradient∇θL . Set L : = 0 ; Set U : = 0s×D1 ; for n = 1 to N do updateProc ( n , False ) ; end for Set∇θL : = 0nparam ; Set G : = 0s×D1 ; for n = N to 1 do updateProc ( n , True ) ; end for Return L , ∇θL . Algorithm 2 updateProc procedure . Input : n ∈ N , binary flag onBP . if onBP then Initialize Φ : = 0 ; end if X : = ( emb ( pAn+l , An + l ) ) Bn l=1 ; for r = 1 to s do Compute T , Γ : = F ( r ) ( X ; θ ) ; if onBP then Update Ur−= ∑Bn l=1 Tl ; end if Set U : = 1BnU > r + PS ( T ) , X : = G ( r ) ( U , Γ ; θ ) ; if onBP then Update Φ+= G > r UBn ; else Update Ur : = UBn ; end if end for Set L ( upd ) : = L ( n ) ( XW ( out ) + b ( out ) ) ; if onBP then Update Φ+= L ( upd ) ; Compute∇θΦ , ∇UΦ through auto-differentiation ; Update∇θL+= ∇θΦ , G : = ∇UΦ ; else Set L+= L ( upd ) ; end if Here θ ∈ Rnparam is a set of all parameters , T ( r−1 ) , U ( r−1 ) ∈ RL×D1 and Γ ( r−1 ) ∈ RL×D2 are the following matrices ( see Figure 1a for an illustration ) . ( a ) T ( r−1 ) is a matrix of representations which are passed into the prefix-sum operator . That is , for each 1 ≤ l ≤ L , T ( r−1 ) l is a concatenation of g ( Kl ) and flattened Vl × g ( Kl ) > for all attention heads computed at the rth step . Consequently , D1 = M ( d+ 1 ) k. ( b ) For each 1 ≤ l ≤ L , U ( r−1 ) l is a concatenation of all corresponding Sl and flattened Rl – results of the prefix sum ( Eq . 2 ) inside each self-attention head at rth layer . ( c ) Γ ( r−1 ) is a matrix of representations which skip the prefix sum . For each 1 ≤ l ≤ L , Γ ( r−1 ) l is a concatenation of X ( r−1 ) l and g ( Ql ) for each attention head 1 ≤ j ≤ k ( 3 ) . Therefore , D2 = Mk + dmodel . F ( r ) and G ( r ) are functionals depending on a set of model ’ s parameters θ . That is , they take subsets of θ corresponding to rth layer weights . F ( r ) is responsible for constructing T ( r−1 ) and Γ ( r−1 ) – representations preceding prefix sum , while G ( r ) finalizes multi-head self-attention and includes the feed-forward block . Importantly , F ( r ) and G ( r ) are applied rowwise , i.e . ( 5 ) can be rewritten as ∀1 ≤ l ≤ L : T ( r−1 ) l , Γ ( r−1 ) l = F ( r ) ( X ( r−1 ) l ; θ ) , X ( r ) l = G ( r ) ( U ( r−1 ) l , Γ ( r−1 ) l ; θ ) . ( 6 ) Hence , the only place where the signal is propagated across the sequence is prefix sum in ( 5 ) . The representation ( 5 ) encapsulates architecture details of Performer inside { F ( 1 ) , G ( 1 ) , . . . , F ( s ) , G ( s ) } . In fact , the representation ( 5 ) holds for various possible modifications , proposed in the literature . This includes , but is not limited by the different positioning of layer normalization [ 41 , 38 ] , adding a stabilizing gating mechanism [ 29 ] , weight sharing across layers [ 22 ] or reversible Transformer layers [ 20 ] . High-level description of the algorithm . The algorithm , proposed in the remainder of the section , iterates over the sequence 1 , . . . , L and only maintains a front of current prefix sum values , thus allowing a substantial memory improvement . The backward pass is implemented similarly : current prefix sums along with their gradients are maintained in a dynamic programming fashion , but the iteration proceeds in a backward direction .

Summary: This paper studies a time and memory trade-off for training Performers with the linear attention mechanism. The goal of the paper is to study this trade-off under memory constrained setting, and propose a memory efficient performer training algorithm at the cost of runtime. The contributions of the paper are:  1. Given a sequence of length L, the authors proposed a dynamic programming algorithm which can process the sequence as a collection of segments of length C (the segments are coupled). This reduce the training memory for performer from O(L) to O(C) without any approximation. The intuition here is that the information across segments only flow through certain partial sums over tensors associated with each token in the sequence.  2. The authors analyzed the memory / time trade-off in the proposed algorithm (including extremal cases).  3. Empirically, the authors shows that 1) by controlling the segment length C, the proposed algorithm indeed present a trade off between runtime and memory; 2) The gradient for performer models only have minimal / negligible difference when using the conventional and the proposed algorithm; 3) A large performer model pretrained with the conventional algorithm can be finetuned via the proposed algorithm with small segment lengths; this achieves accuracy matching the conventional algorithm with less memory.