Sub-Linear Memory: How to Make Performers SLiM
Year: 2021

Abstract:
Transformer architectures have become very popular yet the original implementation requires O(L) in serial time and memory as functions of input length L. Recent works proposed various linear self-attention mechanisms, scaling only as O(L) for serial computation. We conduct a thorough complexity analysis of Performers, a class which includes most recent linear Transformer mechanisms. We note a remarkable computational flexibility: the gradient computation can be performed with no approximations using sublinear memory as a function of L (in addition to negligible storage for the input sequence), at a cost of greater time complexity in the parallel setting. In the extreme case, a Performer consumes only O(1) memory, and still requires O(L) time. Due to complete backwardcompatibility, this discovered time-memory tradeoff can be used for fine-tuning on low-memory devices in a decentralized fashion without any server computations.

heading: 1 Introduction
text: 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.

heading: 2 Background: linear self-attention and Performer
text: 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)

heading: 3 Low-memory back-propagation algorithm

heading: 3.1 Compact notation for Performer
text: 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.

heading: 3.2 Forward computation
text: Suppose the memory budget is not enough to perform a complete forward pass through Performer (Equation 5 for r = 1, . . . , s), because the input sequence length L is too big. We show that instead
we can emulate the full forward computation under the memory needed for a forward pass through the input of length C ≤ L, plus a small addition. 1 ≤ C ≤ L is arbitrary and user-defined.
Split matrices X(r),T(r),Γ(r),U(r), into N slices of size at most C along the vertical axis (N = dL/Ce): for each n, 1 ≤ n ≤ N ,
X(r,n) = (X (r) An+l )Bnl=1 ∈ R Bn×dmodel , T(r,n) = (T (r) An+l )Bnl=1,
U(r,n) = (U (r) An+l )Bnl=1 ∈ R Bn×D1 , Γ(r,n) = (Γ (r) An+l )Bnl=1 ∈ R Bn×D2 ,
where An = (n − 1)C and by Bn, 1 ≤ n ≤ N , we denote the size of the nth slice: Bu = C for u < N , BN ≤ C. Based on (5), we conclude that for each 1 ≤ n ≤ N and 1 ≤ r ≤ s it holds that
T(r−1,n),Γ(r−1,n) = F (r)(X(r,n); θ), U(r−1,n) = 1Bn × (U (r−1,n−1) Bn−1 )> (7)
+PS(T(r−1,n)), X(r,n) = G(r)(U(r−1,n),Γ(r−1,n); θ). (8)
Here, 1Bn ∈ RBn is a vector of Bn ones and we denote U (r−1,0) B0 = 0D1 (a vector of D1 zeros).
Now, instead of iterating over r = 1, . . . s and computing (5) for the whole sequence at once, we first iterate over n = 1, . . . , N and then iterate over r = 1, . . . , s in a nested loop to compute (7-8). As can be deduced from the (7-8), we only need to maintain the current value of (U(r−1,n−1)Bn−1 ) s r=1 ∈ Rs×D1 in the outer iteration over n. Denote U (n) = (U(r−1,n)Bn ) s r=1 ∈ Rs×D1 , 0 ≤ n ≤ N . The memory-efficient algorithm for the forward pass is as follows. First, initialize L = 0 and U (0) = 0s×D1 . Then, iterate over n = 1, . . . , N and maintain the current value of U (n−1). During each iteration, compute X(0,n) = (emb(pAn+l, An + l)) Bn l=1. Then iterate over r = 1, . . . , s, where compute (7-8) and update U (n) r = U (r−1,n) Bn
. Finally, compute X(out,n) = X(s,n)W(out) + b(out) and update L+= L(n)(X(out,n)), where we denote L(n)(X(out,n)) = (L− 1)−1 ∑Bn l=1 LAn+l(X (out,n) l ).
By the end of the iteration over n, the correct loss value (4) is computed. As a result, the forward pass takes O(L) serial time or O((L/C) logC) parallel time and consumes only O(C) memory. This is in addition to the input sequence p ∈ ΣL storage, which is O(L) in principle, however the constant is negligibly small. For instance, if p is a flattened image or an ASCII text string, then it occupies precisely L bytes in memory. The logC term in the parallel time complexity is due to the parallel prefix-sum algorithm taking logarithmic time, as discussed in Subsection 2.

