Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Authors: Sukjun Hwang
Aakash Lahoti
Ratish Puduppully
Tri Dao
Albert Gu
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep learning
sequence models
state space models
mamba
Primary area: deep_learning_architectures

Abstract: A wide array of sequence models are built on a framework modeled after Transformers, comprising alternating sequence mixer and channel mixer layers.  This paper studies a unifying *matrix mixer* view of sequence mixers that can be conceptualized as a linear map on the input sequence. This framework encompasses a broad range of well-known sequence models, including the self-attention of Transformers as well as recent strong alternatives such as structured state space models (SSMs), and allows understanding downstream characteristics such as efficiency and expressivity through properties of their structured matrix class. We identify a key axis of matrix parameterizations termed *sequence alignment*, which increases the flexibility and performance of matrix mixers, providing insights into the strong performance of Transformers and recent SSMs such as Mamba. Furthermore, the matrix mixer framework offers a systematic approach to developing sequence mixers with desired properties, allowing us to develop several new sub-quadratic sequence models. In particular, we propose a natural bidirectional extension of the Mamba model (**Hydra**), parameterized as a *quasiseparable matrix mixer*, which demonstrates superior performance over other sequence models including Transformers on non-causal tasks. As a drop-in replacement for attention layers, \name outperforms BERT by 0.8 points on the GLUE benchmark and ViT by 2% Top-1 accuracy on ImageNet.

Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Sukjun Hwang".1, Aakash Lahoti* Ratish Puduppully 22 Tri Dao and Albert Gu
Machine Learning Department, Carnegie Mellon University 'IT University of Copenhagen
DDepartment of Computer Science, Princeton University
*Cartesia Al [ alahoti)@ca cmu .du, rapu@itu.de du. tri@tridao.mu me, aguêcs. agu@cs.cmu.ed cmu. edu
Abstract
A wide array of sequence models are built framework modeled after Transform- ers, comprising alternating sequence mixer and channel mixer layers. This paper studies a unifying matrix mixer view of sequence mixers that can be conceptualized linear map the input sequence. This framework encompasses broad range well-known sequence models, including the self-attention Transformers as well recent strong alternatives such structured state space models (SSMs), and allows understanding downstream characteristics such as efficiency and expressivity through properties of their structured matrix class. We identify key axis of matrix parameterizations termed sequence alignment, which increases the flexibility and performance of matrix mixers, providing insights into the strong performance of Transformers and recent SSMs such as Mamba. Furthermore, the matrix mixer framework offers a systematic approach to developing sequence mixers with desired properties, allowing us to develop several new sub-quadratic sequence models. In particular, we propose a natural bidirectional extension of the Mamba model (Hydra). parameterized as quasiseparable matrix mixer, which demonstrates supe- rior performance over other sequence models including Transformers on non-causal tasks. As drop-in replacement for attention layers, Hydra outperforms BERT by 0.8 points on the GLUE benchmark and ViT by 2% Top-I accuracy on ImageNet.
Introduction
Large-scale pretrained models such as GPT (35), BERT [10.. and VIT [T]] exhibit saate-of-the--ee perfor- mance across wide range of tasks in multiple domains, including language and vision. A large number of these pretrained models follow a multi multi-layer architectural blueprint: a sequence mixer, such as Self- Attentioo [42] aggregates information across the input sequence, followed by a channel mixer proixerppo-oeo cesses information within each sequence element. Over the years, Attention has been the predominant choice for sequence mixing due to its ability to facilitate direct pairwise interactions between elements of the input sequence single step. However, this capability incurs quadratic cost with respect sequence length, making the training and deployment of these models prohibitively expensive for longer sequences. Although numerous alternatives have been proposed, designing principled sequence models that match the performance and versatility of attention- based system remains a substantial challenge.
One general strategy in designing alternative sequence models involves substituting the Attention matrix with different matrix parameterizations the core sequence mixer. These modifications are motivated by various goals. For instance, simplifying the sequence mixer has led the development of models such as MLP-Mixer [40), which uses dense matrices, and FNet [33 which utilizes the Discrete
'Equal Contributions.
WWrk done at A STTRR, Singapore.
In this paper, Attention exclusively refers Self-Attention.
38th Conference Neural Information Processing Systems (NeurIPS 2024).

---

maz maz iggo mes
mo
mi m,
Channel Mixing
Me
122 Mas myo
mas
my
mf
mi
c/Aiabo ci c/Ai,,b
mg
mi
Ai cfA in by A
Sequence Mixing
Myn Mas
maz Maz
ma
m3
m
ejaann Gj by Mi c/A
Dense
Vandermonde
Semiseparable
Y
M
my
my
my
my
9au Huz
Koo 110 K2 Ka
C Aia by Aiib,
m_1
9:n
Koz kos 121 kss
c Aga
61
C by
10g
m:
42u 421
Ai by
My b
m_- m_2 m_
mg
Ran 9a1
c Ag CA C
6g
Toeplitz
Low Rank
Quasiseparable
Figure 1: (Left) A schematic of the matrix mixer framework. (Right) An overview of matrix mixer classes: dense, Vandermonde, Toeplitz, low-rank, semiseparable, and quasiseparable.
Fourier Transform matrix. Additionally, incorporating inductive biases such as positional information
has resulted in the use of Toeplitz. matrices in models like TNN (35]. Enhancing computational
efficiency has spurred the creation of low-rank structures in models like Linear Attention (L.A) (22
and Linformer [45], as well as Monarch matrices (5) in models such as M2 (13).
Despite the achievements of these approaches, they often lack a systematic analysis of their theoretical foundations and empirical consequences. Moreover, these models typically exhibit lower empirical performance than Attention, which is successfully applied across diverse domains and tasks. In another line of work, structured state space models (SSMs) (17 18 such Mamba (16) have been popularized for its attention-like performance with linear-time computational scaling. Furthermore, recent work [6] has shown that SSMs can also be expressed as semiseparable matrices. However, these models are primarily causal as underscored by their recurrent view of computation, limiting their application to auto-regressive settings. Several concurrent efforts to adapt Mamba into a bidirectional model [48 388 have been made, but these attempts remain ad-hoc.
In this work, we study the matrix mixer framework (see Figure D. which provides a powerful abstraction for enhancing the understanding of sequence models through the analysis of their matrix structures. Specifically, we suggest the following:
(i) Formalization of matrix mixer sequence models. We establish the conceptual foundation of the Matrix Mixer sequence models, delineating key axes of matrix parameterizations that influence model characteristics such efficiency, causality, and expressivity. For instance, we highlight that structured matrix [29) parameterizations underpin efficient sequence models through computationally efficient algorithms (Section[.1]
(ii) Introduction of sequence alignment. Within this framework, we define sequence aligument as novel attribute, which induces sequence models with essential features of data dependeeee and extendability. We introduce Sequence Aligned Matrices (SAM) that exhibit these properties, demonstrating their superior performance in downstream tasks compared to non ool-nnd counterparts (Section 2.2)
(iii) Exploration of structured matrix parameterizations, Through the matrix mixer framework, we sy systematically explore and categorize a broad spectrum of existing sequence models (Sec- tion 2.3). Motivated by sequence alignment, we also present new sequence models with underex- plored structured matrix configurations, such as Vandermonde and Cauchy matrices (Section (..4)
Building our matrix mixer framework, we introduce novel sequence model Hydra, which employs a quasiseparable matrix mixer (Section 3). Quasiseparable matrices are a fundamental matrix structure with several important properties, making them ideal for sequence modeling. For example, they gener- alize both the low-rank matrix mixers found in linear attention models and the semiseparable matrices utilized in state space models. In fact, quasiseparable matrix mixers can be seen as a natural bidirectional extension of semiseparable matrices, addressing a major limitation in state space models by making their use in non-causal settings possible. Additionally, Hydra maintains the strong performance and linear-time computational efficiency of SSMs, thanks structured matrix multiplication algorithms. Unlike prior attempts to make SSMs bidirectional by ad-hoc methods -typically by combining separate models for forward and backward sequence processing through element-wiss addition [5] 48 38), the Hadamard product [44] or concatenation [44 12] 38 the systematic approach of Hydra offers a more coherent and theoretically grounded advancement. Specifically, the definition of quasiseparable mixer matrices generalize both the heuristic bidirectional extensions of SSMs and linear attention, providing mathematical interpretation of the strong expressivity exhibited by Hydra.

---

We provide extensive experimental results that substantiate our claims. Our systematic ablation studies control architectural variables to highlight the impact of matrix parameterization. These careful experi- ments confirm that Sequence Alignment, property we newly identified in certain matrix mixers, signif- icantly enhances downstream performance. Furthermore, our experimental findings demonstrate that novel sequence models like those using Cauchy matrix mixers match the performance of established ma- trix mixers such as low-rank. This indicates that strong performance of low-rank variants is sotsol attributable to their matrix structure, but can be matched by other mixers that share similar properties. In addition, we also validate that the bidirectional extension of the Mamba model, implemented through quasiseparable matrix mixers, outperforms naive bidirectional approaches [15] 48 38 [4 [2]. Impor- tantly, Hydra excels as performant, general-purpose bidirectional sequence model, as evidenced by its strong performance across diverse domains. It achieves state-of-the tt-thee-rttt results on the GLUE bench- mark (43] with an average accuracy of 84.3%, outperforming BERT (10) by 0.8 points, and records
a Top-1 accuracy of 81.0% on the ImageNet- 1K benchmark (9. surpassing ViT [1] by 2.2 points.
We publicly release source code at attthttp//// /tt.//gtgttuuu.tg thub. cttubb.commgge haan/ooomaaabbhydrrz.
2 The Matrix Mixer Framework:
Bridging Sequence Mixers and Structured Matrices
In Section 2.1
we formally define the matrix mixer framework, conceptualizing sequence mixers as linear maps input sequences. In Section 2.2 we introduce sequence alignment, a new axis of variation of matrix structures which controls important characteristics of downstream sequence models such as data-dependent parameterizations and extendability. Section 2.3 leverages these definitions to categorize a wide array of previous works based on their matrix mixers, facilitating the understanding of sequence mixers through their matrix structures. Furthermore, in Section 2.4 we propose novel sequence mixers utilizing Vandermonde and Cauchy matrices, demonstrating the flexibility of our framework in systematically designing new sequence mixers.
2.1 Formalizing the Matrix Mixer Framework
(The matrix mixer Let be the input sequence. consisting of L elements, each with C channels. Let fx: - be the impul preprocessing function that en- capsulates common data transformations such as short convolutions, and linear layers. Lei H and P be the number of heads and head dimension respectively, such that HP=D. Let M underlying class of mixer matrices. For each head he [H), let fly: RRLCC - M be the matrix con- struction frrutioofunction function that maps input data to mixer matrices, where is the space of learnable parameters. We denote M = fi (X,0), when M,0,X are clear from the context. Then, we define matrix mixing as,
Definition 2.1
C R L represent the
Y T -fx
where (fx(X)) y denote the preprocessed input and output slice corresponding head h.
This definition encapsulates the view that existing sequence mixers can be mathematically represented as L L mixer matrices that act the sequence length (see Figure T left). The framework not only incorporates typical data preprocessing steps like projections [2.39 and short convolutions [16] 6.. but it also accommodates data dependent mixer matrices (39 34 Furthermore, it is also powerful enough to capture the concept of head structure (42) by equivalently sharing the mixer matrices within a head.
Moreover, Definition 2.1 offers a different lens to conceptualize the computational cost of sequence mixers. If we assume that both the preprocessing function and the matrix construction function are sub-quadratic in the sequence length L. which generally true in practice, then the computational bottleneck lies in the Mfx (X) operation. For a general matrix M, this multiplication will incur O(L?) cost. The only way to mitigate this restrict M to being structured matrix, which are known to possess sub-quadran quadratic matrix multiplication algorithms. We refer to such sequence mixers as structured matrix mixers. This paves the way to systematically develop new sequence mixers: select an appropriate class of structured matrices from established mathematical literature, devise a data-dependent matrix construction function, and integrate them into the matrix mixer framework.
2.2 Sequence Aligned Matrices
Unstructured mixer matrices, also known as dense matrices, lack two key useful properties: 1) these matrix mixers cannot be casily made data detaddeeenndntt, which has been identified key property of performant sequence models such as Transformers and Mamba, and 2) they can only be applied to fixed-length sequences, a property we call extendability. Due to substantial importance of these
3

---

Table I: Categorization of existing methods matrix mixers. L denotes input sequence length.
Matrix Structure M
Formulation (m.y)
Complexity Sequence Aligned Method Instantiations
Dense
Way
O(12)
MI.e-Mixer (40)
Dense
softmax, kg)
O(L2)
Transformer
(Softmax Attention)
Low-Rank
q. k,
O(L)
Linear Attention
22
(Linear Attention)
Linformer 45
Kaleidoscope
Butterfly
See 7
OLLogg)
Monarch 5 13
Toeplitz
S4
17 18
H3
14
(Convolution)
my-
O(LlogL)
TNN 35
CKConv
Discrete Fourier Transform
w/
O L)
FNet
23
Vandermonde
(m:) y
O(Llog 1.)
Ours (Section
2.4
Cauchy
Enleid kza)
Semiseparable
by (120)
(())
Mamba (S6 16 SSD
Quasiseparable
Equation
O(L)
Ours (Hydra) (Section 3
features for sequence mixers, we formally define the class Sequence Aligned Matrices (SAM) to
systematically explore matrices that are characterized with both properties.
Definition 2.2 (Sequence Aligned Matrices). Let L be the sequence length and let M MMELLX denote matrix with a parameter set P. Then, we say that M is a Sequence Aligned Matrix if there exists a partition of PPP, and + Ø, such tha all sets e II, there exists bijective map fe:[L] and. for each i.j [[L, i j. the sub-matrix Mij +1 (.:jj is composed solely from the parameters in the subset Ug Vggsskij (kkk)) CP.
In simpler terms, this implies that each parameter of a SAM either mapped to a specific element of the sequence or is left data- indeleeddeedt, ensuring that every upper-left sub-matrix upto index i is constructed using only the parameters corresponding sequence segment upto and including index i and/or the data iattaindeeendent ones.
Proposition 2.3 (Data Dependency). Sequence aligned matrices exhibit canonical data-dependent purameterization.
Proof. This property arises from the parameter partition structure guaranteed in the definition. Specifically, for each partition we associate a parametric function that, for any given element i computes the parameter's 's value by treating the element itself an input.
Proposition 2.4 (Extendability). Sequence aligned sequence mixers can be extended beyond their trained length.
Proof. This is a direct consequence of Proposition 2.3 to each partition enable the computation of matrices larger than those encountered during training.
the pretrained parametric functions assigned
We identify data dependency - the property induced by SAM - as key axis of differentiation amongst existing models. Although data ditttdddepddnncy is a popular notion that is widely regarded as being crucial to performance of Attention, it lacked any formal definition the literature. Consequently, various works have adopted different interpretations of this notion: Hyena [3] implements it via data- dependent linear operators over the Toeplitz matrix mixer; GSS (26 adds data-dependent gate post sequence mixing: LA, SSD 22 6 like Attention, directly map input data to the parameters of the matrix mixer. We adopt the third notion of data dependency, where each parameter is function of a particular input token, and under this definition it is clear that models like Hydra, SSD, and LA are data-dependent, whereas MLP-Mixer, FNet, S4, S4D, and TNN, only have data- data-independent parameters.
2.3 Prior Sequence Models as (Structured) Matrix Mixers
Using the formalization of the Matrix Mixer framework, categorize wide array of previous works
MI.-Mixxr [40), Transformer (42). Linear Attention (22) Linformer (45). S4 [18] 17 1, H3 (14),
TNN 33 CKCony (37.. FNet [23) Kaleidoscope 88 7 Monarch (5) 13]. and Mamba 116 6 as
matrix mixers in Table For illustrative purposes, we explicitly show that MLP-Mixer, FNet and LA
matrix mixers (prools in Appendix(),. and leave out normalizing factors for simplicity follows:
Proposition 2.5. (MLP-Mixer a matrix mixer). MLP-Mixe. employs fx as an identity function, and its mixer matrix M has unstructured parameterization with single head (H [H=1) 1).

---

Proposition 2.6. (LA is structured matrix mixer with sequence alignment). LA employs fx as linear projection fx fx(X) XXV, with its low-rank mixer matrix being SAM. The mixer matrix M has H heads with structured
parameterization. parantttrd where each (3.j)-elementt mag is defined as:
mij = (
)())=elu())-
Wq. ERCDD
Additionally, it is important to note that not all structured matrix mixers inherently exhibit sequence alignment. For example, consider the matrix mixer formulation M = W where WQ.Wk< RR are trainable parameters. This formulation yields low-rank matrix mixer, which qualifies structured matrix mixer. However, this structure is independent of the input data, directly contradicting the concept of SAM. Another illustrative example is FNet:
Proposition 2.7. (FNet is a structured mairix mixer without sequence alignment). FNet employs fx as an identity function, functtin, using the canonical Vandermonde matrix- - Discrele Fourier Transform (DFT)
for its mixer matrix. The mixer matrix M has a single head H = =1) with structured paramtteratoan where each (p.q)-elemens myy is defined as mag- = PP, where w = 8
Using these definitions, we present a series of matrix mixers, incorporating various mixer matrices both with and without the SAM property, and provide extensive experimental evaluations to investi the influence of SAM on the expressivity of sequence mixers in Section 4.1.1
2.4 The Matrix Mixer Framework as a Creative Toolbox
Armed with this framework, we illustrate how more classes of performant sequence models can be derived. In particular, we introduce Vandermonde and Cauchy matrix mixers, chosen because they are: well-known families of structured matrices with sub-quadratic matrix multiplication, hence are effi- cient as sequence models, and 2) we show that they have SAM parameterizations. In Section|4. we validate that these properties matrix mixer is on par with LA. This validates the central theme of this paper that developing sequence models can be reduced to choosing structured matrix classes with target properties. We assume an input sequence X LRC is given, and utilize the prevalent query-key concept that we employ Q=XWº.
enough to create efficient and strong sequence models, e.g. our Cauchy
K = =XW k where Wo,WK K E CxD, For simplicity, we show the single-headed (H = 1) case, where D= P. Further details including the implementations of multi-headed extensions are in Appendix
Sequence aligned Vandermonde mixer. In contrast to FNet, which employs a -arammerr Vandermonde mixer based on the Discrete Fourier Transform (DFT), propose trainable, and data- dependent Vandermonde-basee matrix mixer. Then, we can construct mixer matrix from a combina- tion of two separate Vandermonde matrices Mq.Mxc each generated from Q and K. Specifically, each (i.j) (i.e)-element of the resulting mixer matrix M is formulated mij- = [c(ss((2r))
cos cosccottttt (2rk)aa) (a)) This formulation utilizes the cosine function the real part of an imaginary number- - which is well established technique used in SSMs 17 28 to prevent the potential for excessive values resulting from the powers of elements. Thanks to the associative property of matrix multiplications, our sequence aligned Vandermonde matrix mixer enjoys algorithms with efficient computational complexities.
Sequence aligned Cauchy mixer. As indicated by the definition of Cauchy matrices in Table the underlying philosophy of using both the Cauchy and low-rank matrix classes as a sequence mixer is remarkably similar: the relevance of pair of elements is directly proportional to their associated value - the greater the relevance, the higher value assigned. Specifically, each (f,,)) (elemet of the mixer matrix M is constructed by mij Ea 1/aa/ -/auu-ka++), where e is a trainable constant that prevents elements leading to excessive values when the denominator yoa - kya approaches zero. To our knowledge, this is the first introduction of Cauchy- based sequence mixer that achieves performance comparable Attention matrix mixers.
3 Hydra: The Double-Headed Mamba
In this section, we introduce Hydra, novel sequence-to-sequence model through a bidirectional extension of Mamba. We briefly begin with a preliminary background that SSMs are semiseparable matrix mixers (Section 3.T) Then, we motivate the design choice of Hydra under the domain of the matrix mixer framework. Specifically, we opt for quasiseparable matrices as our matrix mixer, choice grounded by solid mathematical foundations (Section 3.2) Additionally, we underline the computational benefits and enhanced parameter efficiency aftorded by adopting quasiseparable matrices (Section 3.3
5

---

8g
cj
dl
2TAt,b, of ofAt, drAt.b.
Ro
dik:
8,
Aiahi
Alab
Mi
dito
qik: 9:4: 4ik
9
Rg
6,
i Gik.
4ikk
954g
45
6g
2b
aito gike gikk
vik.
cZk.
qikk
Q.R.
(SA e dl C. 8,
(a) Semissearrblle
(b) Quasiseparable
(c) Additiooo of Two Separate SSMs
(d) Hydra
Figure 2: (a) A semiseparable (SS) matrix. (b) A quasiseparable (QS) matrix. (c) A mixer matrix of adde bidirectional SSMs. (d) A QS mixer matrix for Hydra. SS and QS matrices are characterized by rank conditions (Definition 3.1 Definition (.2). The rank characterization of SS matrices include the diagonals (e.g,, green submatrices), whereas that of QS matrices hold for off- diagonal submatrices (e.g., yellow submatrices). Because of the similar rank properties, naive addition-basse bidirectional SSM is provably a QS matrix mixer. Hence, QS matrix mixers generalize this common heuristic for bidirectional SSMs. The freedom in the diagonal values of Hydra leads to higher expressivity compared to the mixer matrices of the addition-based bidirectional SSMs, where the diagonal values constrained by the colored vectors.
3.1 Background: State Space Models are Semiseparable Matrix Mixers
SSD (6), the latest advancement in the iterations of SSMs, presents a sub-quadratic sequence mixer that attains language modeling performance on-par with Attention. Crucially, SSD underscores that all SSMs are inherently parameterized by semiseparable matrices, which play pivotal role in their computational efficiency and strong empirical performance. Specifically, the operational essence of selective SSMs such as Mamba the transformation of an input sequence < ERLxC, + Y Y RLxC can be succinctly represented within the matrix mixer framework, as detailed below:
Y1 C. A A,X,
I A
fori>j
Y (= MX May myy C A. A.-AA)
(I)
x 1.
fori=j-
(2)
for
Here, may represents an (ele)-element of the mixer matrix M L with each matrix A, NxN and vector b, E IR Nx as time-varying parameters of selective SSMs. This formulation reveals that the mixer matrices M follow a fundamental class of structured matrices known as semiseparable matrices, defined as follows:
Definition 3.1 (The rank characterization of semiseparable matrices). A lower triangular matrix M is N semiseppaeaelee iff any submatrixfron the lower triangle (on or below the diagonal) has a rank of at most N. See Figure 2 (a).
However, a key limitation of these matrices - and by extension, SSMs - is their inherent causality, which restricts their use in scenarios where bidirectional processing is vital. To circumvent this limitation, previous efforts [44] [5] [2] have explored employing two separate SSMs, one for forward and the other for backward sequence processing, then combine the outputs using strategies like element-wise addition, the Hadamard product, concatenation. Among such heuristics, can be conceptualized within our matrix
addition- based bidirectional extensions of SSMs [15 48 38 mixer framework, illustrated in Figure (c).
3.2 Quasiseparable Matrices: A Principled Bidirectional Matrix Mixer
We fully utilize the matrix mixer framework, which is discussed in Section to explore novel bidirectional sequence mixer and identify quasiseparable matrices as a prime candidate. Our exploration focuses on structured matrix classes that meet the following criteria: 1) they feature upper triangular components for bidirectionality, 2) they possess strong expressivity, and 3) they benefit from sub-quadratic matrix multiplication algorithms.
The structural design of quasiseparable matrices inherently meets the first criterion, which is defined as follows: a matrix M is N quaasssseparrle if each element may satisfies
if
mij
ifi=j,
bjj
ifi<j
where each 8, is a scalar, b E Nx1 and A, E NXN [4). Clearly, this matrix class features non-zero upper triangular components, enabling bidirectional processing.

---

Furthermore, the second requirement - the expressivity of quasiseparable matrices is confirmed by their rank characterization:
Definition 3.2 (The rank characterization of quasiseparable matrices 30 A matrix M is N-quasiseparable iff any submatrix from either the stricily upper or lower triangle (off from the diagonal) has rank of at most N. See Figure
This definition emphasizes the rank constraint inherent in quasiseparable matrices, which is also evident from Equation (3) given that each vector c, b, E R- x1 and matrix A; E NNN has a rank of at most N. This structural flexibility of quasiseparable matrices directly leads to significant generalizations, extending the capabilities of both low-rank and semiseparable matrices.
Corollary 3.3. Quasiseparable matrices generalize low-rank matrices.
Corollary 3.4. Quasiseparable matrices generalize and extend semiseparable matrices.
Additionally, we revisit the previous Iaditiin-basee bidirectional extensions of SSMs 15
48
38
through the lens of our matrix mixer framework. Unlike other elements, the diagonal values in mixer matrix embody a unique concept of residuals, serving as a critical aspect of model expressivity. As demonstrated in Figure (c). the mixer matrices in these bidirectional SSMs exhibit constraints in
their diagonal elements { b, +c b. which are directly influenced by the shared non-diagonal construction vectors { C. b. C b, L. Importantly, the rank characterization of semiseparable matrices includes on-diagonal elements, whereas that of quasiseparable matrices applies only off-diagonal submatrices. This generosity in the rank-based definition suggests that sequence models employing quasiseparable mixers not only offer inherent extendability in handling both causal and bidirectional processing, but also exhibit strong expressivity.
Corollary 3.5. Quasiseparable matrices strictly more expressive than mixer matrices of addition-bass bastdd bidirectional SSMs.
Leveraging this inherent flexibility of quasiseparable matrix mixers, our Hydra in Section 3.3 is defined by incorporating shift operations. Our experimental results strikingly confirm that this nuanced parameterization difference leads to a notable improvement downstream task performance, thereby substantiating our theoretical claims (see Appendix(D.]]
Moreover, with their structural similarity to semiseparable matrices, quasiseparable matrices confirmed as sequence aligned matrices. Given our experimental results that SAM parameterizations are the key to the strong representational power (Section (.11I). we further validate our choice of quasiseparable matrices for the bidirectional sequence mixer.
Proposition 3.6. N N-quasixeparable matrices are sequence aligned matrices.
Proof. quasiseparable matrices, due to their structural similarity semiseparable matrices, belong to the class Sequence Aligned matrices. Specifically, the set of parameters is given by and for each
P=p= (A,bb, C,, r.. We consider the partition II= {{A,} element of the partition set, we choose the bijection that maps token to A;, b; C2, and 6; respectively. Finally, it easy to see that the sub-matrin M[: + 1,: i +-+ indeed only contains parameters in the set (A,.b, b,, 5. thus satisfying the last requirement of SAM matrices.
In Section 3.3 we detail how quasiseparable sequence mixer can be effectively implemented using existing SSM frameworks to achieve sub-quadratic multiplication efficiencies, fulfilling the final criterion of our matrix mixer exploration.
3.3 Taming the Hydra
As direct consequence of the favorable mathematical properties of quasiseparable matrices, present a new sequence mixer Hydra. We adopt quasiseparable matrices as matrix mixers in Hydra, which bring forth three significant advantages: 1) higher representational power compared to its heuristic alternatives, 2) easy to implement sub-quadratic matrix multiplication algorithm, and 3) significant parameter savings.
We exploit the relationship between semiseparable and quasiseparable matrices to develop an seyy-to-impleent aub-quadratic matrix multiplication algorithm. Specifically, we recognize that quasiseparable matrices can be expressed as combination of two semiseparable matrices.
Proposition 3.7. Let R. be the input sequence, and let QS(-) and SS(-) denote the action of quasiseparable and semiseparable matrix respectively. Let the two matrices share the parameters
7

---

Table 2: Matrix mixer ablations. A systematic empirical study of matrix mixers on the GLUE
benchmark by controlling for the architecture and varying only the matrix parameterization. Sequence- aligned matrices dynamically parameterize via input projections, becoming data-dependent (DD) that significantly increases performance. Most DD variants achieve competitive GLUE scores.
Structure
DD #Paams
C4 Pretrain
GLUE Tasks
GLUE
Lee
Acc
MNLI
QNIJ
QQP
RTE
SST2
MRPC
COLA
STS
Avg
Dense
71M
2.05
59.6
73.3
76.2
85.3
64.4
90.8
84.7
45.7
76.8
74.7
Toeplitz
71M
1.97
60.8
74.6
79.6
86.6
66.1
90.9
84.2
45.7
79.1
75.8
Toeplitz
72M
1.91
61.9
77.3
81.8
88.1
67.1
90.7
87.3
45.3
84.1
77.7
DFT
71M
2.46
53.1
70.4
70.8
84.5
59.9
89.8
83.6
44.4
69.8
71.7
Vandermonde
71M
2.46
53.0
55.2
61.3
82.5
66.3
87.4
84.2
45.8
84.2
70,8
Vandermonde
70M
2.04
59.7
74.1
80.0
86.2
67.9
89.3
84.3
46.0
80.1
76.0
Cauchy
71M
2.25
56.2
75.3
81.3
66.6
88.8
84.5
27.4
83.2
74.2
Cauchy
70M
1.94
61.6
77.5
84.4
84.2
68.0
86.7
48.1
85.2
78.2
Low-Rank
71M
2.06
59.4
73.7
76.5
85.1
61.5
90.6
85.8
49.2
76.6
74.9
Low-Rank
V
70M
1.90
62.2
77.6
84.1
88.2
69.1
91.0
85.9
47.6
83.9
78.4
Attention
71M
2.08
59.1
71.0
70.4
83.5
62.3
89.9
83.3
49.6
65.2
71.9
Attention
70M
1.87
62.9
78.5
85.4
88.4
67.9
91.2
86.4
47.8
84.3
78.8
Quasiseparable
72M
2.03
59.8
73.8
78.1
87.1
64.3
90.2
84.4
45.5
77.2
75.1
Quasiseparable
71M
1.84
63.3
79.5
85.5
88.6
69.8
91.9
88.4
48.4
85.6
79.7
(A, bi, C. and define D diag(01, 5L), where 8, are the diagonal parameters of the quasiseparable matrix. Then,
shiftsssi((sSS())) (HSNSSSXXXX)))HHiissSCsCCCCCCCCA
where ( denotes the operation that reverses the input sequence, while shiff(-)) refers to shifting the sequence right by position. padding the beginning with zero. (Proof in Appendix
The above proposition demonstrates that quasiseparable matrix multiplication be decomposed into two operations of semiseparable matrix multiplication with simple functions such as flip and shift. Given that semiseparable matrix structure encompasses SSMs, this flexibility allows for the selection of any SSM variant for implementation. In this paper, we employ SSD (6), chosen for its linear-time and dedicated hardware-efficient implementation. However, the architecture of Hydra is compatible and can also generalized with other recurrent models (47
with a variety SSMs [88] 17
16
Furthermore, Hydra significantly improves parameter over the heuristic approaches to bidirectional modeling using SSMs [1]
44 12]. For example, some approaches utilize two distinct SSMs, which doubles the number of training parameters. In contrast, since we conceptualize the model as quasiseparable matrix mixer (see Figure we naturally share the fxx projection layer, which accounts for bulk of the model size. Empirically, we observe only marginal increase in the total number of parameters compared to a single SSM, and can cut the the number parameters nearly in half compared to the heuristic approaches to bidirectionality.
48
4 Experiments
we begin by analyzing the matrix mixer framework through extensive performance comparisons between different structured matrix classes (Section 4.1.1p. The data-deeendent parameterization of Quasiseparable matrices surpasses the performance of all other matrix classes, validating our selection of it as the sequence mixer. Furthermore, we compare our method against other ad-hoc solutions that extend SSMs to acquire bidirectionality (Appendix [DI),, and show that our Hydra outperforms them, underscoring the utility of the matrix view of sequence mixers.
In Section 4.1
Then, in Section 4.2 we validate the effectiveness of our method by evaluating Hydra on guintessential language and image benchmark tasks: Masked Language Modeling (Section 4.2.1 and Image Classificattin (Section (.22.2). State-of the-te- cene-arrteefformaae performance in both tasks has generally been dominated by wansformer-basel models [10 [T]. Our Hydra serves efficient and powerful replacement to the transformer layer, outperforming it on both tasks.
In the presentation results across all tables, the highest performing scores are highlighted in bold, while the recond-hights scores marked with an underline. Each number is the average of five runs.

---

Table 3: Performance of various approaches that ex- tend Mamba to a bidirectional model. We compare our quasiseparable matrix mixer to element-wise addition (Add), the Hadamard product (Mult), and concatenation (Concat) variants.
2
Concat
Mult
1.9
Add
Quasi
Method
#Params
C4 Pretrain
GLUE Avg
1.8
fee
Acc (%)
1.7
Mamba
68M
3.45
39.5
77.7
Add
70M
1.68
65.6
80.6
10
30
40
50
60
70
Mult
70M
1.72
64.9
79.6
Training Steps (K)
Concat
69M
1.71
65.4
81.1
Figure 3: Cross-entropy loss of various bidirec- tional variants, measured on the C4 validation set.
Quasi
70M
1.66
65.9
81.7
4.1
Analysis of the Matrix Mixer Framework
4.1.1
Effects of Different Structured Matrix Families
Our controlled experimental setting distinctly separates the mixer matrices from other architectural components, enabling a rigorous and focused comparison between different types of mixer matrices. Specifically, utilizing the recent Mamba [6] block, only replace SSD with different mixer matrices M. In 2 we provide experimental results that showcase the expressivity of various matrix mixers that support bidirectional sequence processing, primarily by utilizing off-the-shelf structured matrix families for the mixer matrix. Further details of the experimental setup are provided in Appendix
The results distinctly highlight the substantial advantage in expressivity conferred by the SAM property (Definition 2.2). Previously, the high performance observed in [22 45) was attributed to their low-rank based mixer matrices, particularly emphasized by the query-kep QK interaction. The results show that the structures of mixer matrices affect the capability of sequence models. However, demonstrate that the key factor that primarily contributes significant improvements in the expressivity of sequence models is not the query-key formulation, but rather the SAM property. Through our systematic extension, we adapt six structured matrix families Toeplitz, Vandermonde, Cauchy, low-rank, Attention, and quasiseparable - to include the SAM property (see Appendix []ror implementation details). This adaptation reveals that the seuuence-alignee Cauchy variant performs nearly as well as the sgquence-aligned low-rank variant, which characterizes LA. Moreover, the experimental results consistently indicate that variants equipped with the SAM property outperform those lacking this feature across all tested matrix families.
Results.
4.1.2 Ablating Approaches for Bidirectionality
We compare the quasiseparable matrix mixer approach to prior bidirectional SSMs models that achieve bidirectionality by aggregating forward and backward SSMs using various heuristics including
element-wise addition [15), the Hadamard product (44), and concatenation 12]. We provide
details of the ablation studies in Appendix(D. D..
Results. The results presented in Table the shortcomings of the unidirectional Mamba (16) when applied to bidirectional contexts, as evidenced in C4 and GLUE. This sigmificant performance disparity (-4.0) underscores the essential need for models to be capable of bidirectional sequence processing. Within the comparison four bidirectional variants shown in Table 3, our approach of utilizing quasiseparable matrix achieves the top validation results on the C4 benchmark and records the highest GLUE average score of 81.7. This advantage of our method further validated in Figure where it demonstrates the cross-entrope loss on the C4 validation set throughout training. Considering the gigantic size of the dataset, the expressivity of Hydra using quasiseparable matrices is clearly manifested by consistently achieving the lowest loss, as well as the highest masked token prediction accuracy.
4.2 Evaluation Results Hydra
4.2.1
Bidirectional Masked Language Modeling
We pretrain our models the masked language modeling objective using the Colossal Cleaned Common Crawl (C4) corpus (36), then finetune and evaluate them on the GLUE benchmark (43. We relegate experimental details in Appendix D.2

---

Table 4: GLUE Results. Evaluation of various sequence models that can be formulated as matrix mixers. For maximum performance, all models are trained using established recipes [2] 13..
Method
#Params
C4 Pretrain
GLUE Tasks
GLUU
Lee
Acc (%)
MNLI
QNLI
QQP
RTE
SST2
MRPC
COLA
STS
Avg
BERT
110M
1.59
67.3
84.4
90.3
89.7
77.1
92.3
90.7
54.2
89.1
83.5
MLP-Mixer
112M
1.77
63.5
77.2
82.4
87.6
67.3
90.5
86.5
43.0
85.2
77.5
FNet
112M
1.94
61.3
74.9
82.1
85.7
63.6
87.6
86.4
42.7
83.1
75,8
M2
116M
1.65
65.9
80.5
86.0
87.0
69.3
92.3
89.2
56.0
86.9
80.9
Hydra
112M
1.46
69.1
84.5
90.0
91.3
77.5
93.5
91.2
57.2
88.9
84.3
Results. The results in Table showcase that Hydra outperforms all existing ofate-of-the-art -
methods. Notably, Hydra surpasses the performance of BERT-trained with the latest HuggingFace recipe [46]- -in in both pretraining and GLUE benchmark scores. BERT has shown a noticeable gap in the masked language modeling task compared to other previous methods (40 23 13 Hydra gains 1.8% improvement in accuracy of C4 validation and 0.8% increase in the average GLUE score, illustrating the effectiveness of leveraging matrix mixer view for bidirectional settings.
4.2.2 Image Classification
Table 5: Top 1 & 5 image classification accuracies eval- uated on the ImageNet- IK benchmark. We also report accuracies using the common model ensembling tech- nique: Exponential Moving Average (EMA) weights. (Top) Reported from literature (27 31. (Bottom): Our
We assess Hydra on the renowned ImageNet- IK benchmark [9]. which comprises 1.28M training images and 50k validation images across 1, 000 categories. We use the ViT- Base model as a baseline to facilitate rigorous comparison of various sequence mixers by substituting the Transformer block in ViT with alternative sequence mixer mod- els. specifically S4ND (27). Hyena (31), Mamba [16]. and our proposed Hydra model. Unlike many off-the-shelf models such as CNN-based [19 25] and vision- specializee Transformers [24] that include additional
unidirectional and bidirectional Mamba results.
Top-1 (%)
Top-5 (%)
Method
#Params
Acc
ACCEMA
Ace
ACCEMA
VÎT-B S4-V
87M
78.8
80.6
94.2
95.2
89M
79.4
80.4
94.2
95.1
Hyena-ViT-a ViT-B
88M
78.4
76.4
94.0
93.0
Mamba-ViT ViT-B
89M
79.1
80.0
94.2
94.9
Hydra- ViT-B
91M
81.0
81.6
95.3
95.6
techniques such as hierarchical spatial downsampling to boost accuracy. our approach involves substituting only the sequence mixer layers within the ViT architecture. In addition, as opposed to other baselines in the setting of [3 our method uses no tuning over the default ViT recipe except for droppath. We found that Hydra fits the training data noticeably better than ViT, perhaps due to
better expressivity and inductive bias, so we simply increased droppath from 0.3 0.5 stronger regularization. We relegate further experimental details in Appendix D.2
The results, as presented in Table 5 compare the performance of Hydra with ViT and other variants [27] 31 on ImageNet-IKK IK. Hydra exhibits superior performance in image classification, outperforming ViT by 2.2% in Top-1 accuracy and 1.1% in Top-5 accuracy. Remarkably, even though Hydra simply flattens images without incorporating any specific 2D architectural adjustments, it still surpasses S4ND (27] which is specifically tailored for image processing by a notable margin of 1.6% in Top-1 accuracy. This showcases the versatility and effectiveness of Hydra in handling diverse data types.
Results.
Conclusion
In this work, we have explored common paradigm for sequence models wherein the sequence mixer can be represented by a matrix. This framework encompasses many well-known models such as MLP-Mixer, FNet, convolutions, Transformers (softmax attention), and recent state-space models such as Mamba. By formalizing the matrix mixer framework and exploring additional matrix variants, we have identified a key axis of variation (sequence alignmen() in matrix parameterizations, which enables benefits such as data dependence. This, in turn, provides increased flexibility and stronger performance for sequence models. Furthermore, we have leveraged the matrix mixer framework to motivate natural bidirectional extension of state space models called Hydra, which can be formulated as quasiseparable matrix mixers. Hydra consistently outperforms unidirectional Mamba and other bidirectional sequence models in tasks such as masked language modeling and image classification.
We adjust the learning rate to le to address instabilities observed in the training of FNet COLA.
10

---

Acknowledgements
This research was made possible by the generous support of computational resources provided by Cartesia AI.
References
[1] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. "Zoology: Measuring and improving recall in efficient language models". In: arXiv preprini 2trXix::312.0492.. (2023).
[2] Simran Arora, Sabri Eyuboglu, Michael Zhang. Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. "Simple linear attention language models balance the recall-throughhut tradeoff". In: arXiv preprint arXiv:2402. 2ee::2442..86600 18668 (2024).
[3] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. "Neural machine translation by jointly learning to align and translate". In: arXiv preprint
(2014).
[4] Tom Bella, Yuli Fidelman, Israel Gohberg, and Vadim Olshevsky. "Computaiions with quasisep- arable polynomials and matrices". In: Theoretical Computer Science (2008).
[5] Tri Dao. Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. "Monarch: Expressive structured matrices for efficient and accurate training". In: International Conference on Machine Learning. PMLR.
2022. pp. 4690-44947 4721.
[6] Tri Dao and Albert Gu. "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality". In: International Conference Machine Learning (ICML). 2024.
[7] Tri Dao, Albert Gu, Matthew Fichhorn, Atri Rudra, and Christopher Ré. "Learning fast algo- rithms for linear transforms using butterfly factorizations". In: International conference on machine learning. PMLR. 2019, pp. 1517-1527.
[8] Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. "Kaleidoscope: An efficient, learnable representation for all structured linear maps". In: arXiv preprint arXiv:2012. 14966 (2020).
[9] Jia Deng, Wei Dong. Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. "Imagenet: A large-scale hierarchical image database". In: CVPR. 2009.
[10] Jacob Devlin, Ming-We Chang, Kenton Lee, and Kristina Toutanova. "Bert: Pre-training of deep bidirectional transformers for language understanding". In: NAACL (2019).
[11] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. "An image is worth 16x 16 words: Transformers for image recognition at scale". . In: ICLR (2021).
[12] Yassir Fathullah, Chunyang Wu, Yuan Shangguan, Junteng Jia, Wenhan Xiong, Jay Mahadeokar, Chunxi Liu, Yangyang Shi, Ozlem Kalinli, Mike Seltzer, and Mark J. F. Gales. "Multi-Head State Space Model for Speech Recognition". In: Proc. INTERSPEECH 2023. 2023, pp. 241-245.
DOT:|10 /01100.20 21437/ Interspeech 2023- 1036
[13] Daniel Y Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. "Monarch Mixer: A simple sub-quadratic GEMM-based architecture". In: NeurIPS (2023).
[14] Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christo- pher Ré. "Hungry Hungry Hippos: Towards Language Modeling with State Space Models" In: The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL: UBL::htttp::/opperrrriiee http:/rwwwwiiieeeerrreeeeeeeee net/pafffd COZDyOWYGg
[15] Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. "It's 's raw! audio generation with state-space models". In: ICML. 2022.
[16] Albert Gu and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces".
In: arXiv preprint 23Xiv:23/2..07552 (2023).
[17] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. "On the parameterization and initial- ization of diagonal state space models". In: NewrIPS (2022).
[18] Albert Gu, Karan Goel, and Christopher Ré. "Efficiently modeling long sequences with struc- tured state spaces", In: ICLR (2022).

---

[19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. "Deep residual learning for image recognition". In: Proceedings the IEEE conference on computer vision and pattern recognition. 2016. pp. 770-778.
[20] Peter Izsak, Moshe Berchansky, and Omer Levy. "How to train BERT with an academic budget".
In: EMNLP. 2021.
[21] Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. "Repeat after me:
Transformers are better than state space models at copying". In: arXiv preprint arXiv:2402.01032 (2024).
[22] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. "Transformers are rans: Fast autoregressive transformers with linear attention". In: ICML. 2020.
[23] James Lee-Thhree Lee-Thor,, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. "Fnet: Mixing tokens with fourier transforms". In: NAACI. (2022).
[24] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. "Swin transformer: Hierarchical vision transformer using shifted windows". In: ICCV. 2021. [25] Zhuang Liu, Hanzi Mao, Chao Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. "A convnet for the 2020s" .. In: Proceedings of the IEEE/CVE conference on computer vision and pattern recognition. 2022. pp. 11976-11988 1986.
[26] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. "Long Range Language Modeling via Gated State Spaces". In: The Eleventh International Conference on Learning Representationss 2023, URL: http:://ppereeierreeee
[27] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Precy Shah, Tri Dao, Stephen Baccus, and Christopher Ré. "S4nd: Modeling images and videos as multidimensional signals with state spaces". In: NeurlPS (2022).
[28] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. "Resurrecting recurrent neural networks for long sequences". In: International Conference on Machine Learning. PMLR. 2023, pp. 26670-26698.
[29] Victor Y Pan. Structured matrices and polynomials: unified superfast algorithms. Springer Science & Business Media, 2012.
[30] Clément Pernet, Hippolyte Signargout, and Gilles Villard. "Exact computations with qua- siseparable matrices". In: Proceedings the 2023 International Symposium on Symbolic and Algebraic Computation. 2023. pp. 480-489.
[31] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. "Hyena Hierarchy: Towards Larger Convolutional Language Models". In: JCML. 2023.
[32] Jacob Portes, Alex Trott, Sam Havens, Daniel King, Abhinav Venigalla, Moin Nadeem, Nikhil Sardana, Daya Khudia, and Jonathan Frankle. "MosaicBERT: A Bidirectional Encoder Opti- mized for Fast Pretraining". In: arXiv preprint arXIv:2312.17445 (2023).
[33] Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran Zhong. "Toeplitz Neural Network for Sequence Modeling". In: ICLR. 2023.
[34] Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. "The devil in linear transformer". In: EMNLP (2022).
[35] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. "Language models are unsupervised multitask learners". In: OpenAl blog 1.8 (2019), p. 9.
[36] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena. Yanqi Zhou, Wei Li, and Peter Liu. "Exploring the limits of transfer learning with a unified text-to-text transformer" In: JMLR (2020).
[37] David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomezak, and Mark Hoogen- doom. "Ckconv: Continuous kernel convolution for sequential data". In: arXiv preprint arXiv:2102.0244 (2021).
[38] Yair Schiff, Chia-Hsiang Kao, Aaron Gokaslan, Tri Dao, Albert Gu, and Volodymyr Kuleshov. "Caduceus: Bi-directional equivariant long-range dna sequence modeling". In: arXiv preprint arXiv:2403.03:: 033334 (2024).
[39] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang. and Furu Wei. Retentive Network: A Successor to Transformer for Large Language Models.
2023. arXiv: 2307. 08621 [cs. CL.
12

---

[40] Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai. Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, al. "Mlp-mixer: An all-mlp architecture for vision". In: NeurlPS (2021).
[41] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. "Training data-efficient image transformers & distillation through attention". In:
ICML. 2021.
[42] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. "Attention is all you need". In: NeurIPS (2017).
[43] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. "GLUE: A Multi- Mulki-Task Benchmark and Analysis Platform for Natural Language Understanding".
In: ICLR. 2019.
[44]
Junxiong Wang, Jing Yan, Albert Gu, and Alexander Rush. "Pretraining Without Attention". In:
EMNLP. 2023.
[45]
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. "Linformer: Self-attention with linear complexity". In: arXiv preprint arXiv:2006. 04768 (2020).
[46] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac. Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. "Transformers: State-of- the-Art Natural Language Processing". In: EMNLP. 2020.
[47] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. "Gated linear attention transformers with hardware-efficient training". . In: arXiv preprint arXiv:2312.0663 (2023).
[48] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. "Vision mamba: Efficient visual representation learning with bidirectional state space model".
In: arXiv preprint 24Xiv:2401.094/ (2024).
13

---

A Discussions
A.1 Limitations
In this section, we share two limitations of our work, namely 1) pmprresentationnComaaauat Tradeoff, and 2) Hardware Efficiency.
Tradeoff. While structured matrix mixers are computationally more efficient than their dense matrix mixer counterparts like softmax attention, they are also representa- tionally less expressive, which may be seen a limitation of these methods. For instance, concurrent works [ 21 have begun investigating the representational power SSMs by analyzing their performance on memorization-cce centriic tasks. They report that SSMs with fixed model capacity are eventually outperformed by softmax attention for longer sequences. This can viewed as consequence of the matrix being too structured, and hence too inexpressive for the problem.
pnrresentation-Complun
On the other hand, we remark that the degree of structure of structured matrix is a knob that can be tuned according the specific task, that we can tradeoff the computational efficiency of the method for larger expressivity. For instance, within the structured matrix class of low rank matrices, we tune the rank upto the size of the matrix, which is the sequence length. As the rank of the matrix class increases so does its expressive power; however, at the same time it also diminishes the compute efficiency associated with the matrix being low rank.
As another example, we demonstrate this tradeoff on the performance of SSMs on retrieval-style tasks for SSD, which the modern variant of Mamba. Specifically, we show that SSD able recover the accuracy attained by softmax attention once we control for the compute capacity of the model. In contrast, hardware limitations of the selective scan algorithm make it impractical to match the compute capacity in Mamba, explaining the emerging findings from I 2 21] that SSMs underperform on memorization-centric tasks. This makes it evident that the deve revelopment and analysis of SSMs an active area of research with substantial room for exploration and improvement.
Despite the fact that structured matrices have associated sub--quadratie matrix multiplication algorithms, their implementation may not hardware- fardware-friendly, which can reduce the execution speed in practice. For instance, one of the core advantages of Transformer is the hardware-friendly nature of its architectural design, mainly composed of pure matrix multiplications.
Hardware Efficiency:
A.2 Broader Impact
This paper presents work whose goal is to advance the field of Machine Learning. The primary objective of this work is to introduce novel component designs that are universally applicable across various tasks. Consequently, we believe that this work is unlikely to have significant socictal impacts that necessitate special emphasis within this context.
B Proofs
Proposition 2.5. MLP-Mixer is a Matrix Mixer
Proof. Recall that MLP-Mixer applies a weight-shared per-channel linear projection on the input sequence. Since MI.PMMxee does not pre-process the input sequence, we set fx to be the identity function. Additionally, since the mixer matrix is data independent and learned as a free parameter, we define the class M = R L, and set (X:():)) = e for 0 E θ, where θ R L Since MLP-Mixer shares the matrix across all channels, the number of heads H = 1. The output is given by Y = ()) where is the learned mixer matrix, which matches with the functional form of MLP-Mixer's sequence mixing primitive.
Proposition 2.6. LA is a Structured Matrix Mixer
Proof. Recall that for each head à [ LA computes,
Y = ("=)(\"<<<>>>>
where Q² Khh v and projections of the input. Since LA preprocesses the input sequence via projec- tion, we set fx = Wy ERODD with W, v as learned parameter. The class of structured matrices M
For simplicity, we overload the notation to denote finite ciatte-dmensional linear map with its corresponding matrix.
14

---

is clearly the class of Low Rank matrices. We the set θ (D= Hid to correspond to the learned parameters WQ. W for the (Q",K*) matrices, where is the internal dimension for keys and queries, We set the matrix-generating function l'm to compute M as M = (o (K*) through linear projections, followed by applying element- nonlinearity a, and matrix multiplication. The output Y M (fx (X)) matches the expected functional form of LA's sequence mixing primitive.
Proposition 3.2. Let X E FLXDD be the input sequence, and let QS(-) and SS(-) denote the action a Quasiseparable and Semiseparable matrix respectively. Let the matrices share the param- eters C; L, and define D diag (oog((ó1""" (,1,"1,51)) ,8L), where 8, 's are the diagonal parameters of the Quasiseparable matrix. Then,
ffip (hift(((((ti(()))) +DX,
+DX,
where flpp(-) denotes the operation that reverses the input sequence, while shifi(-) refers to shifting the sequence right by one position, padding the beginning with zero.
Proof. The proof follows by observing that the first term, shift(SS(X)), models the effect of the lower triangular part of the QS matrix. The next term, (ip(shift(SS) flip( X)))). represents the upper triangular part of the QS matrix, and the last term, DX, accounts for the QS matrix's diagonal.
def hydra(
X,
H+P)
A ):
# (8,) Parameter
norm
x.b flip(x, dim-1)
(B,L,
dt_f, dt_b proj_dt(x), proj_dt (x_b)
*
(B,L, y-f = SSD(
Y
fD
discretize_ACA dt.f). diecretize_bc(z, dt_f), (B,L,1)
(B,
)
y.b SSD(
x.b.
discretize_ACk, dt_b), (B,L,H)
M
M
discratiza_bc(x.bb dt_b), (B,1,
y-f ehirt(y_1 I, din-1) flip(shift (y_b, dim-1), dim-1)
A
y_f y_b
repeat(
proj_D(x), B B
conv bc
conv X
conv bc
return
bc
fc
bc
Figure 5: Pseudo code for Hydra. B,L,H,P denote
batch size, sequence length, number of heads, and head dimension respectively. The suffices _f and bb denote forward and backward. shift: Right- shift.
X
Figure 4: Detailed illustration of Hydra.
C Architectural Details
Here, we provide architectural specifics of our method Hydra. As shown in Figure and previously discussed Figure Hydra is largely divided into two componmmpnnnns, each for forward and backward sequence processing respectively. A key architectural difference from the unidirectional Mamba [16) and SSD (6) models is the type of convolutional layers employed. Unlike Mamba and SSD, which utilize causal convolutions due to their unidirectional nature, Hydra incorporates standard convolutional layers, reflecting its capacity for bidirectional sequence processing.
D Additional Details
D.1 Details of Ablation Studies in Section 4.1
Bidirectionality. We provide details of the experiments in Each method under consideration employs 12 layers. As depicted in Table due to shared projection layers for forward and backward
15

---

SSMs, Hydra only requires an additional 2M parameters compared to its unidirectional counterpart. Moreover because of the parameter increase the concatenation variant, we lower its hidden dimension to match parameters, As [CLS] token is typically placed at the start of sequence, we add an additional technique only to Mamba, which substitutes a [CLS] token with global average pool of tokens.
D.2 Details of Main Results in Section 4.2
Training Details of Hydra We align our setup with the BERT-Base [0] architecture, which consists of 110M parameters across 12 transformer encoder layers. To ensure a fair comparison, MLP. Mixer [40). FNet (23), and M2 [13] are configured with 12 layers, with the number of channels adjusted to match the parameter count of BERT. Similarly, Hydra is structured with 23 layers to align with the
total number of parameters of BERT. We follow the well-established BERT training recipes [32 13] for optimal performance of each method. The specific hyperparameters for reproducing the results in aar reported in Table and the settings used for obtaining the results of Hydra in Table are listed in Table
GLUE
We pretrain on the C4 dataset, and follow the recipe from MosaicBERT 32 which has
been widely adopted in recent works- the models are trained for 70k steps with batch size of 4096.
We adopt the same hyperparameters as M2 [3) for pretraining, including the optimizer, learning rate schedule, and the best-base uncased tokenizer. We employ the controlled setting from (20) for finctuning models on the GLUE benchmark. Specifically, we utilize weights pretrained on the C4 corpus for the five tasks Masks-MNLI, MNLI, QNLI, QQP, SST2, and COLA and finetune RTE, STSB, and MRPC from an MNLJ checkpoint. For reliability, the performance metrics reported represent the average scores from five runs for each GLUE task, with seeds selected completely at random.
The well-established BERT and M2 models benefit from highly optimized training and finetuning recipes (20 [3] Therefore, we perform short sweep for the learning rate and number of epochs for finetuning tasks to obtain the best results. To note, we ensure that the number of epochs do not surpass their original values for fairness. In Tarse row, we observe that Hydra of 24 layers trained using the M2 recipe out-of -ot--ofthe-boo bbx outperforms both BERT (83.5) and M2 (80.9) on the GLUE benchmark without any hyperparameter tuning. From the tuned hyperparameter, Hydra with 23 layers (second row) gains additional improvements, outperforming the results obtained by using the M2 recipe.
In Table S
we report extended results for our GLUE comparisons (Table 4). including reference numbers directly reported from the literature. We include these to show that our results are a fair comparison of different models using our anternally-consister training framework, and that our baselines are strong and generally consistent with those found in MosaicBERT 32).
ImageNet-IK. For T we use 12 layers for the Transformer (42)
S4ND 27 and Hyena (3]
blocks, and set to 24 layers for Mamba and Hydra to match the number of parameters. The training involves images of 224 x 224 resolution, tokenized using a 16 x 16 patchifying layer. For Hydra, we adopt a row-major ordering to flatten patches and deliberately exclude positional embeddings. The models are trained over 300 epochs, leveraging hyperparameters primarily sourced from a standard training recipe (24). including an initial learning rate of 1e 4 and a batch size of 1024. We also incorporate regular ImageNet- IK training techniques [41). such as augmentation and regularization strategies. All other configurations remain consistent with the original ViT [1]] setup, ensuring that any observed differences in performance can be directly attributed to the expressivity of the respective sequence mixer layers.
E Additional Details of Ablation Studies in Table 2
To rigorously demonstrate the representational power of various families of matrix mixers, we ensure a fair comparison by meticulously controlling architectural hyperparameters. Specifically, all variants are constructed on the SSD (6) block, with a consistent configuration of 12 layers, an expansion factor of and a hidden dimension (d_oode) of 768. We then adjust the channel dimensions for qk_ ok_dim and headdim to keep the total parameter counts close to 70M. The models are pretrained for 24k steps in C4, each with a 4096 batch size. thus processing approximately 12.5B tokens in total. The finetuning phase adheres to the standardized protocol used for Hydra, ensuring comparability.
The primary differentiating factor among the variants is the mixer matrix M, which has two main configurable attributes: 1) the presence of the SAM property, and 2) the family of a mixer matrix. According to Definition 2.2 mixers with the SAM attribute contain parameters that are dynamically generated from elements of the input sequence, making them input data data-dependent (Proposition[2... In contrast, variants without SAM are data-inde iata-indeepndeet, with shared parameters across all inputs. The
16

---

Table 6: Hyperparameters of different recipes used for training GLUE benchmark tasks.
MNLI
QNLI
QQP
RTE
SST2
MRPC COLA
STS
LR
5e-5
le-5
3e-5
le-5
3e-5
8e-5
5e-5
3e-5
WD
5e-6
le-6
3e-6
le-6
3e-6
8e-6
5e-6
3e-6
epochs
3
10
5
3
3
10
10
10
seq_len
256
256
256
256
256
256
256
256
LR
5e-5
5e-5
3e-5
le-5
3e-5
5e-5
5e-5
7e-5
WD
5e-6
le-6
le-2
le-2
3e-6
le-2
5e-6
le-2
nLepochs
3
10
10
10
seq_len
128
128
128
128
128
128
128
128
LR
1e-4
5e-5
5e-5
1e-5
5e-5
8e-5
le-4
3e-5
Ours
WD
5e-6
le-6
3e-6
le-6
3e-6
8e-6
5e-6
3e-6
eeeeoohh 2
7
3
3
2
10
10
10
seq_len
256
256
256
256
256
256
256
256
Table Evaluation of Hydra on C4 dataset and GLUE Benchmark using different training recipes. Specific hyperparameters for reproducing the results are provided in Table
Recipe
#Params
Pretrain
GLUE Tasks
GLUE
Lee
Acc (%)
MNLI
QNLI
QQP
RTE SST2
MRPC COLA
STS
Avg
13
115M
1.45
69.3
83.7
89.7
89.7
77.4
92.8
91.5
54.7
90.1
83.7
Ours
112M
1.46
69.1
84.5
90.0
91.3
77.5 93.5
91.2
57.2
88.9
84.3
Table 8: Official GLUE benchmark results from the referenced papers [32
13
23
Method
Source
#Params
GLUE Tasks
GLUE
MNLI
QNLI
QQP
RTE
SST2
MRPC
COLA
STS
Avg
Ours
110M
84.4
90.3
89.7
77.1
92.3
90.7
54.2
89.1
83.5
32
110M
84.1
89.8
91.2
77.2
91.2
87.5
54.6
88.9
83.2
BERT
23
110M
82.5
91.0
87.0
69.0
93.0
83.0
73.0
89.0
83.3
131
110M
79.6
MLP-Mixer
Ours
112M
77.2
82.4
87.6
67.3
90.5
86.5
43.0
85.2
77.5
Ours
1 112M
74.9
82.1
85.7
63.6
87.6
86.4
42.7
83.1
75.8
FNet
23
83M
72.5
80.0
83.0
63.0
95.0
76.0
69.0
79.0
76.7
23
238M
77.0
85.0
85.0
69.0
94.0
88.0
78.0
84.0
81.9
M2
13
116M
80.5
86.0
87.0
69.3
92.3
89.2
56.0
86.9
80.9
Hydra
Ours
112M
84.5
90.0
91.3
77.5
93.5
91.2
57.2
88.9
84.3
inclusion of projection layers for data dependency results in a marginal increase in parameter count for models with SAM. Consequently, precisely configure qk_ jk__dim= 16, headdim 128 for SAM variants, and gk_diim- = 64, headdim 64 for non-SAM variants. For dense, Toeplitz (SAM and non-SAM), and Vandermonde DFT, we slightly adjust the hyperparameters to match the number of parameters.
To facilitate understanding and reproducability of these methods. we provide PyTorch codes in Figure Figure 7 Figure Figure 9 Figure I0, and Figure Our primary focus on this ablation the comparison of expressivity between different Matrix Mixers. Therefore, our implementations do not necessarily adopt algorithms with efficient computational complexities for simplicity. The abbreviation di stands for data-in iata-inddepedeent and dd represents data-dependent, in which dd variants are equipped with the SAM property. As the Quasiseparable variant is equivalent to the Hydra blocks, its implementation can be found in the Python file provided in the supplementary. Further specific details are described below.
Dense. While extending Dense matrices incorporate the SAM attribute is not straightforward, (i.e. without SAM) is extremely simple: they are equivalent to
implementing the vanilla Dense mixers MLP-Mixer [40), employing mixer matrix of [ L
17

---

Table ViT settings for ImageNet-IK.
Parameter
Value
Image size
224
Optimizer
AdamW W
Optimizer momentum
1, = 0.9,0.999
Weight init
trunc. normal (std=0.02)
ViT base learning rate
1e-3
VIT weight decay
0.05
Dropout
None
Batch size
1024
Training epochs
300
Learning rate schedule
cosine decay
Warmup epochs
10
Warmup schedule
linear
Randaugment
(yy0,5,
Mixup
0.8
Cutmix
1.0
Random erasing
0.25
Label smoothing
0.1
Stochastic depth Exp. mov. avg (EMA)
0.5
None
Toeplitz matrix mixers without SAM propertics are equivalent to convolution operations [8] 17). formulated with 22 parameters. We also present Toeplitz matrix mixer with SAM properties, thus data-dependent and can handle sequences arbitrary lengths. The core idea is fairly similar to the quasiseparable matrix mixer Hydra that the Toeplitz matrix mixer integrates outputs from two separate forward and reverse sequence convolutions. Specifically, for i E I. -- each token J; generates two convolution parameters m- and m,, Using all parameters, data-dependent dynamic convolutional weight is generated as { {m-L+11 L-+1 m -+2....... m. m...., -.
Toeplitz.
Vandermonde. In Section 2.4 we introduced the single-headed sequence aligned Vandermonde matrix mixer. Leveraging the definition, we present multi-headed implementation by seamless extension.
Cauchy.
The constant for preventing the denominator approaching zero is initialized to 0.0 and 0.5 for the di and the dd variants, respectively.
Definition E.1 (Cauchy Matrix), Given q.k R², matrix M is Cauchy if each (i.j)-entry m.j
satisfies m.j . 4: k #0.
F Ablating the impact of Shift and Diagonal operations in Hydra
Observe that general Quasiseparable (QS) matrices be viewed the Add variant augmented by two operations: an extra diagonal component and a shift operation. In this ablation we assesses the impact of each operation the model performance by individually removing either the diagonal or shift term, as well as removing both (Add) to analyze why QS outperforms Add.
find that using only the shift (80.7) or only the diagonal (80.7) yields the same performance as the Add method (80.6). It is only when both operations are combined that observe clear improvement. This result validates our choice of QS matrices, neither Add+Diag nor Add+Shift alone has the expressivity of QS which is strictly more expressive than the addition of two SS matrices.
Table 10: Analyzing the impact of Shift and Diagonal operations in Hydra.
Method
#Params
Acc (%)
GLUE
Add
70M
1.68
80.6
Add+Diag
70M
1.68
80.7
Add+Shift
70M
1.67
80.7
Quasi
70M
1.66
81.7
18

---

class Danse (nn.Nod
def _.init_..
d_model, nax_ssq__een max_seq_len is necessary for Dense. expand=2, headdim-128.
devica=tione,
atype-None,
):
Iactory_kvarge ('device"; device, "dtype": dtype) super 10....t))..... self d_model d_model solf.max_seq_een len sela.expand expand
self.d_inner self. expand self d_model
self-headdin headdin
assert
self. headdin -- 0 delned // headdin
self.headdin
sold sgrt max_seq_len)
self.N nn Paranater( torch.enptyy (self: abeads, aelf. aen.... self max_seq_len, *=Iactory.
init. Antt. (sel.)
def forvard hidden_states):
residual sidden_atates
Rearrange hidden states to shape [batch, m_hoads, length. headdim]
hidden_atates rearrange (harrage(hidden, (hidden_states, ib (n h) -> b h², n=self nheads)
outpat torch. winsune n h, mit >> h?, bidden_atatea, self.N)
output stddestd output
output rearrange (ontput,
h) residual
output
Figure 6: Py PyTorch code for the Dense variant in Table
19

---

class Toaplitz (ot.ttoonn)
def
_.init_.(
ie_data__ependent,
d_madel,
lax,seq_len, max_seg_ necessary Toeplitz.
expand-2,
headdin=128,
device-lone,
dtype=Nonne,
factory_kvarge ("devicee' device, "dtype"; dtype)
super (o.per()...iit..... "elt..ts_att.ddpeeddee self.d_node. d_model
1s_data_dependeer
self self.nax_seq_len max_seq_lon
expane expand
self d.ined expaad relf.d_nodel
held.ne headdin
assert self.d_inae self. headdin
solf.nhoads self.d_inne // self headdim
self.kernel_stee size self.pad, pale-pad_iize welf.max.eel max_aeq.len salf.std_dev
self welf.max_neq__11
ap. (self (self.max_seq_lan
if self self_ie_datt_ddeeadde
self seli,coav_pe conv_params pa Parameter(
torch. torch.enpty (self aheads, self keffkkeeee_size. aepeettorc..leattte device
device)
ininnit. ravier uniiorr (self.ccov.parrmaa)
def forvard (self, x, converd_ cov=
een
(batch, seglen, theadithe forvard_conve conv: (batch, seqlez, abeads) (batch, seglen, nhoads)
residual -x
rearrange 1b bn b P'
Pad the hidden states Fx,ad(x, psdlf.pad_sizo,
x_fft forch.fft.fftt fft. fleat32),,, n-2*selí max_soq_len-11
aaf.-aa_dataddeeedee.
forvard_comv rearrange (forward_ooo tb 1 b
rearrange (ravarss.rrae
conv parana corch.cat(
[torch. fwwp(reverse...p.. (o, (-1)), forward_coav)) din--1
(fttrrch..loat32)) PPT requires float32.
fft_conv_parane torch. fft. (tto..arcoo. n-self kernel_ unequeeze (1)
ftt_conv_param torch. fft. (self.sll.... conv_parana n=self. kernel_size)
output torch. Sarch.fft.iffte (f. ift_conv. paanv__aaaas... n=self "sself.kermel.ds. real
output self std_ output [i,
iself max__eq_len)
output rearrange (output,
t)))) (ttyep) residual
return output
Figure Py Torch code for the Toeplitz variants Table
20

---

class Vandermonde Module):
def ..init_..
je_data_dependeeut
d_model,
gk_dim,
is_dft- True, nax_saq_1en
Used only when
is Palse.
max_aeq_ len necessary for non-DFT data- -nddpaadnnnd version.
expand=2,
beaddim-128,
device-Nlnze dtype-Mone,
factory_kuarge d"dvicaa": device, "dtype": dtyps)
cuper()...intt...0
deeeepeedattaddpeddeent de_dnta,
self denodd_ d_model
self.qk_dim qk.dim
ieldfis_fft is_dft
solf.max_deq_len nax_seq_lon
self.expand expand
self.d_immer self sela.expand self.d_model
self.headdis headdin
assert aelf.d_inner
headdin
gelf, nheads s.lf.d.inner // gelf
self self.d_rte d_state self nheads qk_din
if sllf.ie_data da salf.atd.dev self.eps 1e-3 Constant to stabilize training.
(2. self mex_ self qq_din)
elae:
self.ie_dft:
indices torch. arange (self nax max_
row_indices_ torch arange (self aax_zog_len) unsqueeze
dft_matrix
torch.pi rox_indices
salumn_indicer
max max_seq_lea)..o to factory_kvarge)
self ref.rrrggsttrrbbeee (owfeerfdfftteeeeeeoooll ('aft_natrii'' dft dft_matrix)
self dev_ 1 ag.agr wwaptttstt.eee........ max_aeq_len)
self. blas nn Paranster (torch. zeros (self nheads self folf.qk_iim, aelf.
max_ factory_kvarge
self k. bias an Paraneter 2eros (self nheads, self seldiqk,din, self.
mox_seq_len, factory. solf.std_der - 1 a2. mali.ax_seqq.ee self gi_din)
def Torvard(self, v, q-Rone, x-flone): batch, seqlen, din = v.shape
residusl
y rearrange 1b (a b) >> nnn h', arself nheads)
if self.
rearrange
arseli.nhest nheads)
rearrange 1 (n bbnd
q_eatrix torch.cos.
2 torch. self. opa torch. bbndt, babaad 1', torch, arange (seqlen, dtype-y.dtype)..o to v.device)
torch. cos(
2 torch. self.eps eps torch. torch. einsun( budt, bab 1 t', k, torch. arange (seglen, (agpien,,, dtspe-v.dtype).
v.devica)
syn, vandermonde (qaatrii matriix) eun (din-2)
output corch.einsun 1, bll n b tn b², syn_ v)
claa:
if self.is_dft
output torch. einsun
t b v, gelf.dft aif.ddft_natrit
self. self.k_bias
q_aatrix torch.cost.
torch.pi torch.einsum(
'a t, - baatll, torch, arange (self, max seq_len, dtype-v. dtype) (e.der)
natrix torch.com cont
2 torch.pi pi torch.einsunl in a',, k, torch arange (self. seq, dtype-v.
ays_ _yn__andernonde (q_matrix k_natrix). (dm(mimm=)
output torch. ( t 1, b btth > beta stah cyn_wandormonde, v)
output self.atd atd_ output
output rearrange (output. "b >> b ())) residual return output
Figure PoTorc code for the Vandermonde variants Table

---

class Cauchy (nn. Module):
def
..init_.
self, le ie_data__ependent,
d_model, gk_dim, aax_seg_len-Kono max_seq_len is necessary for data- data-inneppoddnt version. expand=2, headdin-128,
device=tione,
atype-None,
factory_kvarge ('device"; device, "dsype"; dtype) super (...att).....t...
self. self. ls_data_depandent
self delf.d_nodel d_model
self.gk.dim qk_dim
self.2ar.deq_1en nax_eeq_len
self.expand expand
self.d_inzer self self.expand self d_model
self helf.headdin headdin
assert self.d_inmer self. self.headdin
self.nheads self.d // self.headdim
self.d_star self welf.nheads qk_din
self.tol 1e-8
self.etd_dev sgg mex......_.....11 self gk_dim)
celf.bias Parameter (torch. tensor (0.5)
else:
self qunatrix nn nn. Parameter torch. torch.enptyy empty(sall (tate/aa.l...aeeeg..eeee max_seq_len, self.nheada, self qk_din, s*factory_kaarga
self self.k_nn kanaiix nn Paraneter torch empty (self mex_seq_len, self cheads, self di_ ".actory_kvarge"
an init zavier normall (self (self.q_asta) q_astrix)
wawior_normll (self k_natrix)
def servard(selt Nonoe, k-llane):
v. rearrange 1b (n h) >> b 1 h', asself .neaads
ie_daa dependent: deppeedee
rearrange (q, "b (a d)
-> bbbal d', n-self nheads)
rearrange (k, bb 1 (n d) bbblle d', n=self nheads)
txpch.oxp(q))
k torch.
bias
inv wwtrcccchymaatiixx matrix - + self cauchy_natris torch. inv_ matrix, dim=-1)
output torch. t tnb, bal b b'mb, v, cauchy_natriz)
else:
g. (nhoads, seglen, gkdin) torch torch.exp (aelf (eeltrq.uatrix) torch. (self k_natrix)
inv cauchy. matriz (q. unequezze (1) k. unsqueeze + self.ter cauchy_natria torch. su caucha _atrriyaatttx,.. dia--1)
cauchy_matriz
output torch. einsun
h, itn b 1 h', cauchy natrix))
output self der output output rearrange (output, bb
(a h) residual
retura output
Figure 9; PyTorch code for the Cauchy variants in Table 2
22

---

class Lovšank (nn. (
def _.init_..
_e_data__ependent,
d_model.
gk_dim, aax_seq_len-Kono
max_seq_len is necessary for data- data-indepondent version.
expand=2,
headdin-128,
device=tione,
dtype-None,
factory_kvarg ('deiice"; device, "dsype"; dtype)
(...nt.......t....)
self. delf. is_data_depandent
self delf.d_nodel d_model
self.gk.dim qk_dim
self.2ar.deq_len nax_eeq_len
self.expand expand
self.d_inzer self selffexpand self d_model
heldinn headdin
assert self.d_inmer self. self.headdin
self. self.d // self.headdim
self.d_stat self shladd qk_din
gelf, -1 / /t.pppp.... ag. (sll...el.... aee_ self solf_issdatt_ddeadden
+ self. gk.din
self.q_matriz na Paraneter( torch. enpty mex_seq_len, self nheads, self qk_ wifactorykvaare
)
self kataiii an Parameter(
torch. esach www.aaaa...aaa..aaaaa.a. mex_seq_len, self.nheada, self galf.qk_din, sractory_kuargs
nh. init. normall seelf.q.aatrin init. sawie__avvirnnonnee k.atrix)
def forward(self, v, q-lione, k-llone):
residual
y rearrange 1b (n bblnb b', a-self nheads)
dedaaaddddddddedeettt = rearrange "b (n d)
b 1 d', n=self . nheads)
rearrange (k, "b 1 (n d)
b 1ll1 d', n-self. p-self.nheadh
output torcb.einsun(1)
btab', 9, k.
else:
output cinsuni dd
bluh b self g_natrix, self ___ntrix,
output atdf_atd_dar output
output rearrange Coutput, bb b (a h) residual
output
Figure 10: Py PoTrr code for the low-rank variants in Table
23

---

class Attantion (nn. Mondll)
def
_.init...
self,
_e_data__ependent,
d_model,
gk_dim,
jax_seq_len-none max_seq_len is necessary for data data-indepondent version.
expand=2,
headdin-128,
device=tione,
dtype-None,
factory_kvargs ('device"; device, "dsype"; dtype) super (...nit...))....
self. daa.ita_.ata.deaadden: is_data_depandent
self self.d_nddel d_model
self.gk.dim qk_dim
self.2ar.deq_1e1 max_eeq_len
expane expand
self.d_inzer self selffexpand self d_model
self held.headdin headdin
assert self.d_inner
velf.headdin
self. nalfdnheads self.d // self.headdim
self d_state self helf.uheads qk_din
if not gelf dat_d_ta_ddeeeeeee dependent:
self self q_aatriz Paraneter(
torch.enpty (self max sen_ self nheade, self.qu_dd =ffattoy_kverr'.
k_aatrix Paranater torch. expty carch.capy(eal...aa.ec (salf.max_seq.lott, self nhoads, self qk. #ffactor.kvargr
)
saitt.aavier.oraan (seel..g_sstie nn. xaveer (self (self. k_natrix)
foreard 7, qooone k-None)::
residual
rearrange 1b b) bln b², n-self. nhoadss)
if dependenn rearrange (q, (n
d', n-self nheadd))
rearrange 1b (n d) -> b n d', n=self.. orsalf.nheads)
q torch. sincub ("b b d,
-> 1',
atta_weighte softuaz22 ap. sqr stleeeff.k..iin
qk. dim--1) 1, b n >> t b', seight, v)
output torcb..insun("b else:
gk torch.eineunn'n d,
seli mantrix keatrix output tarch. einsun n'n /ap. sp.rggt(self qindim) + qk, dim--1) 1) seightaigh v)
atta_weights torch, softnaz torch.softnaz(1
output
rearrange (output, "
(n b) residual
retura output
Figure 11: Py Morch code for the Attention variants in Table 2
G Background
the FNet (23 architecture, the sequence mixing module utilizes a Discrete Fourier Transform (DFT) for processing input sequences. One of the approaches they adopt for short sequences is the application the matrix representation of the DFT to the input sequence. This representation, denoted as M, takes the form of Vandermonde matrix, constructed from the roots of unity:
FNet
M = nk
(4)
where indices and range from to L-1. First, we define wn as the n-th root of unity:
wn 2
Then, the matrix M, which represents the DFT, expressed using wi, in the form of a Vandermonde matrix as follows:
3
M=
w
wh-1
24

---

MLP-Mixer In the MLP-Mixe architecture, sequence mixing is implemented using a MLP M. This
MLP operates on each token in the sequence and is formulated follows:
M = (W)
where denotes non-linear activation function (such ReI.U). In this context:
W1 is weight matrix with dimensions Rdss x1, transforming the sequence from its original sequence length L to an intermediate dimension ds.
W2 a weight matrix with dimensions Lxxds, converting the dimensions back from the intermediate dimension dy to the original sequence length L.
This design means that W1 and W2 first compress and then expand the sequence dimensions, respec- tively. The choice of the inner dimension dg made independent of the sequence length L. Such configuration leads to computational complexity that is linear with respect to the sequence length L, contrasting with the quadratic complexity commonly seen in sequence mixing operations in attention mechanisms.
Note that M is a dense matrix, effectively capturing interactions across different tokens in the sequence.
Linear Attention Transformers process sequences of feature vectors, denoted as X E R where N represents the number vectors and C' their dimension. The core component of a Transformer is the self-attention mechanism, which mixes information across the sequence.
In self-attention, the sequence X is projected into queries Q, keys K, and values V using matrices We. Q. W K. Wy ERCD The attention output is calculated as follows:
Q = ===Wo Q+
K K;
(8)
V V V.
(9)
QK T
Attention(X === softmax
(10)
VD
A generalized version of self aelf-attention can be represented as [22]
ンj=
similarity(q::k))
(11)
similarity (q;
TTh
where the standard softmax attention employs similarity(q.k)- VD
as the similarity function.
In the context multi-headed Linear Attention (LA) (22). the input sequence is preprocessed and projected into three matrices for each head h [H] This is achieved through learned parameters and nonlinear transformations. Specifically, for each head, the input is projected into matrices a (Q").. o (K*),, and Vh, where a denotes a non-linearity function. The set learned parameters for these projections, denoted O, includes W W к corresponding to the query and key matrices across all heads, defined θ (RCx Mdj2 Here, H is the number of heads, and d represents the dimension for each head for keys and queries.
The class of structured matrices M used in LA characterized as the class of Low Rank matrices. The matrix- gerrrraingnnrtingggg function for each head h. denoted as fhr computes the matrix M as M = (Q² (o This computation involves linear projections of the queries and keys, application of clement-wise nonlinearity a, and matrix multiplication.
The output for each head h, denoted Y is then computed as Y (fx((XX))* where fx Wv E CDDD represents projection of the input sequence with W v a learned parameter matrix. This projection transforms the input data into a suitable form for processing by the attention mechanism. The LA method, as described, offers computational efficiency with time and memory complexity of O(N), in contrast to the traditional softmax attention mechanism which scales with O(N²).
Discrete Convolution Consider convolution operation between sequence X and the filter sequence to produce the output sequence Y. Such a convolution can be represented by a matrix multiplication with an L x L Toeplitz matrix M 31
Let:
X be the input signal, a column vector of length L, i.c.,, X 4x11
25

---

h be the convolution filter, with a length N where typically N - L.
Y be the output of the convolution, also of length L.
The Toeplitz matrix M for the convolution is constructed as follows:
ho h-
h-1+1
h1
ho
142
M
hg-1
hL-2
ho
Thus, the output Y MX represents the convolution of X with h, with both X and Y having the same length L, achieved through the L L L Toeplitz matrix M as defined above.
26

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and socictal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant
information is Not Available.
Please provide a short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
*
Do not modify the questions and only use the provided macros for your answers.
1.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction are clearly discussed throughout the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed in Section A.
Guidelines:
27

---

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e-g., independence assumptions, noiseless settings, model vell-specification asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e-g.. if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should diseuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification:
Theoretical assumptions are provided with proof's in Section 2, 3, and Section B.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reprodueibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide specific details of our models, and experimental settings. In addition, we provide source code of our method in the supplementary.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
28

---

If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (c.g., with an open-souree dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Ys]]
Justification: In the supplementary, we provide source code with instructions in README file. Moreover, provide implementations of ablations studies in Section E.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
Srriddss/ooeeuumiisse
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7 7/nips..cca hatpsscccpbblccccccc guwdes/ CooeeSooiisssiii wwt.bbmiissoooooooiiioooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
29

---

Answer: [Yes]
Justification: We provide specific experimental details in Section 4, Section C, Section D, and Section E.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Each number is the average of five runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call library function, bootstrap, etc.)
The assumptions made should be given (c.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [No]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider. including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well as estimate the total compute.
The paper should disclose whether the full research project required compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
30

---

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of http::/p/ww.apppppppppapa/aaaaaa
Answer: [Yes]
Justification:
The research conform Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special considera- tion due laws regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We briefly discuss potential impacts in Section
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [No]
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
31

---

We recognize that providing effective safeguards is challenging, and many papers do require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: All assets are open-license.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, fatasets.paperssttt paperswi paperswithoode thcode. com ddtasststsssssssss has
curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
32

---

According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "y5LexPAw2a", "reviewer": "Reviewer_9VST", "summary": "Most of sequence models include the token mixers and the channel mixers, and this paper provides a detailed summary. They also identify the matrix parameterization is crucial for recent SSMs. Therefore, they extend the Mamba model by adding bidirectional quasiseparable matrix mixer. The experiments on GLUE benchmark and ImageNet data verify their method.", "questions": "1. The methods to implement token mixer include not only structured matrix mixer, such as implicit reparameterization of CKConv. So the conclusion in Line 117 is inaccurate. What do you think are the advantages and disadvantages of implicit matrices over structured matrices?\n\n2. How to understand \"Sequence Aligned\" property in Table 1?", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper summarizes the previous relevant methods very well, as illustrated in Table 1. The authors also claim that the following two properties are important for sequence aligned matrices: data dependency and extendability. The former one is a well-known propoerty.\n\n2. The proposed quasiseparable matrix mixer is simple and easy to understand. It is implemented through two Semiseparable matrixes.\n\n3. The provided results show the advancement of their method.", "weaknesses": "1. The authors claim that the sequence model should be extended beyond their trained length, but they didn't provide the corresponding results.\n\n2. I am curious about the computational complexity of Hydra compared with Mamba. Because they use two SSD operations as shown in Figure 4. The ablation studies are also necessary for the different variants.\n\n3. Please compare with the most advanced methods, such as xLSTM. It also expands the hidden state into a matrix form.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for recognizing the **detailed categorization of prior methods using the Matrix Mixer framework and for appreciating Hydra’s simple yet effective design**.\n\n---\n\nWe now address the reviewer’s concerns as follows:\n\n### A1. Ablations on extendability:\nThe extendable property of sequence mixers is implicitly validated throughout our experimental section. Non-SAM matrix mixers, such as the Dense variant (e.g., MLP-Mixer [38]), lack extendability and require retraining from scratch when adapting to different sequence lengths. Therefore, in Table 2, we fixed the sequence length to 128 across all variants to ensure a fair comparison, thus emphasizing data dependence (DD) only. Conversely, Hydra, as a SAM matrix mixer, inherently supports both data dependence and extendability. Specifically, in Table 4, Hydra was pretrained on sequences of length 128 (C4) and then fine-tuned with sequences of length 256 (GLUE).\n\n### A2. Computational complexity of Hydra compared to Mamba:\nHydra shares the projection layers for both forward and backward passes, with the additional computation over Mamba being the backward SSD. However, the computations for forward and backward passes can be parallelized by batching, resulting in only a ~30% reduction in throughput compared to Mamba.\n\n### A3. Comparisons to Advanced Methods such as xLSTM\nWe thank the reviewer for pointing out xLSTM, and indeed there are many recent subquadratic models that have shown to be performant. We specifically chose to focus on Mamba as it is a recent model that introduces an interpretation of recurrent models as **matrix mixers**, which is the main focus of our work. Furthermore, **since xLSTM is language model and does not have canonical bidriectional variants so it is not a suitable for comparison on bidirectional tasks**.\n\n\n### A4. Analysis on Implicit Reparameterization:\nWe appreciate the reviewer’s observation that the same matrix class can be parameterized differently. However, this property is orthogonal to the underlying matrix class and its computational complexity. **Specifically, an implicitly parameterized mixer like CKConv [35] and S4 [18] is indeed a Toeplitz matrix mixer, as indicated in Table 1**. Therefore, the conclusion in Line 117 remains accurate.\n\n\n### A5. Clarification of the ‘Sequence Aligned’ Terminology in Table 1:\nThe formal definition of Sequence Aligned Matrices (SAMs) is provided in Definition 2.2. To summarize informally, SAMs ensure that the parameters for every submatrix $M[: i+1, :i+1]$ are functions only of the tokens up to index $i$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eKWfV10hCC", "reviewer": "Reviewer_onHY", "summary": "The paper introduces the concept of matrix mixer and sequence alignment for explaining recent sequence models including Transformer, linear transformer, and Mamba. It also proposes a quasiseparable matrix mixer (Hydra) as an alternative to the bidirectional SSM model. The experiments show that the quasiseparable matrix structure performs better than others including low rank (linear attention), attention, and dense. Also, Hydra outperforms other naive bidirectional approaches and some baselines for masked language modeling and image classification.", "questions": "1. Is there any speed benefit with the quasiseparable bidirectional model compared to the naive bidirectional approaches? It would be good to include some speed comparisons in Table 3.", "limitations": "The limitations are discussed and are reasonable.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The idea of designing a structure matrix for a bidirectional case is reasonable and effective while the implementation is simple.  \n2. The explanation of different matrix mixers, their relation to other methods, and the advantages/disadvantages are clear. This is very informative. \n3. The ablation study shows how different matrix mixers perform and the benefit of quasiseparable matrix structure.", "weaknesses": "1. The presentation of the paper needs to be improved. \n    1. It contains unnecessary details. The main contributions of the paper are matrix mixer, sequence alignment, and Hydra, but the focus diverges throughout the paper, especially in Section 2. Sections 2.3 and 2.4 can be moved as a side note after introducing Hydra.\n    2. The purpose of introducing Cauchy and Vandermonde matrix mixers is unclear. How can this be useful or helpful in some ways? It's not clearly explained or shown in the experiments. \n\nOverall, the current representation makes it difficult for the readers to understand the core contributions. \n\n2. The experiments shown in the paper are limited. The main comparison table (Table 4) does not include any recent transformers or mamba-based models. Also, the experiments include one example of non-causal language modeling (masked language modeling) and image classification as applications for Hydra (bidirectional model settings). Could authors show at least one more example application that Hydra can be useful and better than other comparable models?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments and we are glad that they **found our Matrix Mixer framework informative and the Hydra model simple and effective**.\n\nThe review focuses on the following concerns:\\\nQ1. The main comparison table (Table 4) does not include any recent transformers or mamba-based models. Could authors show at least one more example application that Hydra can be useful?\\\nQ2. Paper presentation issues\\\nQ3. Speed comparisons between Quasiseparable and other naive approaches\n\n---\nWe now answer these questions in detail below:\n\n### A1. Table 4 does not include recent models and the two domains chosen are limited\n\n> (Table 4) does not include any recent transformers or mamba-based models\n\nWe reiterate that the primary objective of Table 4 is to compare the core **sequence mixers** proposed by earlier works. Many recent works, including Caduceus [36], BiGS [42], and Vision Mamba [46], simply **use the Add, Concat, or Element-wise multiplication heuristic on Mamba. These comparisons have already been included in our experiments** in Table 3.\n\nWe acknowledge that integrating the Quasiseparable matrix mixer into domain-specific backbones presents excellent directions for domain-specific future work; however, these efforts are beyond the scope of this paper.\n\n> Could authors show at least one more example application that Hydra can be useful?\n\nWe chose the MLM and ImageNet-1k classification tasks since **they are canonical tasks used by previous works that proposed new sequence mixers like M2[13] and Hyena[29] and their baselines are well-tuned and established**. We are confident our method will also excel in other domains such as DNA modeling and graphs. We invite domain-specific research communities to explore Quasiseparable in their applications.\n\n---\n\n### A2. Paper presentation issues\n\nWe appreciate the reviewer’s feedback on the structure of our paper. We would like to convey that we structured our paper this way to match the flow of our contributions:\n\n1. **Formalization of the Sequence Mixer:** In Section 2.1, we formalize a sequence mixer as a generalized matrix mixer, identifying the computational bottleneck as the cost associated with the matrix class. This leads us to focus on structured matrices.\n2. **Optimal Structured Matrices:** In Section 2.2, we define SAM matrices, characterizing structured matrices that enjoy data dependence and extendability to determine which classes would have superior empirical performance.\n3. **Framework Generality:** In Section 2.3, we demonstrate the broad applicability of our framework by showing that more than 13 past methods, including Attention [40], Mamba [6, 16], MLP-Mixer [38], and Linear Attention [22], fall under this paradigm, indicating our framework’s generality.\n4. **Validation and Prescriptive Power:** In Section 2.4, we validate our framework's prescriptions by identifying Cauchy, Vandermonde, and Quasiseparable matrices as Sequence Aligned Matrices (SAM) which have not been fully utilized in past works and showing that they exhibit strong empirical performance.\n5. **Hydra:** We chose Quasiseparable matrices due to their hardware properties and connections to Mamba, scaling them up and proposing them as the next sequence mixer of choice. We show that Quasiseparable outperforms other sequence mixers on bidirectional language modeling (+0.8 over BERT) and ImageNet-1k image classification (+2.2 over ViT).\n\nWe request the reviewer to also take a look at the first global response wherein we enlist our core contributions in depth and tie them to the different sections of the paper\n\n---\n\n### A3. Speed comparison between QS and naive approaches\n\nHydra shares the projection layers for both forward and backward passes, with the additional computation over Mamba being the backward SSD. However, the computations for forward and backward passes can be parallelized by batching, resulting in only a ~30% reduction in throughput compared to Mamba.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7MBRc5Mh8G", "reviewer": "Reviewer_5fYM", "summary": "This paper presents Hydra, an innovative framework that builds upon the Mamba model by introducing bidirectional capabilities. Hydra's approach centers on a matrix mixer perspective, which allows it to consolidate various sequence models, including Transformers and structured state space models, into a unified framework. The primary strength of Hydra is its capacity to surpass other sequence models in non-causal tasks while retaining efficiency and expressiveness. The study demonstrates how Hydra's novel bidirectional methodology and matrix parameterizations effectively enhance the performance of sequence models.", "questions": "- The experiments are mainly around ~ a 100 M parameter scale. I am curious about whether the model can be scaled up to ~1.5B parameters and how the model will perform.\n- I am wondering, in such non-causal tasks, how is a comparison between the proposed hydra framework with stacked bidirectional Mamba frameworks (https://arxiv.org/pdf/2401.09417,https://arxiv.org/abs/2404.15772)?\n- Compared to the original Mamba framework, does Hydra require additional computational resources and training time, as the training objective is harder? I am also curious about whether the matrix mix method will harm performance on general tasks.", "limitations": "See Questions.", "rating": 9, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The authors proposed a novel framework, Hydra, which extends mamba with bidirectional capabilities and presents an interesting perspective on improving sequence models.\n- The proposed method of matrix mix offers a cohesive understanding of various sequence model architectures, which also offers valuable insights into how matrix parameterizations and sequence alignment affect model flexibility and performance.\n- The paper is well-structured, making it easy for readers to follow the development of the Hydra framework and its contributions to sequence modeling.\n- Abundant experiments covering both language and vision tasks illustrate the efficacy of the proposed method.", "weaknesses": "I don't identify a specific weakness, but I have a few questions regarding efficiency. Please see the following section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the thorough evaluation and for recognizing the novel contributions of our framework. We are pleased that the reviewer appreciated the **bidirectional capabilities**, the **valuable insights provided by the matrix mixer perspective**, and the **extensive experimental validation** across both language and vision tasks.\n\n---\n\nWe now address the reviewer’s key questions:\n\n### Scalability to larger models (~1.5B parameters)\n\nWhile our experiments primarily focused on models with ~100M parameters due to resource constraints, Hydra is certainly scalable. We anticipate that Hydra will continue to demonstrate effectiveness in non-causal tasks as the model size increases.\n\n### Comparison with stacked bidirectional mamba frameworks:\n\nWe appreciate the suggestion to compare Hydra with previous stacked bidirectional Mamba frameworks (e.g., provided references). Under the matrix mixer perspective, most of **these methods fall into Add, Mult, or Concat categories** listed in Table 3. As demonstrated, the quasiseparable matrix mixer used in **Hydra outperforms other bidirectional variants.** Therefore, we confidently expect that substituting the bidirectional components of previous models with Hydra would lead to a boost in performance.\n\n### Computational resources and training time:\n\nAs similar to all other bidirectional extensions of Mamba [12, 15, 42], Hydra does introduce additional computations due to its bidirectional nature. However, unlike many of the previous extensions that utilize two completely separate unidirectional components, Hydra greatly reduces this overhead by sharing the projection layer for both forward and backward sequence processing and by batching both directions to maximize GPU utilization. This approach results in only a ~30% reduction in throughput compared to Mamba.\n\n### Performance of different matrix mixers on general tasks:\n\nHydra, derived from the matrix mixer framework, has demonstrated significantly stable training and enhanced performance across well-established domains such as vision and NLP. We are confident that models developed using the matrix mixer framework, when logically designed for specific domains, will continue to achieve superior performance across a diverse range of tasks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uUzVIn4WU0", "reviewer": "Reviewer_RwkG", "summary": "The paper introduces a matrix mixer framework for sequence models which linearly applies an LxL matrix M  to a sequence representation X of length L. Popular sequence models can be framed within this context e.g. softmax self-attention or SSMs, according to different properties of M. The authors use their framework to identify desirable properties of M, such as data dependence or extendability (where the parameterisation is such that the sequence length L can change for different sequences), or efficient matmul. The authors then use their framework to design sequence models with Vandemonde or Cauchy mixing matrices that perform comparably to attention. The main contribution of the paper is the introduction of Hydra, a bi-directional variant of SSM like Mamba. To do this, the authors let M instead denote a quasiseparable matrix (which is a bidirectional variant of the causal semiseparable matrices introduced in the Mamba-2 paper [6]). The authors show that Hydra (i.e. using quasiseparable matrix mixers) outperforms other matrix mixers (4.1), outperforms other approaches for bidirectional SSMs (4.2), and outperform other standard sequence model arhchitectures like transformers on Masked Language Modelling and ImageNet classification.", "questions": "- Is there a theoretical argument for why data dependence/extendability are nice properties? Or just intuitive? Are there settings where these properties can hurt e.g long range?\n\nTypo:\n- linear 38 \"they often lack a systematic\" - remove \"and\"", "limitations": "There is a discussion of some limitations in the appendix.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "1. The paper provides a general framework for matrix mixing and introduces desirable properties like Sequence Alignment.\n2. Introduces a bidirectional version of SSMs, called Hydra, which is motivated through their framework.\n3. Hydra seems to work really well, beats unidirectional Mamba and transformers across different settings where bidirectionality is desired.\n4. Shows their framework can motivate new sequence models using e.g. Vandermonde or Cauchy matrices, and provides insights into why existing methods may perform well (like low rank matrices in linear attention).", "weaknesses": "1. The writing is unclear in large parts. This detracts from the flow and readability, and ultimately makes it somewhat a frustrating read, as it seems there are nice ideas here but communicated poorly. For example:\n- What does “common data transformations” mean in definition 2.1.\n- The term “Structured matrices” is not standard terminology and is not formally defined. In my reading the closest thing to a definition is line 117: “... a structured matrix, which are known to possess sub-quadratic matrix multiplication”. Is a structured matrix defined by the ability to do sub-quadratic matrix multiplication, because line 117 suggests there are other properties involved. Also the reference to established “mathematical literature” in line 120 but no citations are provided nor examples of structured matrices.\n - “Data dependence” also doesn’t seem to have a theoretical definition, but is part of the theoretical results (proposition 2.3) so should be defined. What does “canonical” mean in the context of line 138? The authors write \"Although data-dependency is a popular notion that is widely regarded as being crucial to performance of Attention, it lacked any formal definition in the literature.\" but there also isn't a formal provided here in my reading. What does \"the third interpretation\" refer to in line 153?\n- What are $\\phi$ or $\\hat{P}$ in definition 2.2 of sequence alignment, are there any constraints on $\\hat{P}$, why/when is it necessary in the definition? I assume $\\phi$ is the empty set but this is not obvious nor defined.\n- The proof of proposition 2.6 does not prove the sequence alignment property.\n- Are Quasiseparable matrices established in the literature or have you defined them? There is no citation provided afaict.\n- Should the tilde be a hat on P in line 237 (over vice versa in line 131).\n- The 2 paragraphs from line 190 to line 205 make reference to results which are a lot later (e.g. table 2). I would move the table earlier or at least refer to table 2.\n- “Taming the Hydra” isn’t particularly insightful for a subsection title.\n- The advantages of Hydra (lines 246 to 249) are not justified until later in the subsection, so on first reading it seems like overclaiming (also see my 2nd point). It is also not clear what is being compared to with these advantages (“heuristic alternatives” is cited for the first one but is this also true for the second and third benefit?)\n- “It is known that quasiseparable matrices are closed under addition” - cite or prove.\n2. It is not clear to me if the theoretical motivation of Quasiseparable matrices is actually a big factor behind the practical gains: in Figure 4 it seems like you are discretising the forward and backward SSD differently, which seems like it will break the symmetry between the A matrix in the forward and backward SSD which is necessary in for the equivalence in proposition 3.2, which to my understanding states in layman terms: “a QS matrix is just two SS matrices that *share parameters*  with flips/shifts”. I could be wrong, but discretise_A is not defined so it is hard to check and the zipped code is different. But if I am right this seems to go against one of the motivations of the paper that previous bidirectional SSMs are heuristic or ad-hoc.\n3. Likewise, the arguments in line 250-258 that motivate Hydra as opposed to previous approaches for bidirectional SSMs make it seem like the only difference is the diagonal elements, which can be seen as skip connections (see e.g. https://arxiv.org/abs/2302.10322 in the example of diagonal elements of attention matrices replacing standard skip connections). Given that the previous methods presumably also use skip connections, this difference seems quite minimal. If anything, it might be better to pitch this framework as encapsulating and generalising previous works.\n4. Ablations to understand the empirical improvements made in Hydra are missing. In general the results seem great, but as discussed a reader may have unanswered questions why Hydra seems to work so well.  For example, related to above, what to performance happens if you remove the extra diagonal part, which should boil down to the Add method in table 3 right? Also what if you change the convolution (e.g. remove it or make it a causal conv) because it doesn’t seem to be connect to the theory here (of quasiseparable mixers)? Also what happens if you remove the parameter sharing between the two SSDs (which again should get closer to the previous methods)? What if you place the different mixers in a standard transformer block instead of a Mamba block in table 2 (does table 2 give an unfair advantage to QS)?\n5. The authors write that \"Hardware efficiency\" is a limitation in the appendix. Do the authors have throughput results for Hydra compared to transformers? If hardware efficiency is a concern then isn't \"fast implementations on hardware\" a more practical desiderata to design new sequence models as opposed to use \"sub-quadratic matmuls in theory\", which is the line taken by this work?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for recognizing the novelty of our framework, its potential to motivate performant sequence mixers, and superior empirical performance enjoyed by Hydra.\n\n---\n\nThe reviewer's feedback highlights the following key concerns:\\\nQ1. Ablating away the shift and the diagonal operations to understand Hydra's performance\\\nQ2. Do separate parameters conflict with the motivations of the matrix mixer framework?\\\nQ3-1. Clarifications on the definitions of Quasiseparable, and Structured Matrices and Data Dependence, along with further writing style comments.\\\nQ3-2. Writing concerns.\n\nA1) We thank the reviewer for this valuable suggestion; the ablation shows that the GLUE scores with only shift or diagonal operations matches the Add variant ($\\approx 80.6$). **It is only when both operations are used together do we see an improvement** ($81.7$). This **validates our framework**, which explains this ablation by noting that a Quasiseparable (**QS**) matrix is strictly more expressive than adding two SS matrices.\n\nA2) We refer the reviewer to the definition of QS matrices [C]: *a matrix M is N-QS if its any submatrix from either the strictly upper or lower triangle (off-diagonal) has a rank of at most N.* We remark that the matrix **with non-symmetric upper and lower triangular forms is a general QS matrix** and is consistent with the motivations of the Matrix Mixer framework.\n\nA3-1) We refer the reviewer to [B, C] for the definition of QS matrices. Structured Matrices [A, B, C, 6, 8, 13, 16, 29] are a well-studied area of mathematics, defined as matrices that admit a subquadratic matrix multiplication algorithm. By data dependence in the \"third interpretation\" we mean that each parameter of a matrix is **either a free parameter or projected from exactly one token**. By canonical data dependence, we mean that the parameters of SAM matrices **can naturally be made data-dependent using the bijective map** $f_{\\mathcal{E}}(\\cdot)$.\n\n---\n\nWe now provide a detailed response to all of reviewer's comments:\n\n### A1. Ablations to understand Hydra's performance\n\nWe thank the reviewer for suggesting an ablation study to assess the impact of shift and diagonal operations on Hydra's performance. In this study, we removed either the Diagonal or Shift operation, or both (equivalent to the Add method). Our results are tabulated below:\n\n|Method|#Params|$L_{ce}$|Acc (%)|GLUE|\n|-|-|-|-|-|\n|Add|70M|1.68|65.6|80.6|\n|No Diag|70M|1.68| 65.7|*80.7*|\n|No Shift|70M|1.67| 65.8|*80.7*|\n|Quasi|70M|1.66|65.9|**81.7**|\n\nWe observe that performance remains unchanged when only one operation is used. **Only when both operations are used together do we see an improvement** ($81.7$ vs $80.6$). This validates our framework, showing that a QS matrix is strictly more expressive than adding two SS matrices.\n\nThe reviewer has raised important ablations involving convolution operations and backbone architectures. **While these are beyond the scope of this paper, as we focus on sequence mixers, they are valuable directions for future research.** We refer the reviewer to the second global response where we clarify the focus of our paper on sequence mixers.\n\n### A2. Do separate parameters conflict the matrix mixer framework?\n\nYes, the reviewer is correct that the two halves of the matrix are discretized separately. This does not deviate from the motivations of our framework as the matrix **with non-symmetric upper and lower triangular forms is simply a QS matrix**. This follows from the rank characterization of QS matrices [C]: *a matrix $M$ is $N$-QS if any submatrix from either the strictly upper or lower triangle (off-diagonal) has a rank of at most $N$.* We note that this definition applies regardless of parameter sharing between the upper and lower halves, and that the parameter shared matrices referred by the reviewer are a strict subset of QS matrices.\n\nProposition 3.2 does not include this detail for pedagogical simplicity and it can easily be extended to QS matrices. To see this, divide the QS matrix into strict upper ($SS_u$), strict lower ($SS_l$), and diagonal ($D$) components, and then:\n\n$$QS(X) = \\text{shift}(SS_l(X)) + \\text{flip}(\\text{shift}(SS_u(\\text{flip}(X)))) + DX$$\n\n### A3-2. Writing concerns\n- “Common data transformations” refers to standard projects commonly used in the ML community, such as those implemented using convolutions or linear layers.\n- Components in Definition 2.2 are necessary as they allow for multiple sets of parameters like B and C. As assumed, $\\emptyset$ represents the empty set.\n- By the definition of SAM, Linear Attention is sequence aligned as each element $(i,j)$ in the matrix $m_{ij}$ is parameterized by $Q_i$ and $K_j$.\n- Quasiseparable matrix are closed under addition [B]\n- Efficacy of data dependence and extendability is born out of the experience of the community and has been shown to be effective on long range tasks [6, 16]\n\nWe appreciate the reviewers constructive feedback on the writing, such as the typo in Line 237 and the ordering of the paper. We will address all the aforementioned issues in our next revision.\n\n### A4. Is Hydra hardware efficient?\n\nWe agree with the reviewer that hardware efficiency is indeed an important desideratum for designing new sequence models. However, achieving hardware efficiency involves a series of complex factors, making it challenging to directly target hardware-friendly models from the outset. Therefore, computational complexity is usually prioritized first, with subsequent efforts focused on developing efficient implementations.\n\nHydra leverages hardware-friendly kernels from Mamba, maintaining competitive speed with Transformers for short sequences. As sequence lengths increase, Hydra’s sub-quadratic design allows it to surpass Transformers in speed, addressing the quadratic bottleneck. The growing interest in processing long sequences [D, E, F] further underscores the importance of Hydra.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
# Global Response

---

We express our sincere gratitude to the reviewers for their valuable feedback and constructive suggestions. We are glad that they found our Matrix Mixer framework insightful and informative, and that they appreciated Hydra's simplicity and strong empirical performance.

In this shared response, **we aim to contextualize our paper within the existing literature** by highlighting our core contributions and outlining the scope and impact of our findings. We wish to emphasize that the **Matrix Mixer (MM) framework is a robust tool for creating performant sequence mixers**, as evidenced by the new sequence mixers like Quasiseparable (QS), Cauchy, and Vandermonde developed using MM prescriptions. We also note that these **mixers can seamlessly integrate with various domain-specific backbones** developed by the community. Our goal is to **showcase that this framework is effective and to advocate for its adoption in developing new sequence mixers.**

--- 

## Core contributions

### 1. The Matrix Mixer Framework

Unlike traditional paradigms that assess a method in its entirety with all its bells and whistles, our **formal** framework offers a different perspective to study models by focusing on their associated matrix mixer which is the core component of a sequence mixer.

We demonstrate that this framework is effective as it possesses the following desirable properties:

1.1. **A formal treatment**

- We formally define a sequence mixer as a **generalized matrix mixer** and identify the computational bottleneck as the cost associated with the matrix class. This prompts us to **focus on structured matrices** that enjoy a fast matrix multiplication algorithm.
- We then seek to determine whether some structured matrices could exhibit superior empirical performance than others. It is widely acknowledged in the machine learning community that data dependency is crucial for performance, and extendability is practically useful. Using these as our desiderata, **our formal framework allows us to define Sequence Aligned Matrices (SAM)**, which characterize matrices possessing these properties.

1.2. **Broad applicability across past methods**
- Next, we demonstrate that our framework is **highly general** and capable of representing a wide range of previous works. In Section 2.3, we show that **more than 13 past methods**, including Attention [40], Mamba [6, 16], MLP-Mixer [38], and LA [22], can be subsumed under our paradigm.
- Moreover, this broad applicability allows us to **compare these previous works on an equal footing**, as shown in Tables 1 and 2, eliminating extraneous complexities of the backbone model architecture.

1.3. **The framework provides effective prescriptions**

- To **validate our framework and test its predictions**, in Section 2.4, we identify Vandermonde, Cauchy, and QS matrices as SAM matrices which have not been fully utilized in past works. Our experiments demonstrate that all these matrices exhibit strong empirical performance, with **QS outperforming all others** and **Cauchy being competitive with Low Rank.**
- Another important validation of our framework is our success with Vandermonde matrices outperforming DFT matrices (a special case of Vandermonde matrices). **FNet [23] (Appendix A.3) previously attempted and failed to make DFT matrices data-dependent and performant.**

This demonstrates that our framework is not only descriptive but also prescriptive, which is a hallmark of robust generalization.

### 2. Hydra: A comprehensive validation and application of QS matrices

- Our framework and ablations suggest that **QS matrices have the potential to be the sequence mixer of choice** for the next generation of sub-quadratic sequence mixers. This potential is further accentuated by the fact that in addition to strong empirical performance, QS matrices also possess the following desirable properties:
  - *Connections to Mamba*: Since QS matrices generalize Semiseparable matrices, they are a **natural mathematical extension of Mamba** to bidirectional settings.
  - *Hardware Efficiency*: QS matrices can be implemented with Mamba as a subroutine, **enabling the use of Mamba's hardware-efficient Triton kernel.**

  To substantiate our claim, we introduced **Hydra**, a bidirectional model with QS matrix as its sequence mixer. Our evaluations on two canonical domains: language and images demonstrated the superior performance of **Hydra over BERT (110M) and ViT-B (88M), achieving +0.8 and +2.2 improvements, respectively.**
 
- Hydra also addresses an ongoing research question in the ML community: **how to make Mamba bidirectional**. In response, numerous methods, including Caduceus [36], BiGS [42], Vision Mamba [46], and MH-SSM [12], have proposed solutions that, when viewed under our framework, employ one of three heuristics: Add, Concat, or Element-wise multiplication. Our results demonstrate that our method **outperforms all these heuristic approaches** (see Table 3). We provide a new ablation study to support our argument (see Response A1 for Reviewer RwkG).

---

## Clarifying the Focus: Sequence Mixer vs. Model Architecture

Our primary objective is to develop a framework that enables us to **derive performant sub-quadratic sequence mixer alternatives** to attention. Consequently, our methodology focuses on the sequence mixers and **does not address the backbone model architecture**, which is typically tuned to meet domain-specific requirements. Therefore,the appropriate comparisons are with other sequence mixers, including attention. For instance, in vision, ViT [11] is a standard backbone that employs attention, whereas Swin Transformer [24], which adopts a hierarchical structure suitable for images, also uses attention.

We would like to emphasize that these two orthogonal **developments actually complement each other**. We invite domain-specific research communities to consider QS as a potential sequence mixer in their respective domains.

Meta-review: The paper presents a compelling formal framework called the Matrix Mixer Framework, which effectively extends sequence models, including Transformers and SSMs, by incorporating a bidirectional quasiseparable matrix mixer. The introduction of Hydra, a novel extension of the Mamba model, showcases significant advancements in non-causal tasks like masked language modeling and image classification. The paper’s contribution is substantial, demonstrating superior performance over existing approaches, particularly through its strategic use of matrix parameterization and bidirectional capabilities. Nevertheless, the paper’s presentation could be improved for clarity and readability. Despite this, I recommend accepting the paper given the strength of its methodology and results.
Decision: Accept (poster)