Title: Sparse Modular Activation for Efficient Sequence Modeling
Venue: NeurIPS 2023
Accepted: yes
Keywords: Sequence Modeling

Modularity

Sparsity

Attention Mechanism

State Space Model

Mixture of Experts

Neural Network

Transformer
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

Abstract

Recent hybrid models combining Linear State Space Models (SSMs) with selfattention 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 qualityefficiency 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 .

1 Introduction

Recent advance on efficient sequence modeling with State Space Models (SSMs) [GGR21; GDE + 20; GGGR22; GB22; SWL23] has shown impressive performance for a wide range of tasks across modalities, such as text classification, image recognition and speech recognition. SSMs, as first-order linear models, defined by a set of input, output, and state variables connected by firstorder differential equations, can efficiently capture the recurrent structure in sequential data with carefully designed state matrices and the application of convolutional parallelism [GGR21]. However, they still significantly underperform the self-attention [BCB14; VSP + 17] based model in both language modeling and machine translation [VPSP23] tasks. A recent work [FDS + 23] reveals that this is due to its deficiency of modeling the second-order pairwise comparisons between the input tokens, and shows that the augmentation of an additional shifted SSM layer can improve SSM's

associative recalling ability. Furthermore, better quality-efficiency trade-off can be achieved by directly introducing extra self-attention modules to form a hybrid model (e.g. MEGA [MZK + 23] and Hybrid H3 [FDS + 23]) that utilizes both the first and the second order inductive biases, i.e. , SSM and self-attention. However, the current hybrid models apply the attention modules statically and uniformly to each of the input token regardless the property of the task itself. This can lead to suboptimal quality-efficiency trade-offs since not all input tokens require second-order modeling and this computation need can vary substantially depending on both its context and the task difficulty.

In this paper, we aim to answer the following research questions for efficiently combining attention with SSMs:

RQ1: Can neural networks learn to activate their attention modules on demand to achieve better quality-efficiency trade-off?

RQ2: How much extra attention is needed for the SSMs on a task-by-task basis?

To answer these questions, we develop a new general mechanism, Sparse Modular Activation (SMA) , that allows a neural network to sparsely and dynamically activate its sub-modules for each of the input token in a fully differentiable manner. Specifically, we assume a neural model can be composed of multiple heterogeneous sub-modules. For the input sequence, a latent configurator sparsely maps tokens to multiple compressed sequences corresponding to sub-modules. Each sub-module is then only applied on its mapped shorter sequence. Compared with activating all sub-modules on the whole input, Sparse Modular Activation can reduce computation and memory consumption for both the training and inference stages. Notably, SMA is proved to have a full coverage of the combinatorial search space of module activation, which is further explained in Section 3.2.

Efficient learning of dynamic sparsity is notoriously difficult under the constraint of the current parallel hardware [LQC + 22; GZYE20; XM22]. To enable the practical efficiency gains from our module-level sparsity, we provide a simple yet efficient parallel implementation of SMA without any custom fused GPU kernels. Specifically, when compressing a batch of sequences in SMA, our implementation conducts both token selection and the sequence re-padding simultaneously using a single scatter operation that is widely optimized and present in modern deep learning frameworks.

To address RQ1 , we apply SMA to construct a novel neural architecture, SeqBoat , that sparsely activate a Gated Attention Unit (GAU) [HDLL22] based on the state representation learned from an SSM. Both the GAU and the SSM representations are then aggregated through simple addition and activation to form a layer-level representation. Multiple same-sized SeqBoat layers are stacked sequentially to form a full neural model. Inspired by the working memory mechanism [AS68] used in human cognition, we further restrict the GAU to only apply local attention on the compressed sequence, which allows our model to have linear sequence inference complexity but theoretically infinite attention span.

We conduct comprehensive experiments to show that SeqBoat has significantly better qualityefficiency trade-off than state-of-the-art hybrid models on a wide range of tasks, including Long Range Arena (LRA) [TDA + 20], speech classification [War18] and language modeling [Hut06]. On the competitive LRA benchmark, SeqBoat achieves 1.96 higher average accuracy than MEGAchunk [MZK + 23], the previous best hybrid model, with a 10.4 × training speed up and a 95% memory 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, addressing RQ2 . We demonstrate that our working memory mechanism provides substantially better computation-accuracy trade-off than chunking based models, and analyze the relationship between the working memory size and the effective attention span on various long sequence modeling tasks.

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 comparisons between our approach with the related works is also included in Appendix F . We begin by reviewing how the standard sequence modeling is formalized to establish the common ground for our discussion.

2.1 Time-Invariant Sequence Modeling

Given a discrete sequence, x = { x 1 , ..., x n } ∈ R n , consisting of n tokens, a time-invariant sequence model P θ is optimized to maximize the likelihood of the observed sequences by factorizing them as follows:

where x <t = { x 1 , ..., x t -1 } is the sequence history at time step t , and the parameter θ is independent of the time step t . This formulation implies that the full model parameters θ and the full history x <t are both essential for the conditional prediction of each token x t . However, one potential issue is as the prediction difficulty of each token may differ depending on the context and the position, this static model P θ can lead to sub-optimal accuracy-efficiency trade-off by wasting computation on either unimportant context [SJP + 21] or easy-to-predict tokens [Gra16].

3 Learning Sparse Modular Activation

To cover a larger search space that may contain more efficient sequence models, we propose to formulate sequence modeling as a problem of finding an optimal time-variant model that can dynamically activate a subset of modules from a pre-defined function space for each time step.

3.1 Time-Variant Sequence Modeling