heading: 3.3 Back-propagation and the final algorithm
text: The goal of a backward pass is to compute the gradient ∇θL of the loss function with respect to parameters θ. One can just perform automatic differentiation [14] (implemented in Tensorflow [1] and Pytorch [30]) through the computation graph induced by the memory-efficient forward pass algorithm from Subsection 3.2. However, such a backward pass would need to store all intermediate tensors produced during the forward pass, resulting in O(L) memory complexity as a function of L and C. Instead, we propose a back-propagation algorithm which has the same time and memory complexity as the efficient forward pass.
Let θ(1) = · · · = θ(N) = θ be results of a symbolic “identity operation” performed on θ, so that for all 1 ≤ n ≤ N , θ(n) is used instead of θ in (7-8). Then the total gradient of θ has the form ∇θL = ∇θ(1)L + · · · + ∇θ(N)L. In Appendix B we derive an expression for ∇θ(n)L, 1 ≤ n ≤ N . Namely, denote G(n) = ∇U(n)L, then∇θ(n)L = ∇θ(n)Φ(n)(θ(n),U (n−1),G(n)), where Φ(n) : Rnparam × Rs×D1 × Rs×D1 → R,
Φ(n)(θ(n),U (n−1),Z) = L(n)(X(out,n)) + s∑ r=1 Z>r U (n)r .
In Φ(n)’s definition, X(out,n) and U (n) = (U(r−1,n)Bn ) s r=1 are results of (7-8) iteration over r = 1, . . . , s with parameters θ = θ(n) and (U(r−1,n−1)Bn−1 ) s r=1 equal to Φ
(n)’s second argument U (n−1). Gradient ∇θ(n)Φ(n) can be computed by automatic differentiation through Φ(n).
An efficient way to compute and sum up all∇θ(n)L is to iterate in a backward direction n = N, . . . , 1 and to maintain values of U (n),G(n). U (N) is known after the end of the forward pass, and
∀1 ≤ n ≤ N : U (n−1) = U (n) − Bn∑ l=1 (T (r−1,n) l ) s r=1. (9)
Further, in Appendix B we show that G(N) = 0r×D1 and, for each 1 ≤ n ≤ N ,
G(n−1) = ∇U(n−1)Φ(n)(θ(n),U (n−1),G(n)). (10)
By a single auto-differentiation through Φ(n) we can compute ∇θ(n)L = ∇θ(n)Φ(n) and the update (10). If w is some vector of length Bn and h is some scalar function of v = PS(w), then for all 1 ≤ l ≤ Bn : ∇hwl = ∑Bn l′=t∇hvl′ . In other words, the gradient through PS(·) is another prefix sum computed backwards. Hence, auto-differentiation through Φ(n) takes the same parallel time O(logC), serial time O(C) and memory O(C), as the forward computation of Φ(n). Since during the whole back-propagation algorithm, we only store and update tensors U (n),G(n), whose size doesn’t depend on L and C, this results in total O((L/C) logC) parallel time, O(L) serial time and O(C) memory in addition to p storage. A full description of the forward-backward pass is presented in Algorithm 1. Figure 1b is an illustration of the algorithm.

heading: 3.4 Analysis of the running time and memory
text: As we have shown, Performer can be trained in parallel time O((L/C) logC) and O(C) memory in addition to the input p storage. Hence, C is a tradeoff parameter: when C is maximal (C = L), the model is fully-parallelized, therefore resulting in the fastest execution. Whereas minimal C = 1 corresponds to step-by-step processing, i.e. a fully-sequential regime which doesn’t benefit from parallelized computations on GPU or TPU, but consumes O(1) memory as a function of L.
During the forward pass, Algorithm 1 requires as many total FLOPs as the naive forward pass through (7-8). As for the backward pass, for each 1 ≤ n ≤ N , the forward pass through n’s slice is repeated for symbolic construction of Φ(n) (see Algorithm 2), and then back-propagation is run through Φ(n). In addition, a backward update of U (n) (9) is computed, taking precisely BnsM(d+ 1)k “add” operations. Hence, we conclude that Algorithm 1 requires as many FLOPs as two forward and one backward pass through (7-8) for the whole sequence p plus LsM(d + 1)k = LsMdmodel + LsMk FLOPs. To characterize this addition, assuming that typically dff (dimension of the feedforward block) is 4dmodel in practice, observe that applying dense Performer layers (11-15, Appendix A) alone requires 3Lsd2model + 2Lsdmodeldff = 11Lsd 2 model FLOPs. This is much bigger than LsMdmodel +LsMk, since M is much smaller than dmodel in practice [11, 18].
Since the back-propagation takes roughly 5 times more FLOPs than the forward pass [14], we conclude that memory efficiency of Algorithm 1 results in a small constant-time increase in FLOPs. FLOPs affect energy consumption [40], a crucial factor for on-device applications.
Further analysis of Algorithm 1 reveals that the C = 1 regime requires as much memory as if Transformer were applied to a sequence of length 1 plus exactly 2sdmodel(M + 1) floats for storing U ,G. For comparison, the subset of θ corresponding to dense layers in self-attention and feed-forward blocks (11-15), occupies 3sd2model + 2sdmodeldff = 11sd 2 model floats. Again, this is much bigger than 2sdmodel(M + 1), since M is much smaller than dmodel in practice.
Table 1 shows comparison of the proposed algorithm with other architectures such as a conventional Transformer, recurrent neural networks (RNNs, 17, 9) and residual networks (e.g. Neural ODEs, 6).

