SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhanpeng Zeng
zzeng38@wisc.edu
Cole Hawkins
colehawk@amazon.com
University of Minnesota, Minneapolis & AWS AI
Aston Zhang
astonz@amazon.com
Nikolaos Pappas
nppappa@amazon.com
Vikas Singh
vsingh@biostat.wisc.edu
Shuai Zheng
shzheng@amazon.com
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.,  $>16\mathrm{K}$  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.
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
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  $4\mathrm{K}$  (and even  $16\mathrm{K}$ ). 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
Figure 1: Model efficiency of processing one sequence on a NVIDIA A100 as sequence length increases (note logarithm  $x$  axis).
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). We hypothesize/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 128K token lengths, with consistent performance benefits (and no complicated architecture changes).
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  $\mathbf{X} \in \mathbb{R}^{n \times d}$  which gives the  $n$  feature vector inputs for a Transformer layer. The output of this Transformer layer,  $\mathbf{X}_{new}$ , is defined as
$$
\mathbf {X} _ {\text {n e w}} = \beta (\alpha (\mathbf {X}) + \mathbf {X}) + \alpha (\mathbf {X}) + \mathbf {X} \tag {1}
$$
using  $\alpha (\mathbf{X})$  as shorthand for  $\alpha (\cdot ,\cdot ,\cdot)$ , which is a multi-head attention (MHA) with  $\mathbf{X}$  as input for queries, keys, and values, described shortly. Here,  $\beta (\cdot)$  is a feed-forward network (FFN). Layer
norms [1] are omitted to reduce clutter. Let the inputs to  $\alpha (\cdot ,\cdot ,\cdot)$  be  $\mathbf{Q},\mathbf{K},\mathbf{V}\in \mathbb{R}^{n\times d}$  for queries, keys, and values. MHA is defined as:
$$
\alpha (\mathbf {Q}, \mathbf {K}, \mathbf {V}) := \operatorname {c a t} _ {i = 1} ^ {i = g} \left[ \operatorname {s o f t m a x} \left(\mathbf {Q W} _ {Q, i} \mathbf {W} _ {K, i} ^ {\top} \mathbf {K} ^ {\top}\right) \mathbf {V W} _ {V, i} \right] \mathbf {W} \tag {2}
$$
where  $g$  is the number of attention heads,  $\{\mathbf{W}_{Q,i},\mathbf{W}_{K,i},\mathbf{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  $\mathbf{W}_{Q,1},\mathbf{W}_{K,1},\mathbf{W}_{V,1},\mathbf{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  $\alpha (\cdot ,\cdot ,\cdot)$  can be expressed as:
$$
\alpha (\mathbf {Q}, \mathbf {K}, \mathbf {V}) := \exp (\mathbf {Q} \mathbf {K} ^ {\top}) \mathbf {V}. \tag {3}
$$
Let  $\gamma (\cdot)$  be a placeholder for all heavy computations in the Transformer layer above:
$$
\gamma (\mathbf {X}) := \beta (\alpha (\mathbf {X}) + \mathbf {X}) + \alpha (\mathbf {X}). \tag {4}
$$
We can verify that the output of a Transformer block (parameters are suppressed to reduce clutter) is,
$$
\mathbf {X} _ {\text {n e w}} = \gamma (\mathbf {X}) + \mathbf {X}. \tag {5}
$$
A Transformer model consists of many such layers: the input of each layer is the output  $\mathbf{X}_{new}$  from the previous layer. Let  $l$  be the number of layers, then the overall complexity is  $\mathcal{O}(ln^2 d + lnd^2)$ .
Efficient Transformers. Many efficient self-attention methods are available to reduce the  $\mathcal{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 self-attention 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  $\mathcal{O}(lrd^2 +lr^2 d +$ $lr\log (n_c)d + lr n_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\ll r\ll 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  $\mathcal{O}(lr d^2 + lr^2 d)$  pertain to the cost for a Transformer, while the remaining terms are the overhead of compression and decompression. The term  $\mathcal{O}(lr \log (n_c)d + lrn_p d)$  is the overhead of compression and updating our data structure at each layer. The  $\mathcal{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  $\mathcal{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 ( $\S 3.1$ ,  $\S 3.2$ ). Then, we will discuss an instantiation of the compression process, by adapting a multi-resolution analysis technique ( $\S 3.3$ ). Lastly, we will introduce a data structure which allows more efficient compression/decompression ( $\S 3.4$ ).
3.1 Elevating the Importance of a Few Tokens: VIP-Token
Let us start with the simplest compression, which identifies a linear transformation  $\mathbf{S} \in \mathbb{R}^{r \times n}$  which acts on the input, resulting in a smaller representation  $\mathbf{S}\mathbf{X} \in \mathbb{R}^{r \times d}$ . Of course, a smaller  $r$  implies that more information about  $\mathbf{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  $\mathbf{P} \in \mathbb{R}^{n_p \times d}$  be the VIP-tokens and  $\mathbf{C} \in \mathbb{R}^{n_c \times d}$  be the non-VIP/remaining tokens,  $\mathbf{X}$  can be expressed as
$$
\mathbf {X} = \left[ \begin{array}{l} \mathbf {P} \\ \mathbf {C} \end{array} \right] \tag {6}
$$
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  $\mathbf{X}$ , the VIP-tokens must still stay (more or less) the same, and the compression matrix  $\mathbf{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  $\mathbf{X}$  denote its input matrix. Express the output of this layer as follows:
$$
\mathbf {X} _ {\text {n e w}} = \mathbf {S} ^ {\dagger} \gamma (\mathbf {S X}) + \mathbf {X} \tag {7}
$$
where  $\mathbf{S} \in \mathbb{R}^{r \times n}$  is a compression matrix compressing  $\mathbf{X}$  to a smaller representation and  $\mathbf{S}^{\dagger}$  is the pseudo inverse for decompression. With the layout in (6), we can write (7) as
$$
\left[ \begin{array}{l} \mathbf {P} _ {\text {n e w}} \\ \mathbf {C} _ {\text {n e w}} \end{array} \right] = \mathbf {S} ^ {\dagger} \gamma \left(\mathbf {S} \left[ \begin{array}{l} \mathbf {P} \\ \mathbf {C} \end{array} \right]\right) + \left[ \begin{array}{l} \mathbf {P} \\ \mathbf {C} \end{array} \right] \tag {8}
$$
where  $\mathbf{P}_{new}$  and  $\mathbf{C}_{new}$  are the new embeddings for  $\mathbf{P}$  and  $\mathbf{C}$ .
Figure 2: Diagram that illustrates a Transformer layer with VIP-token centric sequence compression.
Always reserve seats for VIP-tokens. What is a useful structure of S? Since  $\mathbf{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  $\mathbf{S}^{\dagger}$ :
$$
\mathbf {S} = \left[ \begin{array}{l l} \mathbf {I} _ {n _ {p}} & 0 \\ 0 & \mathbf {S} _ {c} \end{array} \right] \quad \mathbf {S} ^ {\dagger} = \left[ \begin{array}{l l} \mathbf {I} _ {n _ {p}} & 0 \\ 0 & \mathbf {S} _ {c} ^ {\dagger} \end{array} \right]. \tag {9}
$$
The rearrangement simply says that we will avoid compressing  $\mathbf{P}$ . But rewriting it in this way helps us easily unpack (8) to check the desired functionality of  $\mathbf{S}_c$ .
Prioritize information in VIP-tokens. Our goal is to ensure  $\mathbf{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  $\mathbf{S}$  defined in (9) first. We see that
$$
\left[ \begin{array}{c c} \mathbf {I} _ {n _ {p} \times n _ {p}} & 0 \\ 0 & \mathbf {S} _ {c} ^ {\dagger} \end{array} \right] \gamma \left(\left[ \begin{array}{c} \mathbf {P} \\ \mathbf {S} _ {c} \mathbf {C} \end{array} \right]\right) = \left[ \begin{array}{c} \beta (\alpha (\mathbf {P}, \mathbf {S X}, \mathbf {S X}) + \mathbf {P}) + \alpha (\mathbf {P}, \mathbf {S X}, \mathbf {S X}) \\ \mathbf {S} _ {c} ^ {\dagger} \beta (\alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) + \mathbf {S} _ {c} \mathbf {C}) + \mathbf {S} _ {c} ^ {\dagger} \alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) \end{array} \right]. \tag {10}
$$
The orange color identifies terms where  $\mathbf{P}_{new}$  interacts with other compression-related terms C and/or  $\mathbf{S}_c$ . We primarily care about  $\mathbf{P}_{new}$  in (8), so the first (orange) row in (10) is the main concern. We see that  $\mathbf{P}_{new}$  only depends on the compressed  $\mathbf{S}\mathbf{X}$  via  $\alpha (\mathbf{P},\mathbf{S}\mathbf{X},\mathbf{S}\mathbf{X})$ . We can further unpack,
$$
\alpha (\mathbf {P}, \mathbf {S X}, \mathbf {S X}) = \exp (\mathbf {P} \mathbf {X} ^ {\top} \mathbf {S} ^ {\top}) \mathbf {S X} = \exp (\mathbf {P} \mathbf {P} ^ {\top}) \mathbf {P} + \exp (\mathbf {P} \mathbf {C} ^ {\top} \mathbf {S} _ {c} ^ {\top}) \mathbf {S} _ {c} \mathbf {C}. \tag {11}
$$
Again,  $\alpha (\mathbf{P},\mathbf{S}\mathbf{X},\mathbf{S}\mathbf{X})$  depends on  $\mathbf{C}$  and  $\mathbf{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 (\mathbf{PC}^{\top}\mathbf{S}_{c}^{\top})\mathbf{S}_{c}\mathbf{C}$ . As long as the following approximation using  $\mathbf{S}_c$  is good
$$
\exp \left(\mathbf {P} \mathbf {C} ^ {\top} \mathbf {S} _ {c} ^ {\top}\right) \mathbf {S} _ {c} \approx \exp \left(\mathbf {P} \mathbf {C} ^ {\top}\right), \tag {12}
$$
we will obtain a good approximation of  $\mathbf{P}_{new}$ . Our remaining task is to outline a scheme of finding a compression matrix  $\mathbf{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  $\mathbf{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  $\mathbf{S}_c$  should operate as follows. If some regions of the sequence  $\mathbf{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  $\mathbf{S}_c\mathbf{C}$ .
Specifically, let us define a compressed representation (via averaging) of the  $x$ -th  $s$ -length segment of sequence  $\mathbf{C}$ :  $\mathbf{c}_x^s \in \mathbb{R}^d$
$$
\mathbf {c} _ {x} ^ {s} := \frac {1}{s} \sum_ {s x - s <   i \leq s x} [ \mathbf {C} ] _ {i} \tag {13}
$$