Sparse Modular Activation for Efficient Sequence Modeling
Authors: Liliang Ren
Yang Liu
Shuohang Wang
Yichong Xu
Chenguang Zhu
ChengXiang Zhai
Venue: NeurIPS 2023 poster (2023)
Keywords: Sequence Modeling
Modularity
Sparsity
Attention Mechanism
State Space Model
Mixture of Experts
Neural Network
Transformer
TL;DR: We present Sparse Modular Activation, a general mechanism for sparsely and dynamically activating sub-modules in neural networks, addressing the challenges of combining attention with Linear State Space Models.

Abstract: 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.

Sparse Modular Activation for Efficient Sequence Modeling
Liliang Ren Yang Liu?
Shuohang Wang? Yichong Xu Chenguang Zhu Chengxiang Zhai'
University of Illinois at Urbana-Champaign Miicrosoft
{liliang3, www.iiiiiiiiiiiiiiooo edu com
[yaliu10, shuowa, xuyc118gmail.com
Abstract
Recent hybrid models combining Linear State Space Models (SSMs) with self- attention mechanisms have demonstrated impressive results across a range of se- quence modeling tasks. However, current approaches apply attention modules stat- ically and uniformly to all elements in the input sequences, leading to sub-optimal -effiityefficiency trade-offs. To address this limitation, we introduce Sparse Mod- ular Activation (SMA), general mechanism enabling neural networks to sparsely and dynamically activate sub-modules for sequence elements in differentiable manner. Through allowing each element to skip non-activated sub-modules, SMA reduces computation and memory consumption of neural networks both train- ing and inference stages. To validate the effectiveness of SMA sequence mod- eling, we design a novel neural architecture, SegBoat, which employs SMA to sparsely activate Gated Attention Unit (GAU) based on the state representa- tions learned from an SSM. By constraining the GAU to only conduct local at- tention the activated inputs, SeqBoat achieve linear inference complexity with theoretically infinite attention span, and provide substantially better quality- wide
efficiency trade-off than the chunking- buused models. With experiments range of tasks, including long sequence modeling, speech classification and lan- guage modeling, BeqBoat brings new state-of- the-art results among hybrid models with linear complexity, and reveals the amount attention needed for each task through the learned sparse activation patterns. Our code is publicly available https:/ //tpttgtt.t/.t//////////eeeeeeat
thub. cww Baamm/eenn//eqBoat
Introduction
Recent advance on efficient sequence modeling with State Space Models (SSMs) [GGR21; GDE+20; GGGR22; GB22; SWL23] has shown impressive performance for 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, carefully designed state matrices and the application of convolutional parallelism [GGR21]. How- ever, they still significantly underperform the self-attention [BCB14] 14; VSP+ 17] based model in both language modeling and machine translation [VPSP23] tasks. A recent work [FDS* 233 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
efficiently capture the recurrent structure in sequential data with
WWrk partially done during internship at Microsoft.
Work done at Microsoft.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

associative recalling ability. Furthermore, better rectly introducing extra self-attention modules to form a hybrid model (e.g. MEGA [MZK 223] and Hybrid H3 [FD that utilizes both the first and the second order inductive biases, ie., 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 sub- optimal -effiieye 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.
trade-off can be achieved by di-
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- 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& neural network to sparsely and dynamically activate its sub-modules for each of the input token in fully differentiable manner. Specifically, we assume a neural model can be com- posed 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 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 [LQC 22; GZYE20; XM22]. To enable the practical efficiency gains from our module-level sparsity, we provide 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 RQL, 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 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 wide range of tasks, including Long Range Arena (LRA) ITDA+ TTD, speech classification [War18] and language modeling [Hut06]. On the competitive LRA benchmark, SeqBoat achieves 1.96 higher average accuracy than MEGA- chunk [MZK* + 233, the previous best hybrid model, with a 10.4x training speed up and a 95% mem- ory reduction compared to the Transformer [VSP* 17] 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, ad- dressing RQ2. We demonstrate that our working memory mechanism provides substantially better cemputaaionnaccurae 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.
2 Background
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 compar- isons between our approach with the related works is also included in Appendix F. We begin by reviewing how the standard sequence modeling formalized to establish the common ground for our discussion.

---

2.1 Time-Invariant Sequence Modeling
Given discrete sequence, x = (fr,..., Zn E R", consisting tokens, time inverinvrri sequence
model Pe is optimized to maximize the likelihood of the observed sequences by factorizing them as follows:
max Po(x)
II P(x 0), [=1
where X (.i,,...;" 2e- } is the sequence history at time step I, and the parameter is independent of the time step I. This formulation implies that the full model parameters 0 and the full history X are both essential for the conditional prediction of each token zt. However, one potential issue is the prediction difficulty of each token may differ depending the context and the position, this static model Pe can lead to sub-optimal eccuracy-efficiene trade-off by wasting computation on either unimportant context [SJP*21] 21] easy-to-predice 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 problem of finding an optimal time-variant model that can dy- namically activate subset of modules from pre-defined function space for each time step.
3.1 Variant Sequence Modeling
Formally, a variant sequence model is defined on a compact function space F X H [0, 1|=xV' where V is the size of the vocabulary and X, {xf xf S X E X C R"). con- tains all possible sub-sequences of the sequence history x Then for each the token prediction the time step the model learns apply function ft with the parameters 0, that maximizes the sequence probability, ie..
max P=(x) II Py. (x.lx, s.t. Xx S f..0e.xE
(1)
This formulation generalizes the previous works in pursuing a dynamic and sparse model for se- this work, we quence modeling, where the connections are further explained in Appendix Li o E, where assume the function space F is chain-structured,, i.e., FF H LN H 1dd + [0, 1] is the classification function, & R + Raxdm the embedding function, N is the number of intermediate layers, dm the model size and L Raxdm.. + xdd. is the function space of the intermediate mappings. We further assume that C is the spanning set of a finite number of the function S with its parameters 0% i.e., f span the number of pre-defined functions. These assumptions justify the design of our Sparse Modular Activation mechanism, which further explained the following section.
J'r where M
3.2 Sparse Modular Activation
Sparse Modular Activation (SMA) introduces a latent configurator at each time step and each layer of neural sequence model. The configurator produces a binary decision vector a, {0, 1] to only activate those functions f that have the decision value aj 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-end differentiability, we also require the latent configurator to output the confidence probability of its decisions, e, [0, M] and respectively scale the output from the activated function with its confidence values cf. The function space of a layer equipped with SMA can described as follow.
generate the final vector, yt. of the layer. To enable end-
LSMA (ar, Ct) = Σ G VG E R, a 1,1is -i
where the index set of activated functions. Notably, SMA has the following important property, Sh Theorem (Function Space Coverage of SMA). For any L C C Si.., I'1 } there exists a pair of cl), that (amn(a(, c)) L' other words, SMA has a full coverage of the function space L.
3

---

Latent Configurator
Tempered Softmox
Mas
Linear
Arginax
a
S
Linear SSM
SIUU
GAU
SILU
S
Spans: Activation
Spanc Activation
Linear
Figure 1: The proposed SeqBoat layer. The black lines indicate that the gradients can be back- propagated and the red lines stand for gradients blocking. ④ means the element-wise multiplication, and e is the point-wise addition. The max, argmax and softmax operators are all applied to the projected dimension after the linear layer 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, recovered, ie., LSMA = L. if all the functions are activated, and the model can adaptively se- lect any subspace of L through controlling the binary decision vector a.. During sequential in- ference stage, for each function that requires the contextual information, we keep a record of its activated input representations across the time steps form a memory Hद C H E Ruxdm (where He: is the input sequence history), so that each function can recall its contextual mem- ory for each of the decoding step. During parallel training stage, the activated input representa- tions for each function (or module) are sparsely selected to compose the compressed sequence He for the module to conduct parallel processing, and this process is implemented as com- press operator. The returned representations are then mapped to its original input position in 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 [PGM+ 19] code snippets of the compress and extract operators are provided the Appendix A.I with an efficient support of batched sequences using the scatter operation.
that the original function space L
Configurator
aq
In this work, we explore applying SMA ad- dress our RQs, in which the following two types of functions the function space L are considered: (1) linear State Space Model (SSM) [GGGR22; WZK+23) that efficiently captures the recurrent sequential structure, and (2) a Gated Attention Unit (GAU) [HDLL22] that performs second-order pairwise compar- isons between the sequence elements. We also keep the SSM as always-activated module and only apply SMA to the GAU, and in this case the effective number of modules M 1. This 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 Se- qBoat layer. Its details are further explained in the next section.
f
ha
y/2
h2
h,
f
y/
Lincar Aggregation
Compress
Extract
Figure 2: The proposed parallel implementation of the Sparse Modular Activation (SMA) mech- anism. In this example, we have two modules fi. f and an input sequence H {hI...., h4). We assume only a}, aq. of have values equal to one. The white block means zero vector. The compressed sequences for modules fl and f are
Hi = (h2, h1) and Hद {h,) respectively. The final outputs are aggregated as y4 = clyl ciy; and y2 dbyy.
3.3 Model Architecture
Figure illustrates the proposed architecture for our SeqBoat layer. Given an input sequence repre- sentation S E xRd for the 1-th layer, where n is the length of the sequence and dm is the hidden size. first apply a linear time-invariant State Space Model (SSM) [GGR21] over the sequence

---

followed by SiLU activation [BUDI7] to obtain the state representations H E R=xdm
SSM(S') = K * S' +
(2)
H = SILU(SSM(S))
where is convolution operator which efficiently implemented using FFTs [GGR21], O indicates the element-wise multiplication, D E IR xdm is the learnable parameters, and K I R=xdm is the SSM convolution kernel. In this work, we use the Multi-dimensional Damped Exponential Moving Average (MI--MMA) proposed by MEGA IMZK+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 more advanced SSM such as Liquid-S4 [HLW 22] or S5 [SWL23] can lead to better performance. but we keep using MD-EMA for fair comparison with MEGA. Given the state representation H E [g [nnxdm, a latent configurator with one linear layer is applied to obtain the decision vector (0, 1)" and the confidence probability vector E [0, 1]th with the Tempered Softmax, i.e.,
(Hw,+b)/T
P; (Hwo+bo)/r + e (Hwi+bi)/7 for i E (0,1)
a arg max Pi, C max Pr
where Wo: W1 R.dm bo. bj E are learnable parameters, 7 E R2 a learnable temperature parameter which is initialized as a (where is a hyperparameter denoted as the initial temper- ature scale). Previous works [RSVG20a; GLL ULL++22 22; FZS22; RZW+ 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 Soft- max 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 de- cision confidence calculation. The details of the theorem and the proof are provided in Appendix C.
Based on the decision values a, Gated Attention Unit (GAU) is activated and applied on the compressed sequence He resulted from the compress operator as previously described in Section 3.2,
H Compress(H, a)
R 1=xdm
Y GAU(H))
E R.
Y = Extract(Y" (Y.,tt
E R_ddd
where r - Et ag I n is the total number of the time steps when the GAU module is activated. We use Squared ReLU [SML+21]] 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 Y E Raxdm, the final layer output is computed with a linear aggregation followed by non-linear activation, i.c.
S = SiLU O Y + HW + b b-s'
where W Rh Xdm b E Rdmm are the learnable parameters. We also add normalization layer either before the Linear SSM (i.e., Pre-Norm) the residual branch or after the final layer output s (i.e., Post-Norm) to stabilize the training process. The SeqBoat model is then constructed as stacking N number of SeqBoat layers the same hidden size dm 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 depen- dent on the activation probability of the GAU module p r/n E.= a,/n for each of the SeqBoat layer. This means that the computation complexity O(r + n log(n)) for parallel training and Or²² + n) for regressegee decoding. While it is possible that the number of the GAU- activated time steps far smaller than the length the full sequence for trained model, a random initialized Beqtoat model will still activate the GAU module for n/2 number of times on average, leading to a non-ideal (n²) training complexity at the early stage of training. To solve this prob- lem, we restrict the GAU to only conduct local attention [AOA* + 20] with sliding window of size
5

---

w & on the compressed sequence H 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 bidierectional sequence modeling, and number of past tokens for auto-regressive sequence modeling. This approach keeps a Frrst-In-Frrs--On (FIFO) memory with size w of the compressed sequence history during auto-regressive decoding, which can be understand as 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(ww + n), while maintaining the ability of at- tending 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.
Experiments and Results
We evaluate the SeqBoat model on three representative long-range sequence modeling benchmarks across image, text and speech modalities to have 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.
4.1 Baseline Models
We compare our proposed SeqBoat model with multiple strong baseline models besides the Trans- former [VSP+ 17] and its variant Transformer-Xi XXL [DYY+19].
S4 [GGR21] 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] simpler variant of S4 that uses a diagonal state matrix.
MEGA [MZK* + 23] is hybrid model that combines gated self-attention mechanism with Multi- dimensional 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 S4D. MEGA also has linear complexity variant, MEGA-chunk, that applies attention to each local chunk of fixed length.
H3 [FDS+ 23] introduces hybrid SSM layer that stacks diagonal SSM and shifted SSM to model the comparisons between the elements in a sequence. also provides a hybrid model that includes extra self- self-attention layers for better language modeling performance.
S5 [SWL23] is the state-of-the-ar SSM model that improves over S4 with multi-inpu multi- outti- and time-domain processing. It leverages the parallel scan operator in JAX [BFH 18] for an efficient parrallel implementation recurrence unrolling.
Liquid-S4 [SWL.23] 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 memory-efficiint 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 [CLD+20] uses the Fast Attention via positive Orthogonal Random features algorithm to ap- proximate the self-attention mechanism, reducing the computational complexity. The Linformer [WIKK employs low-rank approximations project the self-attention mechanism into a lower- dimensional space, resulting in more efficient system.

---

Table 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 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
59.24 88.21
1x
1x 0.31
MEGA*
63.14
90.43
91.25
90.44
96.01
97.98
2.9x
Sub-quadratic Inference Complexity
Reformer
37.27
56.10
53.40
38.07
68.50
50.67
0.8
0.24x
BigBird
36.05
64.02
59.29
40.83
74.87
X
55.01
1.1x
0.30x
SeqBoat-full
61.65
89.60
91.67
89.96
95.87
95.28
87.33
6.2x
0.07x
Linear Inference Complexity
Performer
18.01
65.40
$3.82
42.77
77.05
X
51.41
5.7x
0.11x
Linformer
35.70
53.94
52.27
38.56
76.34
51.36
5.5x
0.10x
Luna-256
37.98
65.78
79.56
47.86
78.55
X
61.95
4.9%
0.16x
59.10
86.53
90.94
88.48
94.01
96.07
85.86
4.8x
0.14x
S4D-LegS
60.47
86.18
89.46
88.19
93.06
91.95
84.89
6.1x
0.14x
62.15
89.31
91.40
88.00
95.33
98.58
87.46
6.1x
0.14x
Liquid-S4
62.75
89.02
91.20
89.50
94.8
96.66
87.32
1.2x
0.17x
H3*
57.50
88.20
91.00
87.30
93.00
91.80
84.80
6.0x
0.24x
MEGA-chunk*
58.76
90.19
90.97
85.80
94.41
93.81
85.66
7.0x
0.09x
SeqBoat* 61.70
89.60
91.28
90.10
96.35
96.68
87.62
10.4x
0.05x
Luna IMKW*21 21] is a linear unified nested attention mechanism that approximates Softmax atten- tion with two nested linear attention functions, resulting 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 [TDA* 20]. 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; [MDP- 1 11)), document retrieval (Retrieval; [RMQAJ13]), image classification (Image; 091). Pathfinder [LKV+18] 18] 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 speed up and 95% mem- ory 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.17x speed up and 64% memory reduction. Our model even achieves significantly better average accuracy than state-of-the-ar SSM model, S5 [SWL.23], with a 0. 0616 absolute improvement, resulting in a new high score for models with lin- ear 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 Section 5, Since the leamned module-level sparsity of our model is dynamic and task-specific, we also include the mea- surements for training and inference speedup and training memory allocation for all the tasks of LRA in Appendix D to provide 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 [GGR2I] to the SC10 subset of the dataset. It is a 10-class
7

---

classification task of raw speech which contains sequences of length 16k. As shown in Table 2, SeqBoat achieves competitive performance compared the state-of-the-art S4 model and MEGA- chunk, with x.32x 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 com- pared with MEGA-chunk.
Model
#Param.
Acc. (1)
Speed (1)
Mem. (1)
S4
300K
97.50
MEGA-chunk
300K
96.92
1.00x
1.00x
SeqBoat
293K
97.35
1.32x
0.44x
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 perfor- mance 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, evaluate our model on the test set of enwik8 with a sequence length of 8, 8,192 tokens and each token 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-Xe [DYY 19], Adaptive Span [SGBJ19], and MEGA-chunk. SeqBoat achieves the same performance as the previ- ous best hybrid model, MEGA-chunk, but with significant 1. 1.16x training speed up. This demon- strates 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--hunk are also reported.
Model
#Param.
bpc ())
Speed (t)
Mem. (L)
Transformer-X XL
41M
1.06
Adaptive Span
39M
1.02
MEGA-chunk
39M
1.02
1.00
1.00x
Seq
39M
1.02
1.16x
1.07
Analysis
In this section, analyze experimental results of SMA from several research aspects. We also include comprehensive ablation studies of our model Appendix
100
80
Activation Timet%))
How much attention is needed for different sequence modeling tasks? As shown in Fig- 3, we draw the activation time of the GAU module at different layers of our BeqBoat-ful models for each of the task in the LRA bench- mark. We measure the mean and the standard deviation (plotted as error bars) of the activa- tion time on 100 sequences randomly sampled from the validation set of each task. The lines seem to be truncated due to the fact that dif- ferent 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- based tasks (i.e,, Text, Retrieval and ListOps). We can also see that our model trained on ListOps (in Orange and Square) has much larger variance of GAU activation time than other tasks, which
60.
40
Image
Licops
Putrinder
That
Retrieval
PathiX
Layer Indox
Figure 3: The activation time (with error bars) of the GAU module different layers of the SeqBoat-full model for different tasks in the LRA benchmark.

---

may reflect that the task difficulty per data sample of ListOps has larger variance than other tasks. Another trend we can see is that the bottom layers near the imputs 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 fea- tures while the attention modules are more demanded at higher levels for conducting more complex computations of second-order pairwise comparisons.
Layer
Layer
Layer
Layer
Sample
Sample 2
Sample
Sumple
Sample
Sample
(a) SeqBoat-ful
(h) SeqBoat
Figure 4: The confidence probabilities the GAU modular activation at each time step for 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 reshaped back to 32 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 imput sequences has different modular activation patterns, which indicates that our modular activation is dynamically sparse. We can also see that the Seq activation pattern learned to 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 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.
100
512
32
128
250
256
64.
256
90
36
85.5
128.
85
Accurscy (%)
60
Accuracy
75
70
32
77.5
753
36
16
-M-MEGAA
Segliont
55
16
32
-128
Segilest
735
50
13
15
17
19
28
33
38
43
GPU--ouur
GPU-hours
(a) Image
Pathfinder
Figure 5: Training Speed Validation Accuracy trade-off Image and Pathfinder of the LRA benchmark for different models with varying memory/chunk sizes. SeqBoat keeps a working mem- ory 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 V100 GPUs with 32GB memory.
How does the working memory size affect the speed-quality trade-off? We compare the trade- off ability of our SeqBoat with MEGA-chumt on Image and Pathfinder tasks of the LRA benchmark.

---

1400
250
1200
Average Attention Span
Average Attention Span
200
1000
150
Memory Size
600
100
+
16
32
400
64
Memary
128
+ 512
256
1024
0-
Layer Index
Layer
(a) Pathfinder
(b) enwik8
Figure 6: Average Attention Span Layer Index 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 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 SeqBoat 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 size working memory affect the effective attention length of the GAU module? We draw Figure 6 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 sequence. The average attention span on Pathfinder measured 100 sequences randomly sampled from the validation set. The average attention span on enwik8 is calculated based on a sequence of length 8, 8,192 sampled from the validation set, and each token is equipped with minimum of 7, 7,000 tokens context window. From the figures, can see general trend that the models with small working memory size (Size 16, 32 and 64 for Pathfinder, and Size 512 for enwik8) can have far longer average attention span than its memory size. Notabaly, the layer 3 of the SeqBoat model with only 16 working memory slots surprisingly achieves an average attention span 245 for the Pathfinder task. This partially explains the superior performance of our model over MEGA-chune -eGGA-chunk with small memory sizes and directly proves that our working memory mechanism can efficiently enable long-term interactions between the sequence elements.
Conclusion
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 im- plementation but also provides a theoretical guarantee for 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 novel neural architecture, SeqBoat, that sparsely activates Gated Attention Unit (GAU) based on the state representations from an SSM. The proposed model demonstrates a significantly better across various sequence modeling tasks. To the best of our knowledge, SMA is the first mecha- nism that allows a neural network to obtain practical efficiency and complexity gains from sparsely activating a self-attention--ike module. The strong empirical performance show that even with 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.
trade-off compared to existing hybrid models
10

---

Acknowledgement
We would like 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.
References
[AOA*20 20] J. Ainslie, Santiago Ontañónn Chris Alberti, V. Cvicek, Zachary Kenneth Fisher, Philip Pham, Anirudh Ravula, Sumit K. Sanghai, Qifan Wang, and Li Yang. Ete: 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 transla- tion by jointly learning align and translate. International Conference On Learning Representations, 2014.
[BI20]] D. Barrett and B. Dherin. Implicit gradient regularization. International Conference On Learning Representations, 2020.
[BFH+ 18] 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 pro- grams, 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-1-22460 2020.
(CLD+20) 20] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, An- dreca Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint 20Xv:2009..4494 2020.
[DFE+22] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAtten- tion: 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 net- works find simple solutions: The many regularizers of geometric complexity. In Koyejo, Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Ad- vances in Neural Information Processing Systems, volume 35, pages 2333-2349, Cur- ran Associates, Inc., 2022.
[DYY+19] Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length con- text. ACI, 2019.
[BUD17] Stefan Elfwing, E. Uchibe, and K. Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Newral Networks, 2017.
[FDS+23] Daniel Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christo- pher Re. Hungry hungry hippos: Towards language modeling with state space models. 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):523--5200 2022.

---

[GB22] Ankit Gupta and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. ARXIVORG,, 2022.
[0D Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In H. Larochelle, M. Ranzato, R. Had- sell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1477- Curran Associates, Inc., 2020.
[GGGR22] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. ARXIVORG, 2022.
[GGR21] Albert Gu, Karan Goel, and Christopher R'e. Efficiently modeling long sequences with structured state spaces. International Conference On Learning Representations, 2021.
[GLH* 19] Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, S. Levine, Yoshua Bengio, and B. Scholkopf. Recurrent independent mechanisms. International Confer- ence on Learning Representations, 2019.
[GLL+ Yue Guan, Zhengyi Li, Jingwen Leng, Zhouhan Lin, and Minyi Guo. Transkimmer: Transformer learns to layer-wise skim. In Proceedings of the 60th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7275- 7286, Dublin, Ireland, may 2022. Association for Computational Linguistics.
[Gra16] A. Graves. Adaptive computation time for recurrent neural networks. ARXIV.ORG, 2016.
[GZYE20] Trevor Gale, M. Zaharia, C. Young, and Erich Elsen. Sparse gpu kernels for deep learning. International Conference For High Performance Computing, Networking, Storage And Analysis, 2020.
[HDLL22] Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le. Transformer quality in linear time. International Conference On Machine Learning. 2022.
[2LW+22] Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arXiv:2209. 12951, 2022.
[Hut06] Marcus Hutter. The human knowledge compression contest. http:///pize.huutter... 2006.
[JG Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and João Carreira. Perceiver: General perception with iterative attention. International Conference On Machine Learning, 2021.
[JGP17] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel- sofumax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017.
[JJHH1]] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton, Adap- tive mixtures of local experts, Neural Computation, 3:79-87, 1991.
[KH+09] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[KKL20] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient trans- former. arXiv preprint arXiv:2001.0445) 2020.
[LHI8] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Interna- tional Conference on Learning Representations, 2018.
[LJH+19] 19] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance the adaptive learning rate and beyond. Inter- national Conference on Learning Representations, 2019.
12

---

[LKV 18] Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. Ad- vances in neural information processing systems, 31, 2018.
[LLX* 20] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yan- ping Huang, M. Krikun, Noam M. Shazeer, and Z. Chen. Gshard: Scaling giant mod- els with conditional computation and automatic sharding. International Conference On Learning Representations, 2020.
[LQC*22] Liu Liu, Zheng Qu, Zhaodong Chen, Fengbin Tu. Yufei Ding, and Yuan Xie. Dy- namic sparse attention for scalable transformer accelcration. IEEE Transactions on Computers, 71:3165-311 -178,, 2022.
[MDP+11 11] Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the association for computational linguistics: Human language technologies, pages 142-150 2011.
[MKW+21] Xuezhe Ma. Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. Advances in Neural Information Processing Systems, 14:2441-22555 2021.
[MZK*2]] 23] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated atten- tion. In The Eleventh International Conference on Learning Representations, 2023.
[NB18] Nikita Nangia and Samuel R Bowman. Listops: A diagnostic dataset for latent tree learning, arXiv preprint 18X4::8804006235 2018.
[PGM+ 19] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural in- formation processing systems, 32, 2019.
[RMQAJ13] Dragomir R Radev, Pradeep Muthukrishnan, Vahed Qazvinian, and Amjad Abu-Jbara. The acl anthology network corpus. Language Resources and Evaluation, 77:919-9444 2013.
[RSVG20a] Aurko Roy, M. Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. International Conference On Topology, Algebra And Categories In Logic, 2020.
[RSVG20b] Aurko Roy, M. Saffar, Ashish Vaswani, and David Grangier. Efficient content- based sparse attention with routing transformers. International Conference On Topology, Algebra And Categories In Logic, 2020.
[RZW+22] Liliang Ren, Zixuan Zhang, Han Wang, Clare Voss, ChengXiang Zhai, and Heng Ji. Language model pre-training with sparse latent typing. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pages 1480- 1494, Abu Dhabi, United Arab Emirates, dec 2022. Association for Computational Linguistics.
[SGBJ19] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adap- tive attention span in transformers. arXiv preprint arXiv:190.007799, 2019.
[SJP+21] Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur Szlam, We- ston, and Angela Fan. Not all memories are created equal: Learning to forget by expiring. International Conference On Machine Learning, 2021.
18LK+18] Sihyeon Seong, Yegang Lee, Youngwook Kee, Dongyoon Han, and Junmo Kim. To- wards flatter loss surface via nonmonotonic learning rate scheduling. In Conference on Uncertainty in Artificial Intelligence, 2018.
13

---

[SLP*21] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Ro- former: Enhanced transformer with rotary position embedding, arXiv preprint arXiv: 2 2104.09864 2021.
[SML+21] 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. ARXIVORG, 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: 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 t enough: Machine translation needs attention. ARXIVORG,, 2023.
[VSP+17] Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NIP,, 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 200.:200000.4668, 2020.
[XM22] Canwen Xu and Julian McAuley. A survey on dynamic neural networks for natural language processing. EMNLP Findings, 2022.
[2LL++22] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, An- drew M Dai, Quoe V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems, 35:7103-7116 2022.
14

---

A Implementation Details
A.1 Code for Compress and Extract Operators in SMA
def compress (q, a):
Args:
9: 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:
Compreased sequences with the shape: (Batch Size, Compressed Sequence Length, Hidden Size)
# Max sequence length of the compressed sequences
a.m max) max(()
Calcuate the indices the compressed sequence
index_q torch. cumsum (a, dim=-2)
# Simultaneossly map and pad to build the compressed sequences for temporarily storing non-activated elementa (max_al+1,) q.shape[-1:]
shape -.shape[:-2]
new_q torch. zeros (shape)
new_q. scatter (-2, index_q. q)
new_q new_q(: ;] Non-activated elements discarded
return new_q
Listing 1: Pytorch-like code snippet for the Compress operator.
def extract (h, a):
Arga:
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 non-activated positions: (Batch Size, Input Sequence Length, Hidden Size)
# Calcuate the indices in the compressed sequence index_q torch. cumsum (a, dim=-2)
Pad the sequence dimension account for non-activated positions
F. (h, (0,0,1,0))
h torch. wet.eratherrr (h, -2, "2nixdex_g (-apaa-s,,,,,ssapp([--1]] (-1)))
return
Listing 2: Pytorch-like code snippet for the Extract operator.
A.2 Details of Gated Attention Unit
Our Gated Attention Unit (GAU) mostly follows the original implementation [HDLL22), with slight modifications of the relative position bias. Given the compressed sequence H E rxdd we first compute the query, key and value sequences, Q, K. V, i.e.
Q = w, @ H + b
E Brxx
(3)
K = ④ H +b
E Rrxd
(4)
V W.
xxd
where Wa, Wg. b by E Rd. W,, BR xdv b, E Rdd are the learnable parameters, @ means the pointwise multiplication, and d,, is the expanded hidden dimension which is set as d,, 2d,, in our work. Then the attention output is calculated as,
15

---

O= f (QK' ++BBa)) V
E RR
where f is the attention function, 8 r is the normalizing factor, B R is the relative posi- tional bias which is set as the simple relative position bias [SUV18; MZK 23] by default. We also have the choice for using RoPE [SLP*21] relative position embedding, and this case, we don't apply relative position bias but use RoPE in its original form to transform the Q and K sequences. Squared ReLU is applied [SMLL 21] for the attention function on the image and the speech tasks and Softmax for the text related tasks. By default, the relative positions token positions in the original sentence H instead of the compressed sequence H For working memory mechanism, we conduct local attention between Q, K, V with the sliding window size (or the working memory size) w, and we set the normalizing factor 8 = w. The final output of the GAU is computed with gating mechanism,
calculated based on the
G Si, + ba)
E R
(7)
Y = ( + b
E IR 1R
(8)
where W, E Rn xd. Wr E Xdm
xdm bg Rdd by E Rd are learnable parameters.
A.3 Details of Multi-dimennional Damped EMA
As special kind of SSM, the Multi-dimensional Damped EMA (MD-EMA) operator is first pro- posed in the MEGA [MZK* 233 paper, and we briefly introduce it here for the completeness of our work. In Equation (2), the kernel K Raxdm of MD-EMA is parameterized as following.
h K = Σ (. i=1
h Enlooao)...... Σ
h
( 3 (x0)))
(9)
i=1
where v) Blxxdm, a [0, 10x1, E [0, 101xdm, , E R.xdd.. 1- 6, learnable parameters, @ means the element-wise multiplication and h is hyperparamter denoted as the EMA dimension. With the pre-computed kernel K, the output of MD-EMA can be efficiently computed with FFTs in time complexity of On log During sequential decoding stage at the time step t, the output of MD-EMA is calculated as
Utjij = BijBt.ji 8t = [8t.5] E TR
u, [ue,i]] E hRg
2, = a ④ U, (1
2t-1
h
MD-EMA(s:) = Σ niZ..i +DQ
where zg = S; E Rd is the input representation at time and D R.d. are learnable parameters. We can see that the inference of MD-EMA has O(1) complexity per time step that is independent of the sequence length n.
A.4 Details of Latent Configurator with Gumbel Softmax
To produce this baseline, apply Gumbel Softmax the probability computation in the latent configurator described in Section 3.3. Specifically, the modified probability calculation is as follow,
(Hw)+b++G)/+
P (Hwo +o)) + (Hw,+b,-G.))
for E {0,
where G; Gumbel(O, 1) is the Gumbel noise sampled from standard Gumbel distribution and 7' follows the following temperature annealing schedule as in the previous work [BZMA20],
max(2 0. 0.999995" 0.5),
where T is the number of training steps.
16

---

Table Hyper-pen apper-parameter Settings of our SeqBoat model for the LRA benchmark and the Speech Command (SC) dataset. DP is the dropout rate, BSZ is batch size, LR is learning rate, WD is weight decay, and Pre-N is Pre-normalization. BN, LN and SN refer to Batch Normalization, Layer Normalization and Scale Normalization.
Task
Depth
dm
Atte-FN
Norm
Pre-N
BSZ
LR
DP
WD
Epochs
ListOps
0.3
softmax
FALSE
0.004
Text
128
0.3
256
sofimux
SN
FALSE
0.004
0.1
0.01
50
Retrieval
128
0.3
256
softmax
SN
FALSE
64
0.003
0.1
0.04
40
Image
160
0.4
relu
BN
TRUE
0
0.02
200
Pathfiuder
128
256
relu?
BN
TRUE
128
0
0.01
200
Path-X
128
1.0
64
256
relu?
BN
TRUF
128
0.01
0
0.01
100
SC-Raw
120
relu
BN
TRUE
0.01
200
A.5
Hyper-parameter settings
For Long Range Arena (LRA) and Speech Command tasks, we use the AdamW [LHI8] optimizer and the detailed settings together with our initial temperature scale a for the latent configurator are shown in Table 4. Following S5 [SWL23], we NLI-style fusion of encoder features for the Retrieval task of LRA. For SeqBoat model, we use a working memory size w = 512 for Path-X and = 256 for other tasks in LRA. For Speech Command, working memory size of 256 used. And for language modeling tasks, we use the RAdam [LJH* 19] optimizer and a working memory of size 1,024. We use RoPE relative position embeddings and the Softmax attention function for language modeling, and the detailed hyperparameters are shown in Table 5. For the speech command task and the Text task of LRA, the relative positions are calculated based on the positions in the compressed sequence.
A.6 Details of Efficiency Measurements
Due to the dynamic sparsity of our model, we first measure the training time the average per step wall time across the full training stage. The training speed is then calculated as the in- verse of the training time. The memory con- sumption is measured as the average per step GPU memory allocation across the full training stage. To ensure fair comparisons, the relative training speedup between different models are calculated based on the training time measured on the same hardware using the same batch size settings. All the experiments are conducted on mixed cluster with 8 NVIDIA V100 32GB GPUs and 2 NVIDIA A5000 24GB GPUs.
Table 5: Plyper-paramedd of our SeqBoat model for language modeling.
enwik8
Sequence Length x Batch Size Optiizerr
8192 x 8
RAdam
Learning Rate
0,005
Initial Learning Rate
0,002
RAdam-8
(0.9, 0.98)
Learning Rate Decay
lincar
Weight Decay
0.1
Dropout
0.15
Attention Dropout
0.0
Gradient Clipping
0.25
Warmup steps
24K
An alternative way of efficiency measurement to calculate the total Floating Point Opera- tions (FLOPs) used for training. However, the FLOPs calculation does not take account of the time-consuming memory copying of large ma- trices (which is exactly the speed bottleneck of SMA due to the scatter operation), and thus we choose to follow the previous works measure the speedup based on the actual wall time for stricter and fairer comparison.
Total updates
400K
Decoder Layers
14
Model Size (dm)
616
Query and Key Sizes (d.)
128
Value Size (d.)
1232
EMA dimension (h)
16
Initial Temperature Scale (a)
1.0
Working Memory Size (w)
1024
Total Parameters
39M
B Proof of Theorem 1
Proof. Given any L' C span(st..... fhr), we want to show that there exists a pair (a/s of) such that LSMA(at) ct) L'. Since L' C L, every function in L' can be written as linear combination
17

---

of the functions (fl...., far Let's consider function L' Then, we can write
M
where 33 E R.. Now, let's define ai and df as follows:
a (af) where at if 8; # 0 and a otherwise.
c = (4') where c = 1 if 3 # 0 and o O otherwise.
Then, have
LSMA LsMA(a/, ci) =
EG4 1AG VG ER, ( (i (
iE!
Since at = 1 and ei = -1 if and only if B; + 0, then, I= {i 3; 10,00,1 0,1 M) we have
LSMA (a/, ci)
Σ
Saft
AGe E
B v3, R R -C
This completes the proof.
C Theorem of Implicit Regularization for Latent Decision Making
We first give brief introduction of the implicit regularization to provide some backgrounds. The im- plicit regularization of gradient decent [BD20] is a side effect caused by the discreteness of the gradi- ent descent updates. Specifically, it has been proved that a discrete gradient update 0' = =0 of the parameter 0 is implicitly minimizing a transformed loss, L' L + -1 VoL(O) where h the learning rate and L is the training loss term. Assuming the usage of Stochastic Gradient Decent (SGD) for gradient updates, [DMRB22] have shown that the implicit gradient regularizer for a multi-class classifi- cation cross-entropy loss term, where fo(x) is neural network that takes the input and outputs the logits for classification. A Transfer Theorem [DMRB22] is further derived to characterize the conditional transfer of the implicit regularization pressure from |/ (x)] to the gradient norm [z) where ||. || means the Frobenius norm. The transfer theorem assumes an ordinary neural architecture fo(x) : Rd + R with / layers in the following form,
[()]3 will conditionally put an implicit pressure on || oe [z)]l fz))
(x))) a, (wh, + b,), for i 1,2......
where is the activation function, (.) indicates the sub-network from the input E IR up to the output of layer and (z x. This implies h₂(x) fo(x).
In this work, consider neural architecture that has per-layer latent configurator p; Rib il م [o, 1]. Pz)) = ys (w.z), that maps the hidden representation from h4-1 to a confidence probability of activating a submodule R h + |R R i.e.,
(x) a, (p. (b)-(()))) for i 1,2,....4
where y is the normalizing function to produce the probability. We use a similar technique of perturbation argument [DMRB22: SLK* 18] for proving that the implicit pressure will also transfer to the latent configurator p:. which leads to the implicit regularization of latent activation decisions.
Theorem (Transfer Theorem for Latent Configurator). Let fe IRd + R& represent deep neural network consisting of I layers, i.e., fe (x) = ft o fi o ... fi(x) where fa(z) = a, (p((z)q : R + [0, 1], Pz(z) Y: (w(w)) is latent configurator, IR =1 + R [hul is sub-module function, a; is the activation function, y; the probability normalization function and w; a weight matrix. For =
I,, we have
|| (x))1 (M))
| fo(
(10)
18

---

where rz(x) == p=())-1(z) and d;(x) == qs (z)))((z)), h₂(2) indicates the sub-network from the imput x: E TRd up to the output of layer i with (o) = x and || 2| denotes the 1.2 operator norm for matrix and L2 norm for vector.
Proof. Following the above notations, we first rewrite the model function as
fw, (a) = (x))))))
where gx(2) is the remaining network of the full model fo(z) following the i-th layer. Considering small perturbation + óæ of the imput there exists corresponding perturbation w; u(x) for the weight matrix in the latent configurator p; so that
fw. dx) (m (z).
(11)
For sufficiently small dz, we can identify ri (z &x) with the linear approximation around I, i.e.
ra(x + (x 6x)) = r:(2)d ( + (n)
where Va d2 : TRd RlA. and V ri Rd [0, are the total derivatives. We have ( ôz) = +
= (z) (x)
rd(:)Vvd((()0rr)
and similarly, for rz(x) = yi (, and the perturbation wi + wi + u(6r), we can have the
linear approximation (x) vi((w u
r.(2) Vw.ri (zz((z)) This implies
fw.+u(62) (z) ( +
Therefore, substituting both sides of Equation (11) with the equations above, and we get
Vw, (()))))(rr)((()))((()()())(())))
+ (
Note the fact that V (1) m. || (V, 1 This implies
u(dx)=
(12)
IV. ( (3))
Defining u(dx) as above and taking the derivative of both sides of Equation (11) with respect to dx dæ have
fo, V (
V V (x)
(i) (x)
Applying the squared L2 operator norm || 112, at both sides the equation, we get
VV (z) = 4 (or Vn, (o
I1 (.rr
dd + (c)
V w. fo(x)]
|| w. ( (
Since (zfw( = (of)( this completes the proof.
Remarks Considering the architecture of the latent configurator in the SeqBoat layer, i.e. the activation confidence probability is calculated as
px(2) 2uul// e m//t//
then we can derive the squared L2 norms of its derivatives.
|| < 1/ (2)1
| (v)13 = (p.(=)(1 px(=)))lmo -a/pil//mw///lll////ll w/112/TT)2
19

---

Substituting the above gradient values to Equation (10), we can see that
d V=
I1 (x) [ v +
S VI
Jo)x
(x) m/ 112 (L-P((h-1
This implies that the implicit regularization pressure will tend to maximize the temperature to increase the exploration of the latent decisions, while the same time maximize the probability of modular activation or de-activation through maximizing Ilw! w l1 and px(h;-1 for exploit- ing the decisions. Therefore, there will be an implicit balancing between the exploration and exploitation of the latent decisions, and we conjecture that this kind of balancing is the cause of the dynamic sparse activation patterns we observed from the experiments.
D
Additional Experiment Results
Table 6: Per-task speed up and memory allocation reduction on the LRA benchmark.
Models
ListOps (2.048)
Text
Retrieval
Image (1,024)
Path (1.024)
Path-X (16,384)
(Input Length)
(4.096)
(4,000)
Training Step Speed Up (t)
MEGA
1.00x
1,00x
1.00x
1.00x
1.00
1.00
MEGA-chunk 2.07 x
2.42x
2.66x
1.39x
1.46x
1.36x
SeqBoat-full
1.79x
2.14x
2.44x
1.45x
1.27
1.24x
SeqBoat
2.26x
3,61
2.50x
1.33x
1.31x
1.30x
Training Step Memory Allocation (J)
MEGA
1.00
1,00
1.00
1.00
1,00
1.00
MEGA-chunk 0.34
0.29
0.31x
0.75x
0.69
0.82
SeqBoat-full
0.45x
0,23x
0.27x
0.74x
0.68
0.59
SeqBoat
0,20x
0.16x
0,23x
0.74x
0.70x
0.89
Evaluation Step Speed Up (t)
MEGA
1.00x
1,00x
1.00x
1.00x
1.00
1.00x
MEGA-chunk 2.15x
2.33x
2.57x
1.48x
1.50x
2.52x
SeqBoat-full
1.70x
2,47x
2.59
1,52x
1,30x
1.95x
SeqBoat 2.67x
3.55x
2.69
1.43x
1.38x
1.93
D.1 Per-task Efficiency Measures for LRA Benchmark
As shown in Table 6, compare the efficiency of our SeqBoat-full and SeqBoat models with the Mega and its chunk variants on all six tasks of the LRA benchmark. The evaluation speed is measured as the average inference throughput of the best performing model on the validation set. More details of the measurement are included in Appendix A.6. We can see that our SeqBoat model consistently and substantially more efficient than Mega, and also has comparable efficiency with Mega-chunk but significantly better average accuracy on the LRA benchmark.
E
Additional Analysis
How do different model design choices affect the performance? We do the ablation study of our SeqBoat-ful model on the Image and ListOps tasks of the LRA benchmark, as in Table 7, to have general understanding of the influence the design of the latent configurator and the SSM. "Always Activated" means that we don't i use SMA and always activate GAU for conducting second- -ondr self- attention. We can see that although the validation accuracy slightly improved on both Image and Listops, the training efficiency drops significantly. We also explore different parameterization of the convolution kernel for the linear SSM and apply S4D [GGGR22] to replace the original MD-FMM [MZK+2] 23] kernel. The S4D variant of our model has much better training efficiency but the quality
20

---

Table Ablation study of the SeqBoat-full model on Image and ListOps of the LRA benchmark.
We report the validation accuracy on the Image and ListOps tasks, together with the training speed up for each of the tasks compared with MEGA.
Model
Image
Speed (1)
ListOps
Speed (t)
SeqBoat-full
90.16
1.45x
60.95
1.79x
Always Activated
90.70
1.25x
62.10
1.11x
Gumbel Softmax
89.94
1.29x
59.75
1.76x
SSM with S4D
82.84
1.88x
59.25
2,45x
significantly worse than SeqBoat-full with the default hyperparameter setting. It is possible that a well-tuned S4D kernel can have better performance than the MD-EMA for the use case of SMA but leave a more comprehensive study as in future work. We also explore different differentiable approximation of latent decision making by replacing Tempered Softmax with Gumbel Softmax [JGP17]. However, we don't observe the performance improvements over Tempered Softmax for our setting. More details of the implementation of Gumbel Softmax based latent configurator is shown in Appendix A.4.
Table Ablation study of the SeqBoat model on Image, ListOps and Pathfinder of the LRA bench- mark. We report the validation accuracy on each of the tasks, together with the training speed up compared with MEGA.
Model
Image
Speed (1)
ListOps
Speed (f)
Path.
Speed (t)
SeqBoat
90.36
1.33x
59.05
2.26x
96.34
1.31
SeqBoat-FFN FFN
89.60
1.13x
56.45
2.05
96.62
0.97x
SeqBoat- EiBBoat-GGA,, FFN)
89.74
1.20x
59.10
1.97
96.10
1.06x
SeqBoat-Sigmovd
87.34
1.77
58.20
2.52x
91.35
1.67x
SeqBoat-MoE
85.60
1.66x
58.10
2.19x
91.41
1.80x
SeqBoat-Tocaa
90.I2
T.46x
59.00
1.78x
96.08
T.54%
Seq SeqBoott-chnnk
86.46
2.17x
58.35
2.79x
93.91
2.97x
Seq SeqBooat-mmm32
78.16
1.73x
58.70
2.14x
92.97
1.86x
Seq BoqBolatlocal32
75.58
1.67
53.90
2.19x
82.05
1.80x
We further present additional ablation studies on our SeqBoat model to justify the effectiveness of both the proposed working memory mechanism and the design choices of our architecture. We explain the meaning of the postfixes in the table as following:
"-FFN": We add additional Feed-Forwar layer (with the same architecture MEGA)
after each SeqBoat layer.
"-(GAU, FFN]": We add FFN module in parallel with the GAU module to have M with two modules in the function space F = [GAU, FFN] } for our SMA mechanism to control the activation.
"Sigmoid": We use Sigmoid instead of Tempered Softmax for latent decision probability caleulation.
"-MoE": Instead of using our Tempered Softmax for module activation decisions, adapt the X-MoE [JJNH91] routing mechanism to our use case by only considering two modules, GAU module and module that always outputs zero. We do not apply load balancing loss in this case because there is only one module that needs to do actual computation.
"-local": We don't SMA and always activate the GAU module with local attention.
"-chunk": We don't use SMA and always activate the GAU module to process the equal- sized chunks of the input sequence. The chunk size set to be the same as the memory size of SeqBoat for fair comparison.
"-mem32": We restrict the memory size of SeqBoat to 32.
"=Jocal32": We restrict the sliding window size of SegBoat-local 32.
We first explore the application of SMA with multiple heterogeneous modules and identify the cur- rent engineering challenges. We can see that SeqBoat- (GAU, FFN] obtains better trade-offs than
21

---

SeqBoat-FFN (which simply appends an extra FFN layer after each SeqBoat layer) on both the Image and the ListOps tasks. However, Seq FFN) still falls behind SeqBoat (which does not include FFN anywhere at all). This is because the current implementation of SMA is 1/O bounded due copying large tensors with the scatter operator, and for a fair comparison with MEGA, we do not include any fused kernel to optimize the memory bandwidth utilization. We acknowledge that it needs more engineering efforts to scale up SMA with multiple modules, and advocate future works for the empirical impacts of SMA the MoE community. We also include additional ablation results with X-MoE routing as an alternative design choice for the latent configu- rator. We that SeqBoat-MoE performs much worse than our SeqBoat model which uses the Tempered Softmax for decision confidence calculation.
Comparing SeyBoat-Sigmoid with SeqBoat, we can see that while theoretically using Softmax is equivalent to using Sigmoid under the binary decision scenario, using Softmax for decision confi- dence calculation gives much better results under our aaper-parameter settings. Comparing SeqBoat with SeqBoat-loca and SeqBoat-chunk, we can see that our model obtains better accuracy than the models without SMA. If we constrain the memory size to smaller value (see Seq v.s. SeqBoat-local32), this performance gap becomes more substantial (more than 10 percent absolute difference on Pathfinder). These evidences demonstrate the effectiveness of our SMA mechanism for capturing long-term dependencies. Comparing SeqBoat-FFl with SeqBoat, we can see that adding extra FFN layers does not provide any empirical performance gains for SeqBoat, and this motivates us to not include FFNs anywhere in our architecture.
F Related Work
Dynamic Routing [RSVG20b] aims alleviate the second-order complexity of self-attention by dynamically mapping each of the token to its nearest-neighhor clusters, and only do self-attention inside the clusters. Dynamic routing can be understood as the following form when applied for sequence modeling,
max Pf(x) = II (, x x2, 0) 0,x%a
t=1
where f is the Transformer architecture with self-attention layers, and it can be seen as a special case of our formulation described in Equation (1). Even though dynamic routing learns both the parameters and the corresponding context for each of the token prediction, it has a fixed architecture that is non-adaptive for both training and the inference.
Mixture of Experts (MoE) [JNHH1 FZS22; LLX+ 20; ZLL+ is designed to learn to select different parameters with the same architecture for each of the input elements. Specifically, it has the following form when applied to sequence modeling,
max (x(x) II P)(o,(2e1s0) 9c
where the model parameter 0, is selected from a pool of candidates at each time step i based on input content. We can see that it is special case of our formulation as in Section 3.1 because it has a fixed context Xxt for each of the token prediction and the model architecture non-adaptive.
The motivation behind our Sparse Modular Activation (SMA) mechanism is to enable neural networks to contextually skip modules of any architectures for efficient sequence modeling, while MoE aims to efficiently scale up the models with more pa- rameters (usually by adding homogeneous modules). This difference of motivation results a fundamental mechanism difference:
Comparison between SMA and MoE:
S.A leverages a latent configurator to decide if each module needs to be activated for each sequence element, while MoE designed to choose predefined number of modules from a large group of modules.
This difference on the core design of the mechanisms further leads the different properties of the mechanisms:
22

---

SMA supports dynamic model architecture that can learn to drop a sub-module entirely (for better efficiency) based on the task it is trained on, while MoE only selects the parame- ters of the modules for the same architecture. As shown in Figure 3, when the SeqBoat-full model is trained on the Text task of L.R,, the first three layers learn to have zero activation time of the GAU module. This means that these layers degenerate to pure SSM layers from the initial SSM+GAU architecture after adapting to the target task.
S.M is guaranteed to have a full coverage of the combinatorial search space ranging from zero activation to full activation of modules, while MoE can only cover subset of the space by choosing a fixed number of modules to activate.
To the best of our knowledge, SMA the first mechanism that successfully enables a neural network to obtain practical efficiency and complexity gains from sparsely activating self-attention- like mod- ule, while none of the previous works on MoE ever achieved this.
Comparison with Flash Attention [2F:+ Flash Attention (FA) performs exact self-attention efficiently through utilizing hardware-aware techniques. FA and SMA focus on improving model efficiency on orthogonal levels. SMA is a module-level activation strategy that is perpendicular to how the attention module is actually doing the computation. In fact, we can also apply FA to the GAU module in our SeqBoat architecture for more efficient self-attention computation, but we decide not to use any custom fused kernels in our implementation for a fair comparison with MEGA and its variant.
Comparison with Efficient Attention Strategies [KKL200; CLD+20; 20 MKW+21J: Our SMA mechanism enables a dynamic architecture that can learn to drop a sub-module entirely for better efficiency when it is adapted to the target task, and we empirically observe such behavior for different tasks in LRA. This capability is not possible for the previous efficient attention strategies because their architectures are static and thus an efficient attention layer is always applied before the Feed-Forwwrd layer. Also, SMA is an orthogonal mechanism to efficient attention approaches because it does not modify the attention itself but the input to the attention modules, and thus can be stacked with other efficient attention strategies.
Comparison with GShard [L.XX + 20]: GShard divides tokens into groups of the same size, while our SMA allows different modules to have different group size. This means SMA can support an adaptive neural architecture whose sub-modules can be completely dropped if no tokens are selected for that module. Also, while GShard is a routing mechanism, SMA is an activation mechanism that can still function when there is only one module under consideration.
Comparison with Perceiver [JGB 21]] Perceiver utilizes a pre-defined number of latent rep- resentations to repeatedly attend to the input array, which can be understood conducting soft selection from the input array to a fixed subset. Our SMA operates on a different level to directly conduct hard selection from the input form dynamic subset. Plus, our SMA can be applied to causal language modeling, while it is unclear how Perceiver can be adapted to this important use case in the era of large language model.
Connection with Recurrent Independent Mechanisms [GLH* 19]: Generally, Recurrent Inde- pendent Mechanisms (RIMs) and our work both touch on how to learn sparsely interacting modules. However, our works focus more on the efficiency gains from the modular sparsity, while RIMs em- phasize on the generalization benefits. Plus, while the GAU module in our work can be regarded as modeling the sparse interaction between the state representations of SSM at different time steps, our work does not explicitly try to use attention bottlenecks for modeling the interaction between different SSM modules. Instead, we apply rather simple strategy of using línear aggregation to combine the representations from both the SSM and the GAU modules. It would be interesting to see if the attention bottleneck can help SMA to better scale up with multiple heterogeneous modules, and we leave it as a future work worth exploring.
23

---

G Limitations, Future Works & Broader Impact
Our study has some limitations that should be acknowledged. First, while our Sparse Modular Acti- vation (SMA) mechanism is generally applicable for activating any sub-modules in neural networks, we have narrowed the scope of our paper focus solely on developing more efficient neural ar- chitecture with SSM and GAU for long sequence modeling, due to limited resources. Additionally, the current implementation of SMA is I/O bounded due to copying large tensors with the scalter operator. For fair comparison with the previous models, we did not include any fused kernel to optimize memory bandwidth utilization. This limitation highlights the need for more engineering efforts to scale up SMA with multiple modules.
There are several promising directions for future research building on our work. One interesting avenue is to explore how SMA can boost the downstream performance of a pre-trained large model, given that pre-training a useful large language model from scratch can be prohibitively expensive. Additionally, since SMA can produce discrete and dynamic module activation patterns for each of the data sample, our work opens up new possibilities for interpretability in neural networks. Future research could investigate the correlations between the activation of specific modules and input tokens or model predictions, to gain a better understanding of the properties of tasks the behavior of models. Furthermore, our mechanism may provide extra controllability over model predictions, which can be achieved by manually modifying the module activation during inference time. Finally, promising future direction is investigate how to scale up SMA with large number of modules, potentially leading to empirical impacts of SMA on the research field of Mixture-of-Exxpe of-Experts, Some initial ideas include setting an upper bound for the activation time of each module, or limiting the number modules that can be activated at the same time.
The SMA and the SeqBoat architecture can bring about considerable societal impact by enhancing the efficiency and interpretability of sequence models. These advancements can result in better- performing automated systems, like open-domain chatbot and automatic speech recognition, poten- tially improving communication, fostering inclusivity, and making such services more accessible to various societal sectors. The enhanced interpretability of models can generate trust in Al systems, especially in sensitive sectors like healthcare and finance. Transparent AI models can assist profes- sionals in these fields to make better-informed decisions, potentially leading to improved societal outcomes.
However, these improvements come with a notable limitation: our research was unable to conduct experiments on models with a large number of parameters due to computational resource constraints. As such, while the proposed mechanisms show promising results in the tested scenarios, their scal- ability and applicability to larger, more complex models remain uncertain. This limitation may restrict the broader applicability of our findings, especially for large-scale problems or when fine- grained control over the model's behavior needed. Furthermore, ethical considerations such as data privacy need to be seriously considered as work progresses. As the efficiency of these mod- els increases, there's an increased potential for processing large amounts of personal data, which necessitates robust privacy safeguards. Future work can also focus on addressing these limitations and ethical concerns, to ensure the advantages of these models can be responsibly realized.
24


Reviews:
[{"review_id": "K5TYfQH6gg", "reviewer": "Reviewer_KtV6", "summary": "The authors employ a hybrid model combining linear state space models and attention modules, while incorporating sparse attention within the attention modules to reduce memory usage and improve speed. The proposed method, Sailboat, outperforms previous approaches in terms of quality, speed, and memory efficiency. The sparse activation module in Sailboat utilizes temperature softmax to activate modules and aggregates output from each activated module weighted by softmax values. Additionally, the authors propose a variant called Sailboat-Mem, which performs window-sized attention on tokens assigned to each module, resulting in speed improvements and memory optimization, with a slight reduction in quality.", "questions": "The equation introduced at line 120 lacks defined variables beforehand, making it challenging to interpret. Could you provide more clarity on this equation by defining all the variables involved?\n\nThe improvement of linear state-space models over transformers appears to be significant. Can you explain why this is the case? Are the number of parameters comparable between the two models?\n\nAre the SSM kernels in the hybrid models identical? Maintaining consistency in the SSM kernels would offer clearer insights into the improvements achieved by sparse activation. Could you provide further information on this aspect?\n\nCould you provide more details about the schedule used for the temperature softmax function in the work? \n\nCan MEGA be considered as having only one attention module where all tokens are processed by it? In contrast, Sailboat divides tokens across a set of attention modules and performs attention only on those assigned to each module. As Sailboat's performance still lags behind MEGA, is there still potential for further improvements in routing? Can you provide the associated FLOPs (floating-point operations) for Sailboat and compare them with MEGA? Additionally, could you justify the memory reduction achieved by Sailboat? Is the reduction because the memory bottleneck is bs x num_tokens x num_tokens tensor?\n\nIn Section 3.2, binary decisions in SMA can be accomplished by computing p_i=0 using only w_0 and b_0, and p_i=1 can be deduced as 1-p_i=0. Could you please provide a justification for your implementation, particularly regarding the usage of w_0, b_0, w_1, and b_1, as it seems to deviate from conventional approaches?\n\nThere is a typo in Line 816 with the repeated phrase \"is optimized.\" Could you please correct this error?", "limitations": "There are no explicit limitations beyond those discussed in the paper.", "rating": 5, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper demonstrates a clever utilization of a mixture of experts (MoEs) and routing techniques, resulting in notable speed improvements. \n\nThe proposed method achieves results comparable to state-of-the-art approaches while achieving increased speedups and reduced memory requirements. \n\nThe simplicity of the proposed method is a significant strength, as it only requires tuning a single temperature parameter in the softmax activation.  \n\nThe ablation studies conducted in the research paper providing comparisons with other variants, such as gumbel-softmax and complete activation of modules, highlight the method's trade-off in terms of performance, speed, and memory utilization.\n", "weaknesses": "One notable weakness is the absence of a baseline comparison with Flash attention which performs exact attention but provide greater speedups while utilizing hardware aware techniques [https://arxiv.org/pdf/2205.14135.pdf]. \n\nDespite achieving speedups and reduced memory compared to MEGA, the proposed method does not demonstrate an improvement in performance over MEGA. \n\nThere are other variants of routing, such as Top-k and REINFORCE, which could potentially be compared with the proposed method. The absence of such comparisons limits the scope and comprehensiveness of the work. Please refer Top-k https://arxiv.org/abs/2101.03961,  REINFORCE https://arxiv.org/abs/2202.01169", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive review of our work and the constructive feedback on our manuscript. In the following, we address the remaining concerns to hopefully motivate a clear acceptance score.\n\n\n**Comparison with Flash Attention(FA):** FA and SMA focus on improving model efficiency on orthogonal levels. SMA is a module-level activation strategy that is perpendicular to how the attention module is actually doing the computation. In fact, we can also apply FA to the GAU module in our Sailboat architecture for more efficient self-attention computation, but we decide **not** to use any custom fused kernels in our implementation for a fair comparison with MEGA and its variant.\n\n**Lack performance improvement over MEGA:** The MEGA model has quadratic inference complexity but our Sailboat only has sub-quadratic complexity. It is unfair to solely compare the performance, while ignoring the better scalability provided by Sailboat. The linear complexity variant of our model, Sailboat-mem, does provide significantly better average accuracy over MEGA-chunk on LRA with better speedup and memory consumption. In L598-L607 of **Appendix D**, we further demonstrate that Sailboat-mem can provide much better quality-speed trade-off than MEGA-chunk on the Image and Pathfinder tasks of LRA.\n\n\n**Comparison with other routing variants:**  We want to first point out that our method is **not** a routing mechanism, but instead an activation mechanism for neural modules. For more details on the comparisons with MoEs, please refer to **Global Response**. Moreover, although SMA and MoE are not directly comparable, we add an additional ablation study in **Global Response** by adapting the X-MoE routing mechanism to our use case, and show that it actually provides much worse performance than our SMA mechanism.\n\n---\n\n**Question 1:** *The equation introduced at line 120 lacks defined variables beforehand, making it challenging to interpret. Could you provide more clarity on this equation by defining all the variables involved?*\n\n**Answer:**   $L'$ represents the function space of a layer equipped with SMA and $M$ number of predefined functions. According to Line 117, the function space $L'$ is a linear combination of the activated functions at layer $l$, thus $\\zeta_i$ indicate the scalar of the linear combination. As said in Line 113-120, $a_t^i$ is the decision value for each sub-module, and $c_t^i$ is the confidence value for the decision, and $I$ means the set of the indices of the activated functions.\n\n**Question 2:** *The improvement of linear state-space models over transformers appears to be significant. Can you explain why this is the case? Are the number of parameters comparable between the two models?*\n\n**Answer:**  According to S4 paper, the number of parameters are tied between two models. While the exact reasons behind the performance improvement of linear state-space models over transformers are out of the scope of our paper, we here provide some related papers [1,2] for the reviewer to futher read on. Note that our work is perpendicular to the previous works on SSM, and as said in Line 165-167, we use SSM simply because it can provide an efficient computation of the recurrent states, plus it is interesting to investigate how much extra attention is needed beyond the SSM module.\n\n**Question 3:** *Are the SSM kernels in the hybrid models identical? Maintaining consistency in the SSM kernels would offer clearer insights into the improvements achieved by sparse activation. Could you provide further information on this aspect?*\n\n**Answer:**  Yes, we use the same MH-EMA kernel for all the layers as indicated in Line 164-165 of our paper.\n\n**Question 4:** *Could you provide more details about the schedule used for the temperature softmax function in the work?*\n\n**Answer:** As said in Line 170-172 of our paper, we didn’t apply any scheduling for the temperature parameter and set it as a learnable parameter with a specific initialization.\n\n**Question 5:** *Can MEGA be considered as having only one attention module where all tokens are processed by it? ... ... Is the reduction because the memory bottleneck is bs x num_tokens x num_tokens tensor?*\n\n**Answer:** Yes, MEGA can be considered as that. Yes, there are further potential for the improvements, since our design of the latent configurator is rather simple. Since the major speed bottleneck of SMA lies on copying a large matrix (which is not included in the FLOPs calculation) through the *scatter* operator, we follow the previous works to measure the speedup based on the actual wall time for a stricter and fairer comparison. The memory reduction achieved by Sailboat is because the GAUs are processing a shorter sequence, and for some layers of the Sailboat model, the GAUs are skipped entirely because they deactivated by the latent configurator all the time. This kind of skipping results in both a smaller bs x num_tokens x num_tokens tensor and fewer SiLU activation operations inside the GAUs. \n\n**Question 6:** *In Section 3.2, binary decisions in SMA can be accomplished by ... ... as it seems to deviate from conventional approaches?*\n\n**Answer:**  While theoretically using Softmax is equivalent to using Sigmoid under the binary decision case, we empirically find that using Softmax with both $w_0$ and $w_1$ gives much better results under our hyper-parameter settings. Please refer to the performance of the Sailboat-mem-sigmoid model under the Additional Ablation Study section in **Global Response** for more details.\n\n**Question 7:** *There is a typo in Line 816 ... ... this error?*\n\n**Answer:** Thanks for pointing this out.  We have fixed this typo by deleting the additional \"is optimized\".\n\n---\n\n[1] Resurrecting recurrent neural networks for long sequences. Orvieto, Antonio, et al. arXiv 2023.\n\n[2] Simple hardware-efficient long convolutions for sequence modeling. Fu, Daniel Y., et al. arXiv 2023.\n"}, {"review_id": "C6VCUYixbs", "reviewer": "Reviewer_kQJ9", "summary": "This method introduces a novel architecture for long text modeling, which builds upon traditional linear state space models (SSMs). Since SSMs have shown inferior performance, combining SSMs with self-attention has become a popular approach. In this paper, several efficiency-related questions are considered in such settings, such as the amount of additional attention required for SSMs on a per-task basis and whether neural networks can dynamically activate their attention modules to achieve improved quality and efficiency trade-offs. These questions pose an interesting inquiry.\n\nTo address these questions, the authors propose a new architecture module that sparsely activates a Gated Attention Unit based on the state representations learned from an SSM. Through experiments on diverse tasks, the proposed approach demonstrates competitive results when compared to strong baselines.", "questions": "Could you please provide more information on how the training speedup is computed?\n\nIn line 119, what does \"c\" represent? Additionally, in Figure 1, what does \"h\" denote?\n\nIt seems that Figure 2 is listed before Figure 1. The correct order should be Figure 1 followed by Figure 2.\n\nWhat is the motivation behind the intricate setting of the latent configurator? Could you elaborate on how each function contributes to the final performance?\n\nHow does the method perform when applied to widely-used pre-trained settings?", "limitations": "n/a", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The research question presented is intriguing. The hybrid combination of SSM and attention has gained popularity. However, the extent of additional attention required for SSMs on a per-task basis remains unclear. By reducing the reliance on attention, significant inference speedup can be achieved.\n\nThe selected baseline for comparison appears to be adequate. The authors conduct comprehensive experiments to demonstrate the improvements brought by their proposed approach. Additionally, the ablation study effectively showcases the effectiveness of specific components of the method.", "weaknesses": "The method section is difficult to comprehend. The notation should be clearly defined before its usage. Furthermore, the complex settings depicted in Figure 3 lack sufficient discussion on the motivation behind each function. It is unclear how these functions contribute to the final performance and the necessity of including them. If possible, it would be helpful to have a simplified implementation that still achieves competitive results.\n\nMore evidence is needed to support the performance of the proposed method. While the method shows similar effectiveness compared to other approaches, the claim of higher training speedup requires further explanation. The evaluation of speedup is not adequately discussed, and the observed speedup appears to be marginal when compared to strong baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the constructive feedback on our manuscript. In the following, we address the remaining concerns to hopefully encourage a positive evaluation.\n\n\n**Presentation clarity of the method section:** Please refer to our answers to Question 2 and 3. We will provide more details in the next version and improve the presentation clarity.\n\n**Motivation of the architecture design:** The design priniciple behind our Sailboat model is to allow a faster sytem with weaker capability, (e.g. an SSM), to learn to activate a slower system with stronger capability (e.g. a GAU) on demand. We explain the design motivations of our latent configurator as following.\n\nFor a sequence of length $n$, the design goal of the latent configurator is to generate the binary decision vector and the confidence probability vector, $a^l\\in \\mathbb{R}^n$ and $c^l \\in \\mathbb{R}^n$ respectively, for the GAU module so that we can sparsely activate the module with the decision vector while allowing the gradients to be backpropaged to the latent configurator through the confidence vector. The first Linear layer is applied to project the contextualized representation $\\mathbf{H}^l\\in \\mathbb{R}^{n \\times d_m}$ to a matrix whose last dimension is two.  This matrix is then normalized by a Tempered Softmax function on its last dimension to produce a probability matrix in $\\mathbb{R}^{n \\times 2}$ whose first column is the probability of not activating the module and the second column is the probability of activating. The decision vector $a^l$ and the confidence vector $c^l$ are then calculated by applying the argmax and the max operator respectively on the last dimension. \n\nFrom the above descriptions, we can see that **all the functions are necessary** to produce the final decision and the confidence vectors, so that the latent configurator can work as we expected. In this sense, it is not possible to ablate some functions in the configurator to see how they contribute to the final performance. However, we do try different alternative designs such as using Gumbel Softmax/X-MoE/Tempered Sigmoid instead of Tempered Softmax for probability matrix calculation, which are further explained in L300-304 and the additional ablation study in **Global Response**.\n\n\n**More evidence of model performance:** Please refer to the Updated Results section in **Global Response**. Generally, with updated results, our Sailboat-mem is now the **new state-of-the-art** among models with linear inference complexity, outperforming the pure SSM-based model, S5, with remarkably faster training speed and much less GPU memory consumption. We also add more detailed ablation studies to explain the speedup and performance.\n\n\n**Explanation of training speedup measurement:**  Since the dynamic module-level sparsity of our Sailboat model will affect the training speed throughout the training process, we first measure the training time as the average per step wall time across the full training stage. The training speed is then calculated as the inverse of the training time. To ensure fair comparisons, the relative training speedup between different models are calculated based on the training time measured on the same hardware using the same batch size settings. All the experiments are conducted on a mixed cluster with 8 NVIDIA V100 32GB GPUs and 2 NVIDIA A5000 24GB GPUs.\n\n---\n\n**Question 1:** *Could you please provide more information on how the training speedup is computed?*\n\n**Answer:** Please refer to the **Explanation of training speedup measurement** paragraph above.\n\n\n**Question 2:** *In line 119, what does \"c\" represent? Additionally, in Figure 1, what does \"h\" denote?*\n\n**Answer:**  As explained in line 118, $\\mathbf{c}_t$ is the confidence probability vector which stores the probability of decisions for all the modules at the time step $t$, so the symbol $\\mathbf{c}$ represents a confidence matrix storing the probabilities for all the modules at all the time steps. In Figure 1, $\\mathbf{h}\\in \\mathbb{R}^{d_m}$ is one of the vectors of the sequence representation matrix $\\mathbf{H}\\in \\mathbb{R}^{n \\times d_m}$.\n\n**Question 3:** *It seems that Figure 2 is listed before Figure 1. The correct order should be Figure 1 followed by Figure 2.*\n\n**Answer:** Thanks for pointing this out. We have fixed this by re-ordering the Figures.\n\n\n**Question 4:** *What is the motivation behind the intricate setting of the latent configurator? Could you elaborate on how each function contributes to the final performance?*\n\n**Answer:** Please refer to the **Motivation of the architecture design** paragraph above.\n\n**Question 5:** *How does the method perform when applied to widely-used pre-trained settings?*\n\n**Answer:**  While SMA is generally applicable for activating any sub-modules in neural networks, we narrow the scope of our paper (due to limited resources) to only focus on developing a more efficient neural architecture with SMA for long sequence modeling. It is interesting to see how SMA can boost the downstream performance of a pre-trained model, but we think it is out of the scope of this paper and leave it as a future work that is worth exploring. For the setting of Language Model (LM) pre-training, we do test our model on the competitive LM task, enwiki8, and obtain better results than the baseline MEGA model.\n"}, {"review_id": "81N74UiA91", "reviewer": "Reviewer_8NeG", "summary": "This paper introduces a framework for representation learning, which involves multiple modules that can be applied dynamically on different inputs. The authors implement this framework using a combination of linear state space models (SSMs) and a gated attention unit (GAU), also combining ideas from adaptive computation, and mixture of experts. Combined, the network, dubbed Sailboat, obtains a better speed-accuracy tradeoff compared to multiple strong baselines on a range of tasks. \n\n==============================\n\nPost rebuttal:\n\n\nThank you for your response and for the clarification.\n\nI now have a better understanding of the main contributions, and have raised my score to 6. It would be great to include this discussion in the next version of the paper. \n\nAs to the specific responses, please note that results important enough to be included in the introduction and/or the conclusion should be included in the main text, not the appendix.", "questions": "+ I had a hard time understanding the main method:\n* * The introduction mentions (#57-62) a method for efficient parallel implementation of the method using the scatter operation, but this is never discussed later. \n- - The introductory paragraph of section 2 seems to describe something broader than what is currently below it in this section. E.g., section 2 does not touch on Sparse Modular Activation at all.\n- - The F function in 3.1 is not entirely clear to me. First, if I understand correctly, the mapping assigns each word in the vocabulary a boolean value, which indicates whether it is used or not? Does this mean that the selection process works at the type level (i.e., each word is always selected or ignored, rather than this decision being dependent on the context)? Further, how is the vocabulary defined for non-textual tasks? Also, how does this function relate to the next part (e.g., $a_t$?)\n- - In 3.2, what is $M$?\n– I didn’t fully understand the Sailboat-Mem approach. Do the authors simply attend to the nearest $w$ tokens, and in this sense the idea is similar to standard hybrid approaches, or do they attend to these tokens and then apply SMA? Also, the comment in #207 (“... maintaining the ability of attending to the key tokens …”) wasn’t clear to me as well.\n\n- The results in table 2 (LRA) show that the proposed method is dramatically faster and consumes far less memory than all baselines (including MEGA). However, in tables 3 & 4, the speed and memory values of the two approaches are comparable. Can you please explain the source of these differences? Is this because the context length is larger for these tasks? The speech experiments discuss a large context length of 16K, so that doesn’t seem like the reason.\n\n- why do different models have different numbers of layers (#275)?\n\n- In what sense is the proposed method able to improve interpretability? (conclusion, limitations section)\n\nTypos and such:\n- #52: sparsely map -> sparsely *maps*\n- #86: is optimized is optimized\n- #144: are instantiate -> are *instantiated*\n- #236: LRA consists of five tasks: -> *six* tasks", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The technical and empirical contributions of this work more than justify acceptance to NeurIPS. \nThese include a very interesting and clever approach, combining multiple lines of research, and addressing many technical challenges, as well as an extensive and diverse set of experiments, showing the superiority of the proposed approach over previous approaches.", "weaknesses": "At the same time, the paper gives the impression of having been hastily written. There are some clarity issues regarding the proposed method (see below); the introduction mentions details that are never discussed in the paper, and there are some inconsistencies between the introductory paragraph of section 2 and the rest of the section. Combined, these raise some concern that other problems might have been overlooked.\n\nMoreover, the human working memory is only loosely connected to the Sailboat-Mem idea. In particular, $w$ is within hundreds or thousands, which is far larger than 5-7, which is the human working memory, and thus completely unrelated. \n\nFinally, there are unclear gaps between the superb results on LRA compared to the results in tables 3 & 4, which are comparable and not substantially different from the baselines (e.g., MEGA). \n\nSee more details below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive review of our work and the constructive feedback on our manuscript. In the following, we address the remaining concerns to hopefully motivate a clear acceptance score.\n\n**Presentation clarity:** Please refer to our Answers to the Q1, Q2, Q3, Q4.\n\n\n**Loose connection between human working memory and the Sailboat-Mem idea:** Thanks for pointing this out. We will weaken the claim, “mimic the human working memory”, in the abstract, but we still want to emphasize the effectiveness of our working memory mechanism. In Figure 5 of the **Appendix D**, we do demonstrate that our working memory size can be as small as **16** but still maintain a reasonably good performance (>90%) on the Pathfinder task of LRA. In contrast, the traditional chunking based method simply fails the tasks with random guessing when the chunk size is less than 128. We use a larger working memory size for the final Sailboat-mem model simply because we want to achieve better performance under the settings of the LRA benchmark. \n\n\n**Unclear performance gaps between LRA and Speech Recognition/Language Modeling:** This is because we follow the notations of models in the original MEGA paper for the Table 3 & 4 of our paper, and the so-called MEGA models listed in the tables are actually MEGA-chunk variants. We have fixed this by explicitly denoting the MEGA models as MEGA-chunk for clearer understanding. Also, in the updated results table for Speech Recognition in **Global Response**, we do have much better performance and speed up than MEGA-chunk.\n\n---\n\n**Q1:** *The introduction mentions (#57-62) a method for efficient parallel implementation of the method using the scatter operation, but this is never discussed later.*\n\n**A1:**  As in **Appendix A.1**, we do discuss the detailed parallel implementation using the scatter operator in the Pytorch-like code snippets, and such snippets are well documented with comments for how they work. The existence of such code snippets is also explicitly said as in Line 131-132 of our paper. \n\n**Q2:** *The introductory paragraph of section 2 seems to describe something broader than what is currently below it in this section. E.g., section 2 does not touch on Sparse Modular Activation at all.*\n\n**A2:** Section 2 provides the background and motivates us to propose a general formulation of Time-Variant Sequence Modeling in Section 3.1 that supports a dynamic neural architecture. Sparse Modular Activation (SMA) is then built upon the two assumptions proposed in Line 104-111 of Section 3.1, and we further show that, in Line 121-123, SMA can recover the original function space $L$ mentioned in Line 109 to justify the design of SMA.\n\n**Q3:** *The F function in 3.1 is not entirely clear to me... ... Also, how does this function relate to the next part (e.g., $a_t$?)*\n\n**A3:** Thanks for pointing this out. We acknowledge that $[0,1]^V$ in Line 105 should be corrected as $[0,1]^{n \\times V}$.  $[0,1]^{n \\times V}$ means a space of $n \\times V$ dimensional matrices whose values are between 0 and 1, thus the $\\mathcal{F}$ function means the whole chain-structured model that takes a sequence of tokens as input and outputs the probability distributions over the target vocabulary. This function introduces the intermediate function space $L$, which is then used for showing that our SMA can recover the original function space (as shown in Line 121).\n\n**Q4:** *In 3.2, what is $M$? – I didn’t fully understand the Sailboat-Mem approach. Do the authors simply attend to the nearest $w$ tokens, ... ... Also, the comment in #207 (“...maintaining the ability of attending to the key tokens …”) wasn’t clear to me as well.*\n\n**A4:** As show in Line 109, $M$ is the number of the pre-defined functions. As said in Line 201, we apply local attention on the **compressed sequence** which only contains the activated inputs for the GAU module. In this sense, we first apply SMA and then we attend to the nearest $w$ tokens in the compressed sequence. For Line 207, the unbounded attention span between the query and the key token is because two consecutive tokens in the compressed sequence can be from the positions that are far away from each other in the original sequence.\n\n**Q5:** *The results in table 2 (LRA) show that the proposed method is dramatically faster ... ... so that doesn’t seem like the reason.*\n\n**A5:**  Please refer to the **Unclear performance gaps between LRA and Speech Recognition/Language Modeling** paragraph above.\n\n**Q6:** *why do different models have different numbers of layers (#275)?*\n\n**A6:** This is because prior works (e.g. MEGA, S4) are using different numbers of layers for different tasks in LRA, we follow the setting of prior works to have a fair comparison.\n\n**Q7:** *In what sense is the proposed method able to improve interpretability? (conclusion, limitations section)*\n\n**A7:**  As shown in Figure 4 in **Appendix D**, our method can produce discrete and dynamic module activation patterns for each of the data sample. In this sense, our mechanism allows people to find the correlations between the activation of a specific module and the input tokens/model predictions, so that we can have better understanding of the property of the tasks or the behavior of the models. In fact, as shown in Figure 3, SMA allows us to answer the questions related to task properties such as \"How much attention is needed for different sequence modeling tasks?\" Also, since our latent decisions are discrete, our mechanism may also provide extra controllability over the model predictions. This can be achieved by doing manual modifications of the module activation during the inference time. Generally, our work opens new possibilities for interpretability which are worth exploring in the future work.\n\n**Typos and such:** Thanks for pointing out the typos. Fixed them.\n"}, {"review_id": "EfUX0yh3vu", "reviewer": "Reviewer_oaMm", "summary": "The paper introduces Sailboat which builds upon MEGA. MEGA uses a combination of an Exponential Moving Average (EMA) block (which can be interpreted as a specific parameterization of the kernel from SSM models) and Gated Attention Units (GAU). MEGA explores both full attention and chunked attention. Sailboat instead takes a modular/MoE approach. Each layer in Sailboat linearly combines the outputs of multiple EMA + GAU modules. Moreover, the GAU layer in each module dynamically takes as input only a subset of the original input; thus reducing per module computation cost.  Sailboat-mem is also proposed as a variant of Sailboat that uses local attention with a sliding window in each module. Sailbot-mem performs better than MEGA-chunked in LRA while being more efficient.", "questions": "1. How are c$_i$ and a$_i$ computed for modules? \n2. I did not understand the hyperparameter setup for the modules. For example, how many modules are used per layer (what is $M$?)? Wouldn't adding multiple modules per layer increase the parameters (according to your notation it seems your modules are layer specific as opposed to having a common set of modules for every layer to select/combine or is that an incorrect interpretation)? How do you control for parameter counts then?\n3. Can be good to contrast/compare (even if not empirically), in the paper, the difference between GShard [4] and this approach. As far as I understand, GShard also creates local token groups for the parallel processing of modules. \n4. MEGA uses Feed forward after the GAU block. Is that removed in Sailboat? Any specific reason? \n\n\nMissing references:\n\n[5] uses Gumbel sigmoid as a selector mechanism to dynamically discard input before attention although not in a modular context and without exact discrete selection during training. \n\n\nMinor:\n* Are ζ$_i$ in the eqn. near L120 learnable parameters? \n* Is \\forall ∀ in the eqn. near L120 a typo?\n* Is there any reason why the function/module weights (ζ$_i$c$_i$) are not normalized? \n* Another extra comparison could be to use something like perciever [3] for creating some constant number of latent states per module from the SSM outputs. \n* “second-order self-attention” -- second-order in which sense?\n* Should be \"instantiated\" in L144\n* I assume $a^l ∈ \\mathbb{R}^n$ should be more specifically $a^l \\in \\\\{0,1\\\\}^n$ in L168?\n* Probably better to have an extended related works section in the appendix with more discussions. The current related work is a bit sparse. \n* Note that there are pure SSM-based approaches (they are very recent works so I am not factoring this point into my decision) that show competitive performances too like BigS [1] and Hyena [2].  The cited paper that shows SSMs perform poorly on machine translation doesn't seem to use GLUs or other features that were found to be crucial for better NLP performance from SSMs. \n* It could be good to clarify/discuss/contrast better on your paper the difference between your approach and \"sparse attention\" approaches \n\n\n------\n\n[1] Pretraining Without Attention, Wang et al. ArXiv 2023: https://arxiv.org/abs/2212.10544\n\n[2] Hyena Hierarchy: Towards Larger Convolutional Language Models, Poli et al. ArXiv 2023: https://arxiv.org/abs/2302.10866\n\n[3] Perceiver: General perception with iterative attention. Jaegle et al. ICML 2021\n\n[4] GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, Lepikhin et al. ICLR 2021\n\n[5] How Does Selective Mechanism Improve Self-Attention Networks?, Geng et al. ACL 2020", "limitations": "Limitations and social impacts are discussed in the appendix. It's more or less adequate.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The routing of a subset of inputs for attention per module is an interesting way to simplify the complexity introduced by soft module selection per layer. The method used in the paper is relatively straightforward and can be used as a general strategy for modular networks and modular Transformers.  \n2. The paper explores the space of models hybridizing SSMs and Transformers which is an interesting space to explore. \n3. The empirical performance upshot is decent in some tasks and there is some efficiency gain over MEGA/MEGA-chunk", "weaknesses": "1. There are several technical unclarities - see questions\n2. A few more ablations can be done - what if we use selective attention without modularity? what if we use just modules with just local attention (no softmax-based dynamic subset input selection)? Or just chunks (like MEGA)? --- particularly simpler alternatives of input subset selection seem to be missing. \n3. In my interpretation the main technical contribution seems to be in the space of modular architectures. The fact that the authors use SSM+GAU as modules are just an implementation detail (although the empirical results of the implementation are noteworthy). In that sense, it could have been good to have more comparison/discussion/analysis with respect to other existing strategies involving MoEs or modules. \n4. The value of the strategy in the space of efficient attention strategies is also a bit obscure. For example, what if we use BigBird/Performer etc. attention strategy to replace the attention in GAU in MEGA?\n5. The empirical results are good on 2-3 tasks in LRA compared to MEGA-chunk, but otherwise Sailboat/Sailboat-Mem performs close to MEGA with limited efficiency gain. \n6. It seems MEGA uses a feed forward MLP after the GAU-like block whereas Sailboat-mem. Uncontrolled variables like that can be potential confounders preventing clear comparison. Could be good to show some results in LRA with MEGA-chunk variant (with controlled parameters) with the same kind of block as Sailboat with the same activation and everything. \n\n-------\n**Post-Rebuttal score update:**\n\nI increased the soundness score to 3 (from 2), but kept the overall score still to 5. Reasons below.\n\n* Some of the issues are resolved in the rebuttal and additional experiments like 6, 2. \n\n* Some of the important technical details are more clear mitigating my concern in 1., but hopefully they will be clarified better in the paper. \n\n* I was initially not too convinced by the rebuttal about the lack of comparison against other efficient attention in MEGA-framework (point 5).  But now I think, the contribution of the adaptive input sparsification can be seen as a sort of orthogonal mechanism to efficient attention. It does not modify the attention itself but the input to the attention - and thus can be stacked with other efficient attention. So I think, under that perspective, point 5 is not as critical. Although perceiver is a relatively less orthogonal possible comparison. For this point and the above points, I am increasing the soundness score to 3. \n\n* Given the clarification that $M=1$ (no. of pre-defined modules), I think on the flip side some of my points under the strengths section do not really apply anymore. A modular mechanism with just one module amounts to just lacking any modular mechanism. So I wouldn't say that the method makes much of a move in the space of modular deep learning or MoE. A theoretical framework is introduced but not tested empirically. Having more modules can also bring forth other engineering challenges and increase compute - which can result in a worse trade-off than Mega. \n\n* Overall, the input compression mechanism is still interesting but my main concern at present is the presentation and the overall positioning of the paper. The paper emphasizes quite a bit on the modular setup but it seems that all those parts can be skipped over and only the layer input sparsification via tempered softmax can be kept to be relevant to the experiments. Given that it's hard to find (or missing) that M=1, the paper also feels misleading and confusing. This is mainly a presentation issue, but still seems substantial to me. Without this, the paper is $6$ for me. But for the last two points, I am still keeping the score as $5$ at the moment. \n\n**Update 2**\n\nI increased the score to $6$ based on further discussions and the authors' expressed willingness for increased clarification on the positioning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for the positive review of our work and the constructive feedback on our manuscript. In the following, we address the remaining concerns to hopefully motivate a clear acceptance score.\n\n**Technical clarity:** Please refer to our answers to Q1, Q2 and Minors. We will also update the Related Work section to include extra comparisons with the sparse attention approaches and the missing reference.\n\n**More Ablations:** We add the new ablations in the Additional Ablation Study section in **Global Response**. Comparing Sailboat-mem with Sailboat-local and Sailboat-chunk, we can see that our model obtains better accuracy than the models without SMA. If we constrain the memory size to a smaller value (see Sailboat-mem32 v.s. Sailboat-local32), this performance gap becomes more substantial (more than 10 percent absolute difference on Pathfinder). These evidences demonstrates the effectiveness of our SMA mechanism.\n\n**Comparison with MoEs:** We add this comparison in the Comparison with Mixture of Experts section in **Global Response**. We also include the additional ablation results with X-MoE routing as an alternative design choice of latent configurator. We can see that X-MoE performs much worse than our Tempered Softmax in our use case.\n\n**Comparison with efficient attention strategies:**  As discussed in **Global Response**, our SMA mechanism enables a dynamic architecture that can learn to drop a sub-module entirely for better efficiency when it is adapted to the target task, and we empirically observe such behavior for different tasks in LRA. This capability is not possible for the previous efficient attention strategies because their architectures are static and thus an efficient attention layer is always applied before the feedforward layer.\n\n**Limited performance improvements on LRA compared with MEGA-chunk:** Please refer to the Updated Results for Long Range Arena section in **Global Response**. We do have substantial performance improvement over MEGA-chunk on 5 out of 6 tasks on LRA, and the average accuracy is 1.96 points higher. In Figure 5 of **Appendix D** in our paper, we also prove that our model can provide substantially better speed-quality trade-off than MEGA-chunk when varying different memory sizes.\n\n**Results of Sailboat-mem with feed forward MLP:** As shown in the Additional Ablation Study section in **Global Response**, we can see that both Sailboat-mem-ffn and Sailboat-chunk produce significantly worse performance than Sailboat-mem. The Sailboat-chunk model uses the same block as Sailboat but follow the MEGA-chunk to use the chunking strategy for sequence subset selection. Sailboat-mem-ffn follows the MEGA-chunk to add an additional FFN after each Sailboat layer.\n\n---\n\n**Q1:** *How are $c_i$ and $a_i$ computed for modules?*\n\n**A1:** The latent configurator will produce $c_i$ and $a_i$ for each of the module. Please refer to L167-L172 for how the configurator is implemented in Sailboat.\n\n**Q2:** *I did not understand the hyperparameter setup for the modules ... ... How do you control for parameter counts then?*\n\n**A2:** As said in L109, $M$ is the number of predefined functions (or modules), and we set $M=1$ to include only one GAU module per layer, whose activation is controlled by the latent configurator. Yes, an interesting future direction is to investigate how to scale up SMA with a large number of modules. It is promising to try some primitive ideas such as setting an upper bound for the activation time of each module, or setting a maximum for the number of modules that can be activated at the same time. However, this direction is out of the scope of our paper, and we leave it as the future work.\n\n**Q3:**  *... ... difference between GShard [4] and this approach?*\n\n**A3:** GShard divides tokens into groups of the **same** size, while our SMA allows different modules to have different group size. This means SMA can support an adaptive neural architecture whose sub-modules can be completely dropped if no tokens are selected for that module. Also, while GShard is a routing mechanism, SMA is an activation mechanism that can still function when there is only one module under consideration.\n\n**Q4:** *Is FFN removed in Sailboat? Any specific reason?*\n\n**A4:** Yes, it is removed. This is because it is kind of redundant in its design as the GAU module and it doesn't provide any empirical performance gains for Sailboat according to our ablation study.\n\n\n**Minor:**\n\n- *Are $\\zeta_i$ in the eqn. near L120 learnable parameters?* \n\n  As in L117, they are intended to represent a linear combination over the pre-defined functions. In the implementation of our Sailboat architecture, it is instantiated as the last linear layer in the GAU module.  \n\n- *Is \\forall $\\forall$ in the eqn. near L120 a typo?* \n\n  No, it is not. $L’$ is the function space of a layer equipped with SMA and $M$ number of predefined functions. Thus, $\\forall \\zeta_i $ indicates the coefficients for a linear combination of the predefined functions.\n\n- *Is there any reason why the function/module weights ($\\zeta_i c_i$) are not normalized?*\n\n  As in L121-123, we want to show that $L'$, the function space of a layer equipped with SMA, can recover the original function space $L$ defined in L109, if all the functions are activated. This won't be true if the module weights are normalized.\n\n- *... ... extra comparison could be to use something like perceiver [3] ... ...*\n\n  Nice catch! We will try this idea and see if it works out, but it seems out of the scope of this paper.\n\n- *“second-order self-attention” -- second-order in which sense?*\n\n  In the sense of the interactions between the input elements. Self-attention does pairwise comparison between tokens, which is a second-order interaction.\n\n- *Should be \"instantiated\" in L144*\n\n  Thanks! Fixed it.\n\n- *I assume $a^l \\in \\mathbb{R}^n$ should be more specifically $a^l \\in \\\\{0, 1\\\\}^n$ in L168?*\n\n  Yes, you are right! Fixed it.\n\n\n"}]

Author rebuttal:
# Global Response

## Updated Results 

For the Sailboat-mem model, we find that down-scaling the attention matrix $QK^T$ with the window size $w$ instead of the compressed sequence length $r$ can lead to substantially better results (as shown in Equation (5) near the Line 521 of **Appendix A.2**). Thus, we apply this change to Sailboat-mem for both the Long Range Arena and the Speech Command benchmark. 

The table below shows the updated results for Long Range Arena (LRA), where * indicates a hybrid model. Our Sailboat-mem is now the new **state-of-the-art** among models with linear inference complexity, outperforming the pure SSM-based model, S5, with remarkably faster training speed and much less GPU memory consumption. Note that for a fair comparison with MEGA, we use the same MH-EMA based SSM for Sailboat-mem, and MH-EMA is generally considered as a simpler version of S4 and has worse modeling power than S5. The fact that we can still achieve performance improvements over S5 under this unfair setting of comparison demonstrates the effectiveness of our proposed SMA mechanism. We also include detailed efficiency measures for each task in the PDF. 

| Model    | ListOps | Text  | Retr. | Image | Path. | Path-X | Avg. | Speed | Mem. |
|-|-|-|-|-|-|-|-|-|-|
| 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×  |
| **Sailboat-mem*** | 61.70  | _89.60_  | _91.28_ | **90.10**  | **96.35** | _96.68_ | **87.62** | **10.4×**   | **0.05×**   |
| | | | | | | | | | |

The table below shows the updated results for Speech Command. 

| Model | #Param. | Acc. | Speed | Mem. |
|-|-|-|-|-|
| S4 | 300K | 97.50 |-|-|
| MEGA-chunk | 300K | 96.92 | 1.00× | 1.00× |
| **Sailboat-mem** | 293K | 97.35 | 1.32× | 0.44× |
|||||| 
## Additional Ablation Study

We present additional ablation studies on Sailboat-mem to justify the effectiveness of both the proposed working memory mechanism and the design choices of our architecture. We explain the meaning of the postfixes in the table as following:
 - “-ffn” means we add an additional feedforward layer (with the same architecture as MEGA) after each Sailboat layer
- “-sigmoid” means we use sigmoid instead of Tempered Softmax for latent decision probability calculation.
- “-moe” means that instead of using our Tempered Softmax for module activation decisions, we adapt the X-MoE [1] routing mechanism to our use case by only considering two modules, a GAU module and a module that always outputs zero. We do not apply load balancing loss in this case because there is only one module that needs to do actual computation.
 - “-local” means that we don't use SMA and always activate the GAU module with local attention.
- “-chunk” means that we don't use SMA and always activate the GAU module to process the equal-sized chunks of the input sequence.
- “-mem32” means we restrict the memory size of Sailboat-mem to 32.
- “-local32” means we restrict the sliding window size of Sailboat-local to 32.

| Model| Image| Speed| ListOps | Speed | Path.|Speed|
|-|-|-|-|-|-|-|
| Sailboat-mem| 90.36 | 1.33×| 59.05| 2.26× | 96.34 | 1.31× |
| Sailboat-mem-ffn|89.60| 1.13×| 56.45| 2.05× | 96.62 | 0.97× |
| Sailboat-mem-sigmoid | 87.34 |1.77×| 58.20 | 2.52× | 91.35 | 1.67× |
| Sailboat-mem-moe | 85.60 | 1.66× | 58.10 | 2.19× | 91.41| 1.80× |
| Sailboat-local | 90.12 | 1.46×| 59.00| 1.78× | 96.08 | 1.54× |
| Sailboat-chunk | 86.46 | 2.17× | 58.35| 2.79×| 93.91 | 2.97×|
| Sailboat-mem32 | 78.16 | 1.73×| 58.70| 2.14× | 92.97 | 1.86×|
| Sailboat-local32 | 75.58 | 1.67× | 53.90| 2.19× | 82.05 | 1.80× |
||||||||
## Comparison with Mixture of Experts (MoE)

The motivation behind our Sparse Modular Activation (SMA) mechanism is to enable neural networks to contextually skip modules of any architectures for efficient sequence modeling, while MoE aims to efficiently scale up the models with more parameters (usually by adding homogeneous modules). This difference of motivation results in a fundamental mechanism difference:
- SMA leverages a latent configurator to decide if each module needs to be **activated** for each sequence element, while MoE is designed to choose a predefined number of modules from a large group of modules.

This difference on the core design of the mechanisms further leads to the following consequences:
- SMA supports a dynamic model architecture that can learn to drop a sub-module entirely (for better efficiency) based on the task it is trained on, while MoE only selects the parameters of the modules for the same architecture. As shown in Figure 3 of our paper, when the Sailboat model is trained on the Text task of LRA, the first three layers learn to have **zero** activation time of the GAU module. This means that these layers degenerate to pure SSM layers from the initial SSM+GAU architecture after adapting to the target task.
- SMA is guaranteed to have a full coverage of the combinatorial search space ranging from zero activation to full activation of modules, while MoE can only cover a subset of the space by choosing a fixed number of modules to activate. This is further explained in L113-L123 of our paper.

To the best of our knowledge, SMA is the **first** mechanism that successfully enables a neural network to obtain practical efficiency and complexity gains from sparsely activating a self-attention-like module, while none of the previous works on MoE ever achieved this.

---

[1] On the Representation Collapse of Sparse Mixture of Experts. Chi et al. NeurIPS 2022.


Meta-review: This paper presents Sparse Modular Activation (SMA), a general mechanism for sparsely and dynamically activating sub-modules in neural networks, addressing the challenges of combining attention with Linear State Space Models (SSMs). The authors introduce a novel neural architecture, Sailboat, which employs SMA to sparsely activate a Gated Attention Unit (GAU) based on state representations from an SSM. The proposed model demonstrates a significantly better quality-efficiency trade-off compared to existing hybrid models across various tasks. The reviewers are overall positive towards the contribution of this paper. I would recommend to accept this paper.
Decision: Accept (poster)