IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs
Authors: Yuzhen Mao
Martin Ester
Ke Li
Venue: ICLR 2024 poster (2024)
Keywords: Efficient Transformers
Inference-time Efficiency
CPU
Primary area: generative models

Abstract: 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 greater speedup of $2.73\times$ - $7.63\times$ 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/.

Published as a conference paper at ICLR 2024
ICEFORMER: ACCELERATED INFERENCE WITH LONG- SEQUENCE TRANSFORMERS ON CPUs
Yuzhen Mao, Martin Ester, Ke Li School Computing Science, Simon Fraser University Burnaby, BC V5A 1S6, Canada (yuzhenm, ester, keli Jasfu.ca
ABSTRACT
One limitation of cxisting 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 novel method for accelerating self-attention 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 leading LLaMA 2-based LLM, on various benchmarks and demonstrate a speedup of 2.73 x while retaining 98.6% 99.6% of the accuracy of the original pretrained models. The code is available on our project website at htttps:: //tttps://yuzhermm.. thub. ic http/eeeerrr//////////o///
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 (OpenAl 2023 Anthropic 2023). While LLMs can be feasibly trained with expensive hardware accelerators (c.g. GPUs), they need to be deployed on commodity devices, which may only be equipped with CPUs.
However, it currently challenging to deploy LLMs on CPUs due their high computation cost (Dice & Kogan 2021). A significant computational bottleneck arises from the self-attention mechanism that is integral 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 meet the needs of LLMs and are therefore difficult apply LLMs.
An ideal acceleration method for LLMs should satisfy four criteria: (1) No retraining the method should not require the model 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 emror 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 al. 2021) and LARA (Zheng et al. (022) 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.

---

Published as a conference paper at ICLR 2024
Vanilla Attention Matrix
Transformer
output
4:
k
Phase
Phase 2
N-S
Attention Matrix
loeFormer
output
top-h attertiio--wightts
inner-product
Figure 1: Comparison between Transformer (Vaswani al. 2017) (top row) and the proposed method, IceFormer (bottom row). We illustrate with one query and k 2 in k-NNS. In the two attention matrices presented, the top-2 largest attention weights in each color. The remaining attention weights are shown in a pale color in the vanilla attention matrix, and are set to zero (depicted in white) the sparse attention matrix.
represented by dark
In this paper, we propose acceleration method, which we dub IceFormer due to its ability to be applied directly frozen models without retraining, that simultaneously satisfies the above four criteria. Specifically, IceFormer (I) does not require retraining, (2) can be applied to most LLMs, (3) can approximate vanilla attention accurately, and (4) achieves significantly faster inference speeds compared to existing methods. We illustrate method in comparison the Transformer in Figure] As shown, the Transformer computes the attention weights aij for every possible combination of query y; and key ky (Phase 1) and exhaustively enumerates all value vectors by for each query (Phase 2). In contrast, our method takes advantage of sparsity of the attention matrix and only computes the highest attention weights and enumerates only the value vectors associated with them.
We conduct experiments on CPUs on the LRA (Tay et al. 2020), ZeroSCROLLS (Shaham et al. 2023) and LvngEval (Li al. 2023) benchmarks. Across all three benchmarks, IceFormer demonstrates substantially faster inference speeds than existing methods while attaining almost no accuracy loss compared the Transformer. speedup relative to the Transformer while retaining 98.6% of its accuracy. Compared to the best efficient Transformer with comparable accuracy for each task, IceFormer on average 3.04 faster. On the ZeMSCCROLLS benchmark, IceFormer achieves 2 speedup on average compared to leading LLaMA 2-based LLM while retaining 99.6% of its accuracy.
the LRA benchmark, on average IceFormer achieves 7.63x
2 RELATED WORK
Efficient Transformers can be categorized along two axes: method type and retraining requirement. Along the first axis are sparsity- basedd methods and low-rank methods. Along the second axis are methods that can and cannot be applied to common pretrained Transformers without retraining.
Sparsity-based methods employ a sparsified attention mechanism capture global information and integrate it with local attention results. Some approaches aim to improve the space complexity compared to the vanilla attention mechanism without improving the time complexity, e.g., top-k Attention (Gupta al. 2021). Other approaches aim improve both, e.g., Sparse Transformer (Child jet al. 2019). Longformer (Beltagy et al.) 2020), and ETC (Ainslie et al. 2020). A substantial limitation of these models is that the tokens that are attended which do not adapt to varying imput sequences. Because the original attention operation is permitted
predefined and remain static,
attend any token, these models must be trained with their respective predefined constraints
2

---

