SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 7: Dev set results for encoder-only models pintuning on HotpotQA, QuALITY, and WikiHop.
Method	Size	Length	HotpotQA	QuALITY	WikiHop
Runtime	EM	F1	Runtime	Accuracy	Runtime	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
Ours	base	4k	114.6	60.9	74.6	126.4	39.6	108.0	75.9
Ours-150k	base	4k	114.6	60.7	74.1	126.4	39.4	108.0	76.1
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.\times$  in some entries means the profiler throws errors and cannot estimate the FLOPs numbers.
Table 8: Dev accuracy and FLOPs for encoder-only models. *: some calculations are NOT captured by profiler, so value is underestimated.
Method	Size	Length	Runtime	WikiHop GFLOPs	Accuracy
RoBERTa	base	512	19.6	-	67.6
RoBERTa	base	4K	414.1	1317	75.2
Big Bird	base	4K	293.3	790	74.5
Longformer	base	4K	369.7	×	74.3
MRA Attention	base	4K	199.2	697*	76.1
Ours	base	4K	108.0	526	75.9
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.
Table 9: Dev results and FLOPs for encoder-decoder models.
Method	Size	Length	WikiHop
Runtime	GFLOPs	EM	F1
T5	base	512	25.7 / 20.5	120	66.7	69.1
T5	base	4K	594.3 / 553.7	1449	76.2	78.1
LongT5	base	4K	270.7 / 233.9	948	72.7	74.8
LED	base	4K	236.6 / 222.9	×	70.0	72.4
Ours	base	4K	181.7 / 148.1	663	76.7	78.4
Method	Size	Length	ContractNLI
Runtime	GFLOPs	EM	F1
T5	base	512	24.0 / 20.5	112	73.5	73.5
T5	base	4K	579.0 / 551.6	1437	86.8	86.8
LongT5	base	16K	1564.2 / 1462.5	4442	85.1	85.1
Ours	base	16K	484.2 / 393.1	1933	87.0	87.0
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],
Table 10: Dev results for encoder-decoder models on MultiNews.
Method	Size	#Param	Length	Runtime	MultiNews
R-1	R-2	R-L
T5	base	223M	512	59.2 / 20.5	42.5	15.3	39.0
T5	base	223M	4K	651.2 / 551.8	46.4	18.2	42.6
LongT5	base	248M	8K	721.7 / 550.6	46.7	18.3	42.9
LED	base	162M	8K	526.5 / 454.2	46.6	17.8	42.7
Ours	base	223M	8K	377.0 / 224.6	46.4	18.1	42.7
T5	large	738M	512	180.8 / 67.0	43.4	15.6	39.8
Ours	large	738M	8K	1140.3 / 651.5	48.2	19.2	44.2
Ours	3b	3B	8K	4094.5 / 2696.0	48.9	19.4	44.7
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
Figure 8: Model runtime vs WikiHop dev accuracy when using different model specific hyperparameters
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
Table 11: Hyperparameters for all experiments.
LM
Task	Encoder-Only	Encoder-Decoder
HotpotQA	QuALITY	WikiHop	WikiHop	HotpotQA	CNN/Dailymail	MediaSum	Qasper
Optimizer	Adam	Adam	Adam	Adam	Adam	Adam	Adam	Adam
Weight Decay	0.01	0.01	0.01	0.01	0.01	0.01	0.01	0.01
LR Decay	Linear	Linear	Linear	Linear	Linear	Linear	Linear	Linear
Precision	FP16	FP16	FP16	BF16	BF16	BF16	BF16	BF16
Batch Size	32	16	32	32	32	32	32	{16, 32}
Learning Rate	{3e-5, 5e-5}	{3e-5, 5e-5}	{3e-5, 5e-5}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}
Epochs	10	{10, 20}	10	10	10	10	10	{10, 20}
Warmup Steps	1000	200	1000	1000	1000	1000	1000	200
Max Output Length	-	-	-	32	40	256	256	128
LM
Task	Encoder-Decoder
QuALITY	ContractNLI	NarrativeQA	Arxiv	SummScreenFD	GovReport	QMSum	MultiNews
Optimizer	Adam	Adam	Adam	Adam	Adam	Adam	Adam	Adam
Weight Decay	0.01	0.01	0.01	0.01	0.01	0.01	0.01	0.01
LR Decay	Linear	Linear	Linear	Linear	Linear	Linear	Linear	Linear
Precision	BF16	BF16	BF16	BF16	BF16	BF16	BF16	BF16
Batch Size	{16, 32}	{16, 32}	32	32	{16, 32}	{16, 32}	{16, 32}	32
Learning Rate	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}	{1e-4, 3e-4}
Epochs	{10, 20}	{10, 20}	5	{10, 20}	{10, 20}	{10, 20}	{10, 20}	10
Warmup Steps	200	1000	1000	1000	200	1000	100	1000
Max Output Length	90	4	47	512	512	512	310	512
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.
Table 12: Test set accuracy of LRA tasks.
Method	Listops	Text	Retrieval	Image	Pathfinder	Avg
Transformer	37.1±0.4	65.2±0.6	79.6±1.7	38.5±0.7	72.8±1.1	58.7±0.3
Performer	36.7±0.2	65.2±0.9	79.5±1.4	38.6±0.7	71.4±0.7	58.3±0.1
Linformer	37.4±0.3	57.0±1.1	78.4±0.1	38.1±0.3	67.2±0.1	55.6±0.3
SOFT	36.3±1.4	65.2±0.0	83.3±1.0	35.3±1.3	67.7±1.1	57.5±0.5
SOFT + Conv	37.1±0.4	65.2±0.4	82.9±0.0	37.1±4.7	68.1±0.4	58.1±0.9
Nystromformer	24.7±17.5	65.7±0.1	80.2±0.3	38.8±2.9	73.1±0.1	56.5±2.8
Nystrom + Conv	30.6±8.9	65.7±0.2	78.9±1.2	43.2±3.4	69.1±1.0	57.5±1.5
YOso	37.0±0.3	63.1±0.2	78.3±0.7	40.8±0.8	72.9±0.6	58.4±0.3
YOSo + Conv	37.2±0.5	64.9±1.2	78.5±0.9	44.6±0.7	69.5±3.5	59.0±1.1
Reformer	18.9±2.4	64.9±0.4	78.2±1.6	42.4±0.4	68.9±1.1	54.7±0.2
Longformer	37.2±0.3	64.1±0.1	79.7±1.1	42.6±0.1	70.7±0.8	58.9±0.1
Big Bird	37.4±0.3	64.3±1.1	79.9±0.1	40.9±1.1	72.6±0.7	59.0±0.3
H-Transformer-1D	30.4±8.8	66.0±0.2	80.1±0.4	42.1±0.8	70.7±0.1	57.8±1.8
Scatterbrain	37.5±0.1	64.4±0.3	79.6±0.1	38.0±0.9	54.8±7.8	54.9±1.4
MRA-2	37.2±0.3	65.4±0.1	79.6±0.6	39.5±0.9	73.6±0.4	59.0±0.3
MRA-2-s	37.4±0.5	64.3±0.8	80.3±0.1	41.1±0.4	73.8±0.6	59.4±0.2
Ours	37.3±0.5	65.3±0.2	80.9±0.5	41.3±1.7	74.6±1.5	59.9±0.9
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 pretend 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.