Addressing novelty criticism
We would like to thank all the reviewers for their constructive feedback. We are happy to hear the that you have found the paper easy to read, and the experimental results comprehensive. On the other hand, most have raised objections regarding the novelty of the method, which is what we would like to address.
We agree with the reviewers that the proposed method is simple and can appear incremental over the existing EVA approach. Nevertheless, we would argue that this does not mean that it is not useful and of interest to the community.
First, we propose to apply the EVA framework in a different setting as the original authors. While they have mainly focused on the training from scratch, we focus on adaptation of pretrained transformers into linearized variants in order to improve model inference efficiency. For this, we leverage the EVA framework due to its Softmax attention approximation guarantees.
Futhermore, we would like to highlight two crucial problems that have limited an adoption of efficient attention approaches, especially for the model inference.
1. There is a need for efficient Implementations in order to be able to materialize the theoretically promised speedups of these methods. By showing that EVA attention can be reformulated as a Softmax attention over modified set of keys and values, we can leverage community improvements to obtain a very efficient implementation of the EVA attention. This allows FlashEVA to optimize inference in two ways: for the prompt processing, it leverages the lower computational complexity of the efficient attention, while for the autoregressive generation/decoding, it benefits from the lower memory footprint, which is the main bottleneck in the decoding process.
2. Linearized attention approaches have generally faced a performance degradation compared the transformers using standard softmax attention. EVA attention has solid theoretical grounding and with its hyperparameters we can tune how well we will approximate softmax attention. In fact, our method consistently outperforms Diljang attention over all model scales, showing that the EVA framework is adequate for this task. Nevertheless, as mentioned, all linear attention approaches will suffer a decreased performance on the retrieval focused tasks, as that is a fundamental limitation of the linearized approach [1]
Given the above points, we feel that the FlashEVA method is a contribution of interest to the wider community.
Furthermore, we would like to expand and contextualize further the main experimental results. The main comparison baseline is the Diliang attention, a recent approach showing comparable performance to Softmas attention transformers on inference, with substantial efficiency improvements. As can be seen, our method outperforms consistently the Dijiang attention for all model scales. In fact, at larger model scales, in some tasks we achieve up to 5% better accuracy. Indeed, one of the limitations of the Dillang method we observed was its variability in performance, drastically underperforming the baseline Transformer in certain tasks.
We additionally compared FlashEVA to a Sliding Window attention baseline, with a larger sliding window (to match the additional Keys and Values in the FlashEVA approach). Note, in the main results of FlashEVA we did not use a sliding window approach, which would likely improve performance (the reason was that we observed some numerical instabilities when testing this on larger models, so for the sake of saving resources, we only used the local window approach in the main experiments). Only in the largest 18 model size we observe some performance regression compared to that baseline, however, that is mainly driven by subpar performance on the WSC and Winogrande tasks, which generally contain quite short sequences. For those sequences, the Sliding Window attention basically acts like full attention over the whole sequence (also seen in the fact that it basically achieves the same performance as the full attention variant)
Finally, in terms of downstream task performance, we also do not expect to achieve higher performance as the baseline transformer, as we are merely trying to efficiently approximate the full attention.
We have thus shown a method that matches better baseline Transformer on downstream tasks, which has a principled way of trading off some performance for efficiency, and has practical speedups in the model generation.
Having said all that, we understand that the novelty criticism is valid, and ICLR might thus not be the best venue for this submission. We will take into account all the useful feedback (e.g. improving the paper writing to describe the FlashEVA better, introduce FlashAttention, more concrete comparison to Mamba, etc.) in any future submission.
[3] Jelassi, S., Brandfonbrener, D., Kakade, S. M. & Malach, E. Repeat After Me: Transformers are Better than State Space Models at Copying. (2024) ArXiv: http://arxiv.org/abs/2402.01032