Published as a conference paper at ICLR 2024
on tokens to be attended to. Reformer (Nikita et al. 2020) can attend to different sets of tokens for different input sequences by using Locality Sensitive Hashing (LSH) (Andoni et al.) 2013) to group tokens into chunks and subsequently attending only to tokens within the same chunk each query and adjacent chunks. However, Reformer imposes two constraints that are not in the original attention operation: keys must be normalized and queries and keys must be the same. Therefore, Reformer must be trained with these constraints built-in, As result, these methods cannot be applied to pretrained, non-modified, models directly; instead, the models must be retrained with the required constraints before these methods can be used.
Low-rank methods approximate the attention weight matrix with low-rank matrix reduce the quadratic time and space complexity. Examples include Linformer (Wang et al. (2200) and Per- former (Choromanski et al. 2020), which decompose the attention weight matrix into product of tall and wide matrices consisting of learned linear features or random features of the keys and queries, respectively. However, these Transformers typically introduce significant approximation errors be- cause attention weight matrices produced by the original attention operation, especially in the case of long input sequences, typically have high rank. Consequently, models that use these approaches must be trained with low-rank approximations built-in, in order to learn to be robust to the associated approximation errors. As a result, these approaches cannot be applied to pretrained, non-modified, models directly; instead, the models must be retrained with the required approximations before these methods can be used. Other approaches provide more general methodologies that can leverage weights pretrained with standard Transformers without retraining, These Transformers accelerate the execution of the standard attention operation without altering the underlying architecture. Two examples are Nyströmformer (Xiong et al. 202T) and LARA (Zheng al. 2022), which replace the softmax structure in the self self-attention mechanism with the product of separately activated query and key matrices. Nyströmformer utilizes the Nyström method, while LARA combines randomized attention (RA) and random feature attentions (RFA) (Peng et al., 2021 to reconstruct the attention weight matrix. In another example, H-Transformer-ll ID (Zhu & Soricut 202T) recursively divides the attention weight matrix into blocks and truncates the small singular values of each off- off-diagonal blocks. All these approaches leverage low-rank approximations, as opposed to sparsity.
Other works propose harewace-specific optimizations without aiming to improve the computational complexity. Examples include AttssAtttnnionn (Dao et al., 2022). which optimizes reads and writes between levels of GPU memory, and H2O (Zhang et al.) 2023), which dynamically retains balance of recent and heavy hitters tokens by a KV cache eviction policy. These strategies are dependent on implementation and are specific particular hardware platforms (e.g. GPU).
3 NOTATION AND PRELIMINARIES
Mathematically, the attention operation takes three matrices as input, K E Ruxdd Q E Raxd V E Rmmxd", which denote keys, queries and values respectively, and outputs a matrix O E Bexd Optionally, it may also take in a mask as input, S E _ whose entries are either or 1. The ith rows of K. Q. V and O, denoted as k, qi and o,, represent the ith key, query, value and output respectively. The entry of in the ith row and jth column, denoted as 8t,j. represents whether the ith query is allowed to attend to the jth key- - if it is 1, it would be allowed; if it is O, it would not be. A common masking scheme is the causal mask, where 8t.j is if 2 j and 0 otherwise. Keys and queries have the same dimension d, and each key is associated with value, and so the number of keys and values is the same and denoted as m.
First the attention operation computes the attention weight matrix A Raxm Its entry in the ith row and jthh column, denoted as ai.j, is computed with the following formula:
k Vd
$t.j exp
0ij
Σ m E1=
S exp
9. k2y
Vd
Then the attention operation combines the values with the attention weights in the following way:
m
Vi
(2)
j=1
3

---

Published as a conference paper at ICLR 2024
The attention matrix A is typically sparse (Nikita et al. 2020 Gupta al. 2021 i.e. in each of A, only a few attention weights have sigmificant (large) values, while the majority of the remaining values close zero. Suppose we somehow identify the k unmasked keys that receive the highest attention weights for each query qi without computing the attention weights for all keys. Then, the original attention matrix A can be approximated by only computing the inner product for the identified keys, which can save significant amount of time and computational resource.
4 ICEFORMER: ACCELERATED ALL-AATENNION FOR GENERAL KEYS WITHOUT RETRAINING
To build general-purpose retraining-free acceleration method, our approach must not require modifications to the attention mechanism to change attention patterns or the introduction of new model parameters to capture regularities in the attention patterns. This precludes popular strategies such as attention mechanisms with predefined sparse attention patterns, e.g., (Child et al.) (2019] Beltagy et al., 2020; Ainslie et al. (2020), and learned dimensionality reduction keys and queries, e.g., Wang et al.] 2020 Choromanski et al., 2020).
Consequently, it is difficult to design an acceleration method that exploits known regularities in the attention patterns without imposing the retraining requirement. We therefore aim to design an acceleration method that does not make assumptions on the existence of regularity in the attention patterns. In order to improve the O(mn) complexity of vanilla attention, we need to adaptively identify the most important keys (i.e., those that receive the highest attention weights) without computing all attention weights. This seems like a chicken-ande -ankeenn-nd-ggggg problem: how can know which attention weights are highest without comparing them to all the other attention weights?
Remarkably, in the special be done by leveraging k-nearest neighbour search (k-NNSS to identify the A most important keys for each query. This relies on the following mathematical fact, whose derivation is in included in Sect. B.1 of the appendix: if | = for all j. arg max, aij = arg min, ||q - k
of normalized keys, as proposed in Nikita al. (2020), this can
However, this fact only holds when all the keys have the same norm it is true when different keys differ in their norms. Intuitively, this is because the norms of keys can modulate the attention weights they receive, all else being equal. So key has a larger norm than key B. key can receive higher attention weight than key B even if key A is farther from the query than key B. As result, naïvely applying k-NNS in the general case would fail to identify the most important keys.
In this paper, we develop an acceleration method that does not require retraining or impose any constraints on keys. is both accurate and computationally efficient, and can also work with attention masks that are common in Transformers, such as causal masks. Below we will describe the details.
4.1 GENERAL RETRAINING-FREI ACCELERATED ATTENTION
Instead of applying k-NNS to the original keys directly, we will first embed the keys and queries into higher dimensional space. Inspired by Neyshabur & Srebro (215). we choose the following key and query embedding functions, which we denote as Tk
1
and TQ Rdd + Rd+1;
Tx (k,) kkj/c
VI- 1k,
(3)
(Q(q:) /q//l/qill 0]
(4)
where maxy' ||kr||2 ||2 is at least the maximum norm across all keys.
turns out that the & most important keys can be identified by performing A-NNS on the key embeddings using the query embedding. We will show this below:
({a) ky
arg max d.,j arg max softmax,
(5)
Vdl
5'=1
4

---

Published as a conference paper at ICLR 2024
kj
arg max vd
(6)
arg min 1 2q. k
arg min q. q 2q kj/c|lq:l12 + k3 kg/c +1
(8)
arg min | 4 ky/cll + 1 |/2/1//c²
(9)
= arg min (q:): - Tx (k,))l
(10)
4.2 ACCURATE k-NNS FOR ACCELERATED ATTENTION
The problem of k--NS is one of the most well studied problems in theoretical computer science. Many algorithms have been developed, and often significant speedups can be obtained by allowing for mistakes with some probability. Such algorithms known randomized algorithms.
In the context of L.Mss, the number of attention layers is typically high and so errors from earlier layers compound. Therefore, it is essential for the k-NNS algorithm achieve high accuracy. Choosing an appropriate k-NNS algorithm is therefore crucial.
Most k-NNS algorithms are bucketing-based which places keys into discrete buckets and searches over buckets that contain the query. On the other hand, ranking-baan based algorithms compares the rankings of different keys relative the query and searches over highly ranked keys. A bucketing- based algorithm effectively uses a fixed threshold on similarity, and so a variable number (including zero) keys can meet the threshold; on the other hand, a ranking- algorithm returns fixed number of keys, which effectively amounts to choosing variable threshold similarity based on the distribution of keys, as shown in Figure 2 An example
bucketing-based algorithm is socality-sensittve hashing (LSH) (Indyk & Motwani 1998), , and example of a ranking-basse based algorithm is Prioritized DCI (Li & Malik 2017). As shown in Figure 2 LSH hashes each key into bucket associated with the hash value, whereas Prioritized DCI ranks keys along random directions.
For accelerating attention, we posit that ranking- based algorithms are better suited than bucketing- based algorithms,, because attention weights depend on how different keys compare to another, rather than an absolute evaluation of each key against a fixed threshold. Therefore, ranking- based algorithms can yield better recall of truly important keys.
Ranking-basse based
Bucketing- bused
Runking-bueel
Baccktingg based
Curakite Crunt
Cunulatus Couc Douz
Projection
Projection oh kays
Figure 2: Difference between ranking- based and bucketing-bate based k-NNS. Left: illustration of two k-NNS methods, Prioritized DCI (ranking-basee based) and LSH bbucketing-basede Right: the number of keys whose projections are less than a threshold. Ranking-base algorithms return fixed number of keys are most similar the query under projection (shown as a exed-size row), which effectively filters out points outside variable-sized window on the projections. Bucketing- based algorithms use a fixed-size window (shown as fixed-size column) and return all keys whose projections lie within it.
4.3 FAST K-NNS FOR ACCELERATED ATTENTION
In Transformer, the keys in attention layer depend the output from the preceding attention layer. Therefore, a database needs to be constructed for each attention layer. Therefore, is important to choose a k-NN algorithm that attains both fast construction and querying.
5

---

Published as a conference paper at ICLR 2024
Moreover, in the context of LLMs, many popular models use decoder-only architectures. The attention layers in such architectures use causal masks to prevent the currently generated token to depend future yet-to-be-gee be-geerrtte tokens. Such masked attention is equivalent to excluding the masked out keys from the set of keys the k-NNS algorithm operates over. So each time a token is generated, one key becomes unmasked. Instead of constructing new database each time token is generated, it is more efficient to add keys incrementally to the database for k-NNS.
Fortunately, Prioritized DCI is efficient at both the construction and querying stages. If the number of random projection directions p is nearly as large as the intrinsic dimensionality of the data d' 2 1 and the number of nearest neighbours k look for is small, Prioritized DCI can return the exact k-nearest neighbours for a query with high probability within approximately (ddk%/dm--p// time, where (-) suppresses log factors. Its preprocessing is lightweight, and so only needs (dpm) time. If compare this the computational complexity of vanilla attention of ((dmn), observe that there is no longer a term that depends on mn, and so there is no longer the quadratic dependence on sequence length. Later in section [5.1 we also empirically validate the efficiency of Prioritized DCI and found it to be faster than eleven other leading K-NNS algorithms.
To support causal masking, we extended the implementation of Prioritized DCI to support incremental database updates. This can be done efficiently, since the data structure consists of sorted lists, so insertions and deletions can be done in Olog m) time if they are implemented as binary search trees.
5 EXPERIMENTS
In this section, we will compare the recall-latency trade- off between different k-NNS algorithms and then analyze the performance of IceFormer on the LRA benchmark (Tay
heatl 2op Time Coretaciio Oetot -der todhhigh better
et al.] (000)) which is popular benchmark for long-context Transformers (Zhu & Soricut, 202T; Xiong et al.) 2021
orylni
Zheng et al. 2022). Next we will demonstrate the advan- tages of IceFormer applied to LLMS with long prompts as input on the ZeroSCROLLS benchmark (Shaham et al.) 2023) and the LongEval benchmark (Li et al.] 2023). To ensure robustness of results, we used variety of CPUs for our experiments we used IRt) Core(TM)) i7-6850K 6-Core for the LRA experiments, AMD Ryzen 5950X 16-Core for the ZeroSCROLLS experiments, and AMD Ryzen 9 5900X 12-Core for the LongEval experiments.
Figure 3: Comparison between twelve k-NNS algorithms on fashion-mnist- 784 dataset. There are in total 60,000 keys and 10,000 queries with 784 di- mensions. The task is to find top-10 closest neighbours from the entire set of keys for every query. X-axis: Average recall across all the queries; Y-axis: To- tal latency (seconds) including database construction and querying.
5.1 DIFFERENT K-NNS ALGORITTMS COMPARISON
We compare the recall of true nearest neighbours and total construction and querying time of 12 K-NNS algorithms, in- cluding Prioritized DCI and the best performing algorithms from ANN benchmarks (Aumüller et al. 2017), on the Fashion MNIST dataset in Figure 3 As shown, Prioritized DCI achieves the best recall-latency trade-off compared to other algorithms, which demonstrates its suitability in our setting, which requires fast construction and querying.
5.2 EVALUATION ON LONG RANGE ARENA (LRA) BENCHMARK
Datasets and Metrics. LRA consists of five different tasks: ListOps (Nangia & Bowman 2018).
document retrieval (Retrieval) (Radev et al., 2013), text classification (Text) (Maas et al. 201 l).
CIFAR-10 image classification (Image) (Krizhevsky et al. 2009) and Pathfinder (Linsley et al. 2018
Specifically, all the five tasks consist of sequences with at most 4k tokens. We summarize the dataset information in the appendix [.I for more details. In this experiment, from Tay et al. (2020) and report the test dataset classification accuracy, average running time of the attention module, and CPU memory usage during inference for each task.
follow the train/test splits

---

Published as a conference paper at ICLR 2024
Baselines. In addition the vanilla Transformer, compare with Nyströmformer (Xiong et al. 1021). H- T--Trasffrmmr--- ID (Zhu & Soricut (021). LARA (Zheng et al. 2022), Reformer NNiitaa et al.
2020 Longformer (Beltagy et al. 2020), Performer (Choromanski et al. (2020), and Linformer (Wang et al. 2020p. In order to compare with Reformer, we train a Transformer model with shared Q and K according to Nikita et al. (2020). For Longformer and Linformer, as they introduce additional parameters, we randomly initialize these parameters when loading the pre-trained weight from the vanilla Transformer. For fair comparisons, we use the LRA evaluation benchmark implemented in Py PyTorch by (Xiong et al. 202T), and only replace the self-attention module while making other parts of each model exactly the same as the vanilla Transformer.
Implementation Details. For each task, we begin by training base model using GPU with vanilla Transformer architecture. Then we replace the vanilla attention module with one of the eight efficient attention modules mentioned carlier and directly apply the pre-trained weights for inference. To ensure fair comparison, we adjust the batch size to I, climinating the need for a padding mask since our proposed IceFormer automatically ignores padding masks during inference. Note that because of the additional shared-KQ constraint, for the Pathfinder task, our attempts to train a shared-K K Transformer were unsuccessful. As result, we have excluded the corresponding results from the subsequent analysis. Additionally, during the inference, we utilize a total of 4 CPU threads. For more comprehensive details, please refer to the appendix C.2
Inference Results. Ideally, the accuracy of the vanilla Transformer (non-shared-K3) serves as an upper bound for the approximated accuracy of the other seven models (IceFormer (non-shared-KO)) Nyströmformer, H- H-Transformer- ID, LARA, Longformer, Performer, and Linformer). Similar for the shared-KS KQ Transformer. Also, the attention module inference time of the vanilla Transformer would be the longest, with other efficient Transformers achieving shorter inference times the cost of sacrificing prediction accuracy. Table I the attention module for each method. The "yper-parameter settings are listed in the appendix C.3 In general, our proposed IceFormer consistently outperforms all efficient Transformers, offering the best accuracy approximation while requiring the least inference time across all five tasks. This demonstrates the generalizability and effectiveness of our model.
presents the prediction accuracy and inference time of
Table 1: The performance of vanilla Transformer, and eight approximate attention methods on the LRA benchmarks.
Method
shared KQ
ListOps
Text
Retrieval
Image
Pathfinder
[s]
Time
Acc
(s)
Time
(s)
Transformer
Vaswani
2017
x
0.4255
0.4455
2,9208
0.5986
0.6019
0.6615
0.6933
0.6586
0.6681
6.7946
8.3588
01:132
0.3844
4,9303
5,9804
0.7514
0.9620
Reformer
Nikiin ot al Proh
V
0.4121
1.4281
0.5941
0.2288
0.6467
1.4751
0.3726
3.6927
/
/
LARA Zhes
2022
0.4125
0.6146
0.5831
0.2348
0.6401
1.8605
0.3094
2.6720
0,7380
0.5961
Nystrmffonsrr
Xllogg
0.5838
0.3542
0.6540
0.3754
1.7644
0.7176
11 Transformer ID
Zig & Singou A12
X
0.3265
1,930]
0.5954
0.4811
0.5808
3.5605
0.2286
1.2586
0.5286
0.5708
Longlommnr Be o R A1AU
X
0.1975
0.7406
0.5236
0.9862
0.1918
1.0443
0.1488
0.5451
0.5009
0.5899
Performer
ATA0
0.1975
0.6571
0.5000
0.3327
0.4974
1.2058
0.1345
0.6404
0.5056
0.6395
Linforms
Wang
0.1975
0.4940
0.1064
0.7387
0.5022
IceFormer (Ours)
X
0.4124
0.4153
0.3766
0.4678
0.6001
0.5978
0.0921
0,0913
0.6541
0.6602
0.8337
0.8480
01,0046
0.3752
0.5076
0,9581
0.7442
0.3058
Speed & Accuracy Trade-off. For IceFormer, increasing the extent of approximation generally improves model efficiency but can lead a decrease in prediction performance. Here, we study how the extent of approximation affects inference speed and accuracy by varying the number of returned candidates of IceFormer, k., from 3 to 10 for each task and present the results in Figure 4 From the figure, we observe that across all tasks, when k becomes larger, IceFormer achieves improved prediction accuracy but becomes less efficient.
Memory Complexity Analysis. Table 2 summarizes the maximum memory usage for each method during inference. We employ the same aaper-pprameters as in Table | and maintain a batch size of I to eliminate the need for padding masks. The table reveals that IceFormer consistently exhibits the lowest peak memory usage across all tasks. In comparison to the vanilla Transformer, IceFormer achieves memory savings of up to 0.862 GB.
7

---

Published as a conference paper at ICLR 2024
Asunc
W31 oi
Nell Clock
120 50
WalCllok
Wol Clock Time
005. 010 720 540 100 200 100
Figure 4: Tradcoff between speed and accuracy k varies five LRA tasks. The horizontal axis of each plot is the averaged wall clock time of attention module, and the vertical axis is the model prediction accuracy. Each point corresponds value of & in the following set: (3, 5, 8, 10).
Table 2:
Peak memory usage (GB) on LRA benchmark. The peak memory usage is the total memory usage of the whole program, which includes the memory for the Prioritized DCI database/index.
shared KQ
ListOps
Resrivall
Image
Pathiinder
Transformer
VVaswnni
3.729
4.327
5,031
3.778
3.926
3.631
4.877
3,740
Reformer Nikia 2020 LARA
/
3.623
3.983
4.250
3.687
/
X
4.566
3,772
3.943
Nystömfanier
x
4.375
3.463
3,845
Irinsffemerr11 ID 2 Sarier
X 3,883
4.328
4.543
3.583
3,603
IceFormer (ours)
X
4.169
3.304
3.465
3.306
3.756
4.053
3,286
5.3 EVALUATION ON LARGE LANGUAGE MODEL (LLM)
We evaluate IceFormer in the LLM setting as well. Specifically, we utilize IceFormer accelerate the prompt processing process in LLMs. We pick Vicuna-7b-vl.5-- fine-tuned from LLaMA (Touvron et al. 2023) and is one of the top-performing open-source LLMs with a context length up to 16K tokens, for the following experiment. For more comprehensive details including the choice of & in k-NNS of IceFormer, please refer the appendix E.
(Zheng et al. 2023), which
For the following LLM experiments, we do not compare IceFormer with Reformer, LARA and Nyströmformer for the following reasons: Reformer requires keys and queries to be shared, which is not the case in pre-trained LLMs; Longformer only proposed a way to speed up the encoder part of the Transformer, thus cannot be applied to decoder-only I.MM;; LARA and Nyströmformer group different tokens into different clusters and so cannot handle causal masks in LLMs, which use decoder-only architectures. All basclines that require retraining (Longformer, Performer and Linformer) also excluded from the comparison. More details can be found in the appendix E.2
ZeroSCROLLS Results. We compare IceFormer with the vanilla Vicuna-7b-v 1-66 model and H-Transformer-1D 1D applied to Vicuna-7b-vl. 1.5-1 on the ZemOCCROLLS benchmark (Shaham et al. 2023) which is specifically designed for LLMs and contains ten diverse natural language tasks that require understanding long input contexts, including summarization, question answering, aggregated sentiment classification and information reordering. Each task has a different sequence length varying between 3k and 10k. We measure ZeroSCROLLS scores and latency of the attention module. Table shows that IceFormer achieves up to 3.0x speed-up compared to standard self self-attention while attaining at least 99.0% of the vanilla unaccelerated model performance at the same time.
LongEval Results & Scalability Analysis. To provide more comprehensive analysis of Ice- Former's scalability in the LLM setting, we conducted additional experiments on the LvngEval benchmark (Li et al. 2023), which is designed to measure long-context performance and consists of two tasks: topic retrieval task with prompt length varying from 3k to 16k, and line retrieval task with prompt length varying from 5k to 16k. In Figure [] we present the averaged latency of the attention module corresponding to different input prompt length as well as the inference accuracy using the vanilla Vicuna-bb-y..5-. model and IceForrer. From the figure, IceFormer achieve nearly
8

---

Published as a conference paper at ICLR 2024
Table 3: The performance of the vanilla Vicuna-7b-v1.5- 1.5-- model, H-Transformer-1D and IoeFormer on the ZeroSCROLLS benchmarks, Numbers in parentheses indicate the relative comparison to the vanilla unaccelerated model, denoted as ceuna-7b-vv.5-100 We employ the same abbreviations for metric and task names as specified in the original paper (Shaham et al. 2023
We refer interested
readers to the original paper for the details.
Tisk (fookens)
Metric
Visuma-7bvi..-16k
H- H.Tansfomer--D
IceFormer
GvRp (8R)
Myro
(100()
6.8 (61 8%)
11.0 (100%)
Time
5477 (1.Dx)
(1,2x)
1,89 (2.7x)
SSFD (8K)
Rye
13.5 (100%)
(1.0x)
(112
6.3 (46.7%)
13.5 (100%)
(2.8x)
Ryo Time (6)
16.9 (100%)
10.7 (63.39))
16.8 (%))4%)
Mam (98)
6,77 (1.0)
4.62 (144)
2.51 (2.6%)
SOAL (8k)
Rys Time
18.9 (109())))))
(1.0%)
7.3
(2.2x)
18.9 (100%)
Qspr (5k)
Time (5)
34.2 (100%)
2,03 (1.0x)
6.2 (18.(1)))))
(1.20)
34.0 (99.4%)
(L89 (2.3x)
Nrt (10K)
Fl Time
14.7 (100%)
(13.66))
4.55
14.7 (100%)
2,85 (2.4x)
QAET
Time (s)
18.5 (100%)
3.76 (1.0x)
(.09
(13.9))
1.26 (3.0x)
199
MusQ (5K)
Time (5)
Fl
18.6 (100%)
(1.00)
16,9 (90.9%)
(1.10)
18.5 (99.5%)
0,37 (1.9x)
5pDy
Time (9)
42.5 (100%)
(1.0%)
2.22 (2.0k))
1.47 (3.0x)
(99.5%)
BLSS (7.5k)
Time (8)
Cla. 7
19.5 (1009))
4,52 (1.Dx)
11.7 (6008)
2.26 (2.0x)
19.3 (99.0%
1,55 (2.9x)
Arg. (7.5k)
f1 Time (s)
23.9 (100%)
7.8 (32,546)
23.8 (99.6%)
4,38 (1.0x)
2.92 (1.5%)
1.60 (2.7x)
Icatorer
loafomer Orienyy)
Icatemmnt (itanc))
-0-
1.0
0.8
Latency (oecomcn))
0.8
15
Laiendy (sccon))))
0.
02
0.-
019 QK
10K
GX
0.-
3K Context Langh
15K
-0
Contect Length
Figure 5: Scalability analysis for IceFormer on the LongEval benchmark. The left figure shows the results of the topic retrieval task; the right figure shows the results of the line retrieval task. X-axis: length of the input prompt; Y-axis (Left): retrieval accuracy; Y-axis (Right): averaged process wall clock time (second) of the attention module.
identical inference accuracy compared with the vanilla Vicuna-7b-vI.5-166 Notably, as the prompt length increases, there is corresponding increase in the inference latency for both methods and for both tasks. However, even with very long prompt lengths, IceFormer maintains its scalability and consistently outperforms the vanilla Transformer. Furthermore, the length of the prompt increases, the difference in the latency between IceFormer and the vanilla Transformer becomes larger, demonstrating the superior scalability and efficiency of IceFormer in the context of LLMs.
CONCLUSION
In this paper, we present IceFormer, new method for improving the inference time efficiency of pretrained Transformers on the CPU. Notably, in contrast to other methods, IceFormer does not require retraining, does not require special constraints imposed on the attention mechanism and simultaneously achieves high accuracy and fast inference. These advantages make Iceformer very well-suited to LLM deployment on CPUs, especially when the LLM needs to handle very long sequences input. The experimental findings three benchmarks compellingly illustrate the effectiveness of our approach in reducing the quadratic time and space complexity of Transformers both in cases with bi-directional and causal attention mechanisms.

