SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Liliang Ren $^{1*}$  Yang Liu $^{2}$  Shuohang Wang $^{2}$  Yichong Xu $^{\dagger}$  Chenguang Zhu $^{2}$  Chengxiang Zhai $^{1}$
<sup>1</sup>University of Illinois at Urbana-Champaign 2Microsoft
{liliang3,czhai}@illinois.edu
{yaliu10,shuowa,chezhu}@microsoft.com
xuyc11@gmail.com
Recent hybrid models combining Linear State Space Models (SSMs) with self-attention mechanisms have demonstrated impressive results across a range of sequence modeling tasks. However, current approaches apply attention modules statically and uniformly to all elements in the input sequences, leading to sub-optimal quality-efficiency trade-offs. To address this limitation, we introduce Sparse Modular Activation (SMA), a general mechanism enabling neural networks to sparsely and dynamically activate sub-modules for sequence elements in a differentiable manner. Through allowing each element to skip non-activated sub-modules, SMA reduces computation and memory consumption of neural networks at both training and inference stages. To validate the effectiveness of SMA on sequence modeling, we design a novel neural architecture, SeqBoat, which employs SMA to sparsely activate a Gated Attention Unit (GAU) based on the state representations learned from an SSM. By constraining the GAU to only conduct local attention on the activated inputs, SeqBoat can achieve linear inference complexity with theoretically infinite attention span, and provide substantially better quality-efficiency trade-off than the chunking-based models. With experiments on a wide range of tasks, including long sequence modeling, speech classification and language modeling, SeqBoat brings new state-of-the-art results among hybrid models with linear complexity, and reveals the amount of attention needed for each task through the learned sparse activation patterns. Our code is publicly available at https://github.com/renll/SeqBoat.
Recent advance on efficient sequence modeling with State Space Models (SSMs) [GGR21;  $\mathrm{GDE^{+}20}$ ; GGGR22; GB22; SWL23] has shown impressive performance for a wide range of tasks across modalities, such as text classification, image recognition and speech recognition. SSMs, as first-order linear models, defined by a set of input, output, and state variables connected by first-order differential equations, can efficiently capture the recurrent structure in sequential data with carefully designed state matrices and the application of convolutional parallelism [GGR21]. However, they still significantly underperform the self-attention [BCB14;  $\mathrm{VSP^{+}17}$ ] based model in both language modeling and machine translation [VPSP23] tasks. A recent work  $[\mathrm{FDS^{+}23}]$  reveals that this is due to its deficiency of modeling the second-order pairwise comparisons between the input tokens, and shows that the augmentation of an additional shifted SSM layer can improve SSM's
*Work partially done during internship at Microsoft.
†Work done at Microsoft.
associative recalling ability. Furthermore, better quality-efficiency trade-off can be achieved by directly introducing extra self-attention modules to form a hybrid model (e.g. MEGA  $\mathrm{[MZK^{+}23]}$  and Hybrid H3  $\mathrm{[FDS^{+}23]}$ ) that utilizes both the first and the second order inductive biases, i.e., SSM and self-attention. However, the current hybrid models apply the attention modules statically and uniformly to each of the input token regardless the property of the task itself. This can lead to suboptimal quality-efficiency trade-offs since not all input tokens require second-order modeling and this computation need can vary substantially depending on both its context and the task difficulty.
In this paper, we aim to answer the following research questions for efficiently combining attention with SSMs:
- RQ1: Can neural networks learn to activate their attention modules on demand to achieve better quality-efficiency trade-off?
- RQ2: How much extra attention is needed for the SSMs on a task-by-task basis?
To answer these questions, we develop a new general mechanism, Sparse Modular Activation (SMA), that allows a neural network to sparsely and dynamically activate its sub-modules for each of the input token in a fully differentiable manner. Specifically, we assume a neural model can be composed of multiple heterogeneous sub-modules. For the input sequence, a latent configurator sparsely maps tokens to multiple compressed sequences corresponding to sub-modules. Each sub-module is then only applied on its mapped shorter sequence. Compared with activating all sub-modules on the whole input, Sparse Modular Activation can reduce computation and memory consumption for both the training and inference stages. Notably, SMA is proved to have a full coverage of the combinatorial search space of module activation, which is further explained in Section 3.2.
Efficient learning of dynamic sparsity is notoriously difficult under the constraint of the current parallel hardware  $\mathrm{[LQC^{+}22;GZYE20;XM22]}$ . To enable the practical efficiency gains from our module-level sparsity, we provide a simple yet efficient parallel implementation of SMA without any custom fused GPU kernels. Specifically, when compressing a batch of sequences in SMA, our implementation conducts both token selection and the sequence re-padding simultaneously using a single scatter operation that is widely optimized and present in modern deep learning frameworks.
To address RQ1, we apply SMA to construct a novel neural architecture, SeqBoat, that sparsely activate a Gated Attention Unit (GAU) [HDLL22] based on the state representation learned from an SSM. Both the GAU and the SSM representations are then aggregated through simple addition and activation to form a layer-level representation. Multiple same-sized SeqBoat layers are stacked sequentially to form a full neural model. Inspired by the working memory mechanism [AS68] used in human cognition, we further restrict the GAU to only apply local attention on the compressed sequence, which allows our model to have linear sequence inference complexity but theoretically infinite attention span.
We conduct comprehensive experiments to show that SeqBoat has significantly better quality-efficiency trade-off than state-of-the-art hybrid models on a wide range of tasks, including Long Range Arena (LRA)  $\left[\mathrm{TDA}^{+}20\right]$ , speech classification [War18] and language modeling [Hut06]. On the competitive LRA benchmark, SeqBoat achieves 1.96 higher average accuracy than MEGA-chunk  $\left[\mathrm{MZK}^{+}23\right]$ , the previous best hybrid model, with a  $10.4\times$  training speed up and a  $95\%$  memory reduction compared to the Transformer  $\left[\mathrm{VSP}^{+}17\right]$  on the Text task with 4,096 input length. Thanks to the intrinsic modular sparsity brought by SMA, SeqBoat directly reveals the amount of attention needed for each data sample of each task through its sparse activation patterns of GAU, addressing RQ2. We demonstrate that our working memory mechanism provides substantially better computation-accuracy trade-off than chunking based models, and analyze the relationship between the working memory size and the effective attention span on various long sequence modeling tasks.
To motivate and clarify our proposed techniques, we first present a mathematical formulation of our Sparse Modular Activation mechanism and show how it encompasses and generalizes previous attempts that aimed for module-level dynamic sparsity. A dedicated section for detailed comparisons between our approach with the related works is also included in Appendix F. We begin by reviewing how the standard sequence modeling is formalized to establish the common ground for our discussion.
2.1 Time-Invariant Sequence Modeling
Given a discrete sequence,  $\mathbf{x} = \{x_{1},\dots,x_{n}\} \in \mathbb{R}^{n}$ , consisting of  $n$  tokens, a time-invariant sequence model  $P_{\theta}$  is optimized to maximize the likelihood of the observed sequences by factorizing them as follows:
$$
\max _ {\theta} P _ {\theta} (\mathbf {x}) = \prod_ {t = 1} ^ {n} P (x _ {t} | \mathbf {x} _ {<   t}, \theta),
$$
where  $\mathbf{x}_{< t} = \{x_1, \dots, x_{t-1}\}$  is the sequence history at time step  $t$ , and the parameter  $\theta$  is independent of the time step  $t$ . This formulation implies that the full model parameters  $\theta$  and the full history  $\mathbf{x}_{< t}$  are both essential for the conditional prediction of each token  $x_t$ . However, one potential issue is as the prediction difficulty of each token may differ depending on the context and the position, this static model  $P_\theta$  can lead to sub-optimal accuracy-efficiency trade-off by wasting computation on either unimportant context [SJP+21] or easy-to-predict tokens [Gra16].
3 Learning Sparse Modular Activation
To cover a larger search space that may contain more efficient sequence models, we propose to formulate sequence modeling as a problem of finding an optimal time-variant model that can dynamically activate a subset of modules from a pre-defined function space for each time step.
3.1 Time-Variant Sequence Modeling
Formally, a time-variant sequence model is defined on a compact function space  $\mathcal{F}:\mathcal{X}_t^c\mapsto [0,1]^{n\times V}$ , where  $V$  is the size of the vocabulary and  $\mathcal{X}_t^c = \{\mathbf{x}_t^c:\mathbf{x}_t^c\subseteq \mathbf{x}_{< t}\in \mathcal{X}\subseteq \mathbb{R}^n\}$ , contains all possible sub-sequences of the sequence history  $\mathbf{x}_{< t}$ . Then for each of the token prediction at the time step  $t$ , the model learns to apply a function  $f_{t}\in \mathcal{F}$  with the parameters  $\theta_t$  that maximizes the sequence probability, i.e.,
$$
\max  _ {f _ {t}, \theta_ {t}, \mathbf {x} _ {t} ^ {c}} P _ {\mathcal {F}} (\mathbf {x}) = \prod_ {t = 1} ^ {n} P _ {f _ {t}} \left(x _ {t} \mid \mathbf {x} _ {t} ^ {c}, \theta_ {t}\right) \quad \text {s . t .} \quad \mathbf {x} _ {t} ^ {c} \subseteq \mathbf {x} _ {<   t} \tag {1}
$$
This formulation generalizes the previous works in pursuing a dynamic and sparse model for sequence modeling, where the connections are further explained in Appendix F. In this work, we assume the function space  $\mathcal{F}$  is chain-structured, i.e.,  $\mathcal{F} = \mathcal{H} \circ \mathcal{L}_N \circ \dots \circ \mathcal{L}_1 \circ \mathcal{E}$ , where  $\mathcal{H}: \mathbb{R}^{n \times d_m} \mapsto [0,1]^{n \times V}$  is the classification function,  $\mathcal{E}: \mathbb{R}^n \mapsto \mathbb{R}^{n \times d_m}$  is the embedding function,  $N$  is the number of intermediate layers,  $d_m$  is the model size and  $\mathcal{L}: \mathbb{R}^{n \times d_m} \mapsto \mathbb{R}^{n \times d_m}$  is the function space of the intermediate mappings. We further assume that  $\mathcal{L}$  is the spanning set of a finite number of the function  $f_i^l$  with its parameters  $\theta_i^l$ , i.e.,  $\mathcal{L} = \operatorname{span}\{f_1^l, \ldots, f_M^l\}$ , where  $M$  is the number of pre-defined functions. These assumptions justify the design of our Sparse Modular Activation mechanism, which is further explained in the following section.
3.2 Sparse Modular Activation
Sparse Modular Activation (SMA) introduces a latent configurator at each time step  $t$  and each layer of a neural sequence model. The configurator produces a binary decision vector  $\mathbf{a}_t \in \{0,1\}^M$  to only activate those functions  $f_i^l$  that have the decision value  $a_t^i = 1$ . Only the activated function will be applied to process the input representations and return the outputs. These outputs are then aggregated with a linear combination to generate the final vector,  $\mathbf{y}_t$ , of the layer. To enable end-to-end differentiability, we also require the latent configurator to output the confidence probability of its decisions,  $\mathbf{c}_t \in [0,1]^M$ , and respectively scale the output from the activated function with its confidence values  $c_t^i$ . The function space of a layer equipped with SMA can be described as follow,
$$
\mathcal {L} _ {\mathrm {S M A}} \left(\mathbf {a} _ {t}, \mathbf {c} _ {t}\right) = \left\{\sum_ {i \in I} \zeta_ {i} c _ {t} ^ {i} f _ {i} ^ {l} \mid \forall \zeta_ {i} \in \mathbb {R}, I = \{i \mid a _ {t} ^ {i} = 1, 1 \leq i \leq M \} \right\}
$$
where  $I$  is the index set of activated functions. Notably, SMA has the following important property,
Theorem 1 (Function Space Coverage of SMA). For any  $\mathcal{L}' \subseteq \mathcal{L} = \text{span}\{f_1^l, \dots, f_M^l\}$ , there exists a pair of  $(\mathbf{a}_t', \mathbf{c}_t')$  that  $\mathcal{L}_{SMA}(\mathbf{a}_t', \mathbf{c}_t') = \mathcal{L}'$ . In other words, SMA has a full coverage of the function space  $\mathcal{L}$ .
Figure 1: The proposed SeqBoat layer. The black lines indicate that the gradients can be backpropagated and the red lines stand for gradients blocking.  $\odot$  means the element-wise multiplication, and  $\oplus$  is the point-wise addition. The max, argmax and softmax operators are all applied to the projected dimension after the linear layer in the latent configurator block. The sparse activation operators are respectively instantiated as compress and extract operators for parallel processing.
The proof is provided in Appendix B. Intuitively, we can see that the original function space  $\mathcal{L}$  is recovered, i.e.,  $\mathcal{L}_{\mathrm{SMA}} = \mathcal{L}$ , if all the functions are activated, and the model can adaptively select any subspace of  $\mathcal{L}$  through controlling the binary decision vector  $\mathbf{a}_t$ . During sequential inference stage, for each function that requires the contextual information, we keep a record of its activated input representations across the time steps to form a memory  $\mathbf{H}_t^c \subseteq \mathbf{H}_{< t} \in \mathbb{R}^{n \times d_m}$  (where  $\mathbf{H}_{< t}$  is the input sequence history), so that each function can recall its contextual memory for each of the decoding step. During parallel training stage, the activated input representations for each function (or module) are sparsely selected to compose the compressed sequence  $\mathbf{H}^c$  for the module to conduct parallel processing, and this process is implemented as a compress operator. The returned representations are then mapped to its original input position in a new sequence whose inactivated positions are filled with zero vectors, which is denoted as an extract operator. We illustrate this parallel implementation in Figure 2 for better understanding.
The Pytorch-like  $\left[\mathrm{PGM}^{+}19\right]$  code snippets of the compress and extract operators are provided in the Appendix A.1 with an efficient support of batched sequences using the scatter operation.
In this work, we explore applying SMA to address our RQs, in which the following two types of functions in the function space  $\mathcal{L}$  are considered: (1) a linear State Space Model (SSM) [GGGR22;  $\mathrm{MZK}^{+}23]$  that efficiently captures the recurrent sequential structure, and (2) a Gated Attention Unit (GAU) [HDLL22] that performs second-order pairwise comparisons between the sequence elements. We also keep the SSM as an always-activated module and only apply SMA to the GAU, and in this case the effective number of modules  $M = 1$ . This is because for long sequence modeling, the main computation and memory bottleneck comes from the second-order self-attention and the gating operations inside the GAU module. We modularize the combination of SSM, GAU and SMA as a neural layer, which we call a SeqBoat layer. Its details are further explained in the next section.
Figure 2: The proposed parallel implementation of the Sparse Modular Activation (SMA) mechanism. In this example, we have two modules  $f_1^l, f_2^l$  and an input sequence  $\mathbf{H} = \{\mathbf{h}_1, \dots, \mathbf{h}_4\}$ . We assume only  $a_2^1, a_4^1, a_4^2$  have values equal to one. The white block means a zero vector. The compressed sequences for modules  $f_1^l$  and  $f_2^l$  are  $\mathbf{H}_1^c = \{\mathbf{h}_2, \mathbf{h}_4\}$  and  $\mathbf{H}_2^c = \{\mathbf{h}_4\}$  respectively. The final outputs are aggregated as  $\mathbf{y}_4 = c_4^1 \mathbf{y}_4^1 + c_4^2 \mathbf{y}_4^2$ , and  $\mathbf{y}_2 = c_2^1 \mathbf{y}_2^1$ .
3.3 Model Architecture
Figure 1 illustrates the proposed architecture for our SeqBoat layer. Given an input sequence representation  $\mathbf{S}^l\in \mathbb{R}^{n\times d_m}$  for the  $l$  -th layer, where  $n$  is the length of the sequence and  $d_{m}$  is the hidden size, we first apply a linear time-invariant State Space Model (SSM) [GGR21] over the sequence
followed by a SiLU activation [EUD17] to obtain the state representations  $\mathbf{H} \in \mathbb{R}^{n \times d_m}$ ,
$$
\operatorname {S S M} \left(\mathbf {S} ^ {l}\right) = \mathbf {K} * \mathbf {S} ^ {l} + \mathbf {D} \odot \mathbf {S} ^ {l} \tag {2}
$$
$$
\mathbf {H} = \operatorname {S i L U} (\operatorname {S S M} (\mathbf {S} ^ {l}))
$$
where  $*$  is a convolution operator which is efficiently implemented using FFTs [GGR21],  $\odot$  indicates the element-wise multiplication,  $\mathbf{D} \in \mathbb{R}^{1 \times d_m}$  is the learnable parameters, and  $\mathbf{K} \in \mathbb{R}^{n \times d_m}$  is the SSM convolution kernel. In this work, we use the Multi-dimensional Damped Exponential Moving Average (MD-EMA) proposed by MEGA  $[\mathrm{MZK}^{+}23]$  for the kernel parameterization, whose details are included in Appendix A.3. Note that our work is orthogonal to the SSM kernel design as long as it can provide an efficient computation of the recurrent state representation. It is possible that using a more advanced SSM such as Liquid-S4  $[\mathrm{HLW}^{+}22]$  or S5 [SWL23] can lead to better performance, but we keep using MD-EMA for a fair comparison with MEGA. Given the state representation  $\mathbf{H} \in \mathbb{R}^{n \times d_m}$ , a latent configurator with one linear layer is applied to obtain the decision vector  $\mathbf{a} \in \{0,1\}^n$  and the confidence probability vector  $\mathbf{c} \in [0,1]^n$  with the Tempered Softmax, i.e.,
$$
\mathbf {p} _ {i} = \frac {e ^ {\left(\mathbf {H} \mathbf {w} _ {i} + b _ {i}\right) / \tau}}{e ^ {\left(\mathbf {H} \mathbf {w} _ {0} + b _ {0}\right) / \tau} + e ^ {\left(\mathbf {H} \mathbf {w} _ {1} + b _ {1}\right) / \tau}} \quad \text {f o r} i \in \{0, 1 \}
$$
$$
\mathbf {a} = \underset {i} {\arg \max } \mathbf {p} _ {i}, \quad \mathbf {c} = \underset {i} {\max } \mathbf {p} _ {i}
$$
where  $\mathbf{w}_0, \mathbf{w}_1 \in \mathbb{R}^{d_m}, b_0, b_1 \in \mathbb{R}$  are learnable parameters,  $\tau \in \mathbb{R}_{>0}$  is a learnable temperature parameter which is initialized as  $\alpha \sqrt{d_m}$  (where  $\alpha$  is a hyperparameter denoted as the initial temperature scale). Previous works [RSVG20a;  $\mathrm{GLL}^{+}22$ ; FZS22;  $\mathrm{RZW}^{+}22$ ] often seek applying auxiliary losses to regularize the exploration and the exploitation of the latent decisions. However, in this work, we don't apply any auxiliary losses and empirically observe that solely using Tempered Softmax works well for SMA. To provide an explanation of this phenomenon, we show that the implicit regularization [BD20] of gradient decent will conditionally encourage the exploration of the latent decisions if the Softmax function with learnable temperature (i.e. Tempered Softmax) is used for decision confidence calculation. The details of the theorem and the proof are provided in Appendix C.
Based on the decision values  $\mathbf{a}$ , a Gated Attention Unit (GAU) is activated and applied on the compressed sequence  $\mathbf{H}^c$  resulted from the compress operator as previously described in Section 3.2,
$$
\mathbf {H} ^ {c} = \operatorname {C o m p r e s s} (\mathbf {H}, \mathbf {a}) \quad \in \mathbb {R} ^ {r \times d _ {m}}
$$
$$
\mathbf {Y} ^ {c} = \operatorname {G A U} \left(\mathbf {H} ^ {c}\right) \quad \in \mathbb {R} ^ {r \times d _ {m}}
$$
$$
\mathbf {Y} = \operatorname {E x t r a c t} \left(\mathbf {Y} ^ {c}, \mathbf {a}\right) \quad \in \mathbb {R} ^ {n \times d _ {m}}
$$
where  $r = \sum_{t=1}^{n} a_t \leq n$  is the total number of the time steps when the GAU module is activated. We use Squared ReLU  $\left[\mathrm{SML}^{+}21\right]$  for the attention function on the image tasks and Softmax for the text related tasks. The detailed implementation of GAU is shown in Appendix A.2. With the extracted output  $\mathbf{Y} \in \mathbb{R}^{n \times d_m}$ , the final layer output is computed with a linear aggregation followed by a non-linear activation, i.e.,
$$
\mathbf {S} ^ {l + 1} = \operatorname {S i L U} (\mathbf {c} \odot \mathbf {Y} + \mathbf {H W} + \mathbf {b} + \mathbf {S} ^ {l})
$$
where  $\mathbf{W} \in \mathbb{R}^{d_m \times d_m}$ ,  $\mathbf{b} \in \mathbb{R}^{d_m}$  are the learnable parameters. We also add a normalization layer either before the Linear SSM (i.e., Pre-Norm) in the residual branch or after the final layer output  $\mathbf{S}^{l+1}$  (i.e., Post-Norm) to stabilize the training process. The SeqBoat model is then constructed as stacking  $N$  number of SeqBoat layers of the same hidden size  $d_m$  with the task specific embedding and classification layers. Notably, our model does not have any Feed-Forward Networks (FFN) anywhere in the architecture.
Working Memory Mechanism The computation complexity of our model is dynamically dependent on the activation probability of the GAU module  $p = r / n = \sum_{t = 1}^{n}a_{t} / n$  for each of the SeqBoat layer. This means that the computation complexity is  $O(r^2 +n\log (n))$  for parallel training and  $O(r^{2} + n)$  for auto-regressive decoding. While it is possible that the number of the GAU-activated time steps is far smaller than the length of the full sequence for a trained model, a random initialized SeqBoat model will still activate the GAU module for  $n / 2$  number of times on average, leading to a non-ideal  $O(n^{2})$  training complexity at the early stage of training. To solve this problem, we restrict the GAU to only conduct local attention [AOA  $^{+}20$ ] with a sliding window of size
$w \ll n$  on the compressed sequence  $\mathbf{H}^c$  during the parallel training stage. Each token will attend to the nearest  $w$  tokens in the compressed sequence, i.e.,  $w/2$  past tokens and  $w/2$  future tokens for bidirectional sequence modeling, and  $w$  number of past tokens for auto-regressive sequence modeling. This approach keeps a First-In-First-Out (FIFO) memory with size  $w$  of the compressed sequence history during auto-regressive decoding, which can be understand as a working memory [AS68] for the GAU module. With this mechanism, SeqBoat reduces the training complexity to  $O(rw + n\log(n))$  and the inference complexity to  $O(rw + n)$ , while maintaining the ability of attending to the key tokens whose distances from the query are longer than the working memory size  $w$ . This is because the elapsed time steps between two GAU activations is generally not bounded by the size of the working memory.
We evaluate the SeqBoat model on three representative long-range sequence modeling benchmarks across image, text and speech modalities to have a comprehensive understanding of its capacity. The implementation details of our model for each of the tasks are presented in Appendix A.
Besides model performance on task metrics, we also report training speed and memory consumption as the efficiency of each model. Since the dynamic sparsity of SeqBoat will affect both the speed and the memory consumption throughout the training process, we measure its training speed up as the average per step wall time across the full training stage, and the memory consumption as the average per step GPU memory allocation. More details of the efficiency measurement are included in Appendix A.6.
We compare our proposed SeqBoat model with multiple strong baseline models besides the Transformer  $\left[\mathrm{VSP}^{+}17\right]$  and its variant Transformer-XL  $\left[\mathrm{DYY}^{+}19\right]$ .