heading: 4 Experiments
text: Our main contribution is a new low-memory gradient computation algorithm for the existing Performer architecture. Performers have very competitive performance among other methods for long sequence modelling [11, 18, 36]. Hence, in the experimental section, we aim to answer the following questions about using this algorithm in practice. (a). Does the theoretical time-memory tradeoff, controlled by C, agree with benchmarks of time-memory for varied C? (b) In precise arithmetic, different values of C lead to the same correct gradient ∇θL. Does this hold in practice, when
finite-precision arithmetic is employed? (c) Can a model, pre-trained with a bigger C (e.g. on a server), be fine-tuned with a smaller C (e.g. on an embedded device)?
We address these questions below. We analyse 4 model configurations (L, s, dmodel): I = (8192, 1, 1024), II = (1024, 3, 512), III = (4096, 3, 1024), IV = (16384, 3, 1024). In all configurations, we set dff = 4dmodel, k = dmodel/64 (number of heads). We set M = d and employ g(x) = (x2i ) d i=1 elementwise-quadratic feature mapping in (1), which we find to work well. In all experiments Σ = {0, . . . , 255} and batch size is set to 1, i.e. we analyse a setup where gradient accumulation cannot be used to decrease memory, and therefore our algorithm is crucial. Our code is in PyTorch 1.7 [30]. To ensure that reproduction of experiments is accessible for a wider audience, we use a single NVIDIA Tesla P100 GPU with 16 GB memory for each experiment.

