Sub-Linear Memory: How to Make Performers SLiM
Authors: Valerii Likhosherstov
Krzysztof Marcin Choromanski
Jared Quincy Davis
Xingyou Song
Adrian Weller
Venue: NeurIPS 2021 Poster (2021)
Keywords: transformer
performer
slim-performer
memory efficient
linear transformer

Abstract: Transformer architectures have become very popular yet the original implementation requires  $O(L^2)$ 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 backward-compatibility, this discovered time-memory tradeoff can be used for fine-tuning on low-memory devices in a decentralized fashion without any server computations.

Sub-Linear Memory: How to Make Performers SLiM
Valerii Likhosherstov University of Cambridge v1304@cam.as ac. uk
Krzysztof Choromanski
Google Brain & Columbia University
Jared Davis
Xingyou Song
DeepMind Stanford University
Google Brain
Adrian Weller University of Cambridge & Alan Turing Institute
Abstract
Transformer architectures have become very popular yet the original implemen- tation requires O(L²) in serial time and memory functions of input length L. Recent works proposed various linear self-attention mechanisms, scaling only OL) for serial computation. We conduct thorough complexity analysis of Performers, a class which includes most recent linear Transformer mechanisms. We note remarkable computational flexibility: the gradient computation can be performed with no approximations using sublinear memory as a function L cost of greater time
(in addition to negligible storage for the input sequence), complexity in the parallel setting. In the extreme case, a Performer consumes only O(1) memory, and still requires (L) time. Due to complete backward- compatibility, this discovered time-memory tradeoff can be used for fine-tuning on low-memory devices in decentralized fashion without any server computations.
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) [1 [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 ((L?) time and memory complexity, where is the length of the Transformer's 's input sequence.
recent line of work aimed address this restriction, using either structured sparsity g truncated [12), clustering (20 31 or linear attention methods [18 10
33 23 For
detailed overview of efficient Transformers, see [37] We refer 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 Trans- formers [T] demonstrating strong performance in a systematic comparison efficient Transformers [36.
This recent trend of feeding longer sequences into Transformers, coupled with the use of deeper mod- els, 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 (NcurIPS 2021).

---

Loss
FFN(-)
H
Rayser
X 11.
X
KO
Boyer
Lavver
LaNe
Dayin
R
R
2/133
2
24133
111
PS
PSC-
skip
loes
lnss
Jees
Joss
Joss
Vix a(K-)
SUK
X
Layor
Iaver
Laveer
Jayce
Layne
Jaares
((--1)
Xula
X
XCG.
X H
Layer
X
x
git
X
G
[0.0
4
X
24(0)
X
(1)
(a)
U
(b)
VI
Figure 1: (a) rth layer and its decomposition into T(- p(r=)) U(r-1), (b) Illustration of Al- gorithm 1 when
2. Red color indicates objects stored in memory. I-II) forward passes for 2 respectively. only the loss value and ((n) are stored. III) backward pass start, for- ward computation through the slice 2 to build symbolic ((2) and update (2(2) - u IV) back-propagation through (22) to find V L and g(), V,VI) the same backward iteration for n 1.
operating on long sequences used moderately small batch sizes of 1-8 instances 20 18
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 decrease memory usage at the cost of longer time, but it only be applied when the batch size is bigger than 1. In this paper, we are discussing situation when a batch size of is prohibitive, while longer processing times affordable e.g. when fine- tuning 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 pack-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 can decrease memory consumption at the cost of smaller parallel bandwidth
for batch size of 1, of the model. Notably, no approximations are introduced, so the obtained gradient 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, < < L, controls tradeoff between the memory, scaling as O(C) in addition to negligible input sequence storage, and parallel running time, scaling as O(LL/C) log C). When 1, the algorithm consumes as much memory as a single token were fed into Performer, plus 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 small addition.
We evaluate the proposed tradeoff empirically, and confirm backward-compatbbllt for the synthetic Copying Task and language modelling on Penn Treebank (2 and Enwik8 (4 datasets
models/slim
Bodelss/iiin.peffommet performer

---

2 Background: linear self-attention and Performer
We commence by defining linear self-attention 18 33
23
Consider sequence of length
and three matrices: queries Q E RLxdd keys K E R Lxd and values V E Riddd Then linear
self-attention is defined as functional producing Y Att(Q K, V) E R 1xxdd
Er= Vo (9 (K() " g(Qi)) EE Vv x ((K)) x (
WE {1,
D}: YI {..
(I)
(Ei=1 (Kr)) (T (QQ)) where by Z₁ Rdzx... denote slice Zi:....: of Z E Rdixddx.... g Rd TR positive mapping for a stable division in ① M is typically much smaller than L, e.g. M = =.
Ex=1 ( (Qi)
For needs of autoregressive generative modelling, when each element depends only previous elements of the sequence [8). Y, only depends on inputs at indices {l.... Self-attention of type ① was proposed for processing sequences of the long length L, since the original self-attention from 38 scales as O(L²). ). The second transition in (1. which is due to associativity matrix multiplication, suggests algorithm to compute linear self self-attention efficiently in subquadratic time.
For tensors (1),... ((n) of the same shape, let Z = )Z be tensor such that for all 1 n, Zi Zii....
Z ( By R E L S E LLM denote a tensor and matrix such that
R PS((V, x ())))
())((()))
where PS(Z) = (Ev=1 Ze 1= is prefix sum along the first dimension of Z. Next, compute
L1 S 1 S L: Y, (R, (S))) (QQ)).
Depending on the prefix-sum algorithm used in 2 we can obtain different complexity estimates for linear self-attention. [88] propose to iterate through 1..... maintaining only current R., Sj, and compute and store the result Y,. This way. tensors R. PS(R) E R are not stored in memory, resulting in OL) time complexity and ( + M) + d.M memory complexity, [88 also propose similar computation of gradients through 3 see Appendix for more discussion.
Alternativeyy, employ a parallel prefix-sum algorithm 21 39 which, for a tensor R..... finds PS(Z) Z) in Ollog L) parallel time and ((L) memory. Applying this algorithm for PS(R). PS(S) and then computing 3 results in only Olog L) parallel time complexity and ((LdM) memory consumption. A similar approach was proposed in the context of parallel RNNs [26]
Linear self-attention is part of Performers [11] efficient Transformers. Below outline Performer architecture [1] for language modelling, while our analysis can be applicable in broader setups.
Let Σ be input sequence of length L, where Σ is finite alphabet. By emb(pr, () E Rdmodel 1 s S L, denote a linear combination of the pa token's learned embedding and positional embedding of I's position (sinusoids with different frequencies, as in (38). Then Performer is defined
parametrized mapping from (0) (emb(pi.. 1))/ E IR L &model into X ((ot) E R through
sequence of hidden representations X x(),,..,xX' (() E LBddmddd For each 1 S r < 8, ( is
obtained from ( by applying Performer layer consisting of two parts. The first part is multi-head self-attention which comprises linear self-attentions applied in parallel (dmoodet = kd) plus skip connection. The second part is feedforward network applied independently to each element of the sequence. See formal definition of the architecture in Appendix [A]
For each I, S 1 < L 1, x (oot) are predicted logits of the next token P1+1- Let C (x(ou))) denote the cross-entropy loss with respect to P1+1, zero when I = T.. The loss is defined as
- (L- L=
(x(()())))) (o0t)) + +... Lx(x(
(4)
3 Low-memory back-propagation algorithm
3.1 Compact notation for Performer
We rewrite transformations (0) + (() + ... (() in the following form. For each Srr S 8,
((r-1) p(r-1) = -000))(0))(((()))00)) ((r = x()) x
((())) (5)
3

---

Algorithm Low-memory forward- backward pass. See Algorithm 2 for updateProc. Compared to notation from the text, redundant indices are dropped and tensor names are reused here and Algorithm
Algorithm updateProe procedure.
Input: E N, binary flag onBP. if onBP then Initialize :- end X (emb (cm.(P.... PA.. An m)):
for r = to do
Compute T, T,T:= == (x))(X;00)
Input: p EL, 0 E R'pparmm, CE N. Output: loss C, gradient VoL.
if onBP then Update U, Di= Ty end
Set U := 2, PS((PT X G (U,r:)) if onBP then
Update 2+= G, U
else
Set L : 0;
Update U, = UB
Set U- == D,
end if
for to N do
end for
updateProc(n, False); end for
Set L(upd) = ((w) (XW(out)) blout));
if onBP then
Set VoL :- 0 Oparem
Update ++= L(upd); Compute Vo Vu through auto-differeemiaation;
Set G == Dis
for n = N to do
Update voL+= Vo : VuΦ;
updateProc(n, True);
else
end for Return L, VoL.
Set C+= ((upd);
end if
Here Sparam a set of parameters, T(r11,, U(r-1) RLLDD and (r- RLXD: are the following matrices (see Figure Da for an illustration). (a) ( a matrix of representations which are passed into the prefix-sum operator. That is, for each s / < L, T.T((r ( is concatenation of ((i) and flattened V2 x 9(K)) for all attention heads computed at the rth step. Consequently, Dj ((d k). (b) For each < 1 < T., U(--1) is concatenation of all corresponding S, and flattened R₁ - results of the prefix sum (Eq. 2 inside each self- attention head at rth layer. (c) ((r-)) is
a matrix of representations which skip the prefix sum. For each 1 < l S L, r (r-1) is concatenation of x (r-1) and g(Q:) for each attention head S S.K. Therefore, D₂ = MR: + dmadde-
F((r and ( are functionals depending on a of model's parameters 0. That is, they take subsets of 0 corresponding rth layer weights. F()) is responsible for constructing T(r and ( representations preceding prefix sum, while (+)) finalizes multi-head self-attention and includes the feed-forward block. Importantly. F(r) and G()) are applied rowwise, i.e. 5 can be rewritten
v1 S T. T(T-1) r(r-1)= (2())(((()))
x(r) ( "irii""),r""" -r(r-):)) ;)).
(6)
Hence, the only place where the signal is propagated across the sequence is prefix sum in (5).
The representation
encapsulates architecture details of Performer inside (F (11),, .G..,..., F((), G(s) J. In fact, the representation 5 holds for various possible modi- fications, proposed in the literature. This includes, but is not limited by the different positioning
of layer normalization 41 38 adding stabilizing gating mechanism 229 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 J....,,, and only maintains a front of current prefix sum values, thus allowing substantial memory improvement. The backward pass is implemented similarly: current prefix sums along with their gradients are maintained in dynamic programming fashion, but the iteration proceeds in backward direction.
3.2 Forward computation
Suppose the memory budget not enough to perform complete forward pass through Performer (Equation for 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 forward pass through
the input of length C' S L. plus small addition. 1 1 C' S L is arbitrary and user-defined.
Split matrices ( T(r) p()) U into N slices of size at most C' along the vertical axis (N [L/CT): for each n, 1 5 n - N.
x (mmm) = (X (x) (r)
An + E Bax T(nm) = (T B.,
B.
AntiL=1 An+1l==1 R B
U(ra)= -U -
Bn
E R B p(ra) (T (r) B.
where A,, (n 1)C and by B,.
S N, we denote the size of the nth slice: B, = C for
u < N, BN 2 C. Based on (5), we conclude that for each S n I N and 1 i S it holds that
( ( F ( 0):
U0-1m)) = 1B, x (U Ba-1
(U(r-1n-1))
+ +PS(T(r)m)) I)), x (r.m) = G ( lm)) r (rr-l.n) 0).
(8)
Here, 1B, E IR Rh is a vector of Bn ones and we denote U (fr-10) Bu = Op,, (a vector of D₁ zeros).
Now, instead of iterating over r = L..... and computing (5) for the whole sequence at once, we first
iterate over n 1,.... N and then iterate over r = 1,,..,& 8 in a nested loop to compute 77
As can be deduced from the 7 8
we only need to maintain the current value of (U (U(r-1,n 1-) )==
B
Rax in the outer iteration over n.
Denote 4(()) = (Ug. (US_L E R 0 S n S N. The memory algorithm for the forward pass is as follows. First, initialize L = 0 and u1()) O,x Dr. Then, iterate over
1,...,,,. N and maintain the current value of (1(n1)), 1), During each iteration, compute ((0..)
(emb(pA.+t: An 1)) Then iterate over 1, I,...,&, s, where compute
8 and update un!))
U 1r- Finally, compute X (oot.n) = X (x. (xs)))
((ut) and update ++= L(m) (x (out.)))
where we denote ((r) (oat,))) (L 11 Em
B. ( ((ntt.)))
By the end of the iteration over n, the correct loss value 44) computed. As result, the forward pass takes (L) serial time or O(L/CC log C)) parallel time and consumes only O(C) memory. This is in addition to the input sequence p E Σd storage, which is (L)) in principle, however the constant is negligibly small. For instance, if p is flattened image or an ASCII text string, then it occupies precisely L bytes in memory. The log C' term in the parallel time complexity is due to the parallel prefix-sum algorithm taking logarithmic time, as discussed in Subsection 2
3.3 Back- propaapropggttio and the final algorithm
The goal of backward pass to compute the gradient C of the loss function with respect to parameters 0. One can just perform automatic differentiation [4] (implemented in Tensorflow M and Pytorch (0 through the computation graph induced by the memory-efficieet 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 (L) memory complexity as a function of I. and C. Instead, we propose a back-propagan complexity as the efficient forward pass.
algorithm which has the same time and memory
Let 0(1)
O(N) 0 be results of symbolic "identity operation" performed on 0, so that for all 1 < < N, g(n) is used instead of 0 in (7
8 Then the total gradient of 0 has Appendix B we derive an expression for Vocn C,
the form VoL = V VonyC + ... + VouN
1 S n S N. Namely, denote ((n) Vuin L, then V Voca
Vocn ((m) (B(n), (00m),,((n-1) (n(11),, ((n-1))g(m) g(n)), where
T param x DX x BX + R.
((a) (0(m)) (11111,,2 Z) = (x(m)(x(outnn) (out.n)) + Σ Z (.
In 4(n)'s definition, ()tt,n) and (n) ( are results of iteration over r L...., s with parameters 0 g(n) and (Ulr Gradient V Voon (())) can be computed by automatic differentiation through (n).
p -لا--ي--)) equal to ((x)'s second argument (n(=
5

---

An efficient way to compute and sum up all Vacn; is to iterate in a backward direction n = ......! and to maintain values of (m), g(m). U(M) (
is known after the end of the forward pass, and
Bu (T(r-Lm)
LA S N: u(n-1) = (1(n)
Σ
) )
(9)
Further, in Appendix B we show that G(M) D. and, for each S annan N,
g(n-1) = V un o (x) (B(n), 11(111111000 g()))
(10)
By a single auto-differentiation through (x) we can compute Vi(m) C = V (()) and the update (1O). If is some vector of length B,, and h is some scalar function of v PS((), then for all
I I S Bn V,W; = Ep_ VAVr. In other words, the gradient through PS(() is another prefix sum computed backwards. Hence, auto- iitoodiiferentiationn through ((() takes the same parallel time (log C)2, serial time (CC) and memory O(C), as the forward computation of (("), Since during the whole back-propagation algorithm, we only store and update tensors 2((m), g(v), whose size doesn't depend on L and C, this results in total O(L//C) log C') parallel time, OLL) serial time and O(C") memory in addition to p storage. A full description of the forward-backward pass is presented in Algorithm Figure P is an illustration of the algorithm.
3.4 Analysis of the running time and memory
As we have shown, Performer can be trained in parallel time O((L/C)ogg C) and OCC) memory in addition to the input p storage. Hence, C is tradeoff parameter: when C' is maximal (C L), the model is pally-parallelized therefore resulting the fastest execution. Whereas minimal = 1 corresponds step-by-step processing, i.e. fully-sequential regime which doesn't benefit from parallelized computations on GPU or TPU, but consumes (1)) memory as a function of L.
During the forward pass, Algorithm requires as many total FLOPs as the naive forward pass through 8 As for the backward pass, for each S n S 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, backward update of (1(m)
9 is computed, taking precisely (d + 1)k: "add" operations. Hence, we conclude that Algorithm requuires as many FLOPs as two forward and one backward pass through (78 for the whole sequence plus LsM (d 1)k = ddMdem + L&Mk FL.Ps. To characterize this addition, assuming that typically d,f (dimension of the feedforward block) is 4dmodde in practice, observe that applying dense Performer layers 11 15 Appendix Aj alone requires 3Lsdmode + 2T.ddmade!d = 11Lsd FLOPPs This much bigger than Ls amodet + LsMk,, k. since M is much smaller than dmodel in practice
Since the baok-propag takes roughly 5 times more FLOPs than the forward pass 14, we conclude that memory efficiency of Algorithm I results in a small constant-time increase in FLOPs. FLOPs affect energy consumption 44), a crucial factor for on-device applications.
Further analysis of Algorithm I reveals that the C' = 1 regime requires as much memory as if Transformer were applied
sequence length plus exactly 2sdmodet (M + 1) floats for storing U1 G. For comparison, the subset of 0 corresponding to dense layers in self-attention and feed-forward blocks (11 15), occupies 3sd + 2ndmodedd 11 floats. Again, this much bigger than 2sdmodel [ + 1), since M is much smaller than dmodel in practice.
Table 1 shows comparison of the proposed algorithm with other architectures such as conventional Transformer, recurrent neural networks (RNNs, 17 9 and residual networks (e.g. Neural ODEs, 6).
4 Experiments
Our main contribution is new low-memory gradient computation algorithm for the existing Per- former architecture. Performers have very competitive performance among other methods for long sequence modelling [ 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 VoL. Does this hold in practice, when

---

Gontgg
Contig I, iter.
Contig. PS
10:
UB Config: liter.
Contiy.
Contig,
UB Config.
10a
Cantle IN.
Cantig IV, (ier.
Contig, IV. PS
LBB for Configi IV. Iber.
xt c-m
Candg IV, PS
Contig. Rer.. fill
LB for Conffig PV.P5 C.C.
Retstive ehcoree Mo
Contig, PS, full
mBUmerov G 10
Time, Sec.
Contig,
Confi.ll
Carágg1 Confi III Carfig. IV
Contig.
Canti
Config. Ror.. full
Config. iter.,
10a
10"h
10
10
(ogalC)
(gga(C)
(egal()
Figure 2: Benchmarks of Algorithm T All plots are averaged over 10 seeds. "iter." stands for iterative computation of 2 while "PS" is for explicit prefix sum computation in @ We omit time and memory for big values of C in "Config. IV, PS" and "Config. IV, full" setups. because these led to memory overflow. Analogous results for configuration II can be found in Appendix D (Left) Time dependence on C. Crosses indicate horizontal time levels for corresponding full memory-ineffiiiet methods. The dotted line indicates x C-1 tangent in logarithmic scale. (Middle) Memory dependence on C. Again, crosses are for horizontal levels of full-sequence methods and the dotted line indicates C' tangent, We do not report curves for config. III, because they completely match curves for config. IV, which is natural, since 8, dmoods are the same for both configurations. "L/B" stands for memory lower bound computed by processing input of length C. (Right) Relative gradient discrepancy as a function of C' and standard errors. Evaluated on random inputs over I
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 sL :model): I= (8192, 1, 1024), Il (1024,3 3, 512), III= (4096, 3, 1024), IV = (16384, 3, 1024). In all configurations, we set ر = 4dmodel, k = dmodet/6c (number of heads). We set M d and employ g(x) = (x²) elementwise-quadee
feature mapping in (T. which we find to work well. In all experiments = {o..... 255} and batch size is set to 1, f.e. we analyse a setup where gradient accumulation cannot be used to decrease memory, and therefore our algorithm is crucial. Our code is in Py PyTorch 1.7 (30. To ensure that reproduction of experiments is accessible for wider audience, we use a single NVIDIA Tesla P100 GPU with 16 GB memory for each experiment.
4.1 Empirical benchmarking of the tradeoff
We run Algorithm T for configurations III, IV and different powers of 2 as C. We use in- put strings sampled randomly from Σ4. 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 fasx__emory_____ alccated function to report peak GPU memory.
As discussed in Section 2] there are two methods compute 3 the first (iterative) method doesn't 't compute and store tensors (3) explicitly, resulting in smaller memory consumption at 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 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-optimizee implementation of the native function, and use it in experiments. As for the iterative algorithm, we implement its "block" version, when, instead of iterating one-by-one, iterate through blocks of small size (see details in Appendix
7

---

Table 1: Complexity for the pack-propagation as func- tion of sequence length L and the tradeoff parameter C' S L. The indicated memory complexity is in addi- tion to the input sequence p storage. The serial time complexity for Performer is reported for the version with iterative PS((.) computation (as (18). while the parallel time is reported for the parallel prefix sum (as in (T].. For both methods, memory complexity is the same, but the constant is smaller for the iterative version.
Table 2: Time per iteration (sec., aged over 1000 iterations) and peak GPU memory (Gb). When using small values of C, we relate the remaining memory to storage of parameters 0 and optimizer's state.
Setup, I,, C
Time
Memory
CT, 8192, full CT, 8192, 4096 CT, 8192, 2048
0.3008
0.938
0.5372
0.595
Serial time
Parallel time
Memo
Model
0.6002
0.436
-ry
PTB, 1024, full
0.1377
0.300
RNN
O(L)
O(L)
O
PTB, 1024, 512 PTB, 1024, 256
0.2526
0.257
Residual NN
O(L)
((L)
(
0.3060
0.231
Transformer
O(L²)
O(log L)
O(L²)
Performer
O(L)
O(log L)
((L)
ENW, 4096, full
0.4598
1,513
Our alg.
O(L)
Og C))
O(C)
ENW, 4096, 2048
0.7922
1.085
Ours, C==1
O(L)
O(L)
O(1)
ENW, 4096, 1366
0.8654
0.909
10
Alg. 1, Iter.
Time, sec.
GPU memory.
10a
Alg. PS
109
Checkp., PS
10a
10
10
loga(C)
log:(C)
Figure 3: Algorithm compared to checkpointing of (4(x) 1SSNN.- Time and memory plots.
Ch. This way. the algorithm has a smaller constant in O()) time complexity and bigger constant in "small" (dd) term of the memory complexity (assuming that d, M < L).
For fixed C, in addition to reporting memory of Algorithm I we also report memory of the naive gradient computation on string of length C, sampled uniformly from Σ This is to confirm that memory use of Algorithm Tis is just slightly above the full computation on the input of length C.
Results are reported in Figure (left, middle). We observe significant improvements in memory compared to the full computation, as decreases. As converges to 2 1, the remaining memory can be attributed to storage of the model's parameters 0. 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 OCC), as C increases. We attribute this to details of the Py Torch internal code. We find that iterative version of 2 works only slightly slower than prefix-sum version, while consuming much less memory. Finally, Algorithm consumes only slightly more memory in practice than the the input of length C L ("L/B" plots Figure 2 (middle)).
full method
In addition, we compare Algorithm with checkpointing understand storing (2/(m) JInSSN during the forward pass and reusing them during the backward pass instead of recomputing. This results in small FLOPSs decrease since (((1), while memory scales as OLL/C). We use config. 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 (
15)
By checkpointing in this context we
Effects finite- finite-precision arithmetic
Since the iterative version of 23 results in good balance between time and memory of Algorithm we use it further. To quantify finite-precision effects, we plot relative discrepancy [ C-
(flll) Cl|2/11 V (/ult) Clla between the gradient v(() produced by Algorithm and the gradient

---

Copying tesk
Pern Treebank
Enwike
100
F/T start Full
4.0
T/T start Full
2.00
C- 256
1316
256 FIT
1366 R/T
Acourcyy
2.75
2.50
512 FIT
11 start
60
2.25
2.5
2.00
2.0
4090
1.75
C- 4026 F/T
150
50000
100000
150000
200000
20000
40000
GOOOD
80003
100300
50600
100000
150000 200000
Beration
Ferstion
Hera
Figure 4: Learning curves for three language modelling setups. We report accuracy newly generated data samples for Copying task, and Bits Per Character (BPC) metric on validation examples for Penn Treebank and Enwik8. F/T stands for "fine-tuning". All curves are roughly the same and result in nearly the same performance, confirming correctness and packward-compatiin gradients computed via memory-efficien Algorithm[] Differences can be attributed finite-precision arithmetic effects, accumulating over many iterations.
of
Table 3: Time per iteration (sec., averaged over 1000 iterations) and peak GPU memory (Gb). We also report memory required store parameters 0 only without any computations.
BPC, F/T
BPC, no F/T
Memory, params
F/T mem.,
F/T mem., full comp.
F/T time,
F/T time,
Setup. L, C
Alg.
Alg.
full alg.
PTB, 1024, 16
1.4263 1.5642
1.6544 1.6141
0.0661 0.2610
0.1077
0.1834 0.8049
0.8938
0.0426
ENW 4096, 64
0.4276
1.1142
0.2045
v (full) L produced by full computation. Figure observe a small discrepancy of order 10 10-6-10
shows results for randomly initialized models. We confirming the correctness of Algorithm [
4.3 Training from scratch and fine-tuning
To confirm backward compatibility of AlgorithmmIJunn training, 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 Ow where is drawn uniformly from (E 4 In this setup, we only aggregate cross-entropy loss from the second half of the input, so the task 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 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 | Jusing various values of C. In addition, we include training from scratch equipped with memory-effecient gradient computation via Algorithm Figure demonstrates results: all methods result in almost same performance. This confirms that memory-eflicient gradient computation is backward-compatibe 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
4.4 One-shot fine-tuning under low memory
To analyze the scenario when model is pretrained on server and then fine-tuned (F/T) with 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. this experiment, the first half of the sequence represents the data generated by user on the device, and the second half is 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 memory improvement compared the full computation, while time (Rx sec.) is less crucial, since fine-tuning can run in the background.
5 Related Work and Extensions
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- self-attention 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
also be used reduce the memory consumption for storing ((a)'s graph, though at the cost of longer execution time. The gradient accumulation technique [28] is also compatible with Algorithm D i.c. one can combine both methods to "collapse" batch and sequence dimensions simultancously. Moreover, our algorithm is compatible with distillation [32] since it can be run on distilled model.
Comparison with [88. In 18 authors mention that single self-attention block can be evaluated in ((1) additional memory. However, during propppropgatioon one still needs to store intermediate states, e.g. in the feedforward block. Hence, the full memory complexity is still (L)) In contrast, our method optimizes memory consumption along the sequence dimension for the whole multilayer model.
a popular regularization technique. used with Transformers when the train dataset is small enough to cause overfitting (e.g. it wasn't '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 S n S N, and reuse these seeds two times during the forward and backward pass through the rihh slice.
Extension to Transformers with dropout. Dropout 34
6 Conclusion
We proposed an algorithm for memory-efficient pack-propagation through a Performer. The algorithm reduces memory consumption along the sequence dimension, and can, therefore, be used for long- sequence training. The algorithm: (1) is completely bockward-compatiin 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.
since it computes precise
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 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
Acknowledgments
We thank Tom Weingarten and Tamas Sarlos for many fruitful discussions.
Funding disclosure. V. acknowledges support from the Cambridge Trust and DeepMind. V. L. was part-time employed by Google while PhD student. A.W. acknowledges support from Turing AI Fellowship under grant EP/V025379/1, The Alan Turing Institute, and the Leverhulme Trust via CFI.
"See e.g. CausalFavor class in httpa//wwww......ooooooooooooooa which is compatible with the official Reformer code.
10

---

References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp. Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg. Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large- scale machine learning on heterogeneous systems, 2015. Software available from tensorflow.org
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 2016.
[3] Emily M. Bender, Timnit Gebru, Angelina McMillan- Majorr and Shmargaret Shmitchell, On the dangers of stochastic parrots: Can language models be too big?.. In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Transparency, FAceT '21, page 610-623, New York, NY, USA, 2021. Association for Computing Machinery.
[4] 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 McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
[5] Nicholas Carlini, Florian Tramèr, Eric Wallace, Matthew Jagielski, Ariel Herbert-Vons Kather- ine Lee, Adam Roberts, Tom B. Brown, Dawn Song, Úlfar Erlingsson, Alina Oprea, and Colin Raffel. Extracting training data from large language models. CoRR, abs/2012.07805,. 2020.
[6] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Resa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 6571-6583. Curran Associates, Inc., 2018
[7] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin. Training deep nets with sublinear memory cost. CoRR, abs/ 16011.604444, 2016.
[8] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 2019.
[9] Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1724-1734, Doha, Qatar, October 2014. Association for Computational Linguistics.
[10] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, David Belanger, Lucy Colwell, and Adrian Weller. Masked language modeling for proteins via linearly scalable long-context transformers. arXiv preprint arXiv:2006. 03555, 2020.
[11] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with Performers. In International Conference Learning Representations, 2021.
[12] Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978-2988. 2019.
11

---

[13] Jacob Devlin, Ming-Wor Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[14] A. Griewank and A. Walther. Evaluating Derivatives: Principles and Techniques of Algorithmic Differentiation, Second Edition. Other Titles in Applied Mathematics. Society for Industrial and Applied Mathematics (SIAM, 3600 Market Street, Floor 6, Philadelphia, PA 19104), 2008.
[15] Andreas Griewank. Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation. Optimization Methods and Software, 1(1):35 1992.
[16] Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. CoRR, abs/ 0tls//660.08444 606.08415, 2016.
[17] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8):1335- 1780, November 1997.
[18] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5165. PMLR, 13-18 Jul 2020.
[19] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors. 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
[20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, OpenReview.net, 2020.
[21] Richard E. Ladner and Michael J. Fischer. Parallel prefix computation. J. ACM, 27(4):831-838, October 1980.
[22] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. ALBERT: lite BERT for self-supervised learning of language representations. In International Conference Learning Representations, 2020.
[23] Rui Li, Chenxi Duan, and Shunyi Zheng. Linear attention mechanism: An efficient attention for semantic segmentation. arXiv preprint arXiv:2007. 14902, 2020.
[24] Matt Mahoney, Large text compression benchmark, 2009.
[25] Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building a large annotated corpus of English: The Penn Treebank. Computational Linguistics, 19(2):313-330, 1993.
[26) Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference Learning Representations, 2018.
[27] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: fast, extensible toolkit for sequence modeling. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstrations) pages 48-53, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
[28] Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Thind Conference on Machine Translation: Research Papers, pages 1-9, Brussels, Belgium, October 2018, Association for Computational Linguistics.
12

---

[29] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg. Raphaël Lopez Kaufman, Aidan Clark, Seb Noury, Matthew Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 7487-7498. PMLR, 13-18 Jul 2020.
[30] Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary De Vito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
[31] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. arXiv, 2003.05997, 2020.
[32] Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv, 1910.01108 100,00110,,0000000 2020.
[33] Zhuoran Shen, Mingyuan Zhang, Shuai Yi, Junjie Yan, and Haiyu Zhao. Factorized attention: Self-attention with linear complexities. CoRR, abs/18 18200124,, 2018.
[34] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):1929-1955 1958, 2014.
[35] Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in NLP. CoRR, abs/1906.. 02244302233 2019.
[36] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference on Learning Representations, 2021.
[37] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arXivv, 9.20006732, 2020.
[38] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, ukasz, Kaiser, and Illia Polosukhin. Attention is all you need. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems 30, pages 5998-6008. Curran Associates, Inc., 2017.
[39] Uzi Vishkin. Thinking in parallel: Some basic data-parallel algorithms and techniques. 2010.
[40] 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.
[41] Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tie-Ye Yan Liu. On layer normalization in the transformer architecture, 2020.
[42] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large batch optimization for deep learning: Training bert in 76 minutes. In International Conference on Learning Representations, 2020.
13


Reviews:
[{"review_id": "ggjl73yhBZ", "reviewer": "Reviewer_NcB2", "summary": "The paper studies memory-time tradeoff and complexity analysis of Performers and proposes an algorithm for memory-efficient back-propagation through a Performer. The experiments highlight the tradeoff between time and memory consumption and the backward-compatibility of the algorithm", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper studies an interesting problem and provided through analysis of the memory-time tradeoff in training/finetuning Performer models. It discusses interesting implications on training models on edge and low-memory devices and highlights the backward-compatible nature of the algorithm.\n\nSome comments/questions:\n\nThe tasks studied are fairly limited in number and type. Did the author consider other tasks  (e.g. tasks from the long arena benchmark)?\n\nThe paper cites the competitive performance of Performers among other methods for long sequence modelling and opts to not report tasks performance numbers. It would be useful to report that, especially as relates to the discussion on gradient discrepancy \n\nIn 4.1, the paper states that the memory scales slower than O(C), can you please elaborate on why you think this is the case?\n\nWhile the paper acknowledges the limitation of the work relying on a specific feature of Performers, it would be useful to discuss if any of the findings could be extended to other efficient Transformers architectures", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "XYgsQx3niCP", "reviewer": "Reviewer_Bx7L", "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:\n\n1. 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.\n\n2. The authors analyzed the memory / time trade-off in the proposed algorithm (including extremal cases).\n\n3. 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.\n\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\n\nThe intuition and ideas in the proposed algorithm is novel when applies in reduced the memory of performer training.\n\nQuality:\n\nIn this paper, the proposed algorithm can reduce the performer training memory from O(L) to O(C) where L is the sequence length and C is the segment length (C < L). The value of C present a knob to exploit the time and memory trade-off in this algorithm. When C is small, the algorithm can reduce performer training memory significantly. Here are a few suggestions to further enhance the paper (and I am willing to raise the score if the questions are resolved and incorporated into writing)\n\n1. The derivation and discussion in the paper are mostly focusing on the one-directional casual language modeling setting. In the modern NLP pipelines, the bi-directional masked language modeling setting is another equally important one. Does the trade-off discussion and complexity still applies to bi-directional setting?\n\n2. I saw the memory saving with small C is only ~2X compared to conventional performer training algorithm (Table 3)? It would be great if the authors could explicitly decompose the memory consumption so that practitioners can know where to attach to make it more savings when implemented in lower level (instead of using the coarse grain pytorch or other frameworks)?\n\n3. The exact attention compute in conventional transformers (vs performers) can be approximated with performers using tailored feature mapping functions g(*). E.g. the exp kernel to compute attention score can be approximated with say linear features (e.g. random fourier features or Nystrom approximation). For the conventional transformer MHA via exp kernel, does the proposed approach with a tailored g(*) provide matching accuracy as conventional MHA while significantly save the memory?\n\nClarity: \n\nI think the paper is well-written in the high level. I have one major comments and a few minor comments.\n\nMajor: The discussion in sec 3.3 can be further improved to allow for easier reading. Essentially, the backward algorithm there is just a recursive / iterative to compute \\partial L / partial \\theta_{n} and \\partial L / partial \\U_{n} where \\theta_{n} is the symbolic model param for the n-th segment and \\U_{n} is the partial sum related quantities for the n-th segment. The current discussion is a bit scattered to be digested very quickly. By giving a theorem or algorithm box on the equations to compute \\partial L / partial \\theta_{n} and \\partial L / partial \\U_{n} would be useful for clarity. \n\nMinor: \n1. Line 163, typo. It should be serial time O(L/C log C)\n2. Line 178, curious when saying two forward, does it just mean a full recompute of forward activations for the current segment?\n         \n\nSignificance: This paper can have positive impact on memory efficient implementation for performers. It can also potentially lead to wider application into more transformer variants.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the attached checklist answers the questions adequately. The overhead / limitation of the method is discussed in the conclusion section.", "review_text": "", "extra_scores": {"ethical_concerns": "I do not see anything improper or having ethical issues with this paper.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6 hours"}, "extra_text": {}}, {"review_id": "93rqgytyn02", "reviewer": "Reviewer_M3kX", "summary": "This work proposes a modification of linear self-attention that allows to trade off parallel execution speed for memory usage. Authors provide extensive comparison of their method with baseline approaches in several configurations and on multiple datasets.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1) A simple but effective approach to memory usage reduction for efficient Transformers. Authors clearly demonstrate the advantages of their approach both theoretically and in experiments.\n2) In general, the paper is well-written and supplied with helpful illustrations.\n3) The work addresses a problem of memory efficiency for Transformers, which is quite important to the research community.\n\nWeaknesses:\n1) Currently, the paper is quite hard to understand due to the heavy usage of custom notation. Describing the core idea in a simpler way without the notation can considerably improve the flow of sections 2 and 3.\n2) Though the benchmarks are reasonable in general, they deviate from standard Transformer architectures in two ways. First, the experiments use no more than 3 Transformer layers, and most of today's architectures use at least 6 of them. Second, the vocabulary size of 256 is used very rarely (e.g., for byte-level models without BPE). As a result, it is not entirely clear whether the reported gains will translate into the same improvements in a more practical setting.\n\nTypos:\nL66: \"subqudratic\" -> \"subquadratic\"", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "7ghASMIYgdv", "reviewer": "Reviewer_s7XB", "summary": "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)\nThey 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.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I agree with the authors that memory limitations are an important concern and constraint for training of large-scale machine learning models. Trading memory from computational time allows you to train a model on a smaller machine, with the caveat that you potentially have to wait for a long time to get your results, but it is possible.\n\nOverall I think the paper is well written and the proposed implementation is clearly described. I am not an expert on transformer architectures so I leave the judgement of novelty and significance to other reviewers.  The only concern I have is that the approach is quite niche and I am not sure how transferrable and relevant this algorithm is for state-of-the-art architectures. \n\nSome questions I have are\n-\tYou need to be able to change the order of the nested loop in computing (5). How general is this implementation tick, can it be extended from linear to more sophisticated transformer architectures? \n-\tFor larger batch sizes a common approach to reduce memory is to do gradient accumulation. Do you need to load full sequences to do gradient accumulation or could your approach also be combined with larger batches and gradient accumulation? You would probably have issues with reading data from storage in chunks I suppose.\n-\tAre you planning to make your code publicly available or integrate it into existing transformer implementations?\n-\tYou are running your experiments on a P100 GPU. I would be curious to know the utilization of the device for different values of C to get an idea of how much sequential nature of your approach and the memory constraints are limiting your ability to fully utilize the GPU.\n-\tYou mention that this approach would allow for fine-tuning of transformer models on devices such as smartphones or microcontrollers. But then your experiments are on a P100 GPU which seems quite a bit more powerful. Do you think these experiments could actually be done on a smartphone?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers have reached a consensus in favor of accepting this paper. I concur with this consensus. I expect that the proposed changes suggested in the author response would improve the clarity of the final version of the paper.
Decision: Accept (Poster)