Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper studies a time and memory trade-off for training Performers with the linear attention mechanism. The goal of the paper is to study this trade-off under memory constrained setting, and propose a memory efficient performer training algorithm at the cost of runtime. The contributions of the paper are:

1. Given a sequence of length L, the authors proposed a dynamic programming algorithm which can process the sequence as a collection of segments of length C (the segments are coupled). This reduce the training memory for performer from O(L) to O(C) without any approximation. The intuition here is that the information across segments only flow through certain partial sums over tensors associated with each token in the sequence.

2. The authors analyzed the memory / time trade-off in the proposed algorithm (including extremal cases).

3. Empirically, the authors shows that 1) by controlling the segment length C, the proposed algorithm indeed present a trade off between runtime and memory; 2) The gradient for performer models only have minimal / negligible difference when using the conventional and the proposed algorithm; 3) A large performer model pretrained with the conventional algorithm can be finetuned via the proposed algorithm with small segment lengths; this achieves accuracy matching the conventional algorithm with less memory.


main_review: Originality:

The intuition and ideas in the proposed algorithm is novel when applies in reduced the memory of performer training.

Quality:

In 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)

1. 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?

2. 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)?

3. 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?

Clarity: 

I think the paper is well-written in the high level. I have one major comments and a few minor comments.

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

Minor: 
1. Line 163, typo. It should be serial time O(L/C log C)
2. Line 178, curious when saying two forward, does it just mean a full recompute of forward activations for the current segment?
         

Significance: This paper can have positive impact on memory efficient implementation for performers. It can also potentially lead to wider application into more transformer variants.
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_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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.
main_review: Strengths:
1) 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.
2) In general, the paper is well-written and supplied with helpful illustrations.
3) The work addresses a problem of memory efficiency for Transformers, which is quite important to the research community.

Weaknesses:
1) 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.
2) 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.

Typos:
L66: "subqudratic" -> "subquadratic"
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: 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)
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.

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.

Overall 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. 

Some questions I have are
-	You 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? 
-	For 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.
-	Are you planning to make your code publicly available or integrate it into existing transformer implementations?
-	You 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.
-	You 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?

limitations_and_societal_impact: yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: 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
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.

Some comments/questions:

The tasks studied are fairly limited in number and type. Did the author consider other tasks  (e.g. tasks from the long arena benchmark)?

The 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 

In 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?

While 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
limitations_and_societal_impact: Yes
review_title: null

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)