SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $s \in \{k^0, k^1, k^2, \dots, n_c\}$  assuming  $n_c$  is a power of  $k$  and  $x \in \{1, 2, \dots, n_c / s\}$ .  $[\cdot]_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  $\mathbf{c}_x^s$ . Higher  $s$  (e.g.,  $s = 8$  in  $\mathbf{c}_1^8$  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  $\mathbf{C}$ . In our scheme, we compress the sequence  $\mathbf{C}$  and use a set of  $\mathbf{c}_x^s$  for some selected  $s$ 's and  $x$ 's as the rows of the compressed  $\mathbf{S}_c\mathbf{C}$  as seen in Fig. 3. The sequence  $\mathbf{C}$  is broken into multiple segments of different lengths, then each segment is compressed into a vector  $\mathbf{c}_x^s$ .
Procedurally, as shown in Fig. 3, our scheme starts with the heaviest compression and progressively refines certain segments of  $\mathbf{C}$  guided by the VIP-tokens  $\mathbf{P}$ . The scheme starts with the heaviest compression that treats  $\mathbf{C}$  as a  $n_c$ -length segment and compresses it to a single  $\mathbf{c}_1^{n_c}$ . Then, starting with  $s = n_c$  (root node), the procedure (1) computes the averaged attention scores between VIP-tokens  $\mathbf{P}$  and  $\mathbf{c}_x^s$ 's for different  $x$ 's (averaged over all attention heads and all VIP-tokens; only one  $\mathbf{c}_1^{n_c}$  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  $\mathbf{P}$  as queries and  $\mathbf{c}_x^s$ 's as keys. Then, it (2) splits the  $s$ -length segments corresponding  $\mathbf{c}_x^s$ 's with higher averaged attention scores (one segment is split in Fig. 3 but we might split more segments, again only one  $\mathbf{c}_1^{n_c}$  at level  $s = n_c$ ) into  $(s/2)$ -length sub-segments: the corresponding  $\mathbf{c}_x^{s/2}$  (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
Figure 3: Illustration of multi-resolution compression.  $n_c = 8$ . Purple line: compute attention scores between  $\mathbf{P}$  and different segments. Green arrow: segment with higher attention score is split into two sub-segments. Accordingly,  $\mathbf{S}_c\mathbf{C} = \left[\mathbf{c}_1^2\quad \mathbf{c}_3^1\quad \mathbf{c}_4^1\quad \mathbf{c}_2^4\right]^\top$  is constructed.
1. We note that this procedure is guided by the VIP-tokens  $\mathbf{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  $\mathbf{P}_{\text {new }}$  (8) is well approximated since the approach preserves the relevant components of  $\mathbf{C}$  that have a high impact on the output  $\mathbf{P}_{\text {new }}$ . Further, if the VIP-tokens  $\mathbf{P}$  have high attention weights for some rows of  $\mathbf{C}$ , then the corresponding row in  $\mathbf{C}$  will be approximated with higher frequencies (less compressed). So, the output in
$\mathbf{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
Figure 4: Proposed data structure  $\mathcal{T}(\mathbf{C})$
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  $\mathbf{S}_c\mathbf{C}$  with an  $\mathcal{O}(n_c d + rn_p d)$  cost at each layer. The main cost  $\mathcal{O}(n_c d)$  is due to computing  $\mathbf{c}_x^s$  defined in (13) for all resolution  $s$  and location  $x$  by using recursive relation from the bottom up:
$$
\mathbf {c} _ {x} ^ {2 s} = \frac {1}{2} \mathbf {c} _ {2 x - 1} ^ {s} + \frac {1}{2} \mathbf {c} _ {2 x} ^ {s} \quad \mathbf {c} _ {x} ^ {1} = [ \mathbf {C} ] _ {x} \tag {14}
$$
We find that these steps could introduce a large overhead. Further, note that if we decompress (apply  $\mathbf{S}^{\dagger}$  to) the output of  $\gamma$  for compressed sequence as in (8), the cost is  $\mathcal{O}(nd)$  since the number of nonzero entries in  $\mathbf{S}^{\dagger}$  is  $n$  (more details in Appendix). As a solution, we now introduce a data structure  $\mathcal{T}(\cdot)$  for storing  $\mathbf{C}$  and  $\mathbf{C}_{\text{new}}$ , as shown in 4, which enables efficient computation of  $\mathbf{c}_x^s$  and eliminates explicit decompression. We note that this data structure is only possible due to the specific structure of  $\mathbf{S}_c\mathbf{C}$  constructed in §3.3. Specifically,  $\mathcal{T}(\mathbf{C})$  stores  $\mathbf{c}_1^{n_c}$  and  $\Delta \mathbf{c}_x^s$  defined as
$$
\Delta \mathbf {c} _ {x} ^ {s} := \mathbf {c} _ {\lceil x / 2 \rceil} ^ {2 s} - \mathbf {c} _ {x} ^ {s} \tag {15}
$$
for every resolution  $s \neq n_c$  and location  $x$ . Similarly,  $\mathcal{T}(\mathbf{C}_{new})$  stores  $(\mathbf{c}_{new})_1^{n_c}$  and  $\Delta (\mathbf{c}_{new})_x^s$  where  $(\mathbf{c}_{new})_x^s$  and  $\Delta (\mathbf{c}_{new})_x^s$  are defined similar to (13) and (15) but using  $\mathbf{C}_{new}$  instead of  $\mathbf{C}$ .
Then, given  $\mathcal{T}(\mathbf{C})$ , any  $\mathbf{c}_x^s$  can be retrieved efficiently in  $\mathcal{O}(\log (n_c)d)$  cost via recursion:
$$
\mathbf {c} _ {x} ^ {s} = \mathbf {c} _ {\lceil x / 2 \rceil} ^ {2 s} - \Delta \mathbf {c} _ {x} ^ {s} = \mathbf {c} _ {\lceil x / 4 \rceil} ^ {4 s} - \Delta \mathbf {c} _ {\lceil x / 2 \rceil} ^ {2 s} - \Delta \mathbf {c} _ {x} ^ {s} = \dots \tag {16}
$$
The only reason we need decompression  $\mathbf{S}^{\dagger}$  is that we need to obtain new representation  $\mathbf{C}_{new}$  (no decompression for  $\mathbf{P}_{new}$  since  $\mathbf{P}$  is uncompressed). Suppose we have  $\mathcal{T}(\mathbf{C}_{new})$ , then we have an alternative way of getting  $\mathbf{C}_{new}$  similar to (16) (note  $(\mathbf{c}_{new})_x^1 = [\mathbf{C}_{new}]_x$ ) without explicit decompression. The key benefit of this data structure is that we can obtain  $\mathcal{T}(\mathbf{C}_{new})$  by changing some nodes in  $\mathcal{T}(\mathbf{C})$ . This only needs updating  $\mathcal{O}(r)$  nodes, and each update takes  $\mathcal{O}(d)$  cost.
An example. We show a  $\mathcal{T}(\mathbf{C})$  for  $n_c = 8$  in Fig. 4. Let  $\mathbf{S}_c\mathbf{C} = \left[\mathbf{c}_1^2\quad \mathbf{c}_3^1\quad \mathbf{c}_4^1\quad \mathbf{c}_2^4\right]^\top$  as in Fig. 3. Since the segment  $\mathbf{c}_1^2$  is not split into sub-segments  $\mathbf{c}_1^1$  and  $\mathbf{c}_2^1$ , we have (details in Appendix):
$$
\left(\mathbf {c} _ {\text {n e w}}\right) _ {1} ^ {1} - \mathbf {c} _ {1} ^ {1} = \left(\mathbf {c} _ {\text {n e w}}\right) _ {2} ^ {1} - \mathbf {c} _ {2} ^ {1} = \left(\mathbf {c} _ {\text {n e w}}\right) _ {1} ^ {2} - \mathbf {c} _ {1} ^ {2} \tag {17}
$$
By rearranging (17), we can verify that  $\Delta (\mathbf{c}_{new})_1^1, \Delta (\mathbf{c}_{new})_2^1$  in  $\mathcal{T}(\mathbf{C}_{new})$  stays the same as  $\Delta \mathbf{c}_1^1, \Delta \mathbf{c}_2^1$  in  $\mathcal{T}(\mathbf{C})$  and thus do not need to be updated:
$$
\Delta (\mathbf {c} _ {\text {n e w}}) _ {1} ^ {1} = (\mathbf {c} _ {\text {n e w}}) _ {1} ^ {2} - (\mathbf {c} _ {\text {n e w}}) _ {1} ^ {1} = \mathbf {c} _ {1} ^ {2} - \mathbf {c} _ {1} ^ {1} = \Delta \mathbf {c} _ {1} ^ {1}
$$
$$
\Delta (\mathbf {c} _ {\text {n e w}}) _ {2} ^ {1} = \left(\mathbf {c} _ {\text {n e w}}\right) _ {1} ^ {2} - \left(\mathbf {c} _ {\text {n e w}}\right) _ {2} ^ {1} = \mathbf {c} _ {1} ^ {2} - \mathbf {c} _ {2} ^ {1} = \Delta \mathbf {c} _ {2} ^ {1} \tag {18}
$$
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  $\mathbf{c}_x^s$  of  $\mathbf{S}_c\mathbf{C}$  (a leaf node in Fig. 3), only the node storing  $\Delta (\mathbf{c})_x^s$  and its ancestor nodes in  $\mathcal{T}(\mathbf{C})$  must be updated, so the total number of nodes (including their ancestors) being updated is  $\mathcal{O}(r)$ . Next, we can update the nodes as follows: first, we get representations  $(\mathbf{c}_{new})_1^2, (\mathbf{c}_{new})_3^1, (\mathbf{c}_{new})_4^1, (\mathbf{c}_{new})_2^4$  by feeding  $\mathbf{S}_c\mathbf{C}$  into Transformer layer (details in Appendix). At level  $s = 1$ , given  $(\mathbf{c}_{new})_3^1$  and  $(\mathbf{c}_{new})_4^1$ , we (1) compute  $(\mathbf{c}_{new})_2^2$  via (14), and then (2) compute  $\Delta (\mathbf{c}_{new})_3^1$  and  $\Delta (\mathbf{c}_{new})_4^1$  via (15). The last two values are the new values for  $\Delta \mathbf{c}_3^1$  and  $\Delta \mathbf{c}_4^1$  in  $\mathcal{T}(\mathbf{C})$ . At level  $s = 2$ , given  $(\mathbf{c}_{new})_2^1$  and  $(\mathbf{c}_{new})_2^2$  computed at previous level, we apply similar procedure to obtain  $(\mathbf{c}_{new})_1^4, \Delta (\mathbf{c}_{new})_2^1, \Delta (\mathbf{c}_{new})_2^2$ , and the last two values are used to update two nodes in  $\mathcal{T}(\mathbf{C})$ . It becomes apparent that each node update takes  $\mathcal{O}(d)$  cost. Putting it together: the complexity of modifying  $\mathcal{T}(\mathbf{C})$  to  $\mathcal{T}(\mathbf{C}_{new})$  is  $\mathcal{O}(rd)$ . The detailed algorithm and complexity analysis are described in Appendix.
By maintaining this data structure, we never need to materialize the entire  $\mathbf{C}$  or  $\mathbf{C}_{new}$  in any intermediate layer, but instead we use (16) to construct the rows of  $\mathbf{S}_c\mathbf{C}$  and perform updates to  $\mathcal{T}(\mathbf{C})$  to obtain  $\mathbf{C}_{new}$  (represented as  $\mathcal{T}(\mathbf{C}_{new})$ ) at each intermediate layer. At the output of a Transformer,  $\mathbf{C}_{new}$  is materialized from  $\mathcal{T}(\mathbf{C}_{new})$  at a  $\mathcal{O}(n_c d)$  cost via the recursion (16) from the bottom up.
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  $\mathbf{S}_c\mathbf{C}$  to have exactly two resolutions for experiments. Specifically, for a pre-defined increment ratio  $k$ , we split and refine all segments  $\mathbf{c}_x^s$  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  $\mathbf{S}_c\mathbf{C}$  would consist of  $(n_c / k - h)$  of  $\mathbf{c}_x^k$  and  $hk$  of  $\mathbf{c}_x^1$  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).  $\mathbf{P}_{new}$  is the approximated representation of VIP tokens computed with compression and let  $\mathbf{P}_{new}^{*}$  be the ground truth representation of
Table 1: Approx quality.
Compression	Error	Correlation
Random	0.403	0.919
Lazy	0.528	0.869
VIP centric	0.137	0.991
VIP tokens computed without compression. We measure the relative error (defined as  $||\mathbf{P}_{new} - \mathbf{P}_{new}^{*}||_{F} / ||\mathbf{P}_{new}^{*}||_{F}$ ) and correlation coefficient between  $\mathbf{P}_{new}$  and  $\mathbf{P}_{new}^{*}$ . As shown in Tab. 1, we can verify that the proposed procedure indeed improves 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
Table 2: Dev set results for encoder-only models.
Method	Size	Length	HotpotQA	QuALITY	WikiHop
Time	EM	F1	Time	Accuracy	Time	Accuracy
RoBERTa	base	512	19.9	35.1	44.9	21.2	39.0	19.6	67.6
RoBERTa	base	4K	422.3	62.2	76.1	403.2	39.5	414.1	75.2
Big Bird	base	4K	297.9	59.5	73.2	307.0	38.5	293.3	74.5
Longformer	base	4K	371.0	59.9	73.6	368.0	27.9	369.7	74.3
MRA Attention	base	4K	203.5	63.4	77.0	200.5	38.7	199.2	76.1
Ours	base	4K	114.6	60.9	74.6	126.4	39.6	108.0	75.9
Ours*	base	4K	114.6	61.4	75.0	125.7	39.5	108.0	76.1
Ours*	large	4K	285.8	66.7	80.0	390.8	41.8	394.3	79.6
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
Figure 5: Model runtime vs Wiki-Hop dev accuracy when using different model specific hyperparameters
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
Table 4: Dev results of NarrativeQA on base model when scaling sequence length from 16K to 128K.
Length	Runtime (ms)	k	h	EM	F1
16K	518.2 / 394.4 / 162.4	16	90	5.9	16.6
32K	946.8 / 671.6 / 212.6	32	55	6.6	17.5
32K	1027.9 / 751.0 / 298.0	16	90	6.4	17.5
64K	1848.7 / 1177.2 / 254.8	64	30	7.2	18.4
64K	2244.8 / 1574.2 / 659.4	16	90	7.5	19.3
128K	6267.8 / 5125.9 / 1902.2	16	90	8.0	19.6
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).
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).
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.
7.1 Definition of Notations