Formally, a time-variant sequence model is defined on a compact function space F : X c t ↦→ [0 , 1] n × V , where V is the size of the vocabulary and X c t = { x c t : x c t ⊆ x <t ∈ X ⊆ R n } , contains all possible sub-sequences of the sequence history x <t . Then for each of the token prediction at the time step t , the model learns to apply a function f t ∈ F with the parameters θ t that maximizes the sequence probability, i.e. ,

This formulation generalizes the previous works in pursuing a dynamic and sparse model for sequence modeling, where the connections are further explained in Appendix F. In this work, we assume the function space F is chain-structured, i.e. , F = H ◦ L N ◦ · · · ◦ L 1 ◦ E , where H : R n × d m ↦→ [0 , 1] n × V is the classification function, E : R n ↦→ R n × d m is the embedding function, N is the number of intermediate layers, d m is the model size and L : R n × d m ↦→ R n × d m is the function space of the intermediate mappings. We further assume that L is the spanning set of a finite number of the function f l i with its parameters θ l i , i.e. , L = span { f l 1 , ..., f l M } , where M is the number of pre-defined functions. These assumptions justify the design of our Sparse Modular Activation mechanism, which is further explained in the following section.

3.2 Sparse Modular Activation

Sparse Modular Activation (SMA) introduces a latent configurator at each time step t and each layer of a neural sequence model. The configurator produces a binary decision vector a t ∈ { 0 , 1 } M to only activate those functions f l i that have the decision value a i t = 1 . Only the activated function will be applied to process the input representations and return the outputs. These outputs are then aggregated with a linear combination to generate the final vector, y t , of the layer. To enable endto-end differentiability, we also require the latent configurator to output the confidence probability of its decisions, c t ∈ [0 , 1] M , and respectively scale the output from the activated function with its confidence values c i t . The function space of a layer equipped with SMA can be described as follow, where I is the index set of activated functions. Notably, SMA has the following important property, Theorem 1 (Function Space Coverage of SMA) . For any L ′ ⊆ L = span { f l 1 , ..., f l M } , there exists a pair of ( a ′ t , c ′ t ) that L SMA ( a ′ t , c ′ t ) = L ′ . In other words, SMA has a full coverage of the function space L .

SiLU

𝐇

!

𝐇

Linear

Tempered

Softmax

SiLU

Argmax

GAU

Linear

Sparse Activation

𝐒

#

Latent Configurator

𝐒

#$%

Linear SSM

𝐜

𝐚

Sparse Activation

𝐏

The proof is provided in Appendix B. Intuitively, we can see that the original function space L is recovered, i.e. , L SMA = L , if all the functions are activated, and the model can adaptively select any subspace of L through controlling the binary decision vector a t . During sequential inference stage, for each function that requires the contextual information, we keep a record of its activated input representations across the time steps to form a memory H c t ⊆ H <t ∈ R n × d m (where H <t is the input sequence history), so that each function can recall its contextual memory for each of the decoding step. During parallel training stage, the activated input representations for each function (or module) are sparsely selected to compose the compressed sequence H c for the module to conduct parallel processing, and this process is implemented as a compress operator. The returned representations are then mapped to its original input position in a new sequence whose inactivated positions are filled with zero vectors, which is denoted as an extract operator. We illustrate this parallel implementation in Figure 2 for better understanding.

The Pytorch-like [PGM + 19] code snippets of the compress and extract operators are provided in the Appendix A.1 with an efficient support of batched sequences using the scatter operation.

In this work, we explore applying SMA to address our RQs , in which the following two types of functions in the function space L are considered: (1) a linear State Space Model (SSM) [GGGR22; MZK + 23] that efficiently captures the recurrent sequential structure, and (2) a Gated Attention Unit (GAU) [HDLL22] that performs second-order pairwise comparisons between the sequence elements. We also keep the SSM as an always-activated module and only apply SMA to the GAU, and in this case the effective number of modules M = 1 . This is because for long sequence modeling, the main computation and memory bottleneck comes from the second-order self-attention and the gating operations inside the GAU module. We modularize the combination of SSM, GAU and SMA as a neural layer, which we call a SeqBoat layer. Its details are further explained in the next section.

3.3 Model Architecture

Figure 1 illustrates the proposed architecture for our SeqBoat layer. Given an input sequence representation S l ∈ R n × d m for the l -th layer , where n is the length of the sequence and d m is the hidden size, we first apply a linear time-invariant State Space Model (SSM) [GGR21] over the sequence

Extract

Compress

𝐡

!

𝐡

"

𝐡

#

𝐡

$

𝒚

!

$

𝒚

#

$

𝒚

#

!

f

l

1

f

l

2

𝒚

!

𝒚

#

Linear

Aggregation

𝑐

!

"

𝑐

#

"

𝑐

#

!

𝑎

!

"

𝑎

#

"

𝑎

#

!

𝐡

!

𝐡

#

𝐡

#

Latent

Configurator

<latexi sh

1_b

64="Z

cp gXQ5

ORfV/

D

7wC0BY

>A

H

NS8

EJ3

W

Kkn

j

rG

L

v

F2

Td

M

9

+

P

m

o

z

q

u

I

U

y

<latexi sh

1_b

64="Q

vORcop

G

YAT

gm

2L

M

>

B7H

V

NS8

EJ3Ur q/

9

0n

P

K

Wy

k

Z

I

Cj zX

u

+

F

w

f

5

D

d

followed by a SiLU activation [EUD17] to obtain the state representations H ∈ R n × d m ,

