SlowGuess/NeurIPS_2023_OCR huggingface dataset:
S4 [GGR21] is a linear State Space Model (SSM) that efficiently captures the recurrent structure in sequential data. It employs a new parameterization for the state space model, enabling efficient and principled modeling of long-range dependencies. S4D-LegS [GGGR22] is a simpler variant of S4 that uses a diagonal state matrix.
MEGA  $\left[\mathrm{MZK}^{+}23\right]$  is a hybrid model that combines a gated self-attention mechanism with a Multidimensional Damped Exponential Moving Average (MD-EMA) parameterized SSM. It aims to achieve a better quality-efficiency trade-off by utilizing both the first and second-order inductive biases from SSM and self-attention. MD-EMA can be considered as a simpler variant of S4 that is similar to S4D. MEGA also has a linear complexity variant, MEGA-chunk, that applies attention to each local chunk of fixed length.
H3  $\left[\mathrm{FDS}^{+}23\right]$  introduces a hybrid SSM layer that stacks a diagonal SSM and a shifted SSM to model the comparisons between the elements in a sequence. It also provides a hybrid model that includes extra self-attention layers for better language modeling performance.
S5 [SWL23] is the state-of-the-art SSM model that improves over S4 with multi-input multi-output and time-domain processing. It leverages the parallel scan operator in JAX  $\left[\mathrm{BFH}^{+}18\right]$  for an efficient parallel implementation of recurrence unrolling.
Liquid-S4 [SWL23] extends over S4 through allowing the state transition to be input-dependent. The proposed approach can be cast back to the convolutional framework and re-use the Cauchy kernel for efficient computation.
Reformer [KKL20] is a memory-efficient variant of the Transformer that uses locality-sensitive hashing (LSH) to perform approximate nearest neighbor search for attention computation.
Performer and Linformer are two efficient variants of the Transformer model. The Performer  $\mathrm{[CLD^{+}20]}$  uses the Fast Attention via positive Orthogonal Random features algorithm to approximate the self-attention mechanism, reducing the computational complexity. The Linformer  $\mathrm{[WLK^{+}20]}$  employs low-rank approximations to project the self-attention mechanism into a lower-dimensional space, resulting in a more efficient system.
Table 1: Test accuracy on the LRA benchmark. We report the relative training speed up and memory reductions on the Text task with 4,096 sequence length. * indicates a hybrid model. "Retr." and "Path." are the abbreviations of Retrieval and the Pathfinder tasks respectively. Best results of linear inference complexity models are in bold, second best are underlined.
Models	ListOps	Text	Retr.	Image	Path.	Path-X	Avg.	Speed	Mem.
Quadratic Inference Complexity
Transformer	37.11	65.21	79.14	42.94	71.83	X	59.24	1×	1×
MEGA*	63.14	90.43	91.25	90.44	96.01	97.98	88.21	2.9×	0.31×
Sub-quadratic Inference Complexity
Reformer	37.27	56.10	53.40	38.07	68.50	X	50.67	0.8×	0.24×
BigBird	36.05	64.02	59.29	40.83	74.87	X	55.01	1.1×	0.30×
SeqBoat-full*	61.65	89.60	91.67	89.96	95.87	95.28	87.33	6.2×	0.07×
Linear Inference Complexity
Performer	18.01	65.40	53.82	42.77	77.05	X	51.41	5.7×	0.11×
Linformer	35.70	53.94	52.27	38.56	76.34	X	51.36	5.5×	0.10×
Luna-256	37.98	65.78	79.56	47.86	78.55	X	61.95	4.9×	0.16×
S4	59.10	86.53	90.94	88.48	94.01	96.07	85.86	4.8×	0.14×
S4D-LegS	60.47	86.18	89.46	88.19	93.06	91.95	84.89	6.1×	0.14×
S5	62.15	89.31	91.40	88.00	95.33	98.58	87.46	6.1×	0.14×
Liquid-S4	62.75	89.02	91.20	89.50	94.8	96.66	87.32	1.2×	0.17×
H3*	57.50	88.20	91.00	87.30	93.00	91.80	84.80	6.0×	0.24×
MEGA-chunk*	58.76	90.19	90.97	85.80	94.41	93.81	85.66	7.0×	0.09×
SeqBoat*	61.70	89.60	91.28	90.10	96.35	96.68	87.62	10.4×	0.05×
Luna  $\left[\mathrm{MKW}^{+21}\right]$  is a linear unified nested attention mechanism that approximates Softmax attention with two nested linear attention functions, resulting in linear time and space complexity.
4.2 Long Sequence Modeling
We first evaluate SeqBoat and SeqBoat-full, a variant of SeqBoat conducting the full attention over the compressed sequence without the working memory mechanism, on the Long Range Arena (LRA) benchmark  $\left[\mathrm{TDA}^{+}20\right]$ . LRA is designed to test the ability of neural models to capture long-range dependencies in various modalities with the following six tasks: ListOps [NB18], text classification (Text;  $\left[\mathrm{MDP}^{+}11\right]$ ), document retrieval (Retrieval; [RMQAJ13]), image classification (Image;  $\left[\mathrm{KH}^{+}09\right]$ ), Pathfinder  $\left[\mathrm{LKV}^{+}18\right]$  and its longer-range variant Path-X.
From Table 1, we can see that compared with other hybrid models with linear inference complexity, SeqBoat achieves the best average accuracy across all tasks, with a  $10.4 \times$  speed up and a  $95\%$  memory reduction compared to the Transformer. When compared with a pure SSM-based model, S4, which has more general modeling power than the MD-EMA used in our model, SeqBoat achieves a substantially higher accuracy on average with a  $2.17 \times$  speed up and a  $64\%$  memory reduction. Our model even achieves a significantly better average accuracy than state-of-the-art SSM model, S5 [SWL23], with a 0.16 absolute improvement, resulting in a new high score for models with linear inference complexity. Surprisingly, SeqBoat outperforms SeqBoat-full on Pathfinder and the Path-X tasks. This demonstrates the effectiveness of our proposed working memory mechanism to efficiently capture long-range dependencies, which is further analysed in Section 5. Since the learned module-level sparsity of our model is dynamic and task-specific, we also include the measurements for training and inference speedup and training memory allocation for all the tasks of LRA in Appendix D to provide a holistic view of the efficiency of our model.
4.3 Speech Classification
We also evaluate SeqBoat on the long-range speech classification task using the Speech Commands dataset [War18]. We follow S4 [GGR21] to use the SC10 subset of the dataset. It is a 10-class
classification task of raw speech which contains sequences of length 16k. As shown in Table 2, SeqBoat achieves competitive performance compared to the state-of-the-art S4 model and MEGA-chunk, with a  $1.32 \times$  speed up and a  $56\%$  memory reduction compared to MEGA-chunk.
Table 2: Test accuracy on the Speech Commands 10-way classification task of raw audio (Input length 16,000). We also report the training speed up and the memory consumption reduction compared with MEGA-chunk.
Model	#Param.	Acc. (↑)	Speed (↑)	Mem. (↓)
S4	300K	97.50	-	-
MEGA-chunk	300K	96.92	1.00×	1.00×
SeqBoat	293K	97.35	1.32×	0.44×
4.4 Language Modeling
We evaluate our model on enwik8 [Hut06], a popular language modeling benchmark. It consists of the first 100 million bytes of the English Wikipedia and is commonly used to evaluate the performance of sequence models on long-range dependencies. We follow previous studies to use it as a character-level language modeling task with a vocabulary size of around 200. During inference, we evaluate our model on the test set of enwik8 with a sequence length of 8,192 tokens and each token is equipped with a minimum of 7,000 tokens context window. We report the bits per character (BPC) metric for evaluation. As shown in Table 3, we compare SeqBoat with Transformer-XL  $\left[\mathrm{DYY}^{+}19\right]$ , Adaptive Span [SGBJ19], and MEGA-chunk. SeqBoat achieves the same performance as the previous best hybrid model, MEGA-chunk, but with a significant  $1.16\times$  training speed up. This demonstrates the effectiveness of our proposed model in auto-regressive language modeling of long texts.
Table 3: Model performance on the enwik8 character-level language modeling task with a 8192 training sequence length. Training speed up and memory consumption reduction compared with MEGA-chunk are also reported.
Model	#Param.	bpc (↓)	Speed (↑)	Mem. (↓)
Transformer-XL	41M	1.06	-	-
Adaptive Span	39M	1.02	-	-
MEGA-chunk	39M	1.02	1.00×	1.00×
SeqBoat	39M	1.02	1.16×	1.07×
In this section, we analyze experimental results of SMA from several research aspects. We also include comprehensive ablation studies of our model in Appendix E.
How much attention is needed for different sequence modeling tasks? As shown in Figure 3, we draw the activation time of the GAU module at different layers of our SeqBoat-full models for each of the task in the LRA benchmark. We measure the mean and the standard deviation (plotted as error bars) of the activation time on 100 sequences randomly sampled from the validation set of each task. The lines seem to be truncated due to the fact that different models have different number of layers. Generally, we can see that Image (in Blue and Circle) and Pathfinder (in Green and Triangle) need significantly more attentions than the text
Figure 3: The activation time (with error bars) of the GAU module at different layers of the SeqBoat-full model for different tasks in the LRA benchmark.
based tasks (i.e., Text, Retrieval and ListOps). We can also see that our model trained on ListOps (in Orange and Square) has a much larger variance of GAU activation time than other tasks, which
may reflect that the task difficulty per data sample of ListOps has a larger variance than other tasks. Another trend we can see is that the bottom layers near the inputs need much less attention (and sometimes even drop the attention module entirely) than the high level layers. We conjecture that this is because the SSMs with first-order complexity is already capable to capture the low-level features while the attention modules are more demanded at higher levels for conducting more complex computations of second-order pairwise comparisons.
(a) SeqBoat-full
(b) SeqBoat
Figure 4: The confidence probabilities of the GAU modular activation at each time step for the last two layers of the SeqBoat-full and the SeqBoat model. The results are measured on three input sequences randomly sampled from the validation set of the Pathfinder task. The sequences are reshaped back to  $32 \times 32$  squares for better visualization. Darker the blue color, higher the confidence. The white blocks indicate the time steps when GAUs are not activated.
How is the learned sparse modular activation distributed over the input sequence? As shown in Figure 4, we draw the confidence probabilities of the sparse modular activation for both the SeqBoat-full and the SeqBoat models on the Pathfinder task from the LRA benchmark. Generally, we can see that different input sequences have different modular activation patterns, which indicates that our modular activation is dynamically sparse. We can also see that the SeqBoat activation pattern is learned to be more structured than the SeqBoat-full, which may explain the superior performance of SeqBoat over SeqBoat-full on the Pathfinder task. Another trend we can see is that different layers have their own patterns: e.g., the Layer 5 of the SeqBoat model seems trying to find the curly paths in the latent representation space, while the Layer 6 seems aggregating the results from the contiguous time steps.
(a) Image
Figure 5: Training Speed v.s. Validation Accuracy trade-off on Image and Pathfinder of the LRA benchmark for different models with varying memory/chunk sizes. SeqBoat keeps a working memory of the compressed sequence, while MEGA-chunk splits the input sequence into non-overlapping sub-sequences. The memory/chunk sizes are marked along the lines. The GPU-hours for Image are measured on NVIDIA RTX A5000 GPUs, and Pathfinder on V100 GPUs with 32GB memory.
(b) Pathfinder
How does the working memory size affect the speed-quality trade-off? We compare the trade-off ability of our SeqBoat with MEGA-chunk on Image and Pathfinder tasks of the LRA benchmark,
(a) Pathfinder
(b) enwik8
Figure 6: Average Attention Span v.s. Layer Index on both the Pathfinder and the enwik8 tasks for different SeqBoat models with different working memory sizes. A smaller layer index indicates the layer is closer to the input. For reference, the average attention span of a sliding window based local attention is the half of the window size.
as shown in Figure 5. We draw the Pareto frontiers by respectively varying the working memory size for SeqBoat and the chunk size for MEGA-chunk from the set \{16, 32, 64, 128, 256, 512\}. The results are averaged with three random seeds. We can see that our models have significantly better trade-off than MEGA-chunk models for both the Image and the Pathfinder tasks. Notably, the Mega-chunk models generally fail (achieving around  $50\%$  accuracy as random guessing) on Pathfinder tasks with less than 128 chunk size, while our SeqBoat can still achieve around  $90\%$  accuracy with only a memory size of 16. This indicates the effectiveness of our working memory mechanism in capturing long term interactions.
How does the size of working memory affect the effective attention length of the GAU module?
We draw Figure 6 to investigate the relationship between the working memory size and the average effective attention length on both Pathfinder and the enwik8 datasets. We first measure the average attention distance from the query token to its key tokens for each of the time step, and the results are further averaged for all the time steps in a sequence. The average attention span on Pathfinder is measured on 100 sequences randomly sampled from the validation set. The average attention span on enwik8 is calculated based on a sequence of length 8,192 sampled from the validation set, and each token is equipped with a minimum of 7,000 tokens context window. From the figures, we can see a general trend that the models with a small working memory size (Size 16, 32 and 64 for Pathfinder, and Size 512 for enwik8) can have a far longer average attention span than its memory size. Notably, the layer 3 of the SeqBoat model with only 16 working memory slots surprisingly achieves an average attention span of 245 for the Pathfinder task. This partially explains the superior performance of our model over MEGA-chunk with small memory sizes and directly proves that our working memory mechanism can efficiently enable long-term interactions between the sequence elements.
In this paper, we present Sparse Modular Activation (SMA), a general mechanism for sparsely and dynamically activating sub-modules in neural networks. SMA not only has an efficient parallel implementation but also provides a theoretical guarantee for a full coverage of the function space of multiple modules. To address the challenges of combining attention with Linear State Space Models (SSMs), we apply SMA to develop a novel neural architecture, SeqBoat, that sparsely activates a Gated Attention Unit (GAU) based on the state representations from an SSM. The proposed model demonstrates a significantly better quality-efficiency trade-off compared to existing hybrid models across various sequence modeling tasks. To the best of our knowledge, SMA is the first mechanism that allows a neural network to obtain practical efficiency and complexity gains from sparsely activating a self-attention-like module. The strong empirical performance show that even with a preliminary application of SMA, we could already see substantial efficiency and interpretability benefits. For future work, we mainly want to explore the scaling behavior of SMA and also how to incorporate SMA with the pre-trained large language models.
We would like to thank Canwen Xu, Ziyi Yang, Ruochen Xu and Azure Cognitive Services Research group members for their feedback. We also want to thank Qi Zeng and Zixuan Zhang for the early discussion of the project.
$\left[\mathrm{AOA}^{+}20\right]$  J. Ainslie, Santiago Ontañón, Chris Alberti, V. Cvicek, Zachary Kenneth Fisher, Philip Pham, Anirudh Ravula, Sumit K. Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs in transformers. Conference On Empirical Methods In Natural Language Processing, 2020.
[AS68] Richard C. Atkinson and Richard M. Shiffrin. Human memory: A proposed system and its control processes. In The psychology of learning and motivation, 1968.
[BCB14] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. International Conference On Learning Representations, 2014.
[BD20] D. Barrett and B. Dherin. Implicit gradient regularization. International Conference On Learning Representations, 2020.
$\left[\mathrm{BFH}^{+}18\right]$  James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
[BZMA20] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. Advances in neural information processing systems, 33:12449-12460, 2020.
$\left[\mathrm{CLD}^{+}20\right]$  Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
$\left[\mathrm{DFE}^{+}22\right]$  Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems, 2022.
[DMRB22] Benoit Dherin, Michael Munn, Mihaela Rosca, and David Barrett. Why neural networks find simple solutions: The many regularizers of geometric complexity. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 2333-2349. Curran Associates, Inc., 2022.
$\left[\mathrm{DYY}^{+}19\right]$  Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. ACL, 2019.
[EUD17] Stefan Elfwing, E. Uchibe, and K. Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks, 2017.
[FDS+23] Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023.
[FZS22] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research, 23(1):5232-5270, 2022.
[SLP+21] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Reformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv: 2104.09864, 2021.
$\left[\mathrm{SML}^{+}21\right]$  David R. So, Wojciech Ma'nke, Hanxiao Liu, Zihang Dai, Noam M. Shazeer, and Quoc V. Le. Primer: Searching for efficient transformers for language modeling. ARXIV.ORG, 2021.
[SUV18] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. *NAACL*, 2018.
[SWL23] Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2023.
[TDA+20] Yi Tay, M. Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, J. Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. International Conference On Learning Representations, 2020.
[VPSP23] Ali Vardasbi, Telmo Pires, Robin M. Schmidt, and Stephan Peitz. State spaces aren't enough: Machine translation needs attention. ARXIV.ORG, 2023.
$\left[\mathrm{VSP}^{+}17\right]$  Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NIPS, 2017.
[War18] Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint arXiv: Arxiv-1804.03209, 2018.
[WLK+20] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
[XM22] Canwen Xu and Julian McAuley. A survey on dynamic neural networks for natural language processing. EMNLP Findings, 2022.
[ZLL+22] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7114, 2022.
A.1 Code for Compress and Extract Operators in SMA
Listing 1: Pytorch-like code snippet for the Compress operator.
def compress(q, a):
    # Arguments:
        # q: Input sequences with the shape:
            # (Batch Size, Input Sequence Length, Hidden Size)
        # a: Binary decision values with the shape:
            # (Batch Size, Input Sequence Length, 1)
        # Returns:
            # Compressed sequences with the shape:
            # (Batch Size, Compressed Sequence Length, Hidden Size)
        # Max sequence length of the compressed sequences
        max_sl = a-sum(-2).max()
        # Calculate the indices in the compressed sequence
        index_q = a * torch.cumsum(a, dim=-2)
        # Simultaneously map and pad to build the compressed sequences
        # +1 for temporarily storing non-activated elements
        shape = q.shape[-2] + (max_sl+1) + q.shape[-1:] 
        new_q = torch.zeros(shape)
        new_q.scanter_(-2,index_q, q)
        new_q = new_q[:,1:,]: # Non-activated elements are discarded
    return new_q
Listing 2: Pytorch-like code snippet for the Extract operator.
def extract(h, a):
    # Arguments:
        # h: Compressed sequences with the shape:
            # (Batch Size, Compressed Sequence Length, Hidden Size)
        # a: Binary decision values with the shape:
            # (Batch Size, Input Sequence Length, 1)
        # Returns:
            # Output sequences with zeros for non-activated positions:
            # (Batch Size, Input Sequence Length, Hidden Size)
        # Calculate the indices in the compressed sequence
        index_q = a * torch cumsum(a, dim=-2)
        # Pad the sequence dimension to account for non-activated positions
        h = F_pad(h, (0,0,1,0))
        h = torch.gather(h, -2, index_q.exp(-1, -1, h.shape[-1]))
    return h
A.2 Details of Gated Attention Unit
Our Gated Attention Unit (GAU) mostly follows the original implementation as [HDLL22], with slight modifications of the relative position bias. Given the compressed sequence  $\mathbf{H}^c\in \mathbb{R}^{r\times d_m}$ , we first compute the query, key and value sequences,  $\mathbf{Q},\mathbf{K},\mathbf{V}$ , i.e.,
$$
\mathbf {Q} = \mathbf {w} _ {q} \odot \mathbf {H} ^ {c} + \mathbf {b} _ {q} \quad \in \mathbb {R} ^ {r \times d _ {z}} \tag {3}
$$
$$
\mathbf {K} = \mathbf {w} _ {k} \odot \mathbf {H} ^ {c} + \mathbf {b} _ {k} \quad \in \mathbb {R} ^ {r \times d _ {z}} \tag {4}
$$
$$
\mathbf {V} = \operatorname {S i L U} \left(\mathbf {H} ^ {c} \mathbf {W} _ {v} + \mathbf {b} _ {v}\right) \quad \in \mathbb {R} ^ {r \times d _ {v}} \tag {5}
$$
where  $\mathbf{w}_q, \mathbf{w}_q, \mathbf{b}_q, \mathbf{b}_k \in \mathbb{R}^{d_z}, \mathbf{W}_v \in \mathbb{R}^{d_m \times d_v}, \mathbf{b}_v \in \mathbb{R}^{d_v}$  are the learnable parameters,  $\odot$  means the pointwise multiplication, and  $d_v$  is the expanded hidden dimension which is set as  $d_v = 2d_m$  in our work. Then the attention output is calculated as,
$$
\mathbf {O} = f \left(\mathbf {Q} \mathbf {K} ^ {T} / s + \mathbf {B} _ {\text {r e l}}\right) \mathbf {V} \quad \in \mathbb {R} ^ {r \times d _ {v}} \tag {6}
$$