FlashEVA: Accelerating LLM Inference via Efficient Attention

The paper present FlashEVA, an efficient implementation of EVA and demonstrate how to finetune transformers to adapt to FlashEVA attention. The experimental results show that it can accelerate LLMs up to 6.7x with little performance drop.

The paper introduces FlashEVA, an efficient implementation of EVA to improve Transformer model performance during inference by reducing memory usage. FlashEVA allows fine-tuning of Transformers with few tokens. While it excels in many applications, it has limitations in retrieval-focused tasks. The method also offers adjustable settings for balancing throughput and accuracy, enhancing flexibility for different use cases.

• This paper proposes an efficient attention implementation via control variants method using custom CUDA and Triton kernels
• can be finetuned from transformer models with only 1.5b tokens
• show higher inference throughput and lower memory.
• suffers in retrieval focused tasks

This paper proposes FlashEVA, an efficient implementation of EVA, aimed at addressing the high memory demands of Transformer models during inference. FlashEVA allows for fine-tuning of large language models (LLMs) with minimal performance degradation while significantly reducing memory usage and increasing throughput. The method achieves this by maintaining a compressed cache of past context, which is less memory-intensive than computing the prefix during inference. However, it shows limitations in retrieval-focused tasks. The implementation offers adjustable hyperparameters to balance throughput, accuracy, and memory usage, providing flexibility for various applications.

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

The paper "FlashEVA: Accelerating LLM Inference via Efficient Attention" introduces FlashEVA, an efficient implementation of Efficient Attention via Control Variates (EVA) to reduce the memory and computational demands of Transformer models during inference. FlashEVA achieves up to 6.7x throughput improvement and 5x memory reduction, enabling fine-tuning with minimal data while maintaining strong performance across general NLP tasks. It provides flexibility through hyperparameter tuning, balancing speed, accuracy, and memory usage, and demonstrates compatibility with existing optimized attention mechanisms like FlashAttention. FlashEVA underperforms on retrieval-focused tasks due to the lack of a sliding window mechanism and faces training instabilities in larger models. The paper is well organized and written. However, the technical novelty of this paper is still limited since it looks like a modification of an existing method. While the paper briefly compares FlashEVA with the Mamba model, it does not thoroughly examine their differences or provide a clear rationale for preferring FlashEVA. I think the paper could be further improved according to the suggestions for the future submission.

Reject