Title: VCC: Scaling Transformers to 128K Tokens or More by Prioritizing Important Tokens
Venue: NeurIPS 2023
Accepted: yes
Keywords: efficient

transformer

roberta

T5

language modeling

question answering

summarization
Abstract: Transformers are central in modern natural language processing and computer vision applications. Despite recent works devoted to reducing the quadratic cost of such models with respect to sequence length, dealing with ultra long sequences (e.g., $>$16K tokens) remains challenging. Applications such as answering questions based on a book or summarizing a scientific article are inefficient or infeasible. Here, we propose to significantly improve the efficiency of Transformers for ultra long sequences, by compressing the sequence into a much smaller representation at each layer. Specifically, by exploiting the fact that in many tasks, only a small subset of special tokens, which we call VIP-tokens, are most relevant to the final prediction, we propose a VIP-token centric compression (VCC) scheme which selectively compresses the sequence based on their impact on approximating the representation of the VIP-tokens. Compared with competitive baselines, our algorithm is not only efficient (achieving more than $3\times$ compute efficiency gain compared to baselines on 4K and 16K lengths), but also offers competitive/better performance on a large number of tasks. Further, we show that our algorithm scales to 128K tokens (or more) while consistently offering accuracy improvement. Code is available at https://github.com/mlpen/VCC.

VCC: Scaling Transformers to 128K Tokens or More by Prioritizing Important Tokens

Abstract

Transformers are central in modern natural language processing and computer vision applications. Despite recent works devoted to reducing the quadratic cost of such models with respect to sequence length, dealing with ultra long sequences (e.g., > 16K tokens) remains challenging. Applications such as answering questions based on a book or summarizing a scientific article are inefficient or infeasible. Here, we propose to significantly improve the efficiency of Transformers for ultra long sequences, by compressing the sequence into a much smaller representation at each layer. Specifically, by exploiting the fact that in many tasks, only a small subset of special tokens, which we call VIP-tokens, are most relevant to the final prediction, we propose a VIP-token centric compression (VCC) scheme which selectively compresses the sequence based on their impact on approximating the representation of the VIP-tokens. Compared with competitive baselines, our algorithm is not only efficient (achieving more than 3 × compute efficiency gain compared to baselines on 4K and 16K lengths), but also offers competitive/better performance on a large number of tasks. Further, we show that our algorithm scales to 128K tokens (or more) while consistently offering accuracy improvement. Code is available at https://github.com/mlpen/VCC .

1 Introduction

The Transformer [32] is a fundamental/foundational architecture for natural language processing (NLP) and computer vision. It has shown remarkable performance across NLP applications including machine translation [32], language inference [10], and summarization [14]. Transformers have also been successfully applied to various visual recognition tasks and achieve impressive results [11, 4, 41]. Unfortunately, the runtime/memory needs of Transformers involve an unfavorable dependence on the input sequence length, making the use of Transformers for ultra-long sequence applications difficult. Therefore, many studies on Transformers make use of strategies such as truncation to ensure that the input sentence length is at most 512 , e.g., BERT, T5, and other Transformer-based language models [36, 22, 27]. Unfortunately, such a truncation, and other related strategies, inevitably results in loss of accuracy, the extent of which can vary from one task/dataset to another. Consequently, improving the efficiency for longer input sequence length is a key focus of many proposals. These developments are

AWS AI

AWS AI

important milestones, and they have reduced the quadratic dependency on sequence lengths to linear [6, 25, 35, 2, 38, 40, 39]. Currently, many Transformer models can process samples with sequence lengths of up to 4K (and even 16K). Very recently, results of newer models being able to handle much longer sequences have appeared [3].

Rationale. It is natural to ask whether the ability to process longer sequences is worth the trouble. The short answer is yes. Improved accuracy has been reported on long sequence tasks [2, 38, 14]. So, what is stopping us from harvesting even stronger gains in accuracy by feeding even longer sequences to such models? Models such as Longformer [2] and Big Bird [38] become slow and consume an excessive amount of memory as the sequence length keeps increasing. See Fig. 1 for illustration. Why? The representation update of every token involves computing

10 3

10 4

10 5

Sequence Length

0

500

1000

1500

Runtime (ms)

10 3

10 4

10 5

Sequence Length

0

10

20

30

Memory (GB)

Transformer

Big Bird

Longformer

Ours

efficient attention and feed-forward network at each layer. This incurs a linear cost relative to the sequence length and is expensive for sequences much longer than 4K (or 16K) tokens. To endow the models the ability to learn ultra-long range dependency, we need to lower this cost. What we describe in this paper is a concrete step forward - based on certain task-specific assumptions which appear to generally hold, we outline a formulation that works and delivers the expected improvements.

(1) Focus on what we need for a task: VIP-token centric compression (VCC). Wehypothesize/find that in many tasks where Transformers are effective, only a small subset of tokens, which we refer to as VIP-tokens, are relevant to the final output (and accuracy) of a Transformer. If these tokens had been identified somehow, we could preserve this information in its entirety and only incur a moderate loss in performance. Now, conditioned on these specific VIP-tokens, an aggressive compression on the other non-VIP-tokens , can serve to reduce (and often, fully recover) the loss in performance while dramatically decreasing the sequence length. This compression must leverage information regarding the VIP-tokens, with the goal of improving the approximation of the representation of the VIP-tokens. In other words, a high-fidelity approximation of the entire sequence is unnecessary. Once this 'selectively compressed' input passes through a Transformer layer, the output sequence is decompressed to the original full sequence allowing the subsequent layers to access the full sequence.

(2) Specialized data structure for compression/decompression. A secondary, but important practical issue, is reducing the overhead when compressing/decompressing the input/output sequences internally in the network. Ignoring this problem will impact efficiency. We give a simple but specialized data structure to maintain the hidden states of the intermediate layers, where the compression can be easily accessed from the data structure, and explicit decompression can be avoid by updating the data structure: the sequence is never fully materialized in intermediate layers.

Practical contributions. Apart from the algorithmic modules above, we show that despite an aggressive compression of the input sequences, we achieve better/competitive performance on a broad basket of long sequence experiments. Compared to baselines, we get much better runtime/memory efficiency. We show that it is now practical to run standard Transformer models on sequences of 128 K token lengths, with consistent performance benefits (and no complicated architecture changes).

2 Preliminaries

We review the Transformer layer, related work on efficient Transformers and define notations/simplifications. BOLD uppercase letters denote matrices, bold lower case letters denote vectors, and regular lower case letters denote scalars or functions.

Brief review of the Transformer Model. Fix n to be the sequence length and let d be the embedding dimension. Define an embedding matrix X ∈ R n × d which gives the n feature vector inputs for a Transformer layer. The output of this Transformer layer, X new , is defined as

using α ( X ) as shorthand for α ( · , · , · ) , which is a multi-head attention (MHA) with X as input for queries, keys, and values, described shortly. Here, β ( · ) is a feed-forward network (FFN). Layer

norms [1] are omitted to reduce clutter. Let the inputs to α ( · , · , · ) be Q , K , V ∈ R n × d for queries, keys, and values. MHA is defined as:

where g is the number of attention heads, { W Q,i , W K,i , W V,i } are trainable projections, and the 'cat' concatenates the outputs of multiple self-attention modules. We omit the biases for notational simplicity. For ease of discussion, let us further simplify the above notation by assuming that g = 1 , and suppress W Q, 1 , W K, 1 , W V, 1 , W as well as the normalization in softmax: they will still be estimated within the model (i.e., this module remains unchanged) but are tangential to the description of our idea. With these simplifications, the α ( · , · , · ) can be expressed as:

Let γ ( · ) be a placeholder for all heavy computations in the Transformer layer above:

We can verify that the output of a Transformer block (parameters are suppressed to reduce clutter) is,

A Transformer model consists of many such layers: the input of each layer is the output X new from the previous layer. Let l be the number of layers, then the overall complexity is O ( ln 2 d + lnd 2 ) .

Efficient Transformers. Many efficient self-attention methods are available to reduce the O ( ln 2 d ) cost. We list a few models noting that this list is not exhaustive. Performer [6], Random Feature Attention [25], and Nyströmformer [35] propose different low rank approximations of the selfattention matrices. Longformer [2] and Big Bird [38] describe global + local sparse attention. Reformer [17] and YOSO [40] exploit locality sensitive hashing for approximating the self-attention matrix. MRA attention [39] gives a multi-resolution approximation of the self-attention matrices. Memorizing Transformers [34] and RMT [3] follow a recurrent design and store the past context in an external memory module. By sequentially processing one segment of input sequences at one time, they avoid blowing up the memory when processing long sequences.

Efficient Transformers do not scale well to ultra-long sequences . Existing self-attention mechanisms often reduce the quadratic cost of MHA to linear. But so far, most experiments report sequence lengths of up to 4K, with some exceptions [2, 38, 14]. Beyond 4K, the linear cost (on n ) for both computing efficient attentions and FFN makes the cost prohibitive, especially for large models. For example, although LongT5 [14] can train on sequence lengths of up to 16K tokens with an efficient self-attention and shows promising results for longer sequences, it is slower and needs a sizable amount of compute (for example, see Fig. 1). Similar to efficient self-attention with linear cost attention, [34] and [3] do not try to reduce the linear cost (on sequence length) for FFN, so the computation might still be expensive for processing ultra long sequences. On the other hand, our method seeks to reduce the overall cost (both self-attention and FFN) of processing ultra long sequences and processes the entire sequence simultaneously.