---

Published as a conference paper at ICLR 2024
REFERENCES
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs in transformers. arXiv preprint arXiv:2004.080 08448,, 2020.
Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, llya Razenshteyn, and Ludwig Schmidt. Practical and optimal Ish for angular distance. Advances in neural information processing systems, 28, 2015.
Anthropic. 100k context windows, 2023. URL https: wwwp http:///wwwwanttppc.cccco////// conee/////////////ooooomm
contect windows
Martin Aumüller, Erik Bernhardsson, and Alexander Faithfull. Ann-benchmarks: A benchmarking
tool for approximate nearest neighbor algorithms. In International conference on similarity search and applications, pp. 34 34-49. Springer, 2017.
Iz Beltagy. Matthew E Peters, and Arman Cohan. Longformer: The long- don-mett transformer. arXiv preprint arXiv:2004.05150, 2020.
Rewon Child, Scott Gray, Alec Radford, and Ila Sutskever. Generating long sequences with sparse 2019.
transformers. arXiv preprint
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 11166444-63555 2022.
Dave Dice and Alex Kogan, Optimizing inference performance of transformers on cpus. arXiv preprint arXiv:2 102. 102.06621, 2021.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient
transformers via top-k attention. arXiv preprint (XXi::200600899) 2021.
Piotr Indyk and Rajeev Motwani. Approximate nearest neighbors: towards removing the curse of dimensionality. In Proceedings of the thirtieth annual ACM symposium on Theory of computing. pp. 604-613, 1998.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Dacheng Li, Rulin Shao, Anze Xie, Ying Sheng, Lianmin Zheng, Joseph Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang. How long can context length of open-source LLMs truly promise? In NeurlPS 2023 Workshop on Instruction Tuning and Instruction Following. 2023. https/
/oppenevveww.oo wwt.:eiwwwwweeeeeeeeeeee fnet/forummidd Ltw.ooiiii-----LWWVV
Ke Li and Jitendra Malik. Fast k-nearest neighbour search via prioritized dei. International conference on machine learning, pp. 2081-2090 PMLR, 2017.
Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long- range spatial dependencies with horizontal gated recurrent units. Advances in neural information processing systems, 31, 2018
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng, and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies, pp. 142-150. Portland, Oregon, USA, June 2011. Association for Computational Linguistics. URL http
//www. org/ wnt wwtthoooo////loool/o0.
Nikita Nangia and Samuel R Bowman. Listops: A diagnostic dataset for latent tree learning. arXiv preprint / 060288 2018.
Behnam Neyshabur and Nathan Srebro. On symmetric and asymmetric Ishs for inner product search. In International Conference on Machine Learning, pp. 1926-1934, PMLR, 2015.
10