heading: 4.1 Empirical benchmarking of the tradeoff
text: We run Algorithm 1 for configurations I, III, IV and different powers of 2 as C. We use input strings sampled randomly from ΣL. In order to characterize the time-memory tradeoff, we measure wall-clock time and peak GPU memory for a single gradient evaluation. We use the torch.cuda.max_memory_allocated function to report peak GPU memory.
As discussed in Section 2, there are two methods to compute (2-3): the first (iterative) method doesn’t compute and store tensors (3) explicitly, resulting in smaller memory consumption at a cost of less parallelization, while the second one computes tensors (3) using the parallel prefix sum algorithm, therefore operating faster, but using more memory. The same methods can be applied for the memory-efficient algorithm when computing (7-8) updates. We implement and benchmark both methods as part of the algorithm. For the explicit prefix-sum method, we find that the torch.cumsum function works faster and consumes less memory than our custom implementation of the parallel prefix sum algorithm. We attribute this to hardware-optimized implementation of the native function, and use it in experiments. As for the iterative algorithm, we implement its “block” version, when, instead of iterating l one-by-one, we iterate through blocks of small size (see details in Appendix
Table 1: Complexity for the back-propagation as a function of sequence length L and the tradeoff parameter C ≤ L. The indicated memory complexity is in addition to the input sequence p storage. The serial time complexity for Performer is reported for the version with iterative PS(·) computation (as in [18]), while the parallel time is reported for the parallel prefix sum (as in [11]). For both methods, memory complexity is the same, but the constant is smaller for the iterative version.
Model Serialtime Parallel time Memo -ry
RNN O(L) O(L) O(L) Residual NN O(L) O(L) O(L) Transformer O(L2) O(logL) O(L2) Performer O(L) O(logL) O(L) Our alg. O(L) O(LC logC) O(C) Ours, C=1 O(L) O(L) O(1)
Table 2: Time per iteration (sec., averaged over 1000 iterations) and peak GPU memory (Gb). When using small values of C, we relate the remaining memory to a storage of parameters θ and optimizer’s state.
Setup, L, C Time Memory
CT, 8192, full 0.3008 0.938 CT, 8192, 4096 0.5372 0.595 CT, 8192, 2048 0.6002 0.436 PTB, 1024, full 0.1377 0.300 PTB, 1024, 512 0.2526 0.257 PTB, 1024, 256 0.3060 0.231 ENW, 4096, full 0.4598 1.513 ENW, 4096, 2048 0.7922 1.085 ENW, 4096, 1366 0.8654 0.909
C). This way, the algorithm has a smaller constant in O(L) time complexity and bigger constant in a “small” O(dM) term of the memory complexity (assuming that d,M L). For a fixed C, in addition to reporting memory of Algorithm 1, we also report memory of the naive gradient computation run on a string of length C, sampled uniformly from ΣC . This is to confirm that memory use of Algorithm 1 is just slightly above the full computation on the input of length C.
Results are reported in Figure 2 (left, middle). We observe significant improvements in memory compared to the full computation, as C decreases. As C converges to 20 = 1, the remaining memory can be attributed to storage of the model’s parameters θ. Time follows two regimes: declining fast as C grows (meaning that prefix sums are parallelized) and declining slower for big values of C (meaning that the practical limit of parallelization is reached). Memory scales slower than O(C), as C increases. We attribute this to details of the PyTorch internal code. We find that iterative version of (2-3) works only slightly slower than prefix-sum version, while consuming much less memory. Finally, Algorithm 1 consumes only slightly more memory in practice than the full method run on the input of length C < L (“L/B” plots on Figure 2 (middle)).
In addition, we compare Algorithm 1 with checkpointing [15]. By checkpointing in this context we understand storing {U (n)}1≤n≤N during the forward pass and reusing them during the backward pass instead of recomputing. This results in a small FLOPs decrease since U (n), while memory scales as O(L/C). We use config. I for comparison (Figure 3). While not faster in practice, checkpointing consumes much more memory as C decreases for both iterative and prefix-sum computation of (2-3).

heading: 4.2 Effects of finite-precision arithmetic
text: Since the iterative version of (2-3) results in a good balance between time and memory of Algorithm 1, we use it further. To quantify finite-precision effects, we plot relative discrepancy ‖∇(C)θ L − ∇(full)θ L‖2/‖∇ (full) θ L‖2 between the gradient ∇ (C) θ produced by Algorithm 1, and the gradient
∇(full)θ L produced by full computation. Figure 2 shows results for randomly initialized models. We observe a small discrepancy of order 10−6–10−4, confirming the correctness of Algorithm 1.

heading: 4.3 Training from scratch and fine-tuning
text: To confirm backward compatibility of Algorithm 1 during training, we run three language modelling setups: Copying Task (CT), symbol-level Penn Treebank (PTB) and Enwik8 (ENW).
For the CT, we follow the setup from [20, 18], sampling inputs as 0ω0ω, where ω is drawn uniformly from (Σ \ {0})L/2−1. In this setup, we only aggregate cross-entropy loss from the second half of the input, so the task is to reproduce the first half. We include the CT as a task where long-range signal is crucial, and the heuristic of “chunking” the input into segments would fail to solve the task.
We use model configurations I, II, III for the CT, PTB and ENW, resulting in sequence lengths L = 8192, 1024, 4096 respectively. For each setup, we compare training with full gradient computation, “fine-tuning” regime, when the first half of iterations is run using the full algorithm, and the second half is run using Algorithm 1 using various values of C. In addition, we include training from scratch equipped with memory-efficient gradient computation via Algorithm 1. Figure 4 demonstrates results: all methods result in almost same performance. This confirms that memory-efficient gradient computation is backward-compatible during training. Table 2 quantifies the memory savings and time tradeoff in all setups. Additional details and results, including bigger version of Figure 4, BPC for CT, train set performance, and curve differences in Figure 4, can be found in Appendix D.

