Paper: Few Heads are Enough (ICLR-2024)

Reviewer comment (Novelty):
The author seems to misunderstand the position of flash-attention, see questions below.

Author response:
> 1. The author seems to misunderstand the position of flash-attention, see questions below.

This is a false accusation (we do not understand what caused the reviewer think of this), as we clarify below.

Reviewer comment (Experiments):
The scale of experiment is small, how would this method generalize to larger models such as llama?

Author response:
> 2. the scale of experiment is small, how would this method generalize to larger models such as llama?

If the reviewer is asking us to run llama scale experiments, that is an unreasonable request.
The cost of training llama-scale models is typically estimated as several millions of dollars. This is clearly not feasible for a university lab, nor necessary to show the core principle of an idea: in fact, we already had invested much compute for all these experiments; it is not as if we only reported on small/toy tasks such as Penn TreeBank. We report results on two scales, 47M and 262M parameters, and two different attention mechanisms (Transformer XL and RoPE) on four different datasets, some of them typically used for training much bigger models: Enwik8, Wikitext 103, C4 and peS2o. We think this is more than enough to demonstrate the feasibility of the model, and large companies with access to a sufficient amount of computing power can scale it up if interested.

Reviewer comment (Experiments):
Some experiments have not been finished (Table 4).

Author response:
> 3. Some experiments have not been finished (Table 4).

We are sorry about this (we had some technical problems with our cluster just before the deadline). The experiments have finished, and we updated Table 4 in the updated PDF.
The results confirm that our model outperforms the baseline in every case.

Reviewer comment (Evaluation):
In what sense flash-attention reduces compute. Do you mean FLOP or wall clock time? FA is exact attention and does not reduce FLOP, it is just a series of clever fusion. If it is wall clock time, then this paper should keep the definition consistent, and provide wall clock time analysis instead of MAC.

Author response:
> (1) In what sense flash-attention reduces compute. Do you mean FLOP or wall clock time? FA is exact attention and does not reduce FLOP, it is just a series of clever fusion. If it is wall clock time, then this paper should keep the definition consistent, and provide wall clock time analysis instead of MAC.

We report MACs by consistently following our main baseline paper (MoA, Zhang et al, 2023: Mixture of Attention Heads: Selecting Attention Heads Per Token). MAC is a good metric as it is not implementation-dependent, while the actual wall clock time is highly implementation-dependent. This having said, we suspect that what might have caused the confusion was a sentence in the abstract “...Flash-Attention reduces both compute…”. We admit that this formulation was confusing. We changed it to “...Flash-Attention reduces both run-time…”

Reviewer comment (Writing):
"Unlike FlashAttention (Dao et al., 2022), it is research-friendly, because it does not hide the internal details of the attention mechanism inside a CUDA kernel." is an either arguably wrong or highly subjective judgement to flash-attention.

Author response:
> (2) "Unlike FlashAttention (Dao et al., 2022), it is research-friendly, because it does not hide the internal details of the attention mechanism inside a CUDA kernel." is an either arguably wrong or highly subjective judgement to flash-attention.

This is not at all wrong or subjective.  As we also clarified in our response to Reviewer DLPx (explaining why our method is “research friendly” when we also need a CUDA kernel), any modification to the “attention mechanism” in FlashAttention requires proficiency in GPU programming: this is unquestionable.