---

Published as a conference paper at ICLR 2024
Kitaev Nikita, Kaiser Lukasz, Levskaya Anselm, et al. Reformer: The efficient transformer. In Proceedings of International Conference on Learning Representations (ICLR), 2020.
OpenAl. Openai gpt-4, 2023, URL httttt://///:aeatt com/ cwm/gpttt/tppooooooooo -4
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. arXiv preprint arXiv:2103.02143, 2021.
Dragomir Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The anthology network corpus. Language Resources and Evaluation, 47:919-9444 2013.
Uri Shaham, Maor Ivgi, Avia Efrat, Jonathan Berant, and Omer Levy. Zeroscrolls: A zero-shot benchmark for long text understanding, 2023.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: benchmark for efficient transformers. arXiv preprint (rxXi:2201.040000 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768. 2020.
Yunyang Xiong, Zhampeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer. A nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 44138-14148 2021.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H 2 o: Heavy-hitter oracle for efficient generative inference of large language models. arXiv preprint wrXiv:230.114443 2023.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric. P Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging Ilm-as-a-judge with mt-bench and chatbot arena, 2023.
Lin Zheng, Chong Wang, and Lingpeng Kong. Linear complexity randomized self-attention mecha- nism. In International Conference Machine Learning. pp. 2701 1-241.. PMLR, 2022.
Zhenhai Zhu and Radu Soricut. H-transformer-ID Fast one-dimensional hierarchical attention for sequences. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3801-388 -315,, Online, August 2021. Association for Computational Lin- guistics. doi: 101188653112221 lol-lcogg.294,, URL https: / waweaaaaannnolttoooooooooooooo hgw.oeoooooooooooo000.000000
ac www.nn..... 294

---

Published as a conference paper at ICLR 2024
A PSEUDOCODE FOR ICEFORMER
We provide the pseudocode below for IceFormer. For completeness, we also include the pseudocode of Prioritized DCI, which we adapted from (Li & Malik 2017) and added our modifications. In the following pseudocode below, Tk and TQ are the key and query embedding functions, respectively. In our implementation, we implement recursive version of the algorithm, which offers faster performance in practice.
Algorithm IceFormer attention
Require: set S., n points g', "".....4" _e S₃ keys 1s .'.....,k hn I8" S, of n values , ,..." E RRd the number of simple indices m that constitute a composite index, the number of composite indices L, the number of points to retrieve ka and the number points to visit k in each composite index Sk, Sv, m. L, ko, k1)
function ICEFORMER_ DIEEFORMER_ATTENNIONNSSSS
O f- zero matrix E IR =x with rows o, E I2d" Sp <- CONSTRUCT(S:, m, L)
for
do
+ QUERY (q', Spp ku, k1)
for j = to n do
if E S: then
8,g 9. kg
vd
else
8,y
end
end for
for j
to n do if E St then
exp(tir)
ã., + softmax, (8,,) Exco, cxp(2i.))
else
< 0
end if
end for
o, + Ejes, Us
end for return
end function
Algorithm 2
Data structure construction procedure
k", Require: A dataset D of n points k ....,,", the number of simple indices mn that constitute composite index
and the number of composite indices L
function CONSTRUCT(D, m, L)
{ug (ur))e(m))r + mL random unit vectors in Rd
+- mL empty binary search trees or skip lists
for j to m do
for 1 = 1 to 7, do
for
to do F < (Tx (Kx))"), u5²)
Insert (kyz, i) into T3 with ky being the key and being the value
end for end for end for
return {(Taa, Uzx) (emm).ee1))
end function
12

