IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs

One limitation of existing Transformer-based models is that they cannot handle very long sequences as input since their self-attention operations exhibit quadratic time and space complexity. This problem becomes especially acute when Transformers are deployed on hardware platforms equipped only with CPUs. To address this issue, we propose a novel method for accelerating self-attention at inference time that works with pretrained Transformer models out-of-the-box without requiring retraining. We experiment using our method to accelerate various long-sequence Transformers, including a leading LLaMA 2-based LLM, on various benchmarks and demonstrate a speedup of 2.73 × -7.63× while retaining 98.6% -99.6% of the accuracy of the original pretrained models. The code is available on our project website at https://yuzhenmao.github.io/IceFormer/ . INTRODUCTION Transformers (Vaswani et al., 2017) have powered incredible advances in NLP, as exemplified by large language models (LLMs) such as GPT-4 and LLaMA 2. Increasingly LLMs are applied to exceptionally long input sequences, which enables many exciting applications such as long-form content creation, extended conversations, and large document search and analysis (OpenAI, 2023; Anthropic, 2023). While LLMs can be feasibly trained with expensive hardware accelerators (e.g. GPUs), they need to be deployed on commodity devices, which may only be equipped with CPUs. However, it is currently challenging to deploy LLMs on CPUs due to their high computation cost (Dice & Kogan, 2021) . A significant computational bottleneck arises from the self-attention mechanism that is integral to Transformers -both time and space complexity are quadratic in the sequence length. This problem is exacerbated in the context of LLMs, which are often used on very long sequences. To handle long input sequences, there has been substantial research into reducing the quadratic time complexity of self-attention -these methods are collectively known as efficient Transformers. However, many do not meet the needs of LLMs and are therefore difficult to apply to LLMs. An ideal acceleration method for LLMs should satisfy four criteria: (1) No retraining -the method should not require the model to be retrained, given the enormous computational expense of training LLMs; (2) Generality -the method should be applicable to a variety of LLMs, rather than just those trained with particular constraints built-in; (3) High accuracy -the method should not introduce large approximation errors, since LLMs have many attention layers and so errors from earlier layers can compound; (4) Fast inference -the method should achieve fast test-time performance. Satisfying all these criteria simultaneously is difficult, and to our knowledge no existing methods can do so. For example, Transformers with fixed attention patterns, e.g., Longformer (Beltagy et al., 2020) , require retraining the model before they can be used. Reformer (Nikita et al., 2020) requires keys to be normalized -this requirement is not met in most pretrained models. Nyströmformer (Xiong et al., 2021) and LARA (Zheng et al., 2022) do not support causal masks, which are commonly found in LLMs. Low-rank methods such as Performer (Choromanski et al., 2020) introduce substantial approximation errors, especially when they are not retrained/finetuned.