SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The construction of  $\mathcal{I}$  is described in Alg. 1. This algorithm describes the same procedure as the Figure 3 in the main text. The  $\mathbf{b}_x^s$ 's in  $\mathcal{J}$  are the rows of  $\mathbf{S}_c$ , and the corresponding  $\mathbf{c}_x^s$ 's (40) are the rows of  $\mathbf{S}_c\mathbf{C}$ . The budgets  $h_2, h_4, \dots, h_{n_c}$  required by Alg. 1 is used to determine the number of components at each resolution that will be added to  $\mathcal{J}$ . Specifically, there are  $2h_{2s} - h_s$  number of components  $\mathbf{b}_x^s$  for  $s \neq 1$  based on simple calculations. We can choose budgets such that the final size of  $\mathcal{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  $\mathbf{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^{\dagger}$  contain exactly a 1 as stated in Claim 7.3,  $S^{\dagger}$  can commute with  $\beta$ , so in summary, we can write the approximation of the computation of a Transformer layer as
$$
\begin{array}{l} \alpha (\mathbf {P}, \mathbf {S X}, \mathbf {S X}) = \exp (\mathbf {P} \mathbf {P} ^ {\top}) \mathbf {P} + \exp (\mathbf {P} \mathbf {C} ^ {\top} \mathbf {S} _ {c} ^ {\top}) \mathbf {D} \mathbf {S} _ {c} \mathbf {C} \\ \mathbf {S} _ {c} ^ {\dagger} \alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) = \mathbf {S} _ {c} ^ {\dagger} \exp \left(\mathbf {S} _ {c} \mathbf {C P} ^ {\top}\right) \mathbf {P} + \mathbf {S} _ {c} ^ {\dagger} \exp \left(\mathbf {S} _ {c} \mathbf {C C} ^ {\top} \mathbf {S} _ {c} ^ {\top}\right) \mathbf {D S} _ {c} \mathbf {C} \tag {54} \\ \left[ \begin{array}{c} \mathbf {P} _ {n e w} \\ \mathbf {C} _ {n e w} \end{array} \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}) \\ \beta (\mathbf {S} _ {c} ^ {\dagger} \alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) + \mathbf {S} _ {c} ^ {\dagger} \mathbf {S} _ {c} \mathbf {C}) + \mathbf {S} _ {c} ^ {\dagger} \alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) \end{array} \right] + \left[ \begin{array}{c} \mathbf {P} \\ \mathbf {C} \end{array} \right] \\ \end{array}
$$
Note that  $\mathbf{D}$  is added as discussed in (50).
There are four main approximation components (purple) in (54). Taking the fact that  $\mathbf{DS}_c = (\mathbf{S}_c^\dagger)^\top$ , all of these approximations are row or column space multi-resolution approximations governed by  $\mathbf{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  $\mathbf{C}$  that have higher dependency to  $\mathbf{P}$  are better approximated. The output  $\mathbf{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  $\mathbf{P}_{\text {new }}$ . Further, the output in  $\mathbf{C}_{\text {new }}$  corresponding to the subset of rows of  $\mathbf{C}$  that have higher dependency with the VIP-tokens will have better approximation than the remaining rows of  $\mathbf{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  $\mathbf{S}_c$  which has some desirable properties, namely satisfying (19) which we elaborate further below.
Note that for equation (19), we need a matrix  $\mathbf{S}_c$  such that the approximated attention matrix involving  $\mathbf{P}$  and  $\mathbf{C}$  is similar to the true attention matrix involving  $\mathbf{P}$  and  $\mathbf{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  $\mathbf{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  $\mathbf{S}_c$  we are looking for) to the signal  $\mathbf{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  $\mathbf{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  $\mathbf{c}_x^s$  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  $\mathcal{T}(\cdot)$ .
7.3.1 Why  $(\mathbf{c}_{new})_1^1 -\mathbf{c}_1^1 = (\mathbf{c}_{new})_2^1 -\mathbf{c}_2^1 = (\mathbf{c}_{new})_1^2 -\mathbf{c}_1^2$  if  $\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 ?$
Claim 7.4. Given the set  $\mathcal{J}$  satisfying restriction (25), if  $\mathbf{b}_x^s \in \mathcal{J}$ , then  $(\mathbf{c}_{new})_x^s - \mathbf{c}_x^s = (\mathbf{c}_{new})_{x'}^{s'} - \mathbf{c}_{x'}^{s'}$  for all  $\mathbf{b}_{x'}^{s'}$  satisfying the support of  $\mathbf{b}_{x'}^{s'}$  is contained in the support of  $\mathbf{b}_x^s$  (the  $x'$ -th  $s'$ -length segment of  $\mathbf{C}$  is a sub-segment of the  $x$ -th  $s$ -length segment of  $\mathbf{C}$ ).
Proof. To simplify the notations a bit, without loss of generality, we assume  $x = 1$ . Then, for  $i \leq s$ , consider  $(\mathbf{c}_{new})_i^1$ :
$$
\begin{array}{l} \left. \left(\mathbf {c} _ {\text {n e w}}\right) _ {i} ^ {1} = \left[ \mathbf {C} _ {\text {n e w}} \right] _ {i} = \left[ \mathbf {S} _ {c} ^ {\dagger} \beta \left(\alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}\right) + \mathbf {S} _ {c} ^ {\dagger} \alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) \right] _ {i} + [ \mathbf {C} ] _ {i} \right. \tag {55} \\ = \left[ \mathbf {S} _ {c} ^ {\dagger} \right] _ {i} \beta \left(\alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}\right) + \left[ \mathbf {S} _ {c} ^ {\dagger} \right] _ {i} \alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {c} _ {i} ^ {1} \\ \end{array}
$$
By Claim 7.3,  $\mathbf{S}_c^\dagger = \mathbf{S}_c^\top \mathbf{D}$  and  $i$ -th row of  $\mathbf{S}_c^\dagger$  contains exactly a 1. The column that contains 1 in the  $i$ -th row of  $\mathbf{S}_c^\dagger$  is exactly  $s\mathbf{b}_1^s$  since  $i$  is contained in the support of exactly one components in  $\mathcal{I}$  due to the restriction (25). Denote this column index as  $j$ , then
$$
\left(\mathbf {c} _ {\text {n e w}}\right) _ {i} ^ {1} = \left[ \beta \left(\alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}\right) \right] _ {j} + \left[ \alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) \right] _ {j} + \mathbf {c} _ {i} ^ {1} \tag {56}
$$
Note that this holds for all  $i \leq s$ . As a result, for  $i, i' \leq s$
$$
\left(\mathbf {c} _ {\text {n e w}}\right) _ {i} ^ {1} - \mathbf {c} _ {i} ^ {1} = \left[ \beta \left(\alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}\right) \right] _ {j} + \left[ \alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) \right] _ {j} = \left(\mathbf {c} _ {\text {n e w}}\right) _ {i ^ {\prime}} ^ {1} - \mathbf {c} _ {i ^ {\prime}} ^ {1} \tag {57}
$$
Algorithm 2 Computation of one Transformer layer with  $\mathcal{T}(\mathbf{C})$
Input: VIP-tokens  $\mathbf{P}$  and data structure  $\mathcal{T}(\mathbf{C})$    
Use Algo. 1 to construct  $\mathcal{I}$  but use (63) to retrieve  $\mathbf{c}_x^s$  from  $\mathcal{T}(\mathbf{C})$    
Construct  $\mathbf{S}_c,\mathbf{S}_c\mathbf{C}$  associated with  $\mathcal{I}$  using Claim 7.2   
Compute  $\left[ \begin{array}{l}\mathbf{P}_{new}\\ \mathbf{S}_c\mathbf{C}_{new} \end{array} \right] = \left[ \begin{array}{l}\beta (\alpha (\mathbf{P},\mathbf{S}\mathbf{X},\mathbf{S}\mathbf{X}) + \mathbf{P}) + \alpha (\mathbf{P},\mathbf{S}\mathbf{X},\mathbf{S}\mathbf{X}) + \mathbf{P}\\ \beta (\alpha (\mathbf{S}_c\mathbf{C},\mathbf{S}\mathbf{X},\mathbf{S}\mathbf{X}) + \mathbf{S}_c\mathbf{C}) + \alpha (\mathbf{S}_c\mathbf{C},\mathbf{S}\mathbf{X},\mathbf{S}\mathbf{X}) + \mathbf{S}_c\mathbf{C} \end{array} \right]$    
Set  $\mathcal{T}(\mathbf{C}_{new})\gets \mathcal{T}(\mathbf{C})$    
for  $s\gets 1,2,4,\dots n_c / 2$  do for  $\mathbf{b}_x^s\in \mathcal{I}$  do Locate row location  $\mathbf{b}_x^s$  in  $\mathbf{S}_c$  , refer the index as  $j$  Compute  $(\mathbf{c}_{new})_x^s = [\mathbf{S}_c\mathbf{C}_{new}]_j$  Mark  $\Delta (\mathbf{c}_{new})_x^s$  dirty end for for dirty  $\Delta (\mathbf{c}_{new})_x^s$  do Compute  $(\mathbf{c}_{new})_{[x / 2]}^{2s}$  and update  $\Delta (\mathbf{c}_{new})_x^s$  Mark  $\Delta (\mathbf{c}_{new})_{[x / 2]}^{2s}$  dirty end for   
end for   
Update  $(\mathbf{c}_{new})_1^{n_c}$    
Output: VIP-tokens  $\mathbf{P}_{new}$  and data structure  $\mathcal{T}(\mathbf{C}_{new})$
$$
\begin{array}{l} (\mathbf {c} _ {\text {n e w}}) _ {\lceil i / 2 \rceil} ^ {2} - \mathbf {c} _ {\lceil i / 2 \rceil} ^ {2} = \frac {1}{2} (\mathbf {c} _ {\text {n e w}}) _ {2 \lceil i / 2 \rceil - 1} ^ {1} + \frac {1}{2} (\mathbf {c} _ {\text {n e w}}) _ {2 \lceil i / 2 \rceil} ^ {1} - \frac {1}{2} \mathbf {c} _ {2 \lceil i / 2 \rceil - 1} ^ {1} - \frac {1}{2} \mathbf {c} _ {2 \lceil i / 2 \rceil} ^ {1} \\ = \frac {1}{2} \left(\left(\mathbf {c} _ {\text {n e w}}\right) _ {2 \lceil i / 2 \rceil - 1} ^ {1} - \mathbf {c} _ {2 \lceil i / 2 \rceil - 1} ^ {1}\right) + \frac {1}{2} \left(\left(\mathbf {c} _ {\text {n e w}}\right) _ {2 \lceil i / 2 \rceil} ^ {1} - \mathbf {c} _ {2 \lceil i / 2 \rceil} ^ {1}\right) \tag {58} \\ = (\mathbf {c} _ {n e w}) _ {2 \lceil i / 2 \rceil - 1} ^ {1} - \mathbf {c} _ {2 \lceil i / 2 \rceil - 1} ^ {1} \\ \end{array}
$$
The rest follows from induction.
7.3.2 How do we get  $(\mathbf{c}_{new})_1^2, (\mathbf{c}_{new})_3^1, (\mathbf{c}_{new})_4^1, (\mathbf{c}_{new})_2^4$  if  $\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$ ?
Claim 7.5. We have
$$
\mathbf {S} _ {c} \mathbf {C} _ {\text {n e w}} = \beta \left(\alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}\right) + \alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}. \tag {59}
$$
And the updated representation  $(\mathbf{c}_{new})_x^s$  of the corresponding  $\mathbf{c}_x^s$  (a row of  $\mathbf{S}_c\mathbf{C}$ ) is the corresponding row of  $\mathbf{S}_c\mathbf{C}_{new}$ .
$$
\mathbf {C} _ {\text {n e w}} = \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}) + \mathbf {C} \tag {60}
$$
$$
\begin{array}{l} \mathbf {S} _ {c} \mathbf {C} _ {\text {n e w}} = \mathbf {S} _ {c} \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} \mathbf {S} _ {c} ^ {\dagger} \alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) + \mathbf {S} _ {c} \mathbf {C} \tag {61} \\ = \beta \left(\alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C}\right) + \alpha \left(\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}\right) + \mathbf {S} _ {c} \mathbf {C} \\ \end{array}
$$
Since the  $\mathbf{S}_c$  is the same for  $\mathbf{S}_c\mathbf{C}_{new}$  and  $\mathbf{S}_c\mathbf{C}$ , the second statement follows.
7.3.3 Algorithm for Making  $\mathcal{T}(\mathbf{C})$  into  $\mathcal{T}(\mathbf{C}_{new})$
In this section, we describe the exact algorithm to update  $\mathcal{T}(\mathbf{C})$  into  $\mathcal{T}(\mathbf{C}_{new})$ . The pseudo code is described in Alg. 2 where  $\mathbf{c}_x^s$  is computed via
$$
\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 {63}
$$
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.
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  $\mathcal{O}(lrd^2 + lr^2d + lr\log(n_c)d + lrn_pd + nd)$  when using the proposed efficient data structure.
7.4.1 Preparing Input Sequence to  $\mathcal{T}(\mathbf{C})$ :  $\mathcal{O}(nd)$
At the first layer, we need to permute the rows of  $\mathbf{X}$  into  $[\mathbf{P};\mathbf{C}]$ , which takes  $\mathcal{O}(nd)$  cost. Then, we process  $\mathbf{C}$  into  $\mathcal{T}(\mathbf{C})$ . This requires (1) computing  $\mathbf{c}_x^s$  defined in (41).  $\mathbf{c}_x^1 = [\mathbf{C}]_x$ , so no compute is needed. With all  $\mathbf{c}_x^1$  given, computing all  $\mathbf{c}_x^2$  takes  $\mathcal{O}(n_c d / 2)$ . With all  $\mathbf{c}_x^2$  given, computing all  $\mathbf{c}_x^4$  takes  $\mathcal{O}(n_c d / 4)\ldots$  So, the cost is
$$
\mathcal {O} \left(n _ {c} d / 2 + n _ {c} d / 4 + \dots + d\right) = \mathcal {O} \left(n _ {c} d\right). \tag {64}
$$
Then (2) computing  $\Delta \mathbf{c}_x^s$  for all  $s$  and  $x$ . Computing each  $\Delta \mathbf{c}_x^s$  takes  $\mathcal{O}(d)$  when given  $\mathbf{c}_x^s$  and  $\mathbf{c}_{[x/2]}^{2s}$ . The amount of cost is the same as the number of nodes in the tree  $\mathcal{T}(\mathbf{C})$ , so the cost is  $\mathcal{O}(n_c d)$ . Note that  $n_c < n$ , so the overall complexity of the above operations is  $\mathcal{O}(nd)$ .
7.4.2 Constructing  $\mathcal{J},\mathbf{S}_c,\mathbf{S}_c\mathbf{C}$  ..  $\mathcal{O}(lr\log (n_c)d + lr n_p d)$
We can analyze the complexity of constructing  $\mathcal{I}$  using Algo. 1. There is only one possible  $\mu_x^{n_c}$ . Then for each  $s$ , there are  $2h_s$  number of  $\mu_x^{s/2}$  being computed since there are 2 components  $\mathbf{b}_x^{s/2}$  for each  $\mathbf{b}_{x'}^s$ . As a result, we need to compute  $\mathcal{O}(1 + \sum_s 2h_s)$  number of  $\mu_x^{s/2}$ . When  $\mathbf{c}_x^{s/2}$  is given, the cost of computing a  $\mu_x^{s/2}$  is  $\mathcal{O}(n_p d)$ , so the overall cost of constructing  $\mathcal{I}$  is  $\mathcal{O}((1 + \sum_s 2h_s)n_p d)$ .
Further, at each  $s$ , the size of  $\mathcal{I}$  is increased by  $h_s$  since  $h_s$  segments are split into  $2h_s$  sub-segments, so the size of  $\mathcal{I}$  is  $\mathcal{O}(\sum_{s} h_s)$ . Since  $\mathbf{S}_c \in \mathbb{R}^{(r - n_p) \times n}$  and  $|\mathcal{I}| = r - n_p$  as discussed in §7.2.4,  $\mathcal{O}(r - n_p) = \mathcal{O}(\sum_{s} h_s)$ . We use  $\mathcal{O}(r)$  for simplicity instead of  $\mathcal{O}(r - n_p)$ . As a result, the overall cost of constructing  $\mathcal{I}$  is  $\mathcal{O}(rn_p d)$ .
The above cost assumes  $\mathbf{c}_x^{s/2}$  is given. If we compute all possible  $\mathbf{c}_x^{s/2}$  using (41), the cost will be  $\mathcal{O}(n_c d)$  as analyzed in §7.4.1. However, if we employ the proposed data structure, each  $\mathbf{c}_x^{s/2}$  can be retrieved in at most  $\mathcal{O}(\log(n_c)d)$  by recursively computing (63). Since we need to retrieve  $\mathcal{O}(1 + \sum_s 2h_s) = \mathcal{O}(r)$  number of  $\mathbf{c}_x^s$ , the complexity of computing necessary  $\mathbf{c}_x^s$  is  $\mathcal{O}(r\log(n_c)d)$ .
As a result, the complexity of constructing  $\mathcal{J}$  is  $\mathcal{O}(rn_p d + r\log (n_c)d)$  at each layer. When summing the cost over all layers, the complexity is  $\mathcal{O}(lr n_p d + lr\log (n_c)d)$ .
By Claim 7.2, the rows of  $\mathbf{S}_c$  and  $\mathbf{S}_c\mathbf{C}$  are simply the  $\mathbf{b}_x^s\in \mathcal{I}$  and the corresponding  $\mathbf{c}_x^s$ , which are already computed during the construction of  $\mathcal{I}$ , so we essentially can get these  $\mathbf{S}_c$  and  $\bar{\mathbf{S}}_c\mathbf{C}$  for free.
7.4.3 Feeding Compressed Sequence into a Transformer Layer:  $\mathcal{O}(lrd^2 +lr^2 d)$
At each layer, we need to compute
$$
\left[ \begin{array}{c} \mathbf {P} _ {\text {n e w}} \\ \mathbf {S} _ {c} \mathbf {C} _ {\text {n e w}} \end{array} \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 {P} \\ \beta (\alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) + \mathbf {S} _ {c} \mathbf {C}) + \alpha (\mathbf {S} _ {c} \mathbf {C}, \mathbf {S X}, \mathbf {S X}) + \mathbf {S} _ {c} \mathbf {C} \end{array} \right] \tag {65}
$$
for updating  $\mathcal{T}(\mathbf{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  $\mathcal{O}(nd^2 + n^2d)$  for a input sequence of length  $n$ . Now a compressed sequence of length  $r$  is fed into a Transformer layer, the cost is simply  $\mathcal{O}(rd^2 + r^2d)$ . We note that there is an additional re-scaling to plug  $\mathbf{D}$  into  $\exp(\mathbf{PC}^\top \mathbf{S}_c^\top) \mathbf{DS}_c$  during multi-head attention computation discussed in 50. However, the additional cost of applying  $\mathbf{D}$  is  $\mathcal{O}(rd)$ , which does not change the complexity. When summing the cost of all layers, the overall complexity is  $\mathcal{O}(lrd^2 + lr^2d)$ .
7.4.4 Updating  $\mathcal{T}(\mathbf{C})$  into  $\mathcal{T}(\mathbf{C}_{\text{new}})$ :  $\mathcal{O}(lrd)$
Once (65) is computed, we need to change  $\mathcal{T}(\mathbf{C})$  into  $\mathcal{T}(\mathbf{C}_{new})$ . The cost of change  $\mathcal{T}(\mathbf{C})$  into  $\mathcal{T}(\mathbf{C}_{new})$  is  $\mathcal{O}(rd)$  as analyzed in the main text. For more specific analysis, let us take a look at the first three iterations:
(1) At the first iteration, there are  $\mathcal{O}(2h_2)$  number of  $(\mathbf{c}_{new})_x^1$  to be computed at the first inner for loop, and there are  $\mathcal{O}(2h_2)$  number of  $\Delta (\mathbf{c}_{new})_x^1$  to be updated in the second inner for loop. Additional  $\mathcal{O}(h_2)$  number of  $\Delta (\mathbf{c}_{new})_{[x / 2]}^2$  are masked dirty.
(2) At the second iteration, there are  $\mathcal{O}(2h_4)$  number of  $(\mathbf{c}_{new})_x^2$  to be computed at the first inner for loop, and there are  $\mathcal{O}(2h_4 + h_2)$  number of  $\Delta (\mathbf{c}_{new})_x^2$  to be updated in the second inner for loop. The second term is due to the dirty  $\Delta (\mathbf{c}_{new})_{[x / 2]}$  from the first iteration. Additional  $\mathcal{O}(h_4 + \frac{h_2}{2})$  number of  $\Delta (\mathbf{c}_{new})_{[x / 2]}^4$  are masked dirty.
(3) At the third iteration, there are  $\mathcal{O}(2h_8)$  number of  $(\mathbf{c}_{new})_x^4$  to be computed at the first inner for loop, and there are  $\mathcal{O}(2h_8 + h_4 + \frac{h_2}{2})$  number of  $\Delta (\mathbf{c}_{new})_x^4$  to be updated in the second inner for loop. The second and third term is due to the dirty  $\Delta (\mathbf{c}_{new})_{[x / 2]}^4$  from the second iteration. Additional  $\mathcal{O}(h_8 + \frac{h_4}{2} +\frac{h_2}{4})$  number of  $\Delta (\mathbf{c}_{new})_{[x / 2]}^8$  are masked dirty.
It becomes apparent that if we sum over the number of computes of  $(\mathbf{c}_{new})_x^s$  and updates of  $\Delta (\mathbf{c}_{new})_x^s$ , the total number is  $\mathcal{O}(\sum_s 2h_s + 2\sum_s\sum_{j = 1}^{\log (s)}\frac{h_s}{2^j}) = \mathcal{O}(\sum_sh_s + \sum_sh_s) = \mathcal{O}(r)$ . Since each compute and update takes  $\mathcal{O}(d)$  cost, the overall complexity of changing  $\mathcal{T}(\mathbf{C})$  into  $\mathcal{T}(\mathbf{C}_{new})$  is  $\mathcal{O}(rd)$ . When summing the cost of all layers, the overall complexity is  $\mathcal{O}(lrd)$ .
7.4.5 Materializing  $C_{new}$  from  $\mathcal{T}(\mathbf{C}_{new})$  at the Last Layer:  $\mathcal{O}(nd)$
At the output of the last layer, we can (1) compute all  $(\mathbf{c}_{new})_x^{n_c / 2}$  via (63) at a cost of  $\mathcal{O}(2d)$ , (2) compute  $(\mathbf{c}_{new})_x^{n_c / 4}$  via (63) at a cost of  $\mathcal{O}(4d)$  ... until all  $(\mathbf{c}_{new})_x^1$  are computed. Then,  $[\mathbf{C}_{new}]_x = c_x^1$  is materialized from  $\mathcal{T}(\mathbf{C}_{new})$  at a total cost of
$$
\mathcal {O} (d + 2 d + 4 d + \dots + n _ {c} d) = \mathcal {O} \left(n _ {c} d\right). \tag {66}
$$
Lastly, undoing the permutation so that  $\left[\mathbf{P}_{new};\mathbf{C}_{new}\right]$  are re-ordered to the original positions has a complexity of  $\mathcal{O}(nd)$ . As a result, the overall complexity is  $\mathcal{O}(nd)$ .
7.4.6 Overall Complexity
In summary, the overall complexity of our method is
$$
\mathcal {O} \left(l r d ^ {2} + l r ^ {2} d + l r \log \left(n _ {c}\right) d + l r n _ {p} d + n d\right) \tag {67}
$$
7.5 Experiments
Table 6: Length statistics of each dataset. The values are the percentiles of number of tokens for the specific tokenizers. For T5 tokenizer, the left value of is for sequence lengths of encoder input, and the right value is for sequence lengths of decoder input.
Percentile	HotpotQA	RoBERTa
QuALITY	WikiHop	WikiHop	HotpotQA	T5
Qasper	QuALITY	ContractNLI
75th	1535	7603	2204	2399 / 6	1692 / 6	7029 / 29	7747 / 17	2991 / 4
95th	1928	8495	3861	4206 / 9	2129 / 10	10920 / 71	8603 / 28	5061 / 4
Percentile	NarrativeQA	CNN/Dailymail	MediaSum	T5
Arxiv	SummScreenFD	GovReport	QMSum	MultiNews
75th	90482 / 10	1242 / 87	2621 / 29	13477 / 364	12119 / 188	13304 / 811	19988 / 110	3032 / 379
95th	260533 / 18	1946 / 130	5061 / 64	26024 / 759	16722 / 330	23795 / 983	31749 / 162	6676 / 468
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.