---

Published as a conference paper at ICLR 2024
Algorithm 3 k-nearest neighbour querying procedure
Require:
Query point R², binary search trees/skip lists and their associated projection vectors
{[Tz, (gx)) jemm), the number of points to retrieve ko and the number of points to visit k in each
composite index
function QUERY(q, ((T u, ka, k1)
Ci + array of size n with entries initialized to 0 VI E [4]
9,1 (g), ust) vj E [m].1 [4]
S f 0 E []]
P + emppy priority queue AI E [4]
for = 1 to L. do
for j to do
(1) (1)) + the node in Tia whose key is the closest to 4zt
Insert (p ( with priority -lp (1)
4zel into Pt
end for end for
for 1' = 1 to k1 do
for I = 1 to L do
[S|| ko then
(P
1 < node with the highest priority in
Remove (0 A from Pt and insert the node in T5 whose key is the next closest to Zzt-
which denoted as (p(++1) ((+t)), with priority -lp Izel into P
Cilhf + (h (+) +1
if Ci(h
then
S - St U (1)))
end if end if end for end for
return k points in Ulg(e) Si that have the maximum inner-product value with
end function
B PROOFS
B.
PROOF
Here, we provide the full step-by-step derivation of the mathematical equivalence between conducting k-nearest neighbour search on normalized keys and identifying the keys that obtain the highest attention weight.
(( Vd
arg max di.j arg max max softmax,
(11)
('m))
arg max j
(12)
Vd
arg min | l1
-2q k, +
(13)
Since |ky |l2 1 for all j', |/ 2q, k, 1
| 2q k, + |k,
||q:
11.
arg max o.j = arg min |/q:10: 2q k, +
(14)
arg min ||q. k 11
(15)
B.2 PROOF 2
Here, we provide the full step-by-step derivation of the result in 4. establishing the mathematical equivalence between conducting k-nearest neighbour search transformed keys and identifying the
13

---

Published as a conference paper at ICLR 2024
keys that obtain the highest attention weight.
(1 ky m Vi
arg max 04j = arg max softmax,
(16)
j'=1/
9: kj
arg max 3 vd
(17)
arg max q. k;
(18)
arg j -2q kj
(19)
arg min 2q, //////
(20)
arg min 2q k,/cll q:1l/ -1
(21)
arg mijn 1 1.11//1lq.11 24, k, k 1k/ +1 -- 1k/1//22
(22)
arg mijn q. q: 9 1| -24 -2 /////////lllll + kkjkk// -1 1k/12/2
(23)
arg min(g://ll 22 (4://l/la ky/c) +1 1k/11//0²
(24)
arg min 1 - //1/ + - |ky1///²
(25)
00 v k 1B/2 2
arg min |q. 00//l|q:1: 2| ////ll2 +
(26)
arg min [Tq (q:) Tk (k2)
(27)
arg min [TQ (q:) - Tx (k2) ll2
(28)
C MORE DETAILS ON THE LRA EXPERIMENTAL SETTING
C.I DATASET DETAILS
In our LRA experiments, for Retrieval, Text and Pathfinder (64 64 version), we directly use the dataset from LRA codebase] Because the original datasets for ListOps and Image only contain short sequences, we generate longer samples for ListOps using the same code from the LRA codebase with 4000 as the maximum length; for Image task, we use version of the CIFAR- 10 dataset super-resolved to 64 x 6 instead of the original low-resolution 32 x 32 CIFAR-10 10 dataset. We follow the exact same train/test split as the original LRA paper (Tay et al. 2020). The details of the LRA dataset is listed in Table
Table 4: LRA Dataset Details.
Task
ListOss]
|Retrieall Image|
Puthiinder
Max longth
3,991
1,000
4,000
4,096
4,096
Avg.
2,232
3,917
4,096
number classes
10
Aecurccy by chince
0.100
0.500
0.500
0.100
0.500
C.2 BASE MODEL CONFIGURATION
We follow the experimental setup of prior work (Zhu & Soricut 2021) for training the base model. However, since we were not able to successfully train base Transformer models to satisfactory accuracy on Image and Pathfinder datasets using the original setting, we decreased the number of heads and layers for these two tasks. The details of the base model for each task are outlined in Table 5
14

---

Published as a conference paper at ICLR 2024
Table 5: Configurations of the base models for different tasks.
Task
LintOps
Text
Retrieval
Image|
Futhiinder
head embedding size
512
512
512
512
512
feed forward size
2048
2048
2048
1024
1024
number heads
8
8
4
number layers
C.3 PAPAM-RAR FOR THE BASELINES AND THE PROPOSED METHOD
For LARA and Nyströmformer, we tundd the parameter num num__andiinn Jandmarks by optimizing over the range (64, 128, 256, 512, 1024). For H-Transformer-11 ID, we tuned the parameter block_size by optimizing over the range 164, 128, 256, 512, 1024). For Reformer, we tuned the parameters num hash and bucket_size we considered the values of num_hash in range {1, 2, 4] and the values of bucket_size in range (64, 128, 256, 512, |.o Longformer, Performer, and Linformer which require retraining, because of their poor performance, we choose of approximation. For IceFormer. we tuned the parameter top over the range [3. 5. 8, 10, 15, 20). In general, a larger value for bucket_size, num Jand_nunk,,s, block_size or top_k k indicates less aggressive approximation, meaning that the model performance is closer to that of the vanilla Transformer. We select the values of the syper-paramettrs that lead to the best accuracy- recuracy-time trade-off for each model, and list them in Table 6
values that result in the least amount
Table 6: Hlyper-parament settings for different methods.
Method
hyper parameter
ListOps
Text
Retrieval
Image
Pathfinder
Reformer
Nikita et al,
2020
bucket_size size
num, hash
512
1
128
1
256
1024
LARA Zheng al.
2022
Jandmarkss
256
256
512
1024
1024
Nystrienfommer
Xiong et
2021
num, Jandmarks
256
256
512
512
1024
H- TransCormer-1D ID
Zha & Soricut
2021
block Size
1024
512
1024
256
1024
Longformer
Beliagy al.
2020
allention, window
2048
2048
2048
2048
20-8
Performer CCoromanski
mund features
2048
2048
2048
2048
Linfommer Wang et al. 2020
nam_proj_diim dim
2048
2048
2048
2048
2048
IceFormer
top_k
3
lefforerr (shural-QK)
top k
10
3
10
20
D ADDITIONAL EXPERIMENTS ON LRA
Approximation Quality. In order to assess how well various efficient Transformers approximate the outputs of the vanilla modified attention module, we measure the approximation error by computing the L2-norm of the difference between their attention module outputs and those of the standard vanilla attention module (o; in Equation 2 The averaged approximation errors for different efficient Transformers, utilizing the same hyper- parameter settings of Table are summarized in Table 7 As indicated in the table, IceFormer consistently achieves the lowest approximation errors across all LRA tasks, providing further evidence of its approximation efficacy.
Table 7: Quality of the approximation on LRA benchmark. The approximation error of the attention module output is reported for each method across all the tasks.
Method
shared-KQ ListOps
Retrieval
Image
Pathrinder
Reformer
Nikila 2070
5.452
LARA Prenz
2,395
9.456
22.066
9.261
Nystrmfforme Xinna de
X
5.758
10.269
6.523
18.789
10.442
H. Transfomerr ID 2hu & Sunier
2021
X
6.110
10.605
5.636
53.926
12:228
leeForeer (ours)
/
2.140
1.562
3,891
1,686
2.499
1,825
6,873
2.127
8,749
Visualization of Tables 1&2. Figures 6 and] respectively.
The results from Table 1 and Table 2 are visually represented in
15

---

Published as a conference paper at ICLR 2024
Lu1Op
Irage
Tweborr
23
03
Torehmwees
Reforra
LARA
02
ba Wo
0.5
03
Figure 6: The inference latency of IceFormer and the baselines (vanilla Transformer, Reformer, LARA, Nyströmer, H-Transformer-IDD ID) on the LRA benchmark (the smaller the better).
Bx/0p
haap
Tarekeree
4.2
36
2.
3.
2.5
Myvacntcerer er
Tunahamenth
LARA
Reformn
452
24
Lo FaFore
425
432
33
Figure 7: Peak memory usage (GB) on the LRA benchmark. The peak memory usage is the total memory usage of the whole program, which includes the memory for the Prioritized DCI database/index (the smaller the better).
E MORE DETAILS ON THE LLM EXPERIMENT
E. LLM EXPERIMENT SETTING
the tested L.M in section 5.3 It contains 32 attention layers, each of which contains 32 attention heads with dimensionality equals to 128. Its maximum input sequence length 16,384. We observe varying levels of sparsity across different layers of LLMs, and this sparsity remains consistent across different prompts. Therefore, in all the LLMs experiments in section 5.3 we apply IceFormer to approximate relatively sparse layers ranging from the sixteenth to the first in ricuna-7b-ve 1.5--v1..-1111 1.5-16k. This selection encompassed a total of 16 layers, equivalent to half of the total number of layers in the model.
We 1icuna-7b-vl.5-111
The k in the k-NNS of IceFormer for each task of the ZeroSCROLLS benchmark and the LongEval benchmark defined as:
= max(min((n * aJ, 50),
(29)
where n is the number of input tokens, [x] is the floor function, and a is haper-yyramenne parameter set by the users. In the ZeROSCROLLS henchmark, we set a equals 4e-3 for tasks SSFD and QMsm; 5e-3 for tasks GvRp, SQAL,, Qspr, Nrtv, MuSQ and BkSS; 6e-3 for tasks QALT and SpD2. In the LongEval benchmark, we set a equals 5e-3 for all the settings both two tasks.
E.2 CAUSAL MASKS AND OTHER INFERENCE- TIME EFFICIENT TRANSFORMERS
In the main paper, we did not compare IceFormer with LARA and Nyströmformer on LLMss. In this section, we elaborate on the problems of causal masks for these two methods.
Most aandom-featurebeane based models such LARA and Nyströmformer group different tokens into different clusters, known as landmarks. In order to enable causal masking these models, not only does the masking need to be applied at the landmark level to prevent the leakage of information from future tokens, an additional set of masks is also required to mask out different numbers of tokens within the same landmark for different queries. The latter is not supported natively and is especially difficult to implement. As a result, it is difficult to apply LARA and Nyströmformer to models that have causal masks.
16

---

Published as a conference paper at ICLR 2024
F TEXT OUTPUTS OF ICEFORMER + LLM
In this section, we provide the text outputs of IceFormer when applied to the LLM (vicuna-7b-vl. .3-- 16k) in Figure 8 folder.
We also include the full information of the input prompts in our supplementary
Prompt (4k tokens)
A chat between curious user and artificial intelligence assistant. The assistant gives belpful, detailed, and polite answers the user's questions. USER: Please help me to summarize the following:
There is nothing better than sitting on the couch after long day. Before though, one of the first thoughts running through person's 's head probably is, "I wonder what should cat for supper?" After working for roughly eight hours or longer, no one wants to is stand over a hot stove for an hour when he or she could be spending that time getting caught up on their favorite Netflix show. Instead of firing up the grill, she pushes few buttons on microwave, and voilà! delicious plate frozen food, it pizza bites corndog, households, tradition that facts, Frozen food causes obesity, just American, but around the globe.
cooked perfection seconds. Although this what happens on regular night most late. time for people to face the hard unwelcome
put to a stop hefore
[..]
Diabetes caused by the lack of hormone produced by the pancreas needed regulate the amount of glucose the blood. The reason for why Type diabetes occurs in humans is unknown. Type 2 diabetes frequent in children that are overweight, less uctive, and often, these children have a purent with diabetes (Hubbard, 2017). "In my cureer. Type 1 diabetes was rure disease in children and Type disease didn't exist. And I'm not that old," said Dr. Robin S. Goland said, who has been in practice for almost 25 years and is the co-director of the Naomi Berrie Diabetes Center at Columbia University Medical Center New York (Saint Louis, 2014). Diabetes could be prevented children for the most part but this is challenging to do when almost every frozen and pre-packaged food item is scriously unhealthy and goes against nutritional needs.
ASSISTANT:
TeeFormer Output Example (k=30)
The article discusses the history and popularity frozen foods, including the development of quick freeze methods by Clarence Birdseye and the marketing strategies used by companies like Swanson to promote frozen meals. Frozen foods have become popular due to their convenience, affordability, and availahility in grocery stores worldwide. However, the article frozen meals are often ultra processedee high calories, leading to concerns about their impact on health. Additionally, risc fast food restuurants has contributed the increusing consumption frozen and processed foods in the US.
notes
Figure 8: Output of IceFormer (Top-30) LLM with 4k input tokens. We ask the LLM summarize an article titled "Frozen Food: The World's Favorite Killer""
17

---

Published as a conference paper at ICLR 2024
Prompt tokens)
chat between curious and artificial intelligence assistant. The assistant gives helpful, detailed, and polite the user's questions. USER: Please help summarize the following:
Rapidly growing populution along with digitalization and globalization re-sult in constuntly increusing number of internet (ITU 2019) Closely relat ed to this development is the expansion of ecommerce husiness. Especially in the times of isolation, caused by the Corona crisis, which seems to be great oppor-tunity for the growth of online shopping industry. Lockdowns in the countries across the whole world have had considcrable influence the numbers of online purchases. The statisties compiled by the Organisation for Economic Co-operation and Development confirm this assumption showing that the e-commeree retail in the United States and in the United Kingdom rose between the first quarter and the second quarter of 2020 (first lockdown in the UK the 26th March 2020, in the U.S. on the 19th March 2020 ( from 11.8% to 16.1% 1% respectively from 20.3% to 31.3%. The increasing numbers of online purchases signify higher revenues and ut the same time add the importance to this sector.
currently live in, spend much more time the internet than we used and this
[..]
udvertisers wish to be successful in their OBA practices, they should also allempl to find out more about usery' reasons for being online. According to the interactive advertising model (Rogers and Thorson 2000), advertisements, which address these motives can achieve better results, bocause consumers tond to pay more attention towards them, and they are also considered to be casier to remem-ber. (Boerman et al. 2017) Consumers are more willing to accept OBA practices, if they foel that they might support them in the pursuit of their objectives. (Kirmani and Camphell 2004)
ASSISTANT:
IceFormer Output Example (k=40)
The rapidly growing population and digitalization have led to incressing number of internet users, with the ecommerce business expanding in particular during the Corona crisis. Online purchases have risen, and retail volumes have decreased due to lockdowns, Companics are investing more capital in online marketing and advertising, but users often foel irritated by online advertisements, especially when they ure not relevant. Online behuvioral advertising (OBA) targets consumers bused om collected personal information, but can be effective only anxiety and fear, and there have been several scandals related to misuse of personal data collected online, The importance scandals related online marketing led to need discuss and investigate OBA practices and their influence on consumers.
highly personalized. However, this can also cause foelings of
the ecommerce sector
Figure 9: Output of IceFormer (Top-40) LLM with 8k input tokens. We ask the LLM to summarize an article titled "Research of How Online Behavioral Advertising Influences Consumers".
18


Reviews:
[{"review_id": "0QwvTp50S5", "reviewer": "Reviewer_5Bmg", "summary": "In this paper, the authors present IceFormer, a novel method for improving the inference efficiency of pretrained Transformers on the CPU without requiring re-training. Leveraging the insight that k most important keys can be identified by performing k-NNS on the key\nembeddings using the query embedding, IceFormer employs the Prioritized DCI algorithm to identify the top-k keys. The evaluations on three benchmarks illustrate the effectiveness of this approach in reducing the quadratic time and space complexity of vanilla Transformers with competitive performance compared to other efficient-attention variants.", "questions": "1. 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? \n2. 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? \n3. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper sheds new insights on an important problem -- choosing the top-k keys in sparse attention for inference acceleration. The authors show that the exact kNNS algorithm is critical for the success of sparse attention. \n2. The paper is well-written and easy to follow.\n3. The evaluation section shows strong performance compared to other efficient attention works.", "weaknesses": "1. 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? 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). \n2. Fig.5's y-axis needs to be annotated with units (especially for latency).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fPySskQyW0", "reviewer": "Reviewer_ad3y", "summary": "The paper aims to address the computational challenges associated with the quadratic complexity of self-attention in long sequences. It proposes using k nearest-neighbor search as a method for approximating a sparse attention matrix, thereby reducing the computational cost.\n\nUpdated post-rebuttal rating.", "questions": "1. Are the baseline results for execution time from original papers or are they reproduced results specific to CPU performance?\n2. Could you provide the typical and range of values for k and m used in the experiments?\n3. What specific settings were used for CPU inference time evaluation? Why limit the experiments to only four CPU threads? Were any BLAS kernels like MKL used in the implementation?\n\nWhile the paper is well-motivated and presents a potential approach to an important problem, there are significant issues that undermine its contributions. Specifically, the paper lacks rigorous evidence to support its claims and fails to provide a thorough methodology for evaluating its proposed solution. Further revision is needed to substantiate the claims and provide a more comprehensive understanding of the algorithm and its performance across various hardware.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper is well-motivated, tackling a pertinent issue in the deployment of large language models.\n2. The idea of employing ranking-based algorithms over bucketing-based algorithms presents an interesting potential for complexity reduction.", "weaknesses": "1. The paper does not adequately support its claim that Prioritized DCI outperforms LSH, lacking both theoretical and empirical evidence.\n2. 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.\n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jK99mN5580", "reviewer": "Reviewer_dxGt", "summary": "Transformers, especially large language models (LLMs) like GPT-4, have shown great promise in handling long input sequences. However, their deployment on CPUs is challenging due to the quadratic time and space complexity of the self-attention mechanism. Existing methods to address this either require model retraining, are not generalizable, or compromise accuracy. The proposed method, IceFormer, aims to accelerate inference time without retraining, maintaining accuracy, and offering fast inference across various LLMs. It achieves this by optimizing the self-attention mechanism, leveraging the sparsity of the attention matrix, and using k-Nearest Neighbor Search (k-NNS). Experimental results on multiple benchmarks show that IceFormer efficiently outperforms other methods in both speed and accuracy.", "questions": "included in weaknesses.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 1, "strengths": "- This method can accelerate inference time (only CPU) for pretrained transformers without the need for expensive and time-consuming retraining.\n- Unlike some other methods, IceFormer ensures that there is minimal approximation error, crucial for LLMs where errors in initial layers can cascade through subsequent ones.\n- Beyond just accuracy, the method also guarantees rapid inference times, making it particularly suitable for CPUs.\n- By capitalizing on the sparsity of the attention matrix and utilizing k-Nearest Neighbor Search (k-NNS), IceFormer effectively reduces the computational burden of the self-attention mechanism.", "weaknesses": "This paper has multiple concerns for acceptance at ICLR 2024:\n\n1) 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\n).\n\n2) 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.\n(Honestly, I'm confusing about the exact meaning of 'wall clock time of attention module'.)\n\n3) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PslTmneGqR", "reviewer": "Reviewer_d3mB", "summary": "This paper presents a method for accelerating inference of long-sequence attention-based Transformers. The proposed method works with pretrained Transformers and does not require retraining. The critical point is to use Prioritized DCI as the k-NNS algorithm.", "questions": "1. What may be the future directions and extensions?\n2. 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?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method does not require retraining the model. The retraining may be expensive and unaffordable.\n2. It can be applied to a broad range of LLMs since it has no assumptions on the Transformer architecture.\n3. The proposed method provides a Pareto optimal solution in terms of time and task-related performance (e.g., accuracy), compared with the baselines.", "weaknesses": "### Major issues\n1. 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?\n2. 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.\n3. Could we apply the proposed method into training?\n\n### Minor issues\n1. It is not clear how IceFormer will perform on short and medium length of sequences.\n2. Could the authors visualize Tables 1 and 2?\n3. What are the limitations and potential negative impact of the proposed method?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xucf5Gd6nM", "reviewer": "Reviewer_nr9P", "summary": "The paper presents a novel approach for adaptively identifying the most important keys for each query using a k-NNS approach with a ranking-based algorithm. This method allows mapping keys and queries to a higher-dimensional space, eliminating the constraint that all keys must have the same norm. Furthermore, the authors extend the implementation of prioritized DCI to support incremental updates to the database.", "questions": "* 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper combines several techniques to accelerate the computation of the attention matrix by pruning the k matrix.\n* The proposed approach is comprehensively evaluated on LRA and LLM benchmarks, and it outperforms SOTA approximation methods, demonstrating its effectiveness.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)