heading: 4.4 One-shot fine-tuning under low memory
text: To analyze the scenario when model is pretrained on server and then fine-tuned (F/T) with a small C on a low-memory device, we add the following experiment. We take a pretrained model from either PTB or ENW setup from Section 4.3 and subsample randomly 5000 examples from the corresponding validation set. We perform a one-step gradient descent with 0.01 learning rate (tuned on other random subset) to minimize the loss computed on the first half of each sequence and evaluate Bits Per Character (BPC) on the second half. In this experiment, the first half of the sequence represents the data generated by user on the device, and the second half is a new data to be predicted. Fine-tuning
procedure, therefore, represents “personalization” of the model to the specific user (see Table 3). We observe BPC improvement without any server compute and a memory improvement compared to the full computation, while time (≈ 1 sec.) is less crucial, since fine-tuning can run in the background.

heading: 5 Related Work and Extensions
text: Compatibility with other memory-optimization techniques. Observe that the specification (5) is compatible with the reversible layer design from [20], when the sparse self-attention is replaced with the linear self-attention2. This can bring more memory savings, since one doesn’t need to store the whole symbolic Φ(n) during the backward pass. Checkpointing along the layer dimension [15, 7] can also be used to reduce the memory consumption for storing Φ(n)’s graph, though at the cost of a longer execution time. The gradient accumulation technique [28] is also compatible with Algorithm 1, i.e. one can combine both methods to “collapse” batch and sequence dimensions simultaneously. Moreover, our algorithm is compatible with distillation [32], since it can be run on a distilled model.
Comparison with [18]. In [18], authors mention that a single self-attention block can be evaluated in O(1) additional memory. However, during back-propagation, one still needs to store L intermediate states, e.g. in the feedforward block. Hence, the full memory complexity is still O(L). In contrast, our method optimizes memory consumption along the sequence dimension for the whole multilayer model.
Extension to Transformers with dropout. Dropout [34] is a popular regularization technique. It is used with Transformers when the train dataset is small enough to cause overfitting (e.g. it wasn’t used with GPT-2, trained on a massive dataset). Our algorithm can be extended to stochastic computation graphs with dropout. For that, use separate random seeds to generate dropout masks for each slice 1 ≤ n ≤ N , and reuse these seeds two times during the forward and backward pass through the nth slice.

heading: 6 Conclusion
text: We proposed an algorithm for memory-efficient back-propagation through a Performer. The algorithm reduces memory consumption along the sequence dimension, and can, therefore, be used for longsequence training. The algorithm: (1) is completely backward-compatible, since it computes precise gradients and does not involve approximation, (2) does not require many additional computations, and (3) enables user control over the tradeoff between time and memory consumption.
Limitations. One limitation of this works is that the proposed memory improvements are traded off by longer running time. Another one is that these improvements rely on a specific feature of prefix sum signal propagation appearing in Performers and, also, in parallelized versions of RNNs [26]. Finally, we don’t see an immediate extension of the proposed construction to the bidirectional self-attention case, used in Transformer encoders [38] and for masked language modelling [13]. The reason for that is that no notion of the prefix sum “front” can be introduced in this case.
Negative societal impacts. This work studies Performers – efficient Transformers. Transformers in general can have the following negative societal impacts: large carbon dioxide emissions [35], privacy and data leak vulnerabilities [5], bias and fairness issues and malicious misuse [4, 3].

heading: Acknowledgments
text: We thank Tom Weingarten and Tamas Sarlos for many fruitful discussions.
Funding disclosure. V. L. acknowledges support from the Cambridge Trust and DeepMind. V. L. was part-time employed by Google while a PhD student. A.W. acknowledges support from a Turing AI Fellowship under grant EP/V025379/1, The Alan Turing Institute, and the Leverhulme Trust via CFI.
2See e.g. CausalFavor class in https://github.com/ google/trax/blob/master/trax/layers/research/sparsity.py, which is compatible with the official Reformer code.

Review summaries:
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.  

The authors propose a low-memory gradient computation algorithm for the linear transformer architecture. Their algorithm allows the user to control the trade-off between computational efficiency and memory by specifying how many tokens of the input sequence should be processed at once in the forward pass. At the extreme, the memory consumption can be reduced to a single token (+model parameters +some overhead) They empirically demonstrate that a) the obtained gradient is correct even considered finite-precision arithmetic, b) the actual memory consumption of their implementation aligns with the claims of their analysis, c) the trade-off between memory and compute can be dynamically adjusted during training to allow for fine-tuning on low-memory devices.