Few Heads are Enough

We thank the reviewer for her/his valuable time reviewing our work. Please find our response as follows.

Clarification on the weaknesses

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.

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.

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.

Response to the questions

(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…”

(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).

General comments

We respectfully note that the provided review is very brief and lacks specific comments/criticisms on our method. This makes it impossible for us to understand the justifications for the scores of “Soundness: 1 poor”, “Contribution: 1 poor” and a final score of 3 (while we acknowledge the rather low confidence score of 3). We emphasize that our EPA provides a novel method to significantly reduce both the compute (MACs) and memory of attention in a research-friendly way. As discussed above (under “Clarifications on the weaknesses”), we rigorously tested it across various scales, datasets, and attention variants with positive outcomes. In light of this, we kindly request more detailed feedback on our methodology, or specific areas of concern that led to the current evaluation, or please consider revising the scores. Thank you for your consideration.

Thank you for getting back to the reviews. The reviewer is happy to see all experiments are finished now.

Given the current updates, the main question that affects whether the reviewer can raise the score is: Does the method achieve wall clock time speedup?

MACs is not the correct measure, as reviewer obUv points out. Alternatively, if you don't think the current implementation leads to wall clock time speedup, you should give a reasonable implementation draft/plan that can lead to wall clock time speed up. If you do not have this plan in mind, you should change your position, and write something like "future work can improve on this, given the MAC is theoretically lower." If this is the case, you should reconsider why you are mentioning flash attention, the reviewer is happy to review again if you decide to reposition.

Thank you very much for the score update. Now naturally, we'd like to know the reasons/justifications why the reviewer still votes for borderline reject.

Based on our discussion above, it seems that we thoroughly addressed and resolved all the concerns the reviewer had initially raised. We would like to respond to any remaining concerns. Thank you.

We would like to thank the reviewer for his/her valuable time.

We are surprised by the reviewer’s comment on the quality of writing (“the poor writing and underwhelming evaluation really makes it hard to appreciate it.”, “needs complete re-write”). In contrast, Reviewer DLPx highlights the clarity as one of the strengths of the paper in his/her high-quality review, and even Reviewer dyhc states that "the method is clearly illustrated". We’d like to hear more specific clarifications/suggestions on what confused the reviewer (including the comment regarding the notation).

We would like to point out that the reviewer’s summary does not reflect what our algorithm does. We have a few heads (not 1), and the value and output projections (not “Q, K, and V”) are a mixture of experts.

Regarding the timing results, we report the commonly used metric which is the number of “Multiply-Accumulated (MAC) operations” required for each method. We do this by 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. FlashAttention does not change the number of MAC operations, so it is the same as the baseline dense Transformer. If the reviewer is curious about the wall-clock time of our current implementation, we also measured the resource usage for both the baseline Transformer and EPA: please refer to Appendix 4 of our updated manuscript for more details. In short, EPA with the current (suboptimal) implementation is already capable of accelerating the entire training pipeline by a factor of 1.5-1.6, and reducing the total memory usage by about the same factor.

We kindly ask the reviewer to clarify what “fine-tuning overheads” we should talk about. In fact, because our method accelerates the Transformer architecture, it accelerates both the training and fine-tuning phases.

Overall, we find the current review very brief, and lacks technical comments/criticisms about our method. The reviewer rates our work as 3 (reject) with a high confidence score of 4; this implies responsibility. We respectfully ask for much more specific criticisms and justifications for such a rating.

Thanks for the clarification.

I appreciate the clarification. But it brings back the question. MACs are not a great metric. It depends heavily on the hardware.

"Question regarding fine-tuning" : I am asking a very simple question, to use your method on say a model we would need to first perform some amount extra training or not ?

Despite all three reviewers pointing out how your method is different than the goal of flash attention, i find it interesting that you are not willing to change your position. You have a approximation based method which prunes the number of heads in transformer layer with gating (at a high level), flash attention concerns with build a hardware friendly implementation of exact attention operation. Two very different things in my opinion (of course you can disagree), which is why all reviewers have been pointing out that there is fundamental difference. "Research Friendliness" is a subjective metric. I find it very hard to digest the positioning of your method.

We would like to thank the reviewer for the insightful review, and for positive comments on the clarity and methodology of the paper. Please find our responses as follows:

Clarification on the weaknesses

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.

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:

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.

For the EPA algorithm, the paper mentions that K/Q source experts are not necessary for good results and only output/value experts are required, which seems to contradict the disadvantages shown in 2.2 naive algorithm.

We are not sure to understand what “contradiction” the reviewer refers to. Could you please clarify?

Response to the questions

Can the authors compare against architectures other than TransformerXL? It is not evident from the text why only 1 architecture is chosen for comparison?

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

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?

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.

We believe we’ve addressed all the concerns raised by the reviewer. If you find our response convincing, please consider increasing the score. Otherwise, we’ll be happy to discuss and/or clarify further.