IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs

Thank you for your valuable review. Please find our responses as follows.

Q1: The adaptation of the k-NNS method may potentially make the computation less parallelizable. It is essential to consider the scalability of the proposed method, even on CPU-based platforms. Leveraging multi-threading and distributed computing for scalability is a relevant concern.

A1: Actually k-NNS does not affect the parallelizability. At the level of tokens, neither vanilla self-attention nor k-NNS parallelizes over them in decoder-only models because they are autoregressive. At the level of heads and token dimensions, both vanilla self-attention and k-NNS can parallelize over them in a straightforward way. In our implementation, we do utilize multi-threading within the k-NNS algorithm. To illustrate the scalability of the proposed method, we added the new experiments as per the reviewer’s suggestions and present the results below.

Q2: It would be valuable if the paper provided insights into the scalability of the proposed method in terms of performance (compute time) on different CPU platforms with varying numbers of cores. How well can the approach be parallelized or scaled up on different hardware configurations?

A2: We appreciate the reviewer’s insightful suggestion and tested the proposed IceFormer on a different machine. We reran all the experiments of ZeroSCROLLS benchmark on an AMD Ryzen 9 5950X CPU with 16 CPU cores (where we used 32 threads in total to take advantage of simultaneous multithreading), and compare to the results in the original paper, which were obtained on an AMD Ryzen 9 5900X CPU with 12 cores (where we used24 threads in total):

As shown by the results in the table above, Iceformer can scale effectively to CPUs with more cores. It is notable that the speed-up of Iceformer applied to Vicuna-7b-v1.5-16k compared to the vanilla Vicuna-7b-v1.5-16k model is larger on a CPU with more cores (16) than the CPU used in for the experiments in the original paper, namely a CPU with 12 cores.

Thank you for your valuable review. Please find our responses as follows.

Q1: The proposed method seems orthogonal to the inference platforms. Why do the authors focus on the CPU settings? What about other inference engines, such as GPUs, TPUs, and other accelerators?

A1: Since our main focus is on the deployment of Transformers to end-user devices, we optimized our method for CPUs, since end-user devices are often not equipped with GPUs.

Other hardware platforms such as GPUs and TPUs are fundamentally different from CPUs, with GPUs and TPUs being able to execute a much larger number of parallel threads than CPUs, but also incurring a much higher communication overhead. Therefore, changes need to be made to take advantage of the GPU/TPU features and avoid the performance penalties specific to GPUs/TPUs. Therefore, the optimization approaches that are applicable to CPUs vs. other platforms are very different. On CPUs, the focus is on maximizing cache efficiency and parallelism, both at the instruction level (SIMD) and the thread level (multi-threading). In contrast, on GPUs, the focus is on effectively utilizing the large number of threads without unnecessary communication and synchronization overhead and optimizing accesses to different levels of the memory hierarchy. There isn’t much documentation or literature on TPUs, so effective optimization approaches are somewhat unknown. As a result, performance optimization for different hardware platforms are typically treated separately in the literature, with some papers focusing specifically on CPU optimization, e.g., [1] and others focusing specifically on GPU optimization, e.g., [2].

While it may be possible to extend our method to GPUs, very different optimization approaches will need to be used and changes to the algorithm may be required to minimize communication overhead between different thread blocks. This is beyond the scope of our paper, and we leave it to future work.

Q2: The authors discuss many related methods in Section 2. However, only few of them are used as the baselines in empirical comparisons? Could the authors add more baselines? Some methods may need retraining. It is better to list more results even if some baseline needs retraining or specific architectures.

A2: We added Linformer [3], Performers [4] and Longformer [5] as additional baselines. Please see Table 1 and Table 3 in our updated manuscript for details.

Q3: Could we apply the proposed method into training?

A3: Yes, it is possible to apply the proposed method to training. Specifically, the same procedure can be used to speed up the forward pass. We did not focus on this use case, however, because we focused on deployment to end-user devices, whereas training is typically done centrally in the cloud.