For example, consider combining Transformer XL style relative positional encoding with FlashAttention. It requires loading positional encodings with different offsets for each row of the attention matrix, and combining it with the pre-computed query; constructing non-trivial memory access patterns that enable the reuse of positional encodings already loaded in shared memory… While all of this is doable, it is by no means trivial. Please check the Flash Attention GitHub repository (https://github.com/Dao-AILab/flash-attention) consists of 4500 lines of highly optimized CUDA code (an alternative Triton implementation around 1000 lines also exists).
In contrast, our method can be trivially combined with PyTorch implementations of RoPE , Transformer XL style attention, Alibi, softmax1 activation, (and also with Flash Attention), etc.  None of these require any modification to our CUDA kernel, because the CUDA kernel is only used for the V and O projections (and optionally K and Q).

Reviewer comment (Experiments):
The paper refers to FlashAttention multiple times and compares against their CUDA kernel (SW designed to exploit HW efficiently) optimization vs algorithmic insight in this paper. I am not sure if its an apple-to-apple comparison since there are tons of other literature for transformers which aim to reduce computation/memory cost (like quantization/sparsity methods) and the paper doesn’t compare against these.

Author response:
We’d like to clarify one very important point regarding the goal of this work. The general goal of typical research on efficient attention is to make/train/evaluate Transformer models faster. Here we emphasize that such methods also have another important potential consequence in **accelerating research** on attention based models (by making them resource efficient). We consider this as an extremely important aspect today, potentially allowing for “resource-limited” organizations (e.g., in academia) to still conduct research on reasonably large models.
However, for a “fast attention” method to be effective in accelerating research, it also has to be “research friendly”: certain existing methods do not allow modifications of certain aspects of attention mechanisms, while others may be more flexible. Such flexibility is required for “research friendliness”.

This perspective is at the heart of our comparison with FlashAttention or with any other methods. We refer to FlashAttention simply as an example of a non-flexible acceleration method (and because it is arguably the most popular “fast attention” method today): it does not allow a researcher to modify the attention mechanism itself (which is hard-coded within their CUDA kernel), while our method still leaves such a possibility (since our CUDA kernel only replaces the linear projection layers to obtain query/key/values)---This should also bring an answer to the reviewer’s second point:

Reviewer comment (Reproducibility):
While authors compare against FlashAttention custom kernel implementation and mention that as a drawback, EPA algorithm itself requires a custom CUDA kernel with its own set of restrictions (pointed in the results section).

Author response:
This perspective is at the heart of our comparison with FlashAttention or with any other methods. We refer to FlashAttention simply as an example of a non-flexible acceleration method (and because it is arguably the most popular “fast attention” method today): it does not allow a researcher to modify the attention mechanism itself (which is hard-coded within their CUDA kernel), while our method still leaves such a possibility (since our CUDA kernel only replaces the linear projection layers to obtain query/key/values)---This should also bring an answer to the reviewer’s second point: 
> While authors compare against FlashAttention custom kernel implementation and mention that as a drawback, EPA algorithm itself requires a custom CUDA kernel with its own set of restrictions

For example, FlashAttention can **not** be used with certain positional encoding (e.g. XL-style relative positional encoding) or activation functions (e.g. geometric attention). In contrast, our method can be trivially combined with RoPE, Transformer XL style attention, Alibi, (and also with Flash Attention), etc. None of these require any modification to our CUDA kernel.

Reviewer comment (Experiments):
Can the authors compare against architectures other than TransformerXL? It is not evident from the text why only 1 architecture is chosen for comparison?

Author response:
We’d like to draw the reviewer’s attention to the fact that our experiments already include an architecture other than TransformerXL: standard (pre-layer norm) Transformer with RoPE positional encoding in Sec. 4, page 8.
The main reason why we chose Transformer XL as our main architecture is that it has significantly better performance with a moderately low context window size (we use context windows of 256 for small models and 512 for big). For RoPE, we double the context window size of the XL transformers (512 for small, 1024 for big) to obtain competitive models (and even in this case, they underperform the XL models).

Reviewer comment (Experiments):
It is not evident from the paper how nHead is chosen for a task. Most results demonstrated fixed the nHead to be 2 or 4. Did the authors perform smaller experiments to first search for optimal nHead before scaling up?

Author response:
No, that is not how we chose $n_{head}$. Our search for $n_{head}$ works as follows: we start with $n_{head}=2$, and if the resulting model performs worse than the baseline, we double it to $n_{head}=4$. We never went beyond $n_{head} > 4$, as 2 or 4 was always enough to obtain good models. We added the corresponding explanation in the updated PDF in Appendix 2.1. Thank you for pointing this out.