where ∗ is a convolution operator which is efficiently implemented using FFTs [GGR21], glyph[circledot] indicates the element-wise multiplication, D ∈ R 1 × d m is the learnable parameters, and K ∈ R n × d m is the SSM convolution kernel. In this work, we use the Multi-dimensional Damped Exponential Moving Average (MD-EMA) proposed by MEGA [MZK + 23] for the kernel parameterization, whose details are included in Appendix A.3. Note that our work is orthogonal to the SSM kernel design as long as it can provide an efficient computation of the recurrent state representation. It is possible that using a more advanced SSM such as Liquid-S4 [HLW + 22] or S5 [SWL23] can lead to better performance, but we keep using MD-EMA for a fair comparison with MEGA. Given the state representation H ∈ R n × d m , a latent configurator with one linear layer is applied to obtain the decision vector a ∈ { 0 , 1 } n and the confidence probability vector c ∈ [0 , 1] n with the Tempered Softmax, i.e. ,

where w 0 , w 1 ∈ R d m , b 0 , b 1 ∈ R are learnable parameters, τ ∈ R > 0 is a learnable temperature parameter which is initialized as α √ d m (where α is a hyperparameter denoted as the initial temperature scale). Previous works [RSVG20a; GLL + 22; FZS22; RZW + 22] often seek applying auxiliary losses to regularize the exploration and the exploitation of the latent decisions. However, in this work, we don't apply any auxiliary losses and empirically observe that solely using Tempered Softmax works well for SMA. To provide an explanation of this phenomenon, we show that the implicit regularization [BD20] of gradient decent will conditionally encourage the exploration of the latent decisions if the Softmax function with learnable temperature (i.e. Tempered Softmax) is used for decision confidence calculation. The details of the theorem and the proof are provided in Appendix C.

Based on the decision values a , a Gated Attention Unit (GAU) is activated and applied on the compressed sequence H c resulted from the compress operator as previously described in Section 3.2,

where r = ∑ n t =1 a t ≤ 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 ∈ R n × d m , the final layer output is computed with a linear aggregation followed by a non-linear activation, i.e. ,

where W ∈ R d m × d m , b ∈ R d m are the learnable parameters. We also add a normalization layer either before the Linear SSM ( i.e. , Pre-Norm) in the residual branch or after the final layer output S l +1 ( i.e. , Post-Norm) to stabilize the training process. The SeqBoat model is then constructed as stacking N number of SeqBoat layers of the same hidden size d m with the task specific embedding and classification layers. Notably, our model does not have any Feed-Forward Networks (FFN) anywhere in the architecture.

Working Memory Mechanism The computation complexity of our model is dynamically dependent on the activation probability of the GAU module p = r/n = ∑ n t =1 a t /n for each of the SeqBoat layer. This means that the computation complexity is O ( r 2 + n log( n )) for parallel training and O ( r 2 + n ) for auto-regressive decoding. While it is possible that the number of the GAUactivated time steps is far smaller than the length of the full sequence for a trained model, a random initialized SeqBoat model will still activate the GAU module for n/ 2 number of times on average, leading to a non-ideal O ( n 2 ) training complexity at the early stage of training. To solve this problem, we restrict the GAU to only conduct local attention [AOA + 20] with a sliding window of size

w glyph[lessmuch] n on the compressed sequence H c during the parallel training stage. Each token will attend to the nearest w tokens in the compressed sequence, i.e. , w/ 2 past tokens and w/ 2 future tokens for bidierectional sequence modeling, and w number of past tokens for auto-regressive sequence modeling. This approach keeps a First-In-First-Out (FIFO) memory with size w of the compressed sequence history during auto-regressive decoding, which can be understand as a working memory [AS68] for the GAU module. With this mechanism, SeqBoat reduces the training complexity to O ( rw + n log( n )) and the inference complexity to O ( rw + n ) , while maintaining the ability of attending to the key tokens whose distances from the query are longer than the working memory size w . This is because the elapsed time steps between two GAU activations is generally not bounded by the size of the working memory.

4 Experiments and Results

We evaluate the SeqBoat model on three representative long-range sequence modeling benchmarks across image, text and speech modalities to have a comprehensive understanding of its capacity. The implementation details of our model for each of the tasks are presented in Appendix A.

Besides model performance on task metrics, we also report training speed and memory consumption as the efficiency of each model. Since the dynamic sparsity of SeqBoat will affect both the speed and the memory consumption throughout the training process, we measure its training speed up as the average per step wall time across the full training stage, and the memory consumption as the average per step GPU memory allocation. More details of the efficiency measurement are included in Appendix A.6.

4.1 Baseline Models

We compare our proposed SeqBoat model with multiple strong baseline models besides the Transformer [VSP + 17] and its variant Transformer-XL [DYY + 19].

S4 [GGR21] is a linear State Space Model (SSM) that efficiently captures the recurrent structure in sequential data. It employs a new parameterization for the state space model, enabling efficient and principled modeling of long-range dependencies. S4D-LegS [GGGR22] is a simpler variant of S4 that uses a diagonal state matrix.

MEGA [MZK + 23] is a hybrid model that combines a gated self-attention mechanism with a Multidimensional Damped Exponential Moving Average (MD-EMA) parameterized SSM. It aims to achieve a better quality-efficiency trade-off by utilizing both the first and second-order inductive biases from SSM and self-attention. MD-EMA can be considered as a simpler variant of S4 that is similar to S4D. MEGA also has a linear complexity variant, MEGA-chunk, that applies attention to each local chunk of fixed length.

H3 [FDS + 23] introduces a hybrid SSM layer that stacks a diagonal SSM and a shifted SSM to model the comparisons between the elements in a sequence. It also provides a hybrid model that includes extra self-attention layers for better language modeling performance.

S5 [SWL23] is the state-of-the-art SSM model that improves over S4 with multi-input multi-output and time-domain processing. It leverages the parallel scan operator in JAX [BFH + 18] for an efficient parrallel implementation of recurrence unrolling.

Liquid-S4 [SWL23] extends over S4 through allowing the state transition to be input-dependent. The proposed approach can be cast back to the convolutional framework and re-use the Cauchy kernel for efficient computation.

Reformer [KKL20] is a memory-efficient variant of the Transformer that uses locality-sensitive hashing (LSH) to perform approximate nearest neighbor search for attention computation.

Performer and Linformer are two efficient variants of the Transformer model. The Performer [CLD + 20] uses the Fast Attention via positive Orthogonal Random features algorithm to approximate the self-attention mechanism, reducing the computational complexity. The Linformer [WLK + 20] employs low-rank approximations to project the self-attention mechanism into a lowerdimensional space, resulting in a more efficient system.

Luna [MKW + 21] is a linear unified nested attention mechanism that approximates Softmax attention with two nested linear attention functions, resulting in linear time and space complexity.

4.2 Long Sequence Modeling

We first evaluate SeqBoat and SeqBoat-full, a variant of SeqBoat conducting the full attention over the compressed sequence without the working memory mechanism, on the Long Range Arena (LRA) benchmark [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 + 11]), document retrieval (Retrieval; [RMQAJ13]), image classification (Image; [KH + 09]), Pathfinder [LKV + 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 a 95% memory reduction compared to the Transformer. When compared with a pure SSM-based model, S4, which has more general modeling power than the MD-EMA used in our model, SeqBoat achieves a substantially higher accuracy on average with a 2.17 × speed up and a 64% memory reduction. Our model even achieves a significantly better average accuracy than state-of-the-art SSM model, S5 [SWL23], with a 0.16 absolute improvement, resulting in a new high score for models with linear inference complexity. Surprisingly, SeqBoat outperforms SeqBoat-full on Pathfinder and the Path-X tasks. This demonstrates the effectiveness of our proposed working memory mechanism to efficiently capture long-range dependencies, which is further analysed in Section 5. Since the learned module-level sparsity of our model is dynamic and task-specific, we also include the measurements for training and inference speedup and training memory allocation for all the tasks of LRA in Appendix D to provide a holistic view of the efficiency of our model.

4.3 Speech Classification

We also evaluate SeqBoat on the long-range speech classification task using the Speech Commands dataset [War18]. We follow S4 [GGR21] to use the SC10 subset of the dataset. It is a 10-class

classification task of raw speech which contains sequences of length 16k. As shown in Table 2, SeqBoat achieves competitive performance compared to the state-of-the-art S4 model and MEGAchunk, with a 1.32 × speed up and a 56% memory reduction compared to MEGA-chunk.

4.4 Language Modeling

We evaluate our model on enwik8 [Hut06], a popular language modeling benchmark. It consists of the first 100 million bytes of the English Wikipedia and is commonly used to evaluate the performance of sequence models on long-range dependencies. We follow previous studies to use it as a character-level language modeling task with a vocabulary size of around 200. During inference, we evaluate our model on the test set of enwik8 with a sequence length of 8,192 tokens and each token is equipped with a minimum of 7,000 tokens context window. We report the bits per character (BPC) metric for evaluation. As shown in Table 3, we compare SeqBoat with Transformer-XL [DYY + 19], Adaptive Span [SGBJ19], and MEGA-chunk. SeqBoat achieves the same performance as the previous best hybrid model, MEGA-chunk, but with a significant 1.16 × training speed up. This demonstrates the effectiveness of our proposed model in auto-regressive language modeling of long texts.

5 Analysis

In this section, we analyze experimental results of SMA from several research aspects. We also include comprehensive ablation studies of our model in Appendix E.

How much attention is needed for different sequence modeling tasks? As shown in Figure 3, we draw the activation time of the GAU module at different layers of our SeqBoat-full models for each of the task in the LRA benchmark. We measure the mean and the standard deviation (plotted as error bars) of the activation time on 100 sequences randomly sampled from the validation set of each task. The lines seem to be truncated due to the fact that different models have different number of layers. Generally, we can see that Image (in Blue and Circle) and Pathfinder (in Green and Triangle) need significantly more attentions than the text-

1

2

4

Layerglyph[suppress]Index

0

20

40

60

80

100

Activationglyph[suppress]Time(%)

Image

ListOps

Pathfinder

Text

Retrieval

Path-X

3

5

6

7

8

based tasks ( i.e. , Text, Retrieval and ListOps). We can also see that our model trained on ListOps (in Orange and Square) has a much larger variance of GAU activation time than other tasks, which

may reflect that the task difficulty per data sample of ListOps has a larger variance than other tasks. Another trend we can see is that the bottom layers near the inputs need much less attention (and sometimes even drop the attention module entirely) than the high level layers. We conjecture that this is because the SSMs with first-order complexity is already capable to capture the low-level features while the attention modules are more demanded at higher levels for conducting more complex computations of second-order pairwise comparisons.

Layer 5

Layer 6

Sample 1

Sample 2

Sample 3

(a) SeqBoat-full

Layer 5

Layer 6

Sample 1

Sample 2

Sample 3

(b) SeqBoat

20

57

15

19

19

15

20

20

05

25

How is the learned sparse modular activation distributed over the input sequence? As shown in Figure 4, we draw the confidence probabilities of the sparse modular activation for both the SeqBoat-full and the SeqBoat models on the Pathfinder task from the LRA benchmark. Generally, we can see that different input sequences has different modular activation patterns, which indicates that our modular activation is dynamically sparse. We can also see that the SeqBoat activation pattern is learned to be more structured than the SeqBoat-full, which may explain the superior performance of SeqBoat over SeqBoat-full on the Pathfinder task. Another trend we can see is that different layers have their own patterns: e.g. , the Layer 5 of the SeqBoat model seems trying to find the curly paths in the latent representation space, while the Layer 6 seems aggregating the results from the contiguous time steps.

16

32

64

128

256

512

16

32

64

128

256

512

73.5

75.5

77.5

79.5

81.5

83.5

85.5

87.5

89.5

91.5

13

15

17

19

21

Accuracy (%)

GPU-hours

MEGA-chunk

SeqBoat

(a) Image

16

32

64

128

256

512

16

32

64

128

256

512

50

55

60

65

70

75

80

85

90

95

100

23

28

33

38

43

Accuracy (%)

GPU-hours

MEGA-chunk

SeqBoat

(b) Pathfinder

How does the working memory size affect the speed-quality trade-off? We compare the tradeoff ability of our SeqBoat with MEGA-chunk on Image and Pathfinder tasks of the LRA benchmark,

1

2

3

4

5

6

Layer Index

0

50

100

150

200

250

Average Attention Span

Memory Size

16

32

64

128

256

512

(a) Pathfinder

2

4

6

8

10

12

14

Layer Index

0

200

400

600

800

1000

1200

1400

Average Attention Span

Memory Size

512

1024

2048

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 Megachunk models generally fail (achieving around 50% accuracy as random guessing) on Pathfinder tasks with less than 128 chunk size, while our SeqBoat can still achieve around 90% accuracy with only a memory size of 16. This indicates the effectiveness of our working memory mechanism in capturing long term interactions.

Howdoes the size of working memory affect the effective attention length of the GAU module? We draw Figure 6 to investigate the relationship between the working memory size and the average effective attention length on both Pathfinder and the enwik8 datasets. We first measure the average attention distance from the query token to its key tokens for each of the time step, and the results are further averaged for all the time steps in a sequence. The average attention span on Pathfinder is measured on 100 sequences randomly sampled from the validation set. The average attention span on enwik8 is calculated based on a sequence of length 8,192 sampled from the validation set, and each token is equipped with a minimum of 7,000 tokens context window. From the figures, we can see a general trend that the models with a small working memory size (Size 16, 32 and 64 for Pathfinder, and Size 512 for enwik8) can have a far longer average attention span than its memory size. Notabaly, the layer 3 of the SeqBoat model with only 16 working memory slots surprisingly achieves an average attention span of 245 for the Pathfinder task. This partially explains the superior performance of our model over MEGA-chunk with small memory sizes and directly proves that our working memory mechanism can efficiently enable long-term interactions between the sequence elements.

6 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 implementation but also provides a theoretical guarantee for a full coverage of the function space of multiple modules. To address the challenges of combining attention with Linear State Space Models (SSMs), we apply SMA to develop a novel neural architecture, SeqBoat, that sparsely activates a Gated Attention Unit (GAU) based on the state representations from an SSM. The proposed model demonstrates a significantly better quality-efficiency trade-off compared to existing hybrid models across various sequence modeling tasks. To the best of our knowledge, SMA is the first mechanism that allows a neural network to obtain practical efficiency and complexity gains from sparsely activating a self-attention-like module. The strong empirical performance show that even with a preliminary application of SMA, we could already see substantial efficiency and interpretability benefits. For future work, we mainly want to explore the scaling behavior of SMA and also how to incorporate SMA with the pre-trained large language models.

References

[AOA + 20] J. Ainslie, Santiago Ontañón, Chris Alberti, V. Cvicek, Zachary Kenneth Fisher, Philip Pham, Anirudh Ravula, Sumit K. Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs in transformers. Conference On Empirical Methods In Natural Language Processing , 2020.

[AS68] Richard C. Atkinson and Richard M. Shiffrin. Human memory: A proposed system and its control processes. In The psychology of learning and motivation , 1968.

[BCB14] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. International Conference On Learning Representations , 2014.

[BD20] D. Barrett and B. Dherin. Implicit gradient regularization. International Conference On Learning Representations , 2020.

[BFH + 18] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye WandermanMilne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.

[BZMA20] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. Advances in neural information processing systems , 33:12449-12460, 2020.

[CLD + 20] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794 , 2020.

[DFE + 22] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems , 2022.

[DMRB22] Benoit Dherin, Michael Munn, Mihaela Rosca, and David Barrett. Why neural networks find simple solutions: The many regularizers of geometric complexity. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems , volume 35, pages 2333-2349. Curran Associates, Inc., 2022.

[DYY + 19] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc V. Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. ACL , 2019.

[EUD17] Stefan Elfwing, E. Uchibe, and K. Doya. Sigmoid-weighted linear units for neural network function approximation in reinforcement learning. Neural Networks , 2017.

[FDS + 23] Daniel Y Fu, Tri Dao, Khaled Kamal Saab, Armin W Thomas, Atri Rudra, and Christopher Re. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations , 2023.

[FZS22] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research , 23(1):5232-5270, 2022.

[GB22] Ankit Gupta and Jonathan Berant. Diagonal state spaces are as effective as structured state spaces. ARXIV.ORG , 2022.

[GDE + 20] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 1474-1487. Curran Associates, Inc., 2020.

[GGGR22] Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. ARXIV.ORG , 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 Conference on Learning Representations , 2019.

[GLL + 22] Yue Guan, Zhengyi Li, Jingwen Leng, Zhouhan Lin, and Minyi Guo. Transkimmer: Transformer learns to layer-wise skim. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 72757286, 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.

[HLW + 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://prize.hutter1.net/, 2006.

[JGB + 21] 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 gumbelsoftmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings . OpenReview.net, 2017.

[JJNH91] Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive 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, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 , 2020.

[LH18] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2018.

[LJH + 19] Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. International Conference on Learning Representations , 2019.

[LKV + 18] Drew Linsley, Junkyung Kim, Vijay Veerabadran, Charles Windolf, and Thomas Serre. Learning long-range spatial dependencies with horizontal gated recurrent units. Advances in neural information processing systems , 31, 2018.

[LLX + 20] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, M. Krikun, Noam M. Shazeer, and Z. Chen. Gshard: Scaling giant models 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. Dynamic sparse attention for scalable transformer acceleration. IEEE Transactions on Computers , 71:3165-3178, 2022.

[MDP + 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 , 34:2441-2453, 2021.

[MZK + 23] Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. Mega: Moving average equipped gated attention. 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 arXiv:1804.06028 , 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 information 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 , 47:919-944, 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 on Empirical Methods in Natural Language Processing , pages 14801494, Abu Dhabi, United Arab Emirates, dec 2022. Association for Computational Linguistics.

[SGBJ19] Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers. arXiv preprint arXiv:1905.07799 , 2019.

[SJP + 21] Sainbayar Sukhbaatar, Da Ju, Spencer Poff, Stephen Roller, Arthur D. Szlam, J. Weston, and Angela Fan. Not all memories are created equal: Learning to forget by expiring. International Conference On Machine Learning , 2021.

[SLK + 18] Sihyeon Seong, Yegang Lee, Youngwook Kee, Dongyoon Han, and Junmo Kim. Towards flatter loss surface via nonmonotonic learning rate scheduling. In Conference on Uncertainty in Artificial Intelligence , 2018.

[SLP + 21] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv: 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. ARXIV.ORG , 2021.

[SUV18] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attention with relative position representations. NAACL , 2018.

[SWL23] Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations , 2023.

[TDA + 20] Yi Tay, M. Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, J. Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. International Conference On Learning Representations , 2020.

[VPSP23] Ali Vardasbi, Telmo Pires, Robin M. Schmidt, and Stephan Peitz. State spaces aren't enough: Machine translation needs attention. ARXIV.ORG , 2023.

[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. NIPS , 2017.

[War18] Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint arXiv: Arxiv-1804.03209 , 2018.

[WLK + 20] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

[XM22] Canwen Xu and Julian McAuley. A survey on dynamic neural networks for natural language processing. EMNLP Findings , 2022.

[ZLL + 22] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoc V Le, James Laudon, et al. Mixture-of-experts with expert choice routing. Advances in Neural Information Processing Systems , 35:7103-7114, 2022.

A Implementation Details

A.1 Code for Compress and Extract Operators in SMA

A.2 Details of Gated Attention Unit

Our Gated Attention Unit (GAU) mostly follows the original implementation as [HDLL22], with slight modifications of the relative position bias. Given the compressed sequence H c ∈ R r × d m , we first compute the query, key and value sequences, Q , K , V , i.e. ,

where w q , w q , b q , b k ∈ R d z , W v ∈ R d m × d v , b v ∈ R d v are the learnable parameters, glyph[circledot] means the pointwise multiplication, and d v is the expanded hidden dimension which is set as d v = 2 d m in our work. Then the attention output is calculated as,

where f is the attention function, s = r is the normalizing factor, B rel ∈ R n × n is the relative positional 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 in 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 [SML + 21] for the attention function on the image and the speech tasks and Softmax for the text related tasks. By default, the relative positions are calculated based on the token positions in the original sentence H instead of the compressed sequence H c . 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 s = w . The final output of the GAU is computed with a gating mechanism,

where W g ∈ R d m × d v , W h ∈ R d v × d m , b g ∈ R d v , b h ∈ R d m are learnable parameters.

A.3 Details of Multi-dimensional Damped EMA

As a special kind of SSM, the Multi-dimensional Damped EMA (MD-EMA) operator is first proposed in the MEGA [MZK + 23] paper, and we briefly introduce it here for the completeness of our work. In Equation (2), the kernel K ∈ R n × d m of MD-EMA is parameterized as following,

where η ∈ R h × d m , α ∈ [0 , 1] h × d m , δ ∈ [0 , 1] h × d m , β ∈ R h × d m , ϕ = 1 -α glyph[circledot] δ , are learnable parameters, glyph[circledot] means the element-wise multiplication and h is a hyperparamter denoted as the EMA dimension. With the pre-computed kernel K , the output of MD-EMA can be efficiently computed with FFTs in a time complexity of O ( n log n ) . During sequential decoding stage at the time step t , the output of MD-EMA is calculated as

where z 0 = 0 , s t ∈ R d m is the input representation at time t , and D ∈ R d m 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, we apply Gumbel Softmax to the probability computation in the latent configurator described in Section 3.3. Specifically, the modified probability calculation is as follow,

where G i ∼ Gumbel (0 , 1) is the Gumbel noise sampled from a standard Gumbel distribution and τ ′ follows the following temperature annealing schedule as in the previous work [BZMA20],

where T is the number of training steps.

A.5 Hyper-parameter settings

For Long Range Arena (LRA) and Speech Command tasks, we use the AdamW [LH18] optimizer and the detailed settings together with our initial temperature scale α for the latent configurator are shown in Table 4. Following S5 [SWL23], we use 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 w = 256 for other tasks in LRA. For Speech Command, a working memory size of 256 is 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 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. The memory consumption 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 a mixed cluster with 8 NVIDIA V100 32GB GPUs and 2 NVIDIA A5000 24GB GPUs.

An alternative way of efficiency measurement is to calculate the total Floating Point Operations (FLOPs) used for training. However, the FLOPs calculation does not take account of the time-consuming memory copying of large matrices (which is exactly the speed bottleneck of SMAdue to the scatter operation), and thus we choose to follow the previous works to measure the speedup based on the actual wall time for a stricter and fairer comparison.

B Proof of Theorem 1

Proof. Given any L ′ ⊆ L = span { f l 1 , . . . , f l M } , we want to show that there exists a pair ( a ′ t , c ′ t ) such that L SMA ( a ′ t , c ′ t ) = L ′ . Since L ′ ⊆ L , every function in L ′ can be written as a linear combination

of the functions { f l 1 , . . . , f l M } . Let's consider a function g ∈ L ′ . Then, we can write

where β i ∈ R . Now, let's define a ′ t and c ′ t as follows:

a ′ t = ( a ′ i t ) where a ′ i t = 1 if β i = 0 and a ′ i t = 0 otherwise.

c ′ t = ( c ′ i t ) where c ′ i t = 1 if β i = 0 and c ′ i t = 0 otherwise.

Then, we have glyph[negationslash]

glyph[negationslash]

Since a ′ i t = 1 and c ′ i t = 1 if and only if β i = 0 , then, I = { i | β i = 0 , 1 ≤ i ≤ M } , we have

̸

glyph[negationslash]

glyph[negationslash]

This completes the proof.

C Theorem of Implicit Regularization for Latent Decision Making

Wefirst give a brief introduction of the implicit regularization to provide some backgrounds. The implicit regularization of gradient decent [BD20] is a side effect caused by the discreteness of the gradient descent updates. Specifically, it has been proved that a discrete gradient update θ ′ = θ -h ∇ θ L ( θ ) of the parameter θ is implicitly minimizing a transformed loss, L ′ = L + h 4 ||∇ θ L ( θ ) || 2 where h is 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 ||∇ θ L ( θ ) || 2 will conditionally put an implicit pressure on ||∇ θ f θ ( x ) || 2 F for a multi-class classification cross-entropy loss term, where f θ ( x ) is a neural network that takes the input x and outputs the logits for classification. A Transfer Theorem [DMRB22] is further derived to characterize the conditional transfer of the implicit regularization pressure from ||∇ θ f θ ( x ) || 2 F to the gradient norm ||∇ x f θ ( x ) || 2 F , where || · || F means the Frobenius norm. The transfer theorem assumes an ordinary neural architecture f θ ( x ) : R d ↦→ R k with l layers in the following form,

where a i is the activation function, h i ( x ) indicates the sub-network from the input x ∈ R d up to the output of layer i and h 0 ( x ) = x . This implies h l ( x ) = f θ ( x ) .

In this work, we consider a neural architecture that has a per-layer latent configurator p i : R | h i -1 | ↦→ [0 , 1] , p i ( z ) = y i ( w i z ) , that maps the hidden representation from h i -1 to a confidence probability of activating a submodule q i : R | h i -1 | ↦→ R | h i | , i.e. ,

where y i 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 i , which leads to the implicit regularization of latent activation decisions.

Theorem 2 (Transfer Theorem for Latent Configurator) . Let f θ : R d → R k represent a deep neural network consisting of l layers, i.e., f θ ( x ) = f l ◦ f l -1 ◦ · · · ◦ f 1 ( x ) where f i ( z ) = a i ( p i ( z ) q i ( z )) , p i : R | z | ↦→ [0 , 1] , p i ( z ) = y i ( w i z ) is a latent configurator, q i : R | z | ↦→ R | h i | is a sub-module function, a i is the activation function, y i is the probability normalization function and w i is a weight matrix. For i = 1 , 2 , . . . , l , we have

where r i ( x ) := p i ( h i -1 ( x )) and d i ( x ) := q i ( h i -1 ( x )) , h i ( x ) indicates the sub-network from the input x ∈ R d up to the output of layer i with h 0 ( x ) = x and || · || 2 denotes the L2 operator norm for matrix and L2 norm for vector.

Proof. Following the above notations, we first rewrite the model function as

where g i ( z ) is the remaining network of the full model f θ ( x ) following the i -th layer. Considering a small perturbation x + δx of the input x , there exists a corresponding perturbation w i + u ( x ) for the weight matrix in the latent configurator p i so that

For sufficiently small δx , we can identify r i ( x + δx ) d i ( x + δx ) with the linear approximation around x , i.e. ,

where ∇ x d i : R d ↦→ R | h i | and ∇ x r i : R d ↦→ [0 , 1] are the total derivatives. We have

and similarly, for r i ( x ) = y i ( w i h i -1 ( x )) , and the perturbation w i → w i + u ( δx ) , we can have the linear approximation r ′ i ( x ) = y i (( w i + u ( δx )) h i -1 ( x )) = r i ( x ) + ∇ w i r i ( x ) u ( δx ) . This implies

Therefore, substituting both sides of Equation (11) with the equations above, and we get

Note the fact that ∇ w i r i ( x ) T ∇ w i r i ( x ) = ||∇ w i r i ( x ) || 2 2 . This implies

Defining u ( δx ) as above and taking the derivative of both sides of Equation (11) with respect to δx at δx = 0 , we have

Applying the squared L2 operator norm || · || 2 2 , at both sides of the equation, we get

Since ∇ x f w i ( x ) = ∇ x f θ ( x ) , 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

then we can derive the squared L2 norms of its derivatives,

Substituting the above gradient values to Equation (10), we can see that

This implies that the implicit regularization pressure will tend to maximize the temperature τ i to increase the exploration of the latent decisions, while at the same time maximize the probability of modular activation or de-activation through maximizing || w 0 i -w 1 i || 2 2 and p i ( h i -1 ( x )) for exploiting the decisions. Therefore, there will be an implicit balancing between the exploration and the 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

D.1 Per-task Efficiency Measures for LRA Benchmark

As shown in Table 6, we 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 is consistently and substantially more efficient than Mega, and it also has comparable efficiency with Mega-chunk but significantly better average accuracy on the LRA benchmark.

E Additional Analysis

Howdodifferent model design choices affect the performance? Wedothe ablation study of our SeqBoat-full model on the Image and ListOps tasks of the LRA benchmark, as in Table 7, to have a general understanding of the influence of the design of the latent configurator and the SSM. 'Always Activated' means that we don't use SMA and always activate GAU for conducting second-order selfattention. We can see that although the validation accuracy is 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-EMA [MZK + 23] kernel. The S4D variant of our model has much better training efficiency but the quality

are 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 we 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.

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 an additional Feed-Forward layer (with the same architecture as MEGA) after each SeqBoat layer.

'-{ GAU , FFN } ': We add a FFN module in parallel with the GAU module to have M = 2 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 calculation.

'-MoE': Instead of using our Tempered Softmax for module activation decisions, we adapt the X-MoE [JJNH91] 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': We don't use 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 equalsized chunks of the input sequence. The chunk size is set to be the same as the memory size of SeqBoat for fair comparison.

'-mem32': We restrict the memory size of SeqBoat to 32.

'-local32': We restrict the sliding window size of SeqBoat-local to 32.

We first explore the application of SMA with multiple heterogeneous modules and identify the current engineering challenges. We can see that SeqBoat-{ GAU , FFN } obtains better trade-offs than

SeqBoat-FFN (which simply appends an extra FFN layer after each SeqBoat layer) on both the Image and the ListOps tasks. However, SeqBoat-{ GAU , FFN } still falls behind SeqBoat (which does not include FFN anywhere at all). This is because the current implementation of SMA is I/O bounded due to 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 on the MoE community. We also include additional ablation results with X-MoE routing as an alternative design choice for the latent configurator. We can see that SeqBoat-MoE performs much worse than our SeqBoat model which uses the Tempered Softmax for decision confidence calculation.

Comparing SeqBoat-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 confidence calculation gives much better results under our hyper-parameter settings. Comparing SeqBoat with SeqBoat-local and SeqBoat-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 SeqBoat-mem32 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-FFN 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 to alleviate the second-order complexity of self-attention by dynamically mapping each of the token to its nearest-neighbor clusters, and only do self-attention inside the clusters. Dynamic routing can be understood as the following form when it is applied for sequence modeling,

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 f that is non-adaptive for both training and the inference.

Mixture of Experts (MoE) [JJNH91; FZS22; LLX + 20; ZLL + 22] 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,

where the model parameter θ t is selected from a pool of candidates at each time step t based on the input content. We can see that it is a special case of our formulation as in Section 3.1 because it has a fixed context x <t for each of the token prediction and the model architecture is non-adaptive.

Comparison between SMA and 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:

SMAleverages 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 different properties of the mechanisms:

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, when the SeqBoat-full 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.

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.

Comparison with Flash Attention [DFE + 22]: 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 [KKL20; CLD + 20; WLK + 20; MKW + 21]: 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-Forward 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 [LLX + 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 representations to repeatedly attend to the input array, which can be understood as conducting soft selection from the input array to a fixed subset. Our SMA operates on a different level to directly conduct a hard selection from the input to form a 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 Independent 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 emphasize 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 a rather simple strategy of using linear 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.

G Limitations, Future Works & Broader Impact

Our study has some limitations that should be acknowledged. First, while our Sparse Modular Activation (SMA) mechanism is generally applicable for activating any sub-modules in neural networks, we have narrowed the scope of our paper to focus solely on developing a more efficient neural architecture 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 scatter operator. For a 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 or 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, a promising future direction is to investigate how to scale up SMA with a large number of modules, potentially leading to empirical impacts of SMA on the research field of Mixture-of-Experts. Some initial ideas include setting an upper bound for the activation time of each module, or limiting the number of 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 betterperforming automated systems, like open-domain chatbot and automatic speech recognition, potentially improving communication, fostering inclusivity, and making such services more accessible to various societal sectors. The enhanced interpretability of models can generate trust in AI systems, especially in sensitive sectors like healthcare and finance. Transparent AI models can assist professionals 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 scalability 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 finegrained control over the model's behavior is needed. Furthermore, ethical considerations such as data privacy need to be seriously considered as our work progresses. As the efficiency of these models 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nLimitations and social impacts are discussed in the appendix. It's more or less adequate.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nAt 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.\n\nQUESTIONS:\n+ 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\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nCould 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?\n\nLIMITATIONS:\nn/a", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nOne 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\n\nQUESTIONS:\nThe 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?\n\nLIMITATIONS:\nThere are no explicit limitations beyond those discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-08"}