Q4: It is not clear how IceFormer will perform on short and medium length of sequences.

A4: We showed performance on different input lengths in Figure 5. As shown, IceFormer is consistently faster and achieves comparable accuracy compared to the vanilla Vicuna-7b-v1.5-16k model across all input lengths.

Q5: Could the authors visualize Tables 1 and 2?

A5: Thanks for the suggestion - we added visualizations to the new manuscript. Please refer to the appendix section D.

Q6: What are the limitations and potential negative impact of the proposed method?

A6: One limitation is that the value of $k$ needs to be picked manually for each Transformer to best balance efficiency and accuracy. However, this is a one-time cost paid before deploying the model, so it does not affect the performance of the method on the end-user device.

A potential negative impact is that making Transformer inference faster may encourage the development of more inefficient Transformer architectures, because inference speed using such inefficient Transformer architectures may become more tolerable after applying our method.

Q7: What may be the future directions and extensions?

A7: One future direction is on developing an extension that can handle non-standard attention masks, beyond simple causal masks explored in this paper. To do so, keys must be inserted into and possibly deleted from the k-NNS database in the appropriate order, so that keys that are masked out for a query cannot be selected by that query.

Q8: The authors mentioned that the code will be released upon acceptance. I do not find implementation in the supplementary material. Is it possible to release it for the reviewers' reference?

A8: The process for releasing the code takes some time, and is unfortunately notready by the rebuttal deadline. But we are happy to answer any questions related to the code.