Other alternatives for sequence compression? Compressing input sequences for efficiency reasons in Transformers is not a new idea. For example, [8] and [16] propose pyramid Transformer variants that progressively compress the sequence as the layers grow deeper via pooling or core-set selection. [23] proposes adaptively compressing the sequence based on the predicted semantic boundaries within the sequence. [26] proposes compressing the fine-grained past activations to coarser memories. There are three key differences with our approach. First, all methods listed above are task agnostic . They seek compressed/smaller representations to represent the original sequence well. Our formulation places no emphasis on representing the original sequence, as long as information pertinent to the VIP-tokens is preserved as much as possible. Second, once these methods compress the sequence, the residual information is lost (for the deeper layers or the later time steps). Our entire approach is predicated on avoiding this loss - we maintain access to the full sequence at each layer (via residual connection at least). Lastly, some of these ideas often involve an n 2 dependence on the sequence length in the initial stages of their formulation, making long sequence experiments problematic.

3 VIP-Token Centric Compression (VCC)

Our main goal is to reduce the dependency on n (but not by modifying Transformer internals). To do this, we describe a scheme that compresses the input sequence of a Transformer layer and

decompresses the output sequence, resulting in a model whose complexity is O ( lrd 2 + lr 2 d + lr log( n c ) d + lrn p d + nd ) . Here, r is the length of the compressed sequence, n p is the number of VIP-tokens described shortly, and n c is the size of non-VIP/remaining tokens. So, we have n p + n c = n and assume n p ≪ r ≪ n . (complexity analysis is provided in the Appendix.)

Parsing the complexity term: Let us unpack the term to assess its relevance. The first two terms O ( lrd 2 + lr 2 d ) pertain to the cost for a Transformer, while the remaining terms are the overhead of compression and decompression. The term O ( lr log( n c ) d + lrn p d ) is the overhead of compression and updating our data structure at each layer. The O ( nd ) term corresponds to pre-processing involving converting the hidden states into our data structure and post-processing to recover the hidden states from the data structure. Note that unlike the dependence on n for vanilla Transformers, this O ( nd ) is incurred only at the input/output stage of the Transformer, but not at any intermediate layers.

High level design choices. We use the standard Transformer layers with a standard feed-forward network (which results in d 2 in the first term) and standard quadratic cost self-attention (which gives the r 2 factor in the second term). Why? These choices help isolate the effect of incorporating their efficient counterparts. The proposed algorithm operates on the input/output of each Transformer layer leaving the Transformer module itself unchanged. Therefore, our goals are distinct from the literature investigating efficient self-attentions and efficient feed-forward networks. This is because one can replace these two vanilla modules with any other efficient alternatives to further reduce the r 2 and d 2 terms directly. Despite these quadratic terms, our approach is faster than baselines (§4).

We will first describe our general idea, as shown in Fig. 2, which uses VIP-tokens to guide the compression/decompression of the input/output of a Transformer layer so that it only needs to process the compressed sequence (§3.1, §3.2). Then, we will discuss an instantiation of the compression process, by adapting a multi-resolution analysis technique (§3.3). Lastly, we will introduce a data structure which allows more efficient compression/decompression (§3.4).

3.1 Elevating the Importance of a Few Tokens: VIP-Tokens

Let us start with the simplest compression, which identifies a linear transformation S ∈ R r × n which acts on the input, resulting in a smaller representation SX ∈ R r × d . Of course, a smaller r implies that more information about X is lost. But we find that in many tasks, only the embedding representations of a few tokens drive the final prediction: we refer to these tokens as VIP-tokens .

Examples of VIP-tokens: Observe that only the embedding outputs of masked tokens in masked language modeling [10] and the CLS token in sequence classification [10, 11] are/is used for prediction. In question answering, only the questions and possible answers associated with the questions are used for prediction. It is important to note that the masked tokens, CLS tokens, and question tokens are (1) defined by the tasks and (2) known to the model (although the embedding representation of these tokens are unknown). These VIP-tokens can be viewed as a task or question that is given to the model. The model can process the sequence with a specific goal in mind so that the model can skip/skim less relevant segments. Our general principle involves choosing a set of tokens as the VIP-tokens that (1) are important to the specific task goals and (2) easily pre-identifiable by the user.

Caveats. Not all important tokens can be pre-identified. For example, the tokens in the correct answer span in answer span prediction are also important to the specific goals, but are difficult to pre-identify, so only the question tokens (and not the answer tokens) are used as VIP-tokens. We assume that any other tokens that are relevant for prediction should have high dependency with these VIP-tokens. For example, the answer tokens should have high dependency (in self-attention) to the question tokens.

VIP-tokens occupy the front seats. VIP-tokens can occur anywhere within a sequence. But we can re-order the sequence as well as the positional encodings so that VIP-tokens are always at the head of sequence to make analysis/implementation easier. With this layout, let P ∈ R n p × d be the VIP-tokens and C ∈ R n c × d be the non-VIP/remaining tokens, X can be expressed as

This is possible since Transformer is permutation invariant when permuting positional encodings (embeddings or IDs) along with tokens. This re-ordering is performed only once for the input of the Transformer model, then the outputs generated by the model are rearranged to their original positions.

Re-ordering makes the analysis, implementation and presentation of our method much clearer and simpler. In fact, placing VIP tokens at the end of the sequence can also serve the same purpose.

From the above discussion, it is clear that one needs to make sure that after compressing the input tokens X , the VIP-tokens must still stay (more or less) the same, and the compression matrix S must be VIP-token dependent . We hypothesize that such VIP-token dependent compression matrices require a much smaller dimension r , compared to VIP-token agnostic compression matrices.

3.2 VIP-Token Centric Compression (VCC): An Initial Proposal

For a Transformer layer, let X denote its input matrix. Express the output of this layer as follows:

where S ∈ R r × n is a compression matrix compressing X to a smaller representation and S † is the pseudo inverse for decompression . With the layout in (6), we can write (7) as

where P new and C new are the new embeddings for P and C .

VIP-Token Centric Compression (

)

S

c

Transformer Layer (

)

γ

VIP-tokens

P

Remaining tokens

C

Residual Connection

Residual Connection

Decompression (

)

S

†

c

S

c

C

Always reserve seats for VIP-tokens. What is a useful structure of S ? Since P new is the embedding output for the VIP-tokens P , we want them to be fully preserved. To achieve this, we impose the following structure on S and S † :

The rearrangement simply says that we will avoid compressing P . But rewriting it in this way helps us easily unpack (8) to check the desired functionality of S c .

Prioritize information in VIP-tokens. Our goal is to ensure P new generated from the compressed sequence in (8) will be similar to its counterpart from the uncompressed sequence. Let us check (8) using the compression matrix S defined in (9) first. We see that

The orange color identifies terms where P new interacts with other compression-related terms C and/or S c . We primarily care about P new in (8), so the first (orange) row in (10) is the main concern. We see that P new only depends on the compressed SX via α ( P , SX , SX ) . We can further unpack,

Again, α ( P , SX , SX ) depends on C and S c via the second (orange) term. Normalization in softmax is omitted for simplicity of discussion. This helps us focus on the key term that matters: exp( PC ⊤ S ⊤ c ) S c C . As long as the following approximation using S c is good

we will obtain a good approximation of P new . Our remaining task is to outline a scheme of finding a compression matrix S c such that this criterion can be assured.

3.3 A Specific Instantiation via Multi-Resolution Compression

What should be the mechanics of our compression such that (12) holds? In general, to get S c , we can use any sensible data driven sketching idea which minimizes the error of (12). Doing so efficiently needs a bit of work; we describe the high level idea below and the low-level details are provided in Appendix.

High level idea. Ideally, an efficient scheme for constructing S c should operate as follows. If some regions of the sequence C have a negligible impact on (12) (via the orange terms above), the procedure should compress the regions aggressively. If other regions are identified to have a higher impact on (12) (again due to the orange terms above), the procedure should scan these regions more carefully for a more delicate compression. This suggests that procedurally a coarse-to-fine strategy may work. For example, multi-resolution analysis does help in approximating self-attention matrices in Transformers [39], but the formulation in [39] cannot be easily written in a form similar to (12), making it incompatible with our design. Nonetheless, we derive an analogous form (details in Appendix) that can be represented in a similar form as (12) and gives a strategy for obtaining S c C .

Specifically, let us define a compressed representation (via averaging) of the x -th s -length segment of sequence C : c s x ∈ R d

where s ∈ { k 0 , k 1 , k 2 , · · · , n c } assuming n c is a power of k and x ∈ { 1 , 2 , · · · , n c /s } . [ · ] i refers to the i -th row of the input matrix. We fix the increment ratio k = 2 for simplicity of discussion. The s represents the resolution of the approximation: it represents the number of non-VIP token embeddings being averaged into a vector c s x . Higher s (e.g., s = 8 in c 8 1 in Fig. 3) means lower resolution and heavier compression of the corresponding segment. The x represents the location of the s -length segment within the sequence C . In our scheme, we compress the sequence C and use a set of c s x for some selected s 's and x 's as the rows of the compressed S c C as seen in Fig. 3. The sequence C is broken into multiple segments of different lengths, then each segment is compressed into a vector c s x .

Procedurally, as shown in Fig. 3, our scheme starts with the heaviest compression and progressively refines certain segments of C guided by the VIP-tokens P . The scheme starts with the heaviest compression that treats C as a n c -length segment and compresses it to a single c n c 1 . Then, starting with s = n c (root node), the procedure (1) computes the averaged attention scores between VIP-tokens P and c s x 's for different x 's (averaged over all attention heads and all VIP-tokens; only one c n c 1 at level s = n c ). We note that the attention scores are obtained by extracting attention matrices from MHA module (2) of the current Transformer layer when using P as queries and c s x 's as keys. Then, it (2) splits the s -length segments corresponding c s x 's with higher averaged attention scores (one segment is split in Fig. 3 but we might split more segments, again only one c n c 1 at level s = n c ) into ( s/ 2) -length sub-segments: the corresponding c s/ 2 x (13) of each sub-segment is computed for finer representation. Then, at next level for s = n c / 2 , the same procedure proceeds. This process continues until the sub-segments have length

C

S

c

C

c

8

1

[

C

]

3

[

C

]

4

[

C

]

1

[

C

]

2

[

C

]

7

[

C

]

8

[

C

]

5

[

C

]

6

c

4

1

c

4

2

c

2

1

c

2

2

c

1

3

c

1

4

P

P

c

2

1

c

4

2

c

1

3

c

1

4

Compression

P

thicker = higher

attention score

1 . We note that this procedure is guided by the VIP-tokens P and designed to maximally reduce the error of approximating (12). No additional learnable parameters are introduced for this scheme. The technical details of this algorithm are less relevant for our overall approach, but for interested readers, the details are discussed in the Appendix.

How good is this approximation? The output P new (8) is well approximated since the approach preserves the relevant components of C that have a high impact on the output P new . Further, if the VIP-tokens P have high attention weights for some rows of C , then the corresponding row in C will be approximated with higher frequencies (less compressed). So, the output in

c

1

Δ

c

1

5

Δ

c

1

6

Δ

c

2

3

Δ

c

1

7

Δ

c

1

8

Δ

c

2

4

Δ

c

4

2

Δ

c

1

1

Δ

c

1

2

Δ

c

2

1

Δ

c

1

3

Δ

c

1

4

Δ

c

2

2

Δ

c

4

1

8

C new (8) for a subset of non-VIP tokens that have a higher dependency with the VIP-tokens will have a better approximation than the others, as desired. This property is useful since some tokens with unknown locations but manifesting a high dependency with the VIP-tokens can also relevant

to the final prediction of a Transformer model in some tasks. The answer in span-based question answering tasks is one example, and our construction ensures that they will be approximated well too.

3.4 Efficient Data Structure for Compression/Decompression

By employing the procedure in §3.3 illustrated in Fig. 3, we can find the compressed S c C with an O ( n c d + rn p d ) cost at each layer. The main cost O ( n c d ) is due to computing c s x defined in (13) for all resolution s and location x by using recursive relation from the bottom up:

We find that these steps could introduce a large overhead. Further, note that if we decompress (apply S † to) the output of γ for compressed sequence as in (8), the cost is O ( nd ) since the number of nonzero entries in S † is n (more details in Appendix). As a solution, we now introduce a data structure T ( · ) for storing C and C new , as shown in 4, which enables efficient computation of c s x and eliminates explicit decompression. We note that this data structure is only possible due to the specific structure of S c C constructed in §3.3. Specifically, T ( C ) stores c n c 1 and ∆ c s x defined as

for every resolution s = n c and location x . Similarly, T ( C new ) stores ( c new ) n c 1 and ∆( c new ) s x where ( c new ) s x and ∆( c new ) s x are defined similar to (13) and (15) but using C new instead of C .

̸

Then, given T ( C ) , any c s x can be retrieved efficiently in O (log( n c ) d ) cost via recursion:

The only reason we need decompression S † is that we need to obtain new representation C new (no decompression for P new since P is uncompressed). Suppose we have T ( C new ) , then we have an alternative way of getting C new similar to (16) (note ( c new ) 1 x = [ C new ] x ) without explicit decompression. The key benefit of this data structure is that we can obtain T ( C new ) by changing some nodes in T ( C ) . This only needs updating O ( r ) nodes, and each update takes O ( d ) cost.

An example. We show a T ( C ) for n c = 8 in Fig. 4. Let S c C = [ c 2 1 c 1 3 c 1 4 c 4 2 ] ⊤ as in Fig. 3. Since the segment c 2 1 is not split into sub-segments c 1 1 and c 1 2 , we have (details in Appendix):

By rearranging (17), we can verify that ∆( c new ) 1 1 , ∆( c new ) 1 2 in T ( C new ) stays the same as ∆ c 1 1 , ∆ c 1 2 in T ( C ) and thus do not need to be updated:

Further, we can verify that only the green nodes in Fig. 4 will be updated. These nodes correspond to the nodes in Fig. 3 that have been traversed. In summary, for each row c s x of S c C (a leaf node in Fig. 3), only the node storing ∆( c ) s x and its ancestor nodes in T ( C ) must be updated, so the total number of nodes (including their ancestors) being updated is O ( r ) . Next, we can update the nodes as follows: first, we get representations ( c new ) 2 1 , ( c new ) 1 3 , ( c new ) 1 4 , ( c new ) 4 2 by feeding S c C into Transformer layer (details in Appendix). At level s = 1 , given ( c new ) 1 3 and ( c new ) 1 4 , we (1) compute ( c new ) 2 2 via (14), and then (2) compute ∆( c new ) 1 3 and ∆( c new ) 1 4 via (15). The last two values are the new values for ∆ c 1 3 and ∆ c 1 4 in T ( C ) . At level s = 2 , given ( c new ) 1 2 and ( c new ) 2 2 computed at previous level, we apply similar procedure to obtain ( c new ) 4 1 , ∆( c new ) 1 2 , ∆( c new ) 2 2 , and the last two values are used to update two nodes in T ( C ) . It becomes apparent that each node update takes O ( d ) cost. Putting it togather: the complexity of modifying T ( C ) to T ( C new ) is O ( rd ) . The detailed algorithm and complexity analysis are described in Appendix.

By maintaining this data structure, we never need to materialize the entire C or C new in any intermediate layer, but instead we use (16) to construct the rows of S c C and perform updates to T ( C ) to obtain C new (represented as T ( C new ) ) at each intermediate layer. At the output of a Transformer, C new is materialized from T ( C new ) at a O ( n c d ) cost via the recursion (16) from the bottom up.

4 Experiments

We perform a broad set of experiments to empirically evaluate the performance of our proposed compression. (See hyperparameters/dataset statistics in Appendix.) We evaluate our method on both encoder-only and encoder-decoder architecture types. We compare our method with baselines on a large list of question answering and summarization tasks, where we found long sequences occur most frequently. Then, we study the model performance of scaling to ultra long sequences enabled by our method. Since efficiency is the focus of the efficient baselines and our work, we include runtime efficiency (of a single sequence) in millisecond in each table. (See the procedure for runtime measurement in Appendix.) We also include a discussion on FLOP efficiency in Appendix.

For ease of implementation and hyperparameter selection, we restrict the rows of S c C to have exactly two resolutions for experiments. Specifically, for a pre-defined increment ratio k , we split and refine all segments c s x with s > k to k -length sub-segments, and select h (pre-defined) k -length segments to further split to 1 -length sub-segments. So, the rows of S c C would consist of ( n c /k -h ) of c k x and hk of c 1 x for some x . To simplify the implementation, we only use the proposed compression in the encoder, and use the vanilla computation in the decoder of encoder-decoder models. We note that our method might be applied to the decoder (more details in Appendix).

Further, we found a few layers of standard Transformer layers to pre-process tokens helps the performance. Therefore, in the initial stage of a Transformer, we segment input sequence into multiple 512 -length segments. For each segment, we use vanilla computation in the first 4 layers (for base models and 6 layers for larger models) of a Transformer. Then, for the remaining layers, segments are concatenated back into one sequence and processed using our proposed compression. There is no communication among any segments, so the downstream tasks cannot be solved by these first 4 transformer layers alone, and the initial stage is used just for getting a reasonable representation for the compression to operate on.

Approximation Quality of VIP-Tokens . We empirical measured the approximation quality of our VIP centric strategy compared to random strategy (the tree growth in Fig. 3 is not guided by VIP-tokens, but is random) and lazy strategy (each k-length segment is compressed to a token). P new is the approximated representation of VIP tokens computed with compression and let P ∗ new be the ground truth representation of

VIP tokens computed without compression. We measure the relative error (defined as || P new -P ∗ new || F / || P ∗ new || F ) and correlation coefficient between P new and P ∗ new . As shown in Tab. 1, we can verify that the proposed procedure indeed improve the approximation quality of VIP-tokens.

Encoder-Only Models . For encoder-only architecture, we compare our method with RoBERTa [22] and three strong baselines: Longformer [2], Big Bird [38], and MRA Attention [39]. We first pretrain a RoBERTa model using masked language modeling task, then for each method, we perform continuous pretraining from the RoBERTa checkpoint to expand the positional embeddings to 4K length and adjust model parameters to adapt approximations used in effi-

cient baselines and our method. We verify that our proposed method can be integrated into a pretrained Transformer with some continuous pretraining. But we note that the amount of reduction in log perplexity for our method ( -0 . 114 ) during pre-training is much larger than Longformer ( -0 . 017 ) and Big Bird ( -0 . 025 ) from 50K steps to 250K steps. The continuous pretraining for these baselines might have saturated since only the self-attention is approximated while our method might require more pretraining to adjust the parameters for more aggressive approximation. So, we run a larger scale pretraining for our method; downstream results are in Tab. 2 and Fig. 5, denoted with *. We use HotpotQA [37], QuALITY [24], and WikiHop [33] to assess the language models. HotpotQA is an answer span extraction task, while QuALITY and WikiHop are multi-choice question answering tasks. We set questions and multi-choice answers (for QuALITY and WikiHop) as VIP-tokens.

As shown in Tab. 2, we verify that our method is consistently better compared to Longformer and Big Bird. Our method obtains better accuracy in QuALITY and WikiHop compared to 4K length

RoBERTa model, but it is a bit worse than 4k length RoBERTa model on HotpotQA. More pretraining helps close the gap. We also use WikiHop to experiment with method specific hyperparameters (such as block size in Big Bird, window size in Longformer, and compression size r in our method). As shown in Fig. 5, our runtime efficiency frontier is consistently better than the baselines. The key takeaway is that our method has a much better runtime efficiency than baselines that have the same sequence length without sacrificing its model performance. Further, we note that our method can be scaled to larger models for accuracy improvement.

Encoder-Decoder Models . We compare our method with T5 [27], LongT5 [14], and LED [2]. We use the public pretrained checkpoints for baselines. The pretrained models for our method are obtained by doing continuous pretraining from the public T5 checkpoints using T5's pretraining task [27]. We note that LED-base has 6 encoder/decoder layers compared to 12 encoder/decoder layers in base models of other methods, its accuracy is usually lower. So, LED is evaluated in limited tasks. We use HotpotQA [37], WikiHop [33], CNN/Dailymail [29], MediaSum [43], Arxiv [7], and GovReport [15], SummScreenFD [5], QMSum [42], Nar-

10 2

2 × 10 2

6 × 10 2

Runtime (ms)

69

71

73

75

77

RoBERTa

Big Bird

Longformer

Ours

Ours*

rativeQA [19], Qasper [9], QuALITY [24], ContractNLI [21] from SCROLLS benchmark [30] to assess the language models. For question answering tasks, we set questions and multi-choice answers (for QuALITY and WikiHop) as VIP-tokens in our method. For query-based summarization, such as QMSum, we use the query as VIP-tokens in our method. For general summartization tasks, we prepend a 'summarize:' in each instance and use it as VIP-tokens in our method. Our method achieves matching or better performance in most tasks compared to T5, LongT5, and LED with much higher efficiency (see Tab. 3). Further, the performance monotonically increases with the model size, so our method can scale to larger models.

Scaling to Longer Sequences . The prior experiments limit the sequence length to at most 4K or 16K since the baselines can only be scaled up to these lengths. However, our method can be scaled

to much longer sequences. We note that NarrativeQA [19] is an ideal testbed as shown in dataset statistics in Appendix. The results are shown in Tab. 4. The left / middle / right values of runtime column are for the entire model / the encoder / the last 8 layers (out of 12 layers) that uses our compression. The performance monotonically increases as sequence length increases. We note that for sequence length 64K, the performance of model with k = 64 is lower than the model with k = 16 . We suspect that since the results are finetuned from the same model that is pretrained with k = 16 , the large gap between the two different k 's may have a negative impact on finetuning performance. Nevertheless, the performance is still higher than 32K length models.

Why focus on 4K - 128K lengths? We believe that the computation required by standard Transformers for processing shorter sequences is not an efficiency bottleneck. As a result, we do not profile the performance of our method for smaller length sequences, since the standard Transformers are sufficiently fast in this case. Further, while our model can be applied to shorter sequences, we suspect that for shorter sequences, there may be less irrelevant information for VIP-tokens. So compressing the irrelevant information will not offer a meaningful speed

up. This is a limitation as the compression works better when there is more compressible information. We have only pushed the sequence lengths to 128K since this length was sufficient to cover a majority of sequence lengths encountered in long sequence tasks (for example, our model is able to process an entire book at once).

5 Limitations

Our method assumes that in many tasks, a subset of tokens are disproportionately responsible for the model prediction, the remaining non-VIP-tokens may play a role but are less critical. Our method excels specifically on such tasks by selectively locating relevant information in the sequence for given VIP-tokens. As the experiments show, this choice is effective in many cases but this behavior is not universal. Occasionally, an embedding is pre-computed which must then serve multiple tasks concurrently, e.g., both text retrieval and natural language inference. In this case, if we do not know the tasks beforehand, VIP-token selection cannot be meaningfully performed. Further, VIP-token selection requires some understanding of the tasks. However, we believe that a reasonable selection can be made with some generic knowledge for most tasks or use cases. Please see Appendix for VIP-token selection guidelines.

To reduce the complexity of our implementation, the method is currently setup for the encoder module of the Transformer that assumes full access to the entire sequence. The proposed compression might be extended to approximate the computation in the decoder, but it needs more implementation work, so we leave it as future work. Consequently, the current implementation is less useful for decoder-only models (but in the appendix, we discuss some strategies).

6 Conclusions

We propose a VIP-token centric sequence compression method to compress/decompress the input/output sequences of Transformer layers thereby reducing the complexity dependency on the sequence length n without sacrificing the model accuracy. Our empirical evaluation shows that our method can be directly incorporated into existing pretrained models with some additional training. Also, it often has much higher efficiency compared to baselines with the same sequence length while offering better or competitive model accuracy. For future work, we believe that extending our method to the decoder of the encoder-decoder and decoder-only models will further boost the efficiency of Transformers while maintaining similar model performance.

Acknowledgments. Zeng and Singh were supported in part by funding from the Vilas Board of Trustees and UW-Madison Office of the Vice Chancellor for Research and Graduate Education.

References

Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer, 2022.

Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision , pages 213-229. Springer, 2020.

Mingda Chen, Zewei Chu, Sam Wiseman, and Kevin Gimpel. SummScreen: A dataset for abstractive screenplay summarization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 8602-8615, Dublin, Ireland, May 2022. Association for Computational Linguistics.

Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations (ICLR) , 2021.

Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian. A discourse-aware attention model for abstractive summarization of long documents. Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 2 (Short Papers) , 2018.

Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. Funnel-transformer: Filtering out sequential redundancy for efficient language processing. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 4271-4282. Curran Associates, Inc., 2020.

Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. A dataset of information-seeking questions and answers anchored in research papers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 4599-4610, Online, June 2021. Association for Computational Linguistics.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Conference of the North American Chapter of the Association for Computational Linguistics (NAACL) , 2019.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR) , 2021.

Alexander Fabbri, Irene Li, Tianwei She, Suyi Li, and Dragomir Radev. Multi-news: A largescale multi-document summarization dataset and abstractive hierarchical model. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 1074-1084, Florence, Italy, July 2019. Association for Computational Linguistics.

Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. LongT5: Efficient text-to-text transformer for long sequences. In Findings of the Association for Computational Linguistics: NAACL 2022 , pages 724-736, Seattle, United States, July 2022. Association for Computational Linguistics.

Luyang Huang, Shuyang Cao, Nikolaus Parulian, Heng Ji, and Lu Wang. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 1419-1436, Online, June 2021. Association for Computational Linguistics.

Xin Huang, Ashish Khetan, Rene Bidart, and Zohar Karnin. Pyramid-BERT: Reducing complexity via successive core-set based token selection. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 8798-8817, Dublin, Ireland, May 2022. Association for Computational Linguistics.

Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations (ICLR) , 2020.

Bryan Klimt and Yiming Yang. The enron corpus: A new dataset for email classification research. In Proceedings of the 15th European Conference on Machine Learning , ECML'04, page 217-226, Berlin, Heidelberg, 2004. Springer-Verlag.

Tomáš Koˇ ciský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The NarrativeQA reading comprehension challenge. Transactions of the Association for Computational Linguistics , 6:317-328, 2018.

Philipp Koehn. Europarl: A parallel corpus for statistical machine translation. In Proceedings of Machine Translation Summit X: Papers , pages 79-86, Phuket, Thailand, September 13-15 2005.

Yuta Koreeda and Christopher Manning. ContractNLI: A dataset for document-level natural language inference for contracts. In Findings of the Association for Computational Linguistics: EMNLP 2021 , pages 1907-1919, Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.

Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

Piotr Nawrot, Jan Chorowski, Adrian Ła´ ncucki, and Edoardo M. Ponti. Efficient transformers with dynamic token pooling, 2022.

Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson, He He, and Samuel Bowman. QuALITY: Question answering with long input texts, yes! In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 5336-5358, Seattle, United States, July 2022. Association for Computational Linguistics.

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations (ICLR) , 2021.

Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations , 2020.

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research , 21(140):1-67, 2020.

David Saxton, Edward Grefenstette, Felix Hill, and Pushmeet Kohli. Analysing mathematical reasoning abilities of neural models. In International Conference on Learning Representations , 2019.

Abigail See, Peter J. Liu, and Christopher D. Manning. Get to the point: Summarization with pointer-generator networks. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 1073-1083, Vancouver, Canada, July 2017. Association for Computational Linguistics.

Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. SCROLLS: Standardized comparison over long language sequences. In EMNLP , 2022.

Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena : A benchmark for efficient transformers. In International Conference on Learning Representations (ICLR) , 2021.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems (NeurIPS) , 2017.

Johannes Welbl, Pontus Stenetorp, and Sebastian Riedel. Constructing datasets for multi-hop reading comprehension across documents. Transactions of the Association for Computational Linguistics , 6:287-302, 2018.

Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference on Learning Representations , 2022.

Yunyang Xiong, Zhanpeng 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 on Artificial Intelligence , 2021.

Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019.

Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. HotpotQA: A dataset for diverse, explainable multi-hop question answering. In Conference on Empirical Methods in Natural Language Processing (EMNLP) , 2018.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems (NeurIPS) , 2020.

Zhanpeng Zeng, Sourav Pal, Jeffery Kline, Glenn M Fung, and Vikas Singh. Multi resolution analysis (MRA) for approximate self-attention. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning , volume 162 of Proceedings of Machine Learning Research , pages 25955-25972. PMLR, 17-23 Jul 2022.

Zhanpeng Zeng, Yunyang Xiong, Sathya Ravi, Shailesh Acharya, Glenn M Fung, and Vikas Singh. You only sample (almost) once: Linear cost self-attention via bernoulli sampling. In International Conference on Machine Learning (ICML) , 2021.

Sixiao Zheng, Jiachen Lu, Hengshuang Zhao, Xiatian Zhu, Zekun Luo, Yabiao Wang, Yanwei Fu, Jianfeng Feng, Tao Xiang, Philip HS Torr, et al. Rethinking semantic segmentation from a sequence-to-sequence perspective with transformers. arXiv preprint arXiv:2012.15840 , 2020.

Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, and Dragomir Radev. QMSum: A new benchmark for query-based multi-domain meeting summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies , pages 5905-5921, Online, June 2021. Association for Computational Linguistics.

Chenguang Zhu, Yang Liu, Jie Mei, and Michael Zeng. Mediasum: A large-scale media interview dataset for dialogue summarization. arXiv preprint arXiv:2103.06410 , 2021.

7 Appendix

7.1 Definition of Notations

we treat v i as a row of the matrix.

We provide a table 5 of notations that are used for more than once so that the readers can refer to their definition easily.

7.2 Details of Multi-Resolution Compression

We describe the omitted technical details of a modified formulation of [39] to construct S c C and corresponding S c satisfying good approximation of

Before diving into the technical details of constructing S c , we introduce some notations and tools that will be used later. We use [ · ] i to refer the i -th entry/row of the input vector/matrix and [ · ] i,j to refer the ( i, j ) -th entry of the input matrix. BOLD uppercase letters denote matrices, bold lower case letters denote vectors, and regular lower case letters denote scalars. Let v i be a vector, when we write a matrix of form

we treat v i as a column of the matrix. When we write a matrix of form

7.2.1 Basic Problem Setup

Let b s x ∈ R n c be a multi-resolution component defined as

for s ∈ { k 0 , k 1 , k 2 , · · · , n c } assuming n c is a power of k (we assume k = 2 for simiplicity). Here, s and x represent the scaling and translation (similar to the concepts in wavelet basis) of the component, respectively. Fig. 6 is the visualization of b s x .

Then, any 1D signal f ∈ R n c can be represented as a linear combination of b s x :

sx

-

s

s

1

s

where c s x are the coefficients for the linear combination. For a signal with multi-resolution structure (that is, signal has high frequency in some regions and has low frequency in other regions), we can find an approximation ˆ f ∗ that can be expressed as a sparse linear combination where most coefficients are zeros, as shown in Fig. 7.

We denote J as the set of major components b s x corresponding to the large coefficients, that is, J := { b s x | | c s x | being large } . Since the set of all possible b s x is an over-complete dictionary, there are multiple possible linear combinations. To reduce the search space of the best set J , we place a mild restriction on the set J :

̸

̸

The conditions state that each entry of signal f is included in the support region of exactly one component in J . With these tools, we will first describe the approximation when J is given, then discuss how the approximation connects the set J to our target S c and S c C . Finally, we will discuss how to construct this J .

0.0

0.5

f(x)

signal

0

100

200

300

400

500

0.0

0.5

f(x)

approx

x

7.2.2 Plugging Our Problem into the Setup

A recent result shows that the self-attention matrix exp( XX ⊤ ) has the multi-resolution structure discussed above [39]. Since exp( PC ⊤ ) is a sub-matrix of exp( XX ⊤ ) , we conjecture that the multiresolution structure also holds in exp( PC ⊤ ) . As a result, we can find a sparse combination of b s x to represent rows of exp( PC ⊤ ) .

Claim 7.1. Given the set J satisfying restriction (25) , we can define an approximation of the i -th row of exp( PC ⊤ ) similar to (24) as illustrated in Fig. 7

where c s x is the optimal solution that minimizes

Then, the approximation can be written as:

̸

where b s x ∈ J is the component that is supported on j (a.k.a. [ b s x ] j = 0 and there is exactly one b s x ∈ J satisfy this condition due to restriction (25) ).

Proof. If J is given, let B ∈ R n c ×|J| be a matrix whose columns are elements b s x ∈ J and let c ∈ R |J | be a vector whose entries are the corresponding c s x :

then the approximation can be expressed as

Due to the restriction (25), the columns of B are orthogonal, so B ⊤ B is a diagonal matrix:

We can also write down B ⊤ [ exp( PC ⊤ ) ] i

Putting everything together, we have

We note that s b s x simply re-scale the entry of b s x such that any non-zero entry becomes 1 . Then, let us consider j -th entry of Bc . Due to the restriction (25), we have exactly one b s x ∈ J whose support region contains j , so the j -th row of the first matrix at the right hand side of (35) contains exactly a 1 and the remaining entries are 0 . Therefore, we have

where b s x ∈ J is the component that is supported on j , which concludes our proof.

If we solve for then

7.2.3 Efficient Approximation

We note that computing (28) for all j would require access to the entire [ exp( PC ⊤ ) ] i . We exploit the same strategy as described in [39], so the exponential of inner product is used as an approximation to inner product of exponential.

We note that ⟨ [ PC ⊤ ] i , b s x ⟩ is the local average of the support region of b s x , which is also the x -th s -length segment of sequence [ PC ⊤ ] i :

̸

By using some arithmetic manipulations, (37) can be efficiently computed

̸

̸

where c s x is defined in the main text:

We note that c s x is the local average of the x -th s -length segment of sequence C . The c s x can be efficiently computed via

Claim 7.2. Given the set J satisfying restriction (25) , let S c be be a matrix whose rows are elements b s x ∈ J

where ̂ exp( PC ⊤ ) is defined as (37) .

Proof. Consider i -th row of exp( PC ⊤ S ⊤ c ) ,

Then, we have

We note that s b s x simply re-scales the entry of b s x such that any non-zero entry becomes 1 . Then, let us consider j -th entry of [ exp( PC ⊤ S ⊤ c ) DS c ] i . Due to the restriction (25), we have exactly one

Then,

̸

b s x ∈ J whose support region contains j , so the j -th column of the second matrix in the right hand side of (45) contains exactly a 1 and the remaining entries are 0 . Therefore, we have

where b s x ∈ J is the component that is supported on j . The second equality is based on (39).

Claim 7.3. If S c and D are defined as Claim 7.2, the pseudo inverse of S c is simply S † c = S ⊤ c D , so each row of S † c and S † contain exactly a 1 (so the number of nonzero entries of S † c and S † are n c and n respectively).

̸

Proof. Since each row of S c is some b s x ∈ J , due to the restriction (25), for i = j ,

As a result, S c S ⊤ c D = I . Further, S ⊤ c DS c is a symmetric matrix. So, all Moore-Penrose conditions are verified. S † c = S ⊤ c D .

From the restriction (25), we have every column of § c contains exactly a non-zero entry. Also,

Since the non-zero entry of b s x is simply 1 s by definition, s b s x simply re-scales the entry of b s x such that any non-zero entry becomes 1 . As a result, each row of S † c has exactly a 1 . Also, by the relation between S and S c :

each row of S † has exactly a 1 .

At the end, the approximation

does not look exactly as (19), but we can insert a simple diagonal matrix D to the formulation (19) and make the whole thing work.

7.2.4 How to Construct J for S c and S c C ?

The derivation so far assume access to J , but in practice, we have no knowledge of J and need to construct J that leads to good approximation. With the approximation scheme in place, we can now analyze the approximation error, which will be leveraged later to find a reasonable set of components J . The approximation error of i -th row of exp( PC ⊤ ) can be expressed as

̸

̸

̸

̸

The second equality and fourth equality are due to (25) and (39). The approximation error is governed by two components multiply together: attention score between [ P ] i and the local average c s x of the x -th s -length segment of sequence C and the inner product of [ P ] i with the amount of deviation of [ C ] j from its local average c s x .

When s = 1 , the deviation is simply zero:

It is reasonable to assume that the deviation [ C ] j -c s x is smaller if s is smaller. Therefore, this actually suggests a simple heuristic for selecting J : when exp( ⟨ [ P ] i , c s x ⟩ ) is large, we should approximate the x -th s -length segment of C with higher resolution (by splitting the segment to shorter sub-segments and using finer approximation). This heuristic describes the selection criteria for one row of exp( PC ⊤ ) , which corresponds to a single VIP-token, for multiple rows of exp( PC ⊤ ) (for multiple VIP-tokens), we simply use

as selection criteria since J is shared by all VIP-tokens.

̸

The construction of J is described in Alg. 1. This algorithm describes the same procedure as the Figure 3 in the main text. The b s x 's in J are the rows of S c , and the corresponding c s x 's (40) are the rows of S c C . The budgets h 2 , h 4 , · · · , h n c required by Alg. 1 is used determine the number of components at each resolution that will be added to J . Specifically, there are 2 h 2 s -h s number of components b s x for s = 1 based on simple calculations. We can choose budgets such that the final size of J is r -n p to make the length of compressed sequence to be r .

7.2.5 How Good is This Approximation?

At high level, the compression S c performs more compression on tokens that are not relevant to the VIP-tokens and less compression to tokens that are important to the VIP-tokens. We will discuss it in more details. Since each row of S † contain exactly a 1 as stated in Claim 7.3, S † can commute with β , so in summary, we can write the approximation of the computation of a Transformer layer as

Note that D is added as discussed in (50).

There are four main approximation components (purple) in (54). Taking the fact that DS c = ( S † c ) ⊤ , all of these approximations are row or column space multi-resolution approximations governed by S c matrix. High attention weight implies higher dependency, and the procedure in Alg. 1 refines regions with large attention weights with higher resolutions. Therefore, the token embedding in C that have higher dependency to P are better approximated. The output P new is well approximated by design since the approximation preserves the higher frequency components of the subset of rows of

C that has high impact on the output P new . Further, the output in C new corresponding to the subset of rows of C that have higher dependency with the VIP-tokens will have better approximation than the remaining rows of C . This property addresses the issue that some tokens with unknown locations are also relevant to the final prediction of a Transformer in some tasks. For example, in question answering tasks, candidate answers are usually expected to have large dependency with question tokens (VIP-tokens), so they are approximated well as well. This approximation property is exactly what we need.

7.2.6 Relation to [39] that Inspires Multi-Resolution Compression

Our work and [39] can be viewed as operating at slightly different levels of abstractions. While [39] tries to approximate self-attention computation efficiently, our paper proposes a general framework for performing a VIP-token centric compression on the sequence to efficiently handle extremely long sequences (the self-attention module remains completely unchanged). Our VIP-token centric compression involves a number of steps described in the main text. But one of the key steps involves constructing a compression matrix S c which has some desirable properties, namely satisfying (19) which we elaborate further below.

Note that for equation (19), we need a matrix S c such that the approximated attention matrix involving P and C is similar to the true attention matrix involving P and C . This is precisely where the general idea of [39] can be used. But the formulation in [39] cannot be applied directly in its original form since it cannot give us S c . Why? One reason is that the formulation in [39] cannot be written as matrix form similar to equation (19). This may be a reason why [39] has to use custom CUDA kernels in their implementation. Nonetheless, the properties of [39] are useful. So we derive the analogous form but for 1D instead: this 1D case is expressed as applying a matrix (this is the S c we are looking for) to the signal C .

One bonus of this modification is that it also removes the need for custom CUDA kernels. At a high level, [39] offers a multi-resolution view of the self-attention matrices, and our modified version is best thought of as a similar multi-resolution view of the sequence itself. But we can also substitute in a different means of obtaining S c (which could simply be a sketching matrix). Finally, we note that a naive implementation of the resulting modification still requires a O ( n c d ) cost due to the computation of c s x for all possible scaling s and translation x . There is a similar cost in [39] (second paragraph in section 4.4 in [39]). The data structure we propose reduces this cost.

7.3 Details of Proposed Data Structure

In section, we describe some omitted technical details of the proposed data structure T ( · ) .

Claim 7.4. Given the set J satisfying restriction (25) , if b s x ∈ J , then ( c new ) s x -c s x = ( c new ) s ′ x ′ -c s ′ x ′ for all b s ′ x ′ satisfying the support of b s ′ x ′ is contained in the support of b s x (the x ′ -th s ′ -length segment of C is a sub-segment of the x -th s -length segment of C ).

Proof. To simplify the notations a bit, without loss of generality, we assume x = 1 . Then, for i ≤ s , consider ( c new ) 1 i :

By Claim 7.3, S † c = S ⊤ c D and i -th row of S † c contains exactly a 1 . The column that contains 1 in the i -th row of S † c is exactly s b s 1 since i is contained in the support of exactly one components in J due to the restriction (25). Denote this column index as j , then

Note that this holds for all i ≤ s . As a result, for i, i ′ ≤ s ,

Algorithm 2 Computation of one Transformer layer with T ( C )

Input: VIP-tokens P and data structure T ( C ) Use Algo. 1 to construct J but use (63) to retrieve c s x from T ( C Construct S c , S c C associated with J using Claim 7.2

Compute

Set T ( C new ) ←T ( C ) for s ← 1 , 2 , 4 , · · · n c / 2 do

Locate row location b s x in S c , refer the index as j for b s x ∈ J do

Mark

s

x

end for for

dirty

∆(

c

new

c

)

new

)

s

x

dirty do

Compute

(

c

new

)

2

s

⌈

x/

2

⌉

and update

end for end for

Update

n

c

1 Output: VIP-tokens P new and data structure T ( C new )

Then,

The rest follows from induction.

Claim 7.5. We have

And the updated representation ( c new ) s x of the corresponding c s x (a row of S c C ) is the corresponding row of S c C new .

Proof. By definition,

Then,

Since the S c is the same for S c C new and S c C , the second statement follows.

7.3.3 Algorithm for Making T ( C ) into T ( C new )

In this section, we describe the exact algorithm to update T ( C ) into T ( C new ) . The pseudo code is described in Alg. 2 where c s x is computed via

We use the term 'dirty' in Alg. 2 to indicate the node needs to be handled due to node updates. This term is commonly used in computer cache implementations to indicate that the data of a specific location has been updated and needs to be accounted for.

(

c

new

∆(

)

)

∆(

c

new

s

x

)

7.4 Complexity Analysis

In this section, we will discuss the detailed complexity analysis of our proposed method. The overall complexity of our proposed method is O ( lrd 2 + lr 2 d + lr log( n c ) d + lrn p d + nd ) when using the proposed efficient data structure.

7.4.1 Preparing Input Sequence to T ( C ) : O ( nd )

At the first layer, we need to permute the rows of X into [ P ; C ] , which takes O ( nd ) cost. Then, we process C into T ( C ) . This requires (1) computing c s x defined in (41). c 1 x = [ C ] x , so no compute is needed. With all c 1 x given, computing all c 2 x takes O ( n c d/ 2) . With all c 2 x given, computing all c 4 x takes O ( n c d/ 4) ... So, the cost is

Then (2) computing ∆ c s x for all s and x . Computing each ∆ c s x takes O ( d ) when given c s x and c 2 s ⌈ x/ 2 ⌉ . The amount of cost is the same as the number of nodes in the tree T ( C ) , so the cost is O ( n c d ) . Note that n c < n , so the overall complexity of the above operations is O ( nd ) .

7.4.2 Constructing J , S c , S c C : O ( lr log( n c ) d + lrn p d )

We can analyze the complexity of constructing J using Algo. 1. There is only one possible µ n c x . Then for each s , there are 2 h s number of µ s/ 2 x being computed since there are 2 components b s/ 2 x for each b s x ′ . As a result, we need to compute O (1 + ∑ s 2 h s ) number of µ s/ 2 x . When c s/ 2 x is given, the cost of computing a µ s/ 2 x is O ( n p d ) , so the overall cost of constructing J is O ((1 + ∑ s 2 h s ) n p d ) .

Further, at each s , the size of J is increased by h s since h s segments are split into 2 h s sub-segments, so the size of J is O ( ∑ s h s ) . Since S c ∈ R ( r -n p ) × n and |J | = r -n p as discussed in §7.2.4, O ( r -n p ) = O ( ∑ s h s ) . We use O ( r ) for simplicity instead of O ( r -n p ) . As a result, the overall cost of constructing J is O ( rn p d ) .

The above cost assumes c s/ 2 x is given. If we compute all possible c s/ 2 x using (41), the cost will be O ( n c d ) as analyzed in §7.4.1. However, if we employ the proposed data structure, each c s/ 2 x can be retrieved in at most O (log( n c ) d ) by recursively computing (63). Since we need to retrieve O (1 + ∑ s 2 h s ) = O ( r ) number of c s x , the complexity of computing necessary c s x is O ( r log( n c ) d ) .

As a result, the complexity of constructing J is O ( rn p d + r log( n c ) d ) at each layer. When summing the cost over all layers, the complexity is O ( lrn p d + lr log( n c ) d ) .

By Claim 7.2, the rows of S c and S c C are simply the b s x ∈ J and the corresponding c s x , which are already computed during the construction of J , so we essentially can get these S c and S c C for free.

7.4.3 Feeding Compressed Sequence into a Transformer Layer: O ( lrd 2 + lr 2 d )

At each layer, we need to compute

for updating T ( C ) . This is the part of a Transformer layer that requires heavy computation. It can be verified that the complexity of a Transformer layer is O ( nd 2 + n 2 d ) for a input sequence of length n . Now a compressed sequence of length r is fed into a Transformer layer, the cost is simply O ( rd 2 + r 2 d ) . We note that there is an additional re-scaling to plug D into exp( PC ⊤ S ⊤ c ) DS c during multi-head attention computation discussed in 50. However, the additional cost of applying D is O ( rd ) , which does not change the complexity. When summing the cost of all layers, the overall complexity is O ( lrd 2 + lr 2 d ) .

7.4.4 Updating T ( C ) into T ( C new ) : O ( lrd )

Once (65) is computed, we need to change T ( C ) into T ( C new ) . The cost of change T ( C ) into T ( C new ) is O ( rd ) as analyzed in the main text. For more specific analysis, let us take a look at the first three iterations:

At the first iteration, there are O (2 h 2 ) number of ( c new ) 1 x to be computed at the first inner for loop, and there are O (2 h 2 ) number of ∆( c new ) 1 x to be updated in the second inner for loop. Additional O ( h 2 ) number of ∆( c new ) 2 ⌈ x/ 2 ⌉ are masked dirty.

At the second iteration, there are O (2 h 4 ) number of ( c new ) 2 x to be computed at the first inner for loop, and there are O (2 h 4 + h 2 ) number of ∆( c new ) 2 x to be updated in the second inner for loop. The second term is due to the dirty ∆( c new ) 2 ⌈ x/ 2 ⌉ from the first iteration. Additional O ( h 4 + h 2 2 ) number of ∆( c new ) 4 ⌈ x/ 2 ⌉ are masked dirty.

At the third iteration, there are O (2 h 8 ) number of ( c new ) 4 x to be computed at the first inner for loop, and there are O (2 h 8 + h 4 + h 2 2 ) number of ∆( c new ) 4 x to be updated in the second inner for loop. The second and third term is due to the dirty ∆( c new ) 4 ⌈ x/ 2 ⌉ from the second iteration. Additional O ( h 8 + h 4 2 + h 2 4 ) number of ∆( c new ) 8 ⌈ x/ 2 ⌉ are masked dirty.

It becomes apparent that if we sum over the number of computes of ( c new ) s x and updates of ∆( c new ) s x , the total number is O ( ∑ s 2 h s + 2 ∑ s ∑ log( s ) j =1 h s 2 j ) = O ( ∑ s h s + ∑ s h s ) = O ( r ) . Since each compute and update takes O ( d ) cost, the overall complexity of changing T ( C ) into T ( C new ) is O ( rd ) . When summing the cost of all layers, the overall complexity is O ( lrd ) .

7.4.5 Materializing C new from T ( C new ) at the Last Layer: O ( nd )

At the output of the last layer, we can (1) compute all ( c new ) n c / 2 x via (63) at a cost of O (2 d ) , (2) compute ( c new ) n c / 4 x via (63) at a cost of O (4 d ) ... until all ( c new ) 1 x are computed. Then, [ C new ] x = c 1 x is materialized from T ( C new ) at a total cost of

Lastly, undoing the permutation so that [ P new ; C new ] are re-ordered to the original positions has a complexity of O ( nd ) . As a result, the overall complexity is O ( nd ) .

7.4.6 Overall Complexity

In summary, the overall complexity of our method is

7.5 Experiments

We run all experiments on NVIDIA A100 GPUs. All code is implemented using the standard PyTorch framework. No custom CUDA kernels are needed. As a result, it can be easily deployed to other platforms or ML frameworks. We will publish all code and checkpoints necessary for reproducibility concurrently with the paper publication.

7.5.1 Runtime Measurement

The runtimes presented in the experiment section are measured runtimes of a complete training step (including both forward and backward). For each method, we use the largest batch size that can fit into a 80GB A100 and measure the average latency of 10 steps. Then, the average latency is divided by the batch size to get the estimated runtime for a single instance. Via this procedure, we seek to measure the peak efficiency of each method when the GPU is at the highest possible utilization.

7.5.2 FLOPs

Some might be interested about the FLOP efficiency of each method, so we provide FLOP profiling results for some experiments. The automatic tools that we used for calculating FLOPs is deepspeed's FlopsProfiler. We note that the FLOP profiling tools do not always work and may throw errors or produce incorrect results when the code contains custom CUDA kernels. The results are summarized in Tab. 8 and Tab. 9. × in some entries means the profiler throws errors and cannot estimate the FLOPs numbers.

Also, the profiler throws errors when we just profile the encoders of encoder-decoder models, so the FLOP numbers in Tab. 9 are for the entire models. As shown in the tables, our method has the lowest FLOPs.

However, we note that FLOP numbers do not always reflect practical latency reductions. Memory bandwidth and latency (which are not captured by FLOP numbers) also play an important role in the overall latency. For example, sparse matrix multiplication (with unstructured sparsity) usually has a much lower FLOP count than dense matrix multiplication. But, the latency reduction is only possible when sparsity is at least 95% or more (depending on the implementation) since sparse matrix multiplication is a memory bandwidth bounded operator.

7.5.3 Pretraining

We use a filtered The Pile dataset [13] for all pretrainings. Since we are using public pretrained tokenizers, we want to enable the distribution of pretraining corpus aligns well with the distribution of corpus used to create the tokenizers. As a result, we use tokens per byte as a proxy for alignment of distributions and filter out PubMed Central, ArXiv, Github, StackExchange, DM Mathematics [28],

Ubuntu IRC, EuroParl [20], YoutubeSubtitles, and Enron Emails [18] components, which have tokens per byte greater than 0 . 3 . Then, the remaining corpus of The Pile dataset is used for pretraining.

For encoder-only models, we pretrain RoBERTa for 750K steps. A batch consists of 8,192 sequences of 512 length. The masking ratio for masked language modeling (MLM) is 15%. Then, 4K length models are continuously pretrained from the RoBERTa checkpoints for 300k steps. The positional embeddings are extended by duplicating the pretrained 512 positional embedding multiple times. For 4K length RoBERTa, Longformer, Big Bird and MRA Attention, the batch size is 64, and the masking ratio is 15%. With 15% masking ratio, there are roughly 616 masked tokens scattering in the sequences. We find that using 616 scattered masked tokens as VIP tokens for 4,096 length sequences might not be indicative for VIP-token centric compression, so we use masking ratio 7.5% and

10 2

2 × 10 2

6 × 10 2

Runtime (ms)

69

71

73

75

77

RoBERTa

Big Bird

Longformer

Ours

Ours*

Ours-150k

batch size 128 for our method. The number of masked tokens per sequence is reduced, and the number of total masked token predictions remains the same during pretraining. We note that with larger batch size, the wall clock pretraining runtime for our method is still smaller than baselines. In case that anyone is interested, we also show downstream finetuning on our method pretrained on the same number of tokens but fewer number of masked token predictions in Tab. 7 and Fig. 8, denoted as Ours-150k. The accuracy is consistent with our model pretrained on 300k steps. For the larger scale pretraining denoted with *, we pretrain our method for 250K steps with batch size 512 and masking ratio 7.5%.

For encoder-decoder architecture of our method, we do continuous pretraining from the public checkpoints of T5 for 250K steps with batch size 256 using the masked span prediction. Since each masked span (consists of multiple tokens) is replaced by a single special token, when using masking ratio is 15%, the number of special tokens in a sequence is not too large, we keep masking ratio 15% unchanged.

7.5.4 Downstream Finetuning

The statistics of the sequence lengths of instances in each dataset are summarized in Tab. 6. The hyperparameters of all experiments are summarized in Tab 11. When there are multiple values in an entry, it means we perform a hyperparameter search on these values. The amount of search is determined by the size of datasets. If a dataset is relatively large, we only search the learning rate. If a dataset is small, we include batch size and the number of epochs in search. For all tasks, if the sequence lengths are longer than the model length m , the sequences will be truncated and only the first m tokens will be used. For encoder-decoder models, we use greedy decoding in sequence generations for simplicity. The maximal decoder output length, specified in Tab. 11, is set such that the maximal length covers the output lengths of more than 99% of instances. When the length of covering 99% of instances is greater than 512, we just set the maximal decoder output length to 512. Additionally, we show one extra experiment on MultiNews [12] in Tab. 10, which is not in the main text due to space limit.

7.5.5 Non-Language Tasks

While the focus of this work is on language tasks, the overall design does not limit its application to other tasks as long as the assumption ('a subset of tokens are disproportionately responsible for

the model prediction' and VIP-tokens can be reasonably selected) holds or partially holds for the tasks. As a result, we also tested our method on the Long Range Arena (LRA) benchmark [31] and obtained very promising performance among all baselines compared in [39]. We get slightly better performance than the top performing baselines presented in MRA-attention, but note that we are not trying to show that our method outperforms other baselines but to verify that our method can indeed be applied to other tasks.

Note that there are multiple implementations and hyperparameters used in the LRA benchmark, and comparisons across different implementations and hyperparameters is awkward. We use the same implementation and same hyperparameters as [39]. The VIP-token selection of our method in LRA experiments is quite easy. We simply use the prepended CLS token as the only VIP token since it is responsible for the final model classification. The results are shown in Tab. 12. All results except for ours are directly cited from [39]. Since LRA consists of a synthetic task (ListOps), language tasks (Text, Retrieval), and vision tasks (Image, Pathfinder), these results can serve as preliminary evidence indicating the potential applications of our method on other non-language tasks.

7.6 Potential Application to Decoders

The goal of this work is to reduce the cost of attention as well as the feedforward network (FFN) block without changing the internals of the transformer block when dealing with ultra long sequences (which differentiates our work from the few recent ideas that deal with ultra long sequences). Based on the best feasibility of this goal, the scope of application is focused on the encoder and encoder-decoder settings. For standard auto-regressive decoder settings, it might be

difficult to compress the token that is being generated with neighboring tokens. It is in fact challenging to reduce the cost of FFN since every generated token will most likely need a full calculation of FFN when it is being generated.

Having said that, there are still clear opportunities for decoder-only models, which are left as our future work. We briefly describe three possible options to do so. (1) We can use the input tokens of the decoder as VIP-tokens to compress the representations of context sequence generated by the encoder before Cross Attention computation to reduce the cost of Cross Attention. (2) Auto-regressive decoding operates using Causal Attention at each step. This Causal Attention operation requires memory and computation that is linear in the length of the prefix. We can keep the same Causal Attention VIP-token (the representation of the token currently being generated) and apply our method to compress the representations of the previously generated tokens. This reduces the linear complexity of the Causal Attention operation to sublinear. This is useful for reducing the cost of inference. For training, we can break the sequence into two segments: prefix segment and decoding segment. Then, we can use the proposed compression in prefix segment and vanilla computation in decoding segment. To prevent look ahead to the future tokens, we might only use the first token in the decoding segment as VIP-token. (3) In many cases, the current large language models (LLMs) are not used directly to generate an ultra long text. Rather, the requirements for processing ultra longer context manifests due to the need to incorporate user input and previous LLM responses (like ChatGPT) or to incorporate search results (such as New Bing). In these cases, there is a prefix context, and LLMs will generate text based on the user prompt and prefix context. We note that our method can indeed be applied to compress the prefix context, and the user prompt and currently generated tokens will be the VIP-tokens.

7.7 Practical Questions

How is the T5's relative positional encoding handled in our method?

To compute the relative positional bias for self-attention in T5, the position indices of queries and keys are needed (to calculate the position distance between each query and key). After applying our method, the input to the T5 block is the compressed sequence (some tokens are compressed while some tokens remain uncompressed). For each uncompressed token, the position index remains unchanged as its true position in the original sequence. For each new token that represents a compressed segment, the position index is the floored average of position indices of tokens in the segment. In this way, we can minimize the amount of modification needed to the internals of the T5 block, and the relative attention bias for the compressed sequence is an approximation of the attention bias for the original sequence.

VIP-token selection is task dependent, how to do this selection?

VIP-token selection requires some understanding of the tasks. However, we believe that a reasonable selection can be made with some generic knowledge for most tasks or use cases. For example, for question answering tasks, we just use questions as VIP-tokens. For classification tasks (for example, Long Range Arane benchmark shown in Appendix), we simply use CLS token as the VIP-token since only CLS token is used for final prediction. For masked language modeling, we use the masked tokens as the VIP-tokens since only these masked tokens are used for final prediction. For question answering, the question tokens (and candidate tokens for multi-choice QA) are used as the VIP-tokens since 'they (1) are important to the specific task goals and (2) easily pre-identifiable by the user.' In the worst case, if there is no obvious token to be selected, we can prepend some learnable 'latent' tokens or certain user commands (such as 'summarize' for summarization tasks as we used in our experiments) and use them as VIP-tokens (in fact, CLS tokens can be thought of as these learnable tokens).

Why is the performance of our method is better than standard models?

Our method is an approximation of the standard models, which should be inferior to the standard models, but in some cases, the performance of our method is better than standard models. We believe the reason is that the correct inductive bias improves the performance for tasks with limited amounts of data. Our approach is forced to compress irrelevant information and the attention is carried out on the compressed sequences, but in standard model with standard attention, each token has access to the entire sequence, which enables a larger degree of freedom. As a result, more training data might be required for the model to learn the correct pattern or bias.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper aims at designing efficient Transformers for ultra long sequences. The design is motivated from the intuition that only a small subset of special tokens, which are referred to as VIP tokens, are most relevant to the final prediction. Based on the intuition, a multi-resolution compression is proposed to reduce the computational cost. Experiments are conducted to show the efficiency and accuracy of the proposed model.\n\nSTRENGTHS:\n+ Nice presentation.\n\n+ The observation that only a small subset of special tokens are most relevant to the final prediction is insightful. The proposed task-driven multi-resolution compression technique is well-motivated and technically sound.\n\n+ The experiments are extensive and show reasonable performance gain.\n\nWEAKNESSES:\n+ The major concern for the proposed approach is that it **only applies to Transformer encoders**. Many important LLMs today are decoder-only, e.g., GPT, PaLM, LLaMA, etc. The authors formulate their method in an encoder layer (Sections 2 & 3), and show their method can lead to decent performance gains on encoder-only models and encoder-decoder models. However it seems non-trivial to generalize the approach to decoder-only models.\n\n+ The compression also seems to **impede the use of relative positional encoding**. Although there are experiments on T5-based models, no details are provided on how the relative positional information is handled.\n\n+ As discussed at the beginning of Section 3.3, the multi-resolution compression technique is proposed to minimize the approximation error of Eq. (12). However, **neither theoretical guarantees nor numerical simulations on the approximation error are provided**. There seems to be some theoretical results in the appendix but I can't find any analysis on the approximation error (defined in Eq. (33)). Since the multi-resolution compression is the main technical contribution of this paper, I encourage the authors to provide some theoretical analysis and/or numerical experiments on the approximation error in the main body of the paper. \n\n+ **Minor issues**.\n   + Line 121: $n_p$ is the size of the VIP-tokens: \"size\" seems ambiguous here and can be misunderstood as \"dimensionality\". I suggest writing _$n_p$ is the number of the VIP-tokens_.\n   + Eq. (9): $\\mathbf I_{n_p\\times n_p}$ can be simplified to $\\mathbf I_{n_p}$.\n   + Line 223, \"$s\\in\\cdots$\" should be \"_where_ $s\\in\\cdots$\".\n\nQUESTIONS:\n+ Is the proposed approach generalizable to decoder-only models?\n\n+ Is the proposed approach compatible with relative positional encoding? How is the T5's positional encoding handled in the experiments?\n\n+ Can you provide more illustrations on the approximation error of the compression technique?\n\nLIMITATIONS:\nLimitations are discussed in Section 4. However, I recommend the authors to explicitly discuss the applicability to decoder-only models, if the proposed approach only applies to Transformer encoders.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

comments: SUMMARY:
This paper proposes VIP-token centric compression (VCC) as a method to improve the efficiency of Transformers for ultra long sequences. Based on the impact on approximating the representation of the VIP-tokens, VCC can compress the long sequence into a much smaller representation. Thorough experiments validated the effectiveness of the proposed method.

STRENGTHS:
* Long sequences input is a practical problem for current LLM-based applications. And the efficient technology for long sequences that this paper focused on has a high application value, in my opinion.
* The paper is well-written, the method is generally novel and the effectiveness is validated by thorough experiments.
* The authors plan to release the code and checkpoints.

WEAKNESSES:
* FLOPs is missing. Runtime is a hardware-dependent metric, and it would be better to see also hardware-independent metrics FLOPs.
* Why only experiments on encoder-only and encoder-decoder architectures? It would be better to see experiments or discussions on recently more popular decoder-only architectures.
* KV Cache is a useful technique for accelerating Transformer inference and has been already integrated into widely used packages such as huggingface transformer and timm. It would be better to see the discussion on whether KV Cache and VCC are orthogonal and whether they can be used simultaneously. 
* Missing reference. It would be better to have a discussion on the difference with the previous [1] and concurrent [2] work that scales Transformer to 262k and one million, respectively. 

1. Memorizing Transformers. In ICLR 2022.
2. Scaling Transformer to 1M tokens and beyond with RMT. In arXiv.

QUESTIONS:
See Weaknesses.

LIMITATIONS:
Adequately addressed.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work introduces a method to scale up the supported number of tokens for transformer-based models. It prioritize the important/dominant tokens and compress those unimportant tokens during the processing. The conducted experiments demonstrate the effectiveness of the proposed method, and show that the proposed method can be directly incorporated into existing pretrained models with additional finetuning.\n\nSTRENGTHS:\nI think this work is trying to target an important problem, i.e., how to break through the limitation on the number of tokens to be processed. The proposed VIP-token centric compression is overall reasonable in terms of its high-level idea. The coducted experiments provide evidences for its efficiency benefits and performance superority when compared to the baselines with the same sequence length.\n\nWEAKNESSES:\nThe targeted problem and the main story are quite attractive at the first glance. However, the proposed methods are not that solid and insightful as expected, making I feel some parts somewhat overclaimed. To name a few, for the high-level idea of this work, evaluating the importance of VIP-tokens should be one of major parts. However, the proposed method for evaluating the importance of difference tokens are based on some intuitive assumptions. These assumptions are task-specific or case-specific, affecting the generalization capacity of the proposed method in fact.\n\nBesides the proposed method for evaluating the importance of tokens, other parts of the proposed method only introduces high-level idea deas and illustrates them with examples, without providing a rigorous modeling.\n\nThe experiments are conducted mainly on ultra long sequences. This matches with the major story, but in fact, prioritizing important tokens should be generally benefitial for a wide range of tasks using transformer-based models. The authors have not provides discussion or analysis on these.\n\nQUESTIONS:\n1. Are the proposed methods, which include but are not limited to the method for evaluating the importance of different tokens, generally applicable (e.g., for other tasks, or for other modalities)? If yes, pls explain more. If not, please analyze the limitations in detail and in-depthly.\n2. The proposed compression and decompression are in fact implemented by matrix multiplication. How can you ensure the tokens of different importances (not referring to splitting as VIP or non-VIP) are compressed or decomposed as expected?\n3. Why can placing VIP tokens in the front seats by re-ordering play the role of prioritizing them? How about repeating the VIP tokens (this is a more commonly used way for highlighting important tokens)? Pls compare these two, if possible.\n\nLIMITATIONS:\nThe biggest problem is that this paper is quite attractive at the beginning, but the main parts, especially for the proposed methods, are somewhat disappointing since they are not that solid and insightful as expected. Pls see the weakness part and my questions for more detailed comments.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-06"}

comments: SUMMARY:
VIP-token centric compression (VCC) aims to solve the quadratic attention complexity bottleneck in the self-attention component and the transformer encoder block. Roughly, it denotes some tokens as very important (VIP) and compresses the remaining sequence using a low-rank approximation. They benchmark on natural language data and do experiments to verify the time and space complexity from their theoretical analysis.

STRENGTHS:
They achieve impressive accuracy and performance numbers. Strong results on question-answering tasks that require long context and small runtime. They scale to 128k tokens and are able to tackle NarrativeQA, although the results are not as strong.

The method is intuitive and the analysis and description of the method is clear.

They were able to get strong results with pretrained weights from existing transformer encoder models.

WEAKNESSES:
The idea of VIP tokens is similar to the global memory component of efficient transformers like ETC, BigBird, and Longformer. It suffers a similar problem that it's not clear which tokens are VIP.

Experiments are only done on encoders. They mention in passing that the method can be adapted to decoding, which seems to be more in the vogue these days.

Some aspects of the experiments like doing a different computation in the first 4 transformer layers make me think this method may be finicky and hard to apply in general.

QUESTIONS:
How does this method perform in other large context non-language tasks like images or audio or even synthetic tasks like Long Range Arena where the VIP token selection is less obvious?

It would be great to see decoder experiments?

LIMITATIONS:
Authors should be commended for addressing weaknesses of the method like the selection of VIP tokens and encoder-only aspect.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a new method to reduce the computational cost of transformer models on long input sequences. It proposes to use certain tokens (VIP tokens) that will be responsible for compressing the rest of the sequence into fewer tokens that can then be processed efficiently. In detail, a tree is constructed from the input features to the transformer block such that at each node we have the mean of the children of that node. To construct the feature sequence that will be processed, the averaged attention of the nodes of the tree with the VIP tokens is used to walk the tree until the leaves are reached. The reached nodes form the input sequence to the transformer block. Using the tree and the stored differences between the nodes and their children the new sequence is constructed after the transformer block. Thorough experiments show that the proposed attention achieves results on par with dense attention while being significantly faster to compute.\n\nSTRENGTHS:\n- Both using a set of tokens to generate the \"clustering\" of the input sequence as well as using the tree of differences to recreate the full sequence are very interesting and original methods.\n- The experimental analysis is thorough and clearly shows that the proposed method works as well as dense attention in the examined benchmarks\n\nWEAKNESSES:\n- The main benefit of the method is the runtime efficiency, however the experiments do not describe adequately what this measured runtime is.\n  1. It is mentioned that it is the time for a single sequence? How does the method scale wrt batch size?\n  2. It is not clear whether it is the forward pass time or the forward/backward. How does this method scale for training?\n- There is a lot of theoretical analysis that seems slightly unnecessary as it doesn't provide much clarity but rather confuses given that the main approximation that allows for the method is taken from [6] and is not analyzed further in this work.\n- The proposed method cannot trivially be used in an autoregressive setting as it is really hard to efficiently compress N tokens while information flows only from the first to the N-th.\n\nQUESTIONS:\nMy main questions have to do with the calculation of the runtime of the presented method.\n\nAre the numbers for a single sequence? In that case how does the method scale with the batch size given that different samples will have different paths in the tree which will result in many out of order copies that may make the method slower.\n\nIn addition, is the calculated runtime for a forward pass or for a full training step?\n\nLIMITATIONS:\nThe authors adequately discuss the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}