Q3: Presently, in the realm of cloud-based LLMs, MQA/GQA (https://arxiv.org/pdf/2305.13245.pdf) emerges as the leading approach for attention layer optimization. I would be interested in hearing your views on this.

A3: MQA/GQA is a promising approach for attention layer optimization and works by reducing the number of query and key heads. MQA uses a single set of key and value heads for all query heads, whereas GQA utilizes a single set of key and value heads for each group of query heads. Our method is general-purpose and can be applied independently of the architecture – in particular, it can be applied regardless of the specific number of key, value and query heads. Hence, it can be applied on top of MQA/GQA to achieve a further speedup on top of the benefits offered by MQA/GQA.As of today, MQA/GQA are still relatively new, so their effectiveness in other Transformer models have not been fully demonstrated. As they gain greater adoption, our method can be applied on top of MQA/GQA to further enhance their efficiency.

References:

[1]: Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. 2022. Efficient Transformers: A Survey. ACM Comput. Surv. 55, 6, Article 109 (June 2023), 28 pages. https://doi.org/10.1145/3530811

[2]: Quentin Fournier, Gaétan Marceau Caron, and Daniel Aloise. 2023. A Practical Survey on Faster and Lighter Transformers. ACM Comput. Surv. 55, 14s, Article 304 (December 2023), 40 pages. https://doi.org/10.1145/3586074

[3:]{https://lilianweng.github.io/posts/2023-01-27-the-transformer-family-v2/}

[4:]{https://lilianweng.github.io/posts/2023-01-10-inference-optimization/}

[5]: Touvron, Hugo, et al. "Llama 2: Open foundation and fine-tuned chat models." arXiv preprint arXiv:2307.09288 (2023).

[6:]{https://www.mosaicml.com/blog/mpt-7b}

[7:]{https://www.together.ai/blog/redpajama-7b}

[8:]{https://huggingface.co/tiiuae/falcon-7b}

Thank you for your valuable review. Please find our responses as follows.

Q1: The most glaring issue is its reliance on outdated methods from 2020 and 2021, with some even referencing the 2017 vanilla transformer. This dated focus suggests a lack of recent advancements in the field. One must ponder why the topic of 'efficient transformers' isn't garnering contemporary attention. Historically, efforts to enhance transformer efficiency via attention layer optimization waned with the introduction of Large Language Models (LLM). Moreover, these methods face challenges in adhering to the scaling-law (https://arxiv.org/abs/2207.10551).

A1:

1. It is not true that we only compared to outdated methods. In fact, we compared to LARA, which is from 2022. To our knowledge LARA is the latest method in the literature. We only included the vanilla Transformer as a baseline because prior work did - it is not the comparison of primary interest.

2. Actually efficient Transformers are garnering contemporary attention – there have been multiple survey papers on the topic in 2022 and 2023 [1,2] and leading researchers from OpenAI working on LLMs have written multiple tutorials on the topic in 2023 [3,4]. So efficient Transformers are actually quite topical and are increasingly important due to the development of LLMs. The reason why research efforts waned with the introduction of LLMs is not due to the lack of importance, but because existing methods require changing the architecture and retraining, which is not feasible for LLMs due to the extremely high computational expense of doing so. Devising a way of improving the efficiency of attention layers without changing the architecture has historically been difficult, since it requires developing a new computational technique to perform the same attention operation used in vanilla Transformers, as opposed to devising a different attention operation. This is precisely the problem we tackle – our method does not need architectural changes or retraining and approximates the same attention operation used in vanilla Transformers without introducing extra parameters, and yet it can still achieve higher efficiency.

3. We note that our paper does not propose a new efficient Transformer architecture – instead it proposes a general-purpose technique for accelerating Transformers without changing the architecture. The referenced paper (https://arxiv.org/abs/2207.10551) studied the scaling behaviour of different Transformer architectures. It found that only switch Transformer, GLUTransformer and Funnel Transformer achieve comparable results to vanilla Transformer in terms of scaling. This in fact reinforces the need for our method – our method can be applied to speed up the vanilla Transformer without needing to change the architecture, and so the efficiency gain we achieve is on top of the scaling benefits offered by the vanilla Transformer, rather than in place of them. Our method can also be applied to speed up the switch Transformer, GLUTransformer and Funnel Transformer, since the changes they make relative to the vanilla Transformer do not affect the self-attention layer, which remains unchanged. Therefore, our method is orthogonal to particular Transformer architectures and can be combined with the Transformer architectures with good scaling behaviour to realize the benefits offered by each.

Q2: Although the paper touts the efficiency of LLM, it primarily showcases small-scale toy examples or, at best, the 7B LLaMa model. Notably, the latter, when applied to a long context, demands an extensive duration, upwards of several seconds, merely for the attention computation. The practical relevance of such scenarios is debatable. (Honestly, I'm confusing about the exact meaning of 'wall clock time of attention module'.)

A2:

We respectfully disagree with the reviewer on the characterization of the 7B LLaMa 2 model as "small-scale" or a "toy example". The 7B LLaMa 2 model is in fact considered a large language model (LLM) by the original authors [5] and is commonly built upon by other LLMs [6,7,8], so it is far from being a "toy example".

In general, it is not meaningful to speak of latency outside of the context of the latency of the original model that our method is used to accelerate. Originally the model took ~5 seconds, and our method reduced it to ~2 seconds. The original model used Intel MKL to perform the attention operation, which is a highly optimized state-of-the-art implementation of BLAS, so the fact that our method is able to achieve a speedup of ~2.5x relative to MKL is quite significant. Furthermore, our work represents just one step towards this direction – future work may build on our work to achieve further speedups.

“wall clock time of attention module” refers to the length of time it takes in the real world to perform attention.

Thank you for your valuable review. Please find our responses as follows.

Q1: The paper does not adequately support its claim that Prioritized DCI outperforms LSH, lacking both theoretical and empirical evidence.

A1: From a theoretical perspective, the query time complexity of Prioritized DCI can be found in the original Prioritized DCI paper [1], which is $O\left(d \max \left(\log n, n^{1-m / d^{\prime}}\right)\right.$. The query time complexity of LSH can be found in the original LSH paper for Euclidean spaces [2], which is $\approx O\left(d n^{1 /(1+\epsilon)^2}\right)$. The epsilon denotes the maximum amount of error that is tolerable, and since it is important to tell apart similar keys that are ranked differently in our use case, epsilon needs to be very small, which makes LSH very slow.

From an empirical perspective, we incorporated LSH into the kNNS comparison experiment detailed in section 5.1, and consequently updated Figure 3 in the new manuscript. As shown, Prioritized DCI substantially outperforms LSH in terms of both efficiency and recall. The original Prioritized DCI paper also included an empirical comparison with LSH in its experiments section. For more details, we kindly refer the reviewer to the literature mentioned above.

Q2: There is insufficient clarity in the algorithm's implementation details, making it difficult to understand the actual complexity and the mechanics of the proposed method.

A2: We thank the reviewer for pointing this out. To improve the clarity, we have included the pseudocode for our method in the updated appendix section A in the new manuscript.

Q3: The evaluation methodology for measuring inference time is not comprehensive. It appears the method is optimized for CPUs but lacks evidence of similar efficacy on GPUs.

A3: CPU code cannot be run on the GPU and vice versa, so one cannot simply evaluate performance of a CPU implementation on the GPU without creating a separate implementation optimized specifically for GPUs. CPUs and GPUs are fundamentally different hardware platforms, with GPUs being able to execute a much larger number of parallel threads than CPUs, but also incurring a much higher communication overhead. Therefore, changes need to be made to take advantage of the GPU features and avoid the performance penalties specific to GPUs. Therefore, the optimization approaches that are applicable to one platform vs. the other are very different. On CPUs, the focus is on maximizing cache efficiency and parallelism, both at the instruction level (SIMD) and the thread level (multi-threading). In contrast, on GPUs, the focus is on effectively utilizing the large number of threads without unnecessary communication and synchronization overhead and optimizing accesses to different levels of the memory hierarchy. As a result, performance optimization for CPUs vs. GPUs are typically treated separately in the literature, with some papers focusing specifically on CPU optimization, e.g., [3] and others focusing specifically on GPU optimization, e.g., [4].

Since our main focus is on the deployment of Transformers to end-user devices, we optimized our method for CPUs, since end-user devices are often not equipped with GPUs. While it may be possible to extend our method to GPUs, very different optimization approaches will need to be used and changes to the algorithm may be required to minimize communication overhead between thread blocks. This is beyond the scope of our paper, and we leave it to future work.

Q4: Are the baseline results for execution time from original papers or are they reproduced results specific to CPU performance?

A4: As reflected in our paper’s title, i.e., “Accelerated Inference with Long-Sequence Transformers on CPUs”, we focus on deployment of Transformers to end-user devices equipped only with CPUs, so the baseline are the results of running the official code accompanying the original papers or widely-used HuggingFace implementations on the CPU.

Q5: Could you provide the typical and range of values for k and m used in the experiments?

A5: We list the range of values for $k$ in Table 6 in the appendix for the LRA benchmark, and in Eq. 29 under Section E of the appendix, for the LLM experiments. The sequence length $m$ of the LRA benchmark can be found in Table 4 in the appendix. The $m$ of the LLM experiments can be found in Table 3 and Figure 5.

Thank the authors for the rebuttal. With the clarified scope of this paper being CPU-oriented, I am okay with the revision and will raise my score.

Thank you for your valuable review. Please find our responses as follows.

Q1: The use of Prioritized DCI k-NNS algorithm needs more experimental or theoretical justification. The authors claim "ranking-based algorithms is better aligned with how attention weights", if so, how would other ranking-based algorithms perform?

A1: There have only been two ranking-based algorithms, vanilla DCI [1] and Prioritized DCI [2]. To respond to the reviewer’s question, we have incorporated both algorithms into the k-NNS algorithm comparison detailed in Section 5.1, and updated Figure 3 accordingly. As shown in the figure, Prioritized DCI is significantly more efficient than vanilla DCI.

Q2: On top of that, the authors show an evaluation of different kNNS algorithms on fashion-mnist-784 dataset in Section 5.1. It would be better to show the exact setup (eg. model architecture) they used, and compare them on a few more tasks (for instance, text generation, token classification).

A2: For the evaluation of different k-NNS algorithms, we followed the evaluation protocol of the ANN benchmarks (https://github.com/erikbern/ann-benchmarks). More details can be found on the website of ANN benchmarks. ANN benchmarks is a general benchmark for k-NNS algorithms, and is not specific to a particular model architecture. It also includes the references for each k-NNS algorithm tested.

As per the reviewer’s request, we compare the second best k-NNS algorithm as measured by ANN benchmarks in Figure 3, on two tasks, the text classification task in the LRA benchmark and the GvRp (text generation) task in the ZeroSCROLLS. The second best k-NNS algorithm is the hnsw(nmslib) algorithm. The latency for hnsw(nmslib) is 1.06 seconds on the text classification task and 15.76 seconds on the GvRp task, significantly slower than the vanilla self-attention (0.69 seconds and 5.39 seconds) and Iceformer (0.09 seconds, and 2.24 seconds). This is because evaluating each attention layer requires constructing a new k-NNS database, which takes a long time forhnsw(nmslib) and many other k-NNS algorithms. On the other hand, constructing a new k-NNS database is much faster with Prioritized DCI. This result is consistent with our k-NNS comparison results in Figure 3 and the literature [3, 4].

Q3: Fig.5's y-axis needs to be annotated with units (especially for latency).

A3: We followed the reviewer’s suggestion and added the units for latencyin Figure 5 of the updated manuscript.

Q4: Why is CPU the target platform for IceFormer? Are there CPU-specific architectural optimizations, on top of the time/space saved due to sparse attention?

A4: Since our main focus is on the deployment of Transformers to end-user devices, we optimized our method for CPUs, since end-user devices are often not equipped with GPUs. We used various CPU-specific optimizations to minimize the wall-clock time, namely instruction-level parallelism (SIMD), bitwise operations and cache-optimized data structures.

Q5: I hope to clarify with the authors: in the long-context evaluation (Fig.5), is the baseline (vanilla Transformer) referring to the Vicuna-7b-v1.5-16k model, or the original Transformer model?

A5: We thank the reviewer for pointing this out. The baseline (vanilla Transformer) refers to the vanilla Vicuna-7b-v1.5-16k model. We have fixed the labels in Table 3 and Figure 5.

Q6: This paper describes a novel method for improving the inference efficiency and evaluates on CPU. But it seems that this method could potentially also apply to GPU, which is more often used. Are there specific reasons/constraints to use CPU?

A6: While there is certainly potential to extend our method to GPUs, very different optimization approaches will need to be used and changes to the algorithm may be required to minimize communication overhead between different thread blocks. This is beyond the scope of our paper, and we leave it to future work. As mentioned earlier, our implementation uses CPU-specific optimizations, such as instruction-level parallelism (SIMD), bitwise operations and cache-optimized data structures. Different optimizations will need to be used to minimize the wall-clock time on GPUs.

References:

[1]: Li, Ke, and Jitendra Malik. "Fast k-nearest neighbour search via dynamic continuous indexing." International conference on machine learning. PMLR, 2016.

[2]: Li, Ke, and Jitendra Malik. "Fast k-nearest neighbour search via prioritized DCI." International conference on machine learning. PMLR, 2017.

[3]: Echihabi, Karima, et al. "The lernaean hydra of data series similarity search: An experimental evaluation of the state of the art." arXiv preprint arXiv:2006.11454 (2020).

[4]: Echihabi, Karima, et al. "Return of the lernaean hydra: Experimental evaluation of data series approximate similarity search." arXiv preprint arXiv:2006.11459 (2020).