Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers

Authors: sukjun  hwang
aakash  lahoti
ratish  puduppully
tri  dao
albert  gu
low  rank
semiseparable  quasiseparable
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-09 07:18:11

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.

Main paper: 
 introduction: 
 Large-scale pretrained models such as GPT [35] , BERT [10] , and ViT [11] exhibit state-of-the-art performance across a wide range of tasks in multiple domains, including language and vision. A large number of these pretrained models follow a multi-layer architectural blueprint: a sequence mixer, such as Self-Attention 1 [42] , aggregates information across the input sequence, followed by a channel mixer processes 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 in a single step. However, this capability incurs a quadratic cost with respect to 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 systems remains a substantial challenge.One general strategy in designing alternative sequence models involves substituting the Attention matrix with different matrix parameterizations as the core sequence mixer. These modifications are motivated by various goals. For instance, simplifying the sequence mixer has led to the development of models such as MLP-Mixer [40] , which uses dense matrices, and FNet [23] , which utilizes the Discrete Fourier Transform matrix. Additionally, incorporating inductive biases such as positional information has resulted in the use of Toeplitz matrices in models like TNN [33] . Enhancing computational efficiency has spurred the creation of low-rank structures in models like Linear Attention (LA) [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 as 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, 38] have been made, but these attempts remain ad-hoc.In this work, we study the matrix mixer framework (see Figure 1 ), 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 as efficiency, causality, and expressivity. For instance, we highlight that structured matrix [29] parameterizations underpin efficient sequence models through computationally efficient algorithms (Section 2.1). (ii) Introduction of sequence alignment. Within this framework, we define sequence alignment as a novel attribute, which induces sequence models with essential features of data-dependence and extendability. We introduce Sequence Aligned Matrices (SAM) that exhibit these properties, demonstrating their superior performance in downstream tasks compared to non-aligned counterparts (Section 2.2). (iii) Exploration of structured matrix parameterizations. Through the matrix mixer framework, we systematically explore and categorize a broad spectrum of existing sequence models (Section 2.3). Motivated by sequence alignment, we also present new sequence models with underexplored structured matrix configurations, such as Vandermonde and Cauchy matrices (Section 2.4).Building on our matrix mixer framework, we introduce a 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 generalize 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 to 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-wise addition [15, 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 a 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 experiments confirm that Sequence Alignment, a property we newly identified in certain matrix mixers, significantly enhances downstream performance. Furthermore, our experimental findings demonstrate that novel sequence models like those using Cauchy matrix mixers match the performance of established matrix mixers such as low-rank. This indicates that the strong performance of low-rank variants is not solely 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, 44, 12] . Importantly, Hydra excels as a performant, general-purpose bidirectional sequence model, as evidenced by its strong performance across diverse domains. It achieves state-of-the-art results on the GLUE benchmark [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 [11] by 2.2 points.We publicly release source code at https://github.com/goombalab/hydra. 
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 on 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.Definition 2.1 (The matrix mixer framework). Let X ∈ R L×C be the input sequence, consisting of L elements, each with C channels. Let f X : R L×C → R L×D be the input preprocessing function that encapsulates common data transformations such as short convolutions, and linear layers. Let H and P be the number of heads and head dimension respectively, such that HP = D. Let M ⊆ R L×L represent the underlying class of mixer matrices. For each head h ∈ [H], let f h M : R L×C ×Θ → M be the matrix construction function that maps input data to mixer matrices, where Θ is the space of learnable parameters. We denote M h = f h M (X,θ), when M,θ,X are clear from the context. Then, we define matrix mixing as,Y h = M h (f X (X)) h ,where (f X (X)) h ,Y h denote the preprocessed input and output slice corresponding to head h.This definition encapsulates the view that existing sequence mixers can be mathematically represented as L×L mixer matrices that act on the sequence length (see Figure 1 , left). The framework not only incorporates typical data preprocessing steps like projections [22, 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 is generally true in practice, then the computational bottleneck lies in the Mf X (X) operation. For a general matrix M, this multiplication will incur a O(L 2 ) cost. The only way to mitigate this to restrict M to being a structured matrix, which are known to possess sub-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.Unstructured mixer matrices, also known as dense matrices, lack two key useful properties: 1) these matrix mixers cannot be easily made data-dependent, which has been identified as a 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 T i kj) O(L 2 ) ✓ Transformer [42] (Softmax Attention) Low-Rank q T i kj O(L) ✓Linear Attention [22] , (Linear Attention)Linformer [45] Butterfly See [7, 5] O(LlogL) Kaleidoscope [8, 7] , Monarch [5, 13] Toeplitz mj-i O(LlogL) S4 [17, 18] , H3 [14] (Convolution) TNN [33] , CKConv [37] Discrete Fourier Transformw ij O(Llog 2 L) FNet [23] Vandermonde (mi) j O(Llog 2 L) ✓ Ours (Section 2.4) Cauchy d (qid -kjd) -1 Semiseparable c T i A × i:j bj1 {i≥j} (1), (2) O(L) ✓ Mamba (S6 [16], SSD [6]) Quasiseparable Equation (3) 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 ∈ R L×L denote a matrix with a parameter set P. Then, we say that M is a Sequence Aligned Matrix if there exists a partition Π of P ⊆ P, and P ̸ = ∅, such that for all sets E ∈ Π, there exists a bijective map f E : [L] → E, and, for each i,j ∈ [L], i ≤ j, the sub-matrix M[i : j +1,i : j +1] is composed solely from the parameters in the subset ∪ E,i≤k≤j f E (k) ⊆ P.In simpler terms, this implies that each parameter of a SAM is either mapped to a specific element of the sequence or is left data-independent, 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-independent ones. Proposition 2.3 (Data Dependency). Sequence aligned matrices exhibit canonical data-dependent parameterization.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 value by treating the element itself as 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: the pretrained parametric functions assigned to each partition enable the computation of matrices larger than those encountered during training.We identify data dependency -the property induced by SAM -as a key axis of differentiation amongst existing models. 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. Consequently, various works have adopted different interpretations of this notion: Hyena [31] implements it via a datadependent linear operators over the Toeplitz matrix mixer; GSS [26] adds a 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 a 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-independent parameters.Using the formalization of the Matrix Mixer framework, we categorize a wide array of previous works -MLP-Mixer [40] , Transformer [42] , Linear Attention [22] , Linformer [45] , S4 [18, 17] , H3 [14] , TNN [33] , CKConv [37] , FNet [23] , Kaleidoscope [8, 7] , Monarch [5, 13] , and Mamba [16, 6] -as matrix mixers in Table 1 . For illustrative purposes, we explicitly show that MLP-Mixer, FNet and LA are matrix mixers (proofs in Appendix B), and leave out normalizing factors for simplicity as follows: Proposition 2.5. (MLP-Mixer is a matrix mixer). MLP-Mixer employs f X as an identity function, and its mixer matrix M has an unstructured parameterization with a single head (H = 1).Proposition 2.6. (LA is a structured matrix mixer with sequence alignment). LA employs f X as a linear projection f X (X,W V ) = XW V , with its low-rank mixer matrix being SAM. The mixer matrix M has H heads with a structured parameterization, where each (i,j)-element m ij is defined as:m ij = ϕ(x i W Q )ϕ(x j W K ) T , ϕ(•) = elu(•)+1, W Q ,W K ∈ R C×D .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 Q W T K where W Q ,W K ∈ R L×C are trainable parameters. This formulation yields a low-rank matrix mixer, which qualifies as a 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 matrix mixer without sequence alignment). FNet employs f X as an identity function, using the canonical Vandermonde matrix -Discrete Fourier Transform (DFT)for its mixer matrix. The mixer matrix M has a single head (H = 1) with a structured parameterization, where each (p,q)-element m pq is defined as m pq = w pq , where w = e -2πi N .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 investigate the influence of SAM on the expressivity of sequence mixers in Section 4.1.1.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: 1) well-known families of structured matrices with sub-quadratic matrix multiplication, hence are efficient as sequence models, and 2) we show that they have SAM parameterizations. In Section 4.1.1, we validate that these properties are enough to create efficient and strong sequence models, e.g. our Cauchy 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 ∈ R L×C is given, and utilize the prevalent query-key concept that we employQ = XW Q , K = XW K where W Q ,W K ∈ R C×D .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 E.Sequence aligned Vandermonde mixer. In contrast to FNet, which employs a fixed-parameter Vandermonde mixer based on the Discrete Fourier Transform (DFT), we propose a trainable, and datadependent Vandermonde-based matrix mixer. Then, we can construct a mixer matrix from a combination of two separate Vandermonde matrices M Q ,M K , each generated from Q and K. Specifically, each (i,j)-element of the resulting mixer matrix M is formulated as m ij = d (cos(2πq j id )-cos(2πk i jd )). This formulation utilizes the cosine function -the real part of an imaginary number -which is a 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 1 , the underlying philosophy of using both the Cauchy and low-rank matrix classes as a sequence mixer is remarkably similar: the relevance of a pair of elements is directly proportional to their associated value -the greater the relevance, the higher value assigned. Specifically, each (i, j)-element of the mixer matrix M is constructed by m ij = d 1 /(q id -k jd +c), where c is a trainable constant that prevents elements leading to excessive values when the denominator q id -k jd approaches zero. To our knowledge, this is the first introduction of a Cauchy-based sequence mixer that achieves performance comparable to Attention matrix mixers.3 Hydra: The Double-Headed MambaIn this section, we introduce Hydra, a 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.1). 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, a choice grounded by solid mathematical foundations (Section 3.2). Additionally, we underline the computational benefits and enhanced parameter efficiency afforded by adopting quasiseparable matrices (Section 3.3). 2). The rank characterization of SS matrices include the diagonals (e.g., green submatrices), whereas that of QS matrices hold for offdiagonal submatrices (e.g., yellow submatrices). Because of the similar rank properties, a naive addition-based 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 a higher expressivity compared to the mixer matrices of the addition-based bidirectional SSMs, where the diagonal values are constrained by the colored vectors.𝐜 ! " 𝐀 !:$ ×" 𝐛 $ " 𝐜 ! " 𝐀 !:& ×" 𝐛 & " 𝐜 ! " 𝐀 !:' ×" 𝐛 ' " 𝐜 $ " 𝐀 $:! ×" 𝐛 ! " 𝐜 $ " 𝐀 $:& ×" 𝐛 & " 𝐜 $ " 𝐀 $:' ×" 𝐛 ' " 𝐜 & " 𝐀 &:! ×" 𝐛 ! " 𝐜 & " 𝐀 &:$ ×" 𝐛 $ " 𝐜 & " 𝐀 &:' ×" 𝐛 ' " 𝐜 ' " 𝐀 ':! ×" 𝐛 ! " 𝐜 ' " 𝐀 ':$ ×" 𝐛 $ " 𝐜 ' " 𝐀 ':& ×" 𝐛 & " (c) Addition of Two Separate SSMs (d) Hydra δ ! 𝐜 $ " 𝐀 $:$ ×" 𝐛 $ " 𝐜 $ " 𝐀 $:& ×" 𝐛 & " 𝐜 $ " 𝐀 $:' ×" 𝐛 ' " 𝐜 ! " 𝐀 !:! ×" 𝐛 ! " δ $ 𝐜 & " 𝐀 &:& ×" 𝐛 & " 𝐜 & " 𝐀 &:' ×" 𝐛 ' " 𝐜 $ " 𝐀 $:! ×" 𝐛 ! " 𝐜 $ " 𝐀 $:$ ×" 𝐛 $ " δ & 𝐜 ' " 𝐀 ':' ×" 𝐛 ' " 𝐜 & " 𝐀 &:! ×" 𝐛 ! " 𝐜 & " 𝐀 &:$ ×" 𝐛 $ " 𝐜 & " 𝐀 &:& ×" 𝐛 & " δ ' 𝐜 ! " 𝐛 ! " + 𝐜 ! " 𝐛 ! " 𝐜 # " 𝐛 # " + 𝐜 # " 𝐛 # " 𝐜 $ " 𝐛 $ " + 𝐜 $ " 𝐛 $ " 𝐜 % " 𝐛 % " + 𝐜 % " 𝐛 % " (a) Semiseparable 𝒒 ! " 𝒌 ! 𝒒 $ " 𝒌 ! 𝒒 $ " 𝒌 $ 𝒒 & " 𝒌 ! 𝒒 & " 𝒌 $ 𝒒 & " 𝒌 & 𝒒 ' " 𝒌 ! 𝒒 ' " 𝒌 $ 𝒒 ' " 𝒌 & 𝒒 ' " 𝒌 ' 𝒒 + " 𝒌 ! 𝒒 + " 𝒌 $ 𝒒 + " 𝒌 & 𝒒 + " 𝒌 ' 𝒒 + " 𝒌 + δ ! 𝒄 ! " 𝒃 $ 𝒄 ! " 𝒃 & 𝒄 ! " 𝒃 ' 𝒄 ! " 𝒃 + 𝒒 $ " 𝒌 ! δ $ 𝒄 $ " 𝒃 & 𝒄 $ " 𝒃 ' 𝒄 $ " 𝒃 + 𝒒 & " 𝒌 ! 𝒒 & " 𝒌 $ δ & 𝒄 & " 𝒃 ' 𝒄 & " 𝒃 + 𝒒 ' " 𝒌 ! 𝒒 ' " 𝒌 $ 𝒒 ' " 𝒌 & δ ' 𝒄 ' " 𝒃 + 𝒒 + " 𝒌 ! 𝒒 + " 𝒌 $ 𝒒 + " 𝒌 & 𝒒 + " 𝒌 ' δ + (b) Quasiseparable3.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 a 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 X ∈ R L×C → Y ∈ R L×C -can be succinctly represented within the matrix mixer framework, as detailed below:EQUATIONEQUATIONHere, m ij represents an (i,j)-element of the mixer matrix M ∈ R L×L , with each matrix A i ∈ R N ×N and vector c i ,b i ∈ R N ×1 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 -semiseparable iff any submatrix from 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, 15, 12] 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, or concatenation. Among such heuristics, addition-based bidirectional extensions of SSMs [15, 48, 38] can be conceptualized within our matrix mixer framework, as illustrated in Figure 2 (c). 
quasiseparable matrices: a principled bidirectional matrix mixer: 
 We fully utilize the matrix mixer framework, which is discussed in Section 2, to explore a 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:EQUATIONwhere each δ i is a scalar, b i ,c i ∈ R N ×1 , and A i ∈ R N ×N [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 strictly upper or lower triangle (off from the diagonal) has a rank of at most N . See Figure 2 (b) .This definition emphasizes the rank constraint inherent in quasiseparable matrices, which is also evident from Equation (3) given that each vector c i ,b i ∈ R N ×1 and matrix A i ∈ R N ×N 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 addition-based bidirectional extensions of SSMs [15, 48, 38] through the lens of our matrix mixer framework. Unlike other elements, the diagonal values in a mixer matrix embody a unique concept of residuals, serving as a critical aspect of model expressivity. As demonstrated in Figure 2 (c), the mixer matrices in these bidirectional SSMs exhibit constraints intheir diagonal elements { -→ c T i - → b i + ← - c T i ← - b i } L , which are directly influenced by the shared non-diagonal construction vectors { - → c i , - → b i , ← - c i , ← - b i } L .Importantly, the rank characterization of semiseparable matrices includes on-diagonal elements, whereas that of quasiseparable matrices applies only to 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 are strictly more expressive than mixer matrices of addition-based 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 in downstream task performance, thereby substantiating our theoretical claims (see Appendix D.1).Moreover, with their structural similarity to semiseparable matrices, quasiseparable matrices are confirmed as sequence aligned matrices. Given our experimental results that SAM parameterizations are the key to the strong representational power (Section 4.1.1), we further validate our choice of quasiseparable matrices for the bidirectional sequence mixer. Proposition 3.6. N -quasiseparable matrices are sequence aligned matrices. Proof. quasiseparable matrices, due to their structural similarity to semiseparable matrices, belong to the class of Sequence Aligned matrices. Specifically, the set of parameters is given byP = P = {A i ,b i ,c i ,δ i } L . We consider the partition Π = {{A i } L ,{b i } L ,{c i } L ,{δ i } L },and for each element of the partition set, we choose the bijection that maps token i to A i , b i , c i , and δ i respectively. Finally, it is easy to see that the sub-matrix M [: i+1,: i+1] indeed only contains parameters in the set {A j ,b j ,c j ,δ j } i , 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.As a direct consequence of the favorable mathematical properties of quasiseparable matrices, we 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 easy-to-implement, sub-quadratic matrix multiplication algorithm. Specifically, we recognize that quasiseparable matrices can be expressed as a combination of two semiseparable matrices. Proposition 3.7. Let X ∈ R L×D be the input sequence, and let QS(•) and SS(•) denote the action of a quasiseparable and semiseparable matrix respectively. Let the two matrices share the parameters {A i , b i , c i } L , and define D = diag(δ 1 , ••• , δ L ), where δ i 's are the diagonal parameters of the quasiseparable matrix. Then,QS(X) = shift(SS(X))+flip(shift(SS(flip(X))))+DX,where flip(•) denotes the operation that reverses the input sequence, while shift(•) refers to shifting the sequence right by one position, padding the beginning with zero. (Proof in Appendix B)The above proposition demonstrates that quasiseparable matrix multiplication can 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 with a variety of SSMs [18, 17, 16, 6] and can also be generalized with other recurrent models [47, 28] .Furthermore, Hydra significantly improves parameter over the heuristic approaches to bidirectional modeling using SSMs [15, 48, 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 a quasiseparable matrix mixer (see Figure 4 ), we naturally share the f X projection layer, which accounts for a bulk of the model size. Empirically, we observe only a marginal increase in the total number of parameters compared to a single SSM, and can cut the the number of parameters nearly in half compared to the heuristic approaches to bidirectionality. 
experiments: 
 In Section 4.1, we begin by analyzing the matrix mixer framework through extensive performance comparisons between different structured matrix classes (Section 4.1.1). The data-dependent 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 D.1), and show that our Hydra outperforms them, underscoring the utility of the matrix view of sequence mixers.Then, in Section 4.2, we validate the effectiveness of our method by evaluating Hydra on quintessential language and image benchmark tasks: Masked Language Modeling (Section 4.2.1) and Image Classification (Section 4.2.2). State-of-the-art performance in both tasks has generally been dominated by transformer-based models [10, 11] . Our Hydra serves as an efficient and powerful replacement to the transformer layer, outperforming it on both tasks.In the presentation of results across all tables, the highest performing scores are highlighted in bold, while the second-highest scores are marked with an underline. Each number is the average of five runs. Results. 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-key QK T interaction.The results show that the structures of mixer matrices affect the capability of sequence models. However, we demonstrate that the key factor that primarily contributes to 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 E for implementation details). This adaptation reveals that the sequence-aligned Cauchy variant performs nearly as well as the sequence-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.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 [44, 12] . We provide details of the ablation studies in Appendix D.1.Results. The results presented in Table 3 show the shortcomings of the unidirectional Mamba [16] when applied to bidirectional contexts, as evidenced in C4 and GLUE. This significant performance disparity (-4.0) underscores the essential need for models to be capable of bidirectional sequence processing. Within the comparison of four bidirectional variants shown in Table 3 , our approach of utilizing a 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 is further validated in Figure 3 , where it demonstrates the cross-entropy 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.We pretrain our models on 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. Results. The results in Table 4 , showcase that Hydra outperforms all existing state-of-the-art methods. Notably, Hydra surpasses the performance of BERT -trained with the latest HuggingFace recipe [46] -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 a 1.8% improvement in accuracy of C4 validation and a 0.8% increase in the average GLUE score, illustrating the effectiveness of leveraging matrix mixer view for bidirectional settings. We assess Hydra on the renowned ImageNet-1K benchmark [9] , which comprises 1.28M training images and 50k validation images across 1, 000 categories. We use the ViT-Base [11] model as a baseline to facilitate a rigorous comparison of various sequence mixers by substituting the Transformer block in ViT with alternative sequence mixer models, 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-specialized Transformers [24] that include additional 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 [31] , 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 to 0.5 as stronger regularization. We relegate further experimental details in Appendix D.2.Results. The results, as presented in Table 5 , compare the performance of Hydra with ViT [11] and other variants [27, 31] on ImageNet-1K. 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. 
conclusion: 
 In this work, we have explored a 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 alignment) 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 a 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.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: [Yes] Justification: In the supplementary, we provide source code with instructions in a README file. Moreover, we provide implementations of ablations studies in Section E. Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a 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, or as supplemental material. 
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, or 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 to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is 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 a 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 or tables in the text. 
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 to 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 of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
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 A.Guidelines:• The answer NA means that there is 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 or 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 a direct path to any negative applications, the authors should point it out. For example, it is 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, it 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 are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [No] Justification:Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: 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?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 or dataset.• The authors should state which version of the asset is used and, if possible, include a 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. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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. 
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 as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of 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 contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: A.1 LimitationsIn this section, we share two limitations of our work, namely 1) Representation-Computation Tradeoff, and 2) Hardware Efficiency.Representation-Computation Tradeoff. While structured matrix mixers are computationally more efficient than their dense matrix mixer counterparts like softmax attention, they are also representationally less expressive, which may be seen as a limitation of these methods. For instance, concurrent works [1, 2, 21] have begun investigating the representational power of SSMs by analyzing their performance on memorization-centric tasks. They report that SSMs with a fixed model capacity are eventually outperformed by softmax attention for longer sequences. This can be viewed as a consequence of the matrix being too structured, and hence too inexpressive for the problem.On the other hand, we remark that the degree of structure of a structured matrix is a knob that can be tuned according to the specific task, that is we can tradeoff the computational efficiency of the method for larger expressivity. For instance, within the structured matrix class of low rank matrices, we can 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 is the modern variant of Mamba. Specifically, we show that SSD is able to 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 [1, 2, 21] that SSMs underperform on memorization-centric tasks. This makes it evident that the development and analysis of SSMs is an active area of research with substantial room for exploration and improvement.Hardware Efficiency: Despite the fact that structured matrices have associated sub-quadratic matrix multiplication algorithms, their implementation may not be hardware-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.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 societal impacts that necessitate special emphasis within this context.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 MLP-Mixer does not pre-process the input sequence, we set f X to be the identity function. Additionally, since the mixer matrix is data independent and is learned as a free parameter, we define the class M := R L×L , and set f M (X;θ) := θ3 , for θ ∈ Θ, where Θ := R L×L . Since MLP-Mixer shares the matrix across all channels, the number of heads H = 1. The output is given by Y (0) = θX (0) , where θ is the learned mixer matrix, which matches with the functional form of MLP-Mixer's sequence mixing primitive. Y h = σ(Q h )σ(K h ) T V h, where Q h ,K h ,V h and projections of the input. Since LA preprocesses the input sequence via projection, we set f X = W V ∈ R C×D , with W V as a learned parameter. The class of structured matrices M is clearly the class of Low Rank matrices. We define the set Θ = (R C×Hd ) 2 to correspond to the learned parameters W Q ,W K for the {Q h ,K h } h matrices, where d is the internal dimension for keys and queries. We set the matrix-generating function f h M to compute M h as M h = σ(Q h )σ(K h ) T through linear projections, followed by applying element-wise nonlinearity σ, and a matrix multiplication. The output Y h = M h (f X (X)) h matches the expected functional form of LA's sequence mixing primitive. Here, we provide architectural specifics of our method Hydra. As shown in Figure 4 and previously discussed Figure 5 , Hydra is largely divided into two sub-components, 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. Bidirectionality. We provide details of the experiments in . Each method under consideration employs 12 layers. As depicted in Table 3 , due to shared projection layers for forward and backward SSMs, Hydra only requires an additional 2M parameters compared to its unidirectional counterpart. Moreover because of the parameter increase in the concatenation variant, we lower its hidden dimension to match parameters. As [CLS] token is typically placed at the start of a sequence, we add an additional technique only to Mamba, which substitutes a [CLS] token with a global average pool of tokens.Training Details of Hydra We align our setup with the BERT-Base [10] 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 Table 4 are reported in Table 6 , and the settings used for obtaining the results of Hydra in Table 5 are listed in Table 9 .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 a batch size of 4096. We adopt the same hyperparameters as M2 [13] for pretraining, including the optimizer, learning rate schedule, and the bert-base uncased tokenizer. We employ the controlled setting from [20] for finetuning models on the GLUE benchmark. Specifically, we utilize weights pretrained on the C4 corpus for the five tasks -MNLI, QNLI, QQP, SST2, and COLA -and finetune RTE, STSB, and MRPC from an MNLI 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, 13] . Therefore, we perform a 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 Table 7 first row, we observe that Hydra of 24 layers trained using the M2 recipe out-of-the-box 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 8 , 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 internally-consistent training framework, and that our baselines are strong and generally consistent with those found in MosaicBERT [32] .ImageNet-1K. For Table 5 , we use 12 layers for the Transformer [42] , S4ND [27] , and Hyena [31] blocks, and set to 24 layers for Mamba and Hydra to match the number of parameters. The training involves images of 224×224 resolution, tokenized using a 16×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-1K training techniques [41] , such as augmentation and regularization strategies. All other configurations remain consistent with the original ViT [11] setup, ensuring that any observed differences in performance can be directly attributed to the expressivity of the respective sequence mixer layers.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 2, and a hidden dimension (d_model) of 768. We then adjust the channel dimensions for qk_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-dependent (Proposition 2.3). In contrast, variants without SAM are data-independent, with shared parameters across all inputs. The To facilitate understanding and reproducability of these methods, we provide PyTorch codes in Figure 6 , Figure 7 , Figure 8 , Figure 9 , Figure 10 , and Figure 11 . Our primary focus on this ablation is 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-independent 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 to incorporate the SAM attribute is not straightforward, implementing the vanilla Dense mixers (i.e. without SAM) is extremely simple: they are equivalent to MLP-Mixer [40] , employing a mixer matrix of R L×L . [18, 17] , formulated with R 2L-1 parameters. We also present a Toeplitz matrix mixer with SAM properties, thus data-dependent and can handle sequences of 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 ∈ 0,...L-1, each token x i generates two convolution parameters m -i and m i . Using all m parameters, a data-dependent dynamic convolutional weight is generated as{m -L+1 ,m -L+2 ,...,m -1 ,m 0 ,m 1 ,...m L-2 ,m L-1 }.Vandermonde. In Section 2.4, we introduced the single-headed sequence aligned Vandermonde matrix mixer. Leveraging the definition, we present multi-headed implementation by a 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 L , a matrix M is Cauchy if each (i,j)-entry m ij satisfies m ij = 1 qi-kj ; q i -k j ̸ = 0.Observe that general Quasiseparable (QS) matrices can be viewed as 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 on the model performance by individually removing either the diagonal or shift term, as well as removing both (Add) to analyze why QS outperforms Add.We 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 we observe a clear improvement. This result validates our choice of QS matrices, as neither Add+Diag nor Add+Shift alone has the expressivity of QS which is strictly more expressive than the addition of two SS matrices. def forward ( self , x , forward_conv = None , reverse_conv = None ) :""" # q , k : ( nheads , seqlen , qkdim ) q = torch . exp ( self . q_matrix ) k = torch . exp ( self . k_matrix ) i n v _ c a u c h y _ m a t r i x = ( q . unsqueeze ( 1 ) + k . unsqueeze ( 0 ) ) + self . tol cauchy_matrix = torch . sum ( 1 / inv_cauchy_matrix , dim = -1 ) 2 .def forward ( self , v , q = None , k = None ) : residual = v = rearrange (v , 'b l ( n h ) -> b l n h ' , n = self . nheads ) if self . i s _ doutput = torch . einsum ( 'b t n h , l t n -> b l n h ' , v , cauchy_matrix ) output = self . std_dev * output output = rearrange ( output , 'b l n h -> b l ( n h ) ') + residual return outputFNet In 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 of the matrix representation of the DFT to the input sequence. This representation, denoted as M , takes the form of a Vandermonde matrix, constructed from the roots of unity:EQUATIONwhere indices n and k range from 0 to L-1. First, we define ω n as the n-th root of unity:EQUATIONThen, the matrix M , which represents the DFT, is expressed using ω n in the form of a Vandermonde matrix as follows:EQUATIONMLP-Mixer In the MLP-Mixer architecture, sequence mixing is implemented using a MLP M . This MLP operates on each token in the sequence and is formulated as follows: M = W 2 σ(W 1 ) where σ denotes a non-linear activation function (such as ReLU). In this context:• W 1 is a weight matrix with dimensions R d S ×L , transforming the sequence from its original sequence length L to an intermediate dimension d S . • W 2 is a weight matrix with dimensions R L×d S , converting the dimensions back from the intermediate dimension d S to the original sequence length L.This design means that W 1 and W 2 first compress and then expand the sequence dimensions, respectively. The choice of the inner dimension d S is made independent of the sequence length L. Such a configuration leads to a 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 ∈ R N ×C , where N represents the number of 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 matricesW Q ,W K ,W V ∈ R C×D .The attention output is calculated as follows:EQUATIONEQUATIONEQUATIONEQUATIONA generalized version of self-attention can be represented as [22] :EQUATIONwhere the standard softmax attention employs similarity(q,k) = e q T k √ D as the similarity function. In the context of 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 σ(Q h ), σ(K h ), and V h , where σ denotes a non-linearity function. The set of learned parameters for these projections, denoted as Θ, includes W Q ,W K corresponding to the query and key matrices across all heads, defined as Θ = (R C×Hd ) 2 . 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 is characterized as the class of Low Rank matrices. The matrix-generating function for each head h, denoted as f h M , computes the matrix M h as M h = σ(Q h )σ(K h ) T . This computation involves linear projections of the queries and keys, application of element-wise nonlinearity σ, and matrix multiplication.The output for each head h, denoted as Y h , is then computed as Y h = M h (f X (X)) h , where f X = W V ∈ R C×D represents a projection of the input sequence with W V as a learned parameter matrix. This projection transforms the input data X 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 2 ).Discrete Convolution Consider a convolution operation between sequence X and the filter sequence h to produce the output sequence Y . Such a convolution can be represented by a matrix multiplication with an L×L Toeplitz matrix M [31] .Let:• X be the input signal, a column vector of length L, i.e., X ∈ R L×1 .• 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:M =     h 0 h -1 ••• h -L+1 h 1 h 0 ••• h -L+2 . . . . . . . . . . . . h L-1 h L-2 ••• h 0    Thus, the output Y = M X represents the convolution of X with h, with both X and Y having the same length L, achieved through the L×L Toeplitz matrix M as defined above.• If the contribution is a dataset and/or model, the authors should describe the steps taken to 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, or 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source 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.Justification: The research conform Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).We adjust the learning rate to 1e-5 to address instabilities observed in the training of FNet on COLA.For simplicity, we overload the notation to denote a finite-dimensional linear map with its corresponding matrix. 


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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 1. The authors claim that the sequence model should be extended beyond their trained length, but they didn't provide the corresponding results.

2. 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.

3. Please compare with the most advanced methods, such as xLSTM. It also expands the hidden state into a matrix form.
questions: 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?

2. How to understand "Sequence Aligned" property in Table 1?
limitations: limitations: N/A
paper_summary: paper_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.
strength_weakness: strength_weakness: 1. The presentation of the paper needs to be improved. 
    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.
    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. 

Overall, the current representation makes it difficult for the readers to understand the core contributions. 

2. 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?
questions: 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: limitations: The limitations are discussed and are reasonable.
paper_summary: paper_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.
strength_weakness: strength_weakness: I don't identify a specific weakness, but I have a few questions regarding efficiency. Please see the following section.
questions: 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.
- 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)?
- 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: limitations: See Questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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:
- What does “common data transformations” mean in definition 2.1.
- 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.
 - “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?
- 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.
- The proof of proposition 2.6 does not prove the sequence alignment property.
- Are Quasiseparable matrices established in the literature or have you defined them? There is no citation provided afaict.
- Should the tilde be a hat on P in line 237 (over vice versa in line 131).
- 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.
- “Taming the Hydra” isn’t particularly insightful for a subsection title.
- 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?)
- “It is known that quasiseparable matrices are closed under addition” - cite or prove.
2. 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.
3. 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.
4. 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)?
5. 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?
questions: 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?

Typo:
- linear 38 "they often lack a systematic" - remove "and"
limitations: limitations: There is a discussion of some limitations in the appendix.

Comments: title: Author response
Thank you for your response
Thanks again
response to the rebuttal
Statement of Thanks
Thanks for the response
Acknowledgement
References
comment: I thank the authors for their efforts in responding to my questions.

## No-Shift
The response has not clarified my question. In the previous reply, the authors stated that "We would first like to clarify that there is no data-dependent $DX$ term in the No Shift variant". But now in the Lemma the $DX$ term appears in the "No-shift" variant. Does the No-Shift result in the ablation table have the $DX$ term included or not?

Moreover, in the Lemma (part 2) you only prove $\text{No-Shift} \subseteq \text{QS}$ but do not show that there is a strict subset. I maintain that the expressivity is the same so that the Lemma part 2 is an equality; please correct me if I am wrong. But if so then the expressivity argument you have made for the ablation table doesn't hold. I agree that the Lemma would make a nice addition to the paper, but only if there is a clear and consistent message for the readers to take away.

## Convolution
Thank you for this ablation. I would include this to strengthen the arguments of the paper.
### Further ablations on the impact of convolution

We appreciate the reviewer's interest in understanding the impact of non-matrix mixer backbone components, such as convolution, on the model's performance. Below, we provide an ablation study to quantify the impact of removing the short convolution from the backbone. Due to the time constraints, we are currently comparing the variants on their pretraining validation accuracy and cross-entropy loss.


| Method | #Params | Conv | $L_{ce}$ | Acc (%) |
| ------ | ------- | ---- | -------- | ------- |
| Add    | 70M     | ❌   | 1.70     | 65.3    |
| Quasi  | 70M     | ❌   | *1.68*     | *65.6*    |
| Add    | 70M     | ✅   | *1.68*     | *65.6*    |
| Quasi  | 70M     | ✅   | **1.66**     | **65.9**    |

We make the following observations: 1) QS consistently outperforms the Add variant, both with and without convolution. 2) Adding a short convolution consistently improves the performance of both Add and QS.

This indicates that the established advantages of Quasi over heuristics like Add, Concat, and Multiply persist across changes to the underlying backbone. At the same time, we fully agree with the reviewer that further studies on backbone components are valuable and can indeed help enhance the performance of sequence mixers in general.
We genuinely appreciate and thank the reviewer for their continued interest and detailed consideration of our work.

### 1. Expressivity of QS ablation matrices


We would like to clarify our previous response: we commented on the expressivity of ([Figure, c](https://photos.app.goo.gl/xAsegE2NcLAwrMi96)) which corresponds to the "Add" ablation and not the "No-Shift" ablation. We apologize for the misunderstanding and would like to take the opportunity to clarify the expressivity of all variants of the quasiseparable ablations that the reviewer suggested.

For convenience, we reproduce the ablation table below:

| Method   | #Params | $L_{ce}$ | Acc (%) | GLUE |
|----------|---------|----------|---------|------|
| Add      | 70M     | 1.68     | 65.6    | 80.6 |
| No Diag  | 70M     | 1.68     | 65.7    | *80.7* |
| No Shift | 70M     | *1.67*     | *65.8*    | *80.7* |
| QS    | 70M     | **1.66**     | **65.9**    | **81.7** |

---
---

**Lemma:** The matrix classes satisfy the following inclusions:
$$\text{Add} \subseteq \text{No-Shift} \subseteq \text{QS}, \quad \text{and} \quad \text{No-Diag} \subseteq \text{QS}.$$

**Proof:**
1. ### $\text{Add} \subseteq \text{No-Shift}$

To see this, observe that we can characterize the set of No-Shift matrices as: 
$$ \text{No-Shift} = \\{ W + D \\: | \\: \forall W \in \text{Add}, \\: \forall D \in \text{diag}(\mathbb{R}^d)\\}.$$Choosing $D$ to be the zero matrix, we have the result:
$$ \text{No-Shift} \supseteq  \\{ W + \mathbf{0} \\: | \\: \forall W \in \text{Add}\\} = \text{Add}$$

---

2. ### $\text{No-Shift} \subseteq \text{QS}$

Let $\tilde{SS}_u$ and $\tilde{SS}_l$ denote the (the tilde represents "not strict") upper-triangular and lower-triangular N-Semiseparable matrices. Then,
$$ \text{No-Shift} = \\{ M = \tilde{SS}_u + \tilde{SS}_l + D \\: | \\: \forall \tilde{SS}_u,\\: \tilde{SS}_l, \\: \forall D \in \text{diag}(\mathbb{R}^d)\\}.$$ We apply the definition of N-QS matrices [C]. Without loss of generality, consider any submatrix $L$ in the strict lower triangular half of $M$. Observe that $L$ is also a submatrix of $\tilde{SS}_l$. By the definition of N-Semiseparable matrices, $\text{Rank}(L) \le \text{N}$, which implies $M \in \text{QS}$, and hence $\text{No-Shift} \subseteq \text{QS}$

We also kindly refer the reviewer to [H, §4.1.5], which discusses this result in detail.

---

3. ### $\text{No-Diag} \subseteq \text{QS}$

Let $SS_u$ and $SS_l$ denote strict upper-triangular and strict lower-triangular parts of a N-Quasiseparable matrices. Then,
$$ \text{No-Diag} =  \\{ SS_u + SS_l \\: | \\: \forall SS_u,SS_l \\},$$

$$ \text{QS} =  \\{ SS_u + SS_l + D\\: | \\: \forall SS_u,SS_l, \\: \\: \forall D \in \text{diag}(\mathbb{R^d}) \\}.$$

Restricting $D$ to be the zero matrix in the definition of $\text{QS}$, we have the required result. 

---
---

We now tie the expressivity of different variants with their empirical performance: higher the expressivity, higher the GLUE score. That is the performance follows the same order: $\text{QS}(81.7) > \text{No-Shift}(80.7) > \text{Add}(80.6)$, and that $\text{QS}(81.7) > \text{No-Diag}(80.7)$. This validates our premise of principally deriving QS from the Matrix Mixer framework over using its heuristical counterparts. 

We again thank the reviewer for suggesting this ablation. We have now added it to our paper and incorporated our discussions on expressivity into the appendix.

---
[H]: Eidelman, Y., Gohberg, I., & Haimovici, I. (2013). *Separable Type Representations of Matrices and Fast Algorithms: Volume 1 Basics. Completion Problems. Multiplication and Inversion Algorithms*. Springer, Basel, Switzerland. [Link](https://link.springer.com/10.1007/978-3-0348-0606-0).
Thanks again for the discussion. 

Regarding 1., thank you for the response and figure. My understanding of the "No shift" ablation to Hydra was that only the shift would be removed and you would still have the data-dependent diagonal component of Hydra. More concretely, you would have the data-dependent diagonal $\delta$ added to Figure c in the linked figure, which should have the same expressivity because the $\delta$s are free to "unlearn" the existing diagonals in Figure c. Is that correct or have I missed something?

Regarding 2. and 3., I see where the authors are coming from and don't have major disagreements, but will clarify my points of view which differ from that of the authors. Regarding 2., I maintain that these ablations would strengthen the current work, and encourage the authors to do so. Different architectural components interact in complicated ways, and it would make the matrix mixers proposed here (like Hydra) much more compelling in my view if they can perform well across different backbones. Regarding 3., the accessibility of the contributions to readers is very important in my view, alongside the contributions themselves.
Thank you for the answers. 

I understand that the paper focuses on introducing the matrix mixer framework and Hydra as an application of the Quasiseparable Matrix (the best matrix mixer). Therefore, the authors think my suggestions about adding recent SOTA comparisons and including another application for Hydra are out of scope. I believe that improving the experiments provides a better understanding of the framework (quality) and the practical use of Hydra. This will enhance the overall quality of the paper. 

Also, regarding my concern about the presentation of the paper, the authors believe that the current structure better matches the flow of the contributions. I found that the current flow distracts from understanding the core ideas (the reviewer RwkG09 also pointed out). It's because the paper introduces many terms/methods, but the connections between them are unclear. This makes the paper difficult to read. I suggest simplifying the paper, so the reliability can improve. 

Overall, the paper introduces a novel framework with great insight. However, due to the weaknesses mentioned above, I retain my rating.
We sincerely thank the reviewer for their very positive feedback and for taking the time to engage with our work. We are glad that we have addressed all their concerns and we appreciate their recognition of the significance and merit of our work.
We are grateful to the reviewer for their continued engagement and thoughtful questions. And we appreciate this opportunity to provide further clarification on their concerns:

---

### 1. Expressivity of the "No Shift/Add" ([Figure, c](https://photos.app.goo.gl/xAsegE2NcLAwrMi96)) variant :

We would first like to clarify that there is no data-dependent $DX$ term in the "No Shift" variant. Intuitively, the "No Shift" variant exhibits reduced expressivity compared to general QS matrices ([Figure, d](https://photos.app.goo.gl/xAsegE2NcLAwrMi96)) because the diagonals in $SS_l$ and $SS_u$, specifically $\overrightarrow{c}^T \overrightarrow{b} + \overleftarrow{c}^T \overleftarrow{b}$, share parameters with the non-diagonal elements. This **parameter tying** limits the model’s expressivity, as *it cannot independently assign values to the diagonal elements*.

Mathematically, observe that the "No Shift" variant satisfies the definition of an N-QS matrix [C] and therefore is a **subset** of the N-QS matrices ([Figure, d](https://photos.app.goo.gl/xAsegE2NcLAwrMi96)). Furthermore, the "No Shift" variant is a **strict subset** of general N-QS matrices because the lack of shift forces additional constraints on the matrix class. For instance, in general QS matrices, the diagonal is freely parameterized, while in the "No Shift" variant, it is a function of non-diagonal elements, $\overrightarrow{c}, \overrightarrow{b}, \overleftarrow{c}, \overleftarrow{b}$.

Our additional ablations reported in our first response (A.1) at the reviewer's suggestion validate this increase in expressivity by demonstrating an improved performance of N-QS matrices compared to the "No Shift" and the "No Diagonal" variants, which are strict subsets of N-QS matrices.

---

### 2. Role of Convolution and Block Structure

We would like to begin by noting that **a short convolution within the block is a standard element** used across a vast majority of sub-quadratic models like Mamba2[6], H3[14], Mamba[16], Hyena[29]. We fully agree with the reviewer's insight that investigating the impact of block structure on the performance of the matrix mixer is important, and this is indeed being explored in parallel works like [G]. However, in our work, we focused specifically on **fixing a sub-quadratic block [6] to control for the impact of the matrix mixer**. We simply selected the block architecture from Mamba2 [6] as it is one of the latest iterations of sub-quadratic models. With this view, we sincerely believe that the reviewer’s suggestion of understanding the impact of these architectural choices outside of the matrix mixer would be an important follow-up work.

---

### 3. Scope of Changes

We are truly grateful for the reviewer's valuable suggestions to improve the clarity and accessibility of our work. However, we would like to emphasize that our changes primarily involve adding more precise and mathematically complete definitions, and this would not alter the core results or contributions of the paper. 

---


[G] Yang, Songlin, et al. "Parallelizing Linear Transformers with the Delta Rule over Sequence Length." arXiv preprint arXiv:2406.06484 (2024).
Thank you for the clarifications, and ablation on diagonal/shift operations. My concerns are largely addressed, although I view the further ablations on the convolution and block structure as within the scope of the submission and would encourage the authors to perform these ablations. Thanks in advance also for updating the presentation in the next revision, though the promised updates to the writing represent significant changes to the original submission's presentation which make it hard to assess the original submission.

I have an additional question following the diagonal/shift ablation:
- Does the "No shift" row in the ablation have the same expressivity as a QS matrix in the definition of [C]? The diagonal terms should be the sum of the $SS_l$ and $SS_u$ diagonal terms along with the $DX$ term, which should have the same expressivity as the $DX$ term I believe.
Thank you for your response. Most of my concerns have been addressed and I will keep my very positive score. I appreciate the significant contribution and merit of this work as an important follow-up to mamba.
[A] Xia, Jianlin, et al. “Fast algorithms for hierarchically semiseparable matrices.” Numerical Linear Algebra with Applications, 17.6 (2010): 953-976.\
[B] Boito, Paola. “Matrix Structures and Matrix Functions.” Proceedings of the 2023 International Symposium on Symbolic and Algebraic Computation. 2023.\
[C] Pernet, Clément, Hippolyte Signargout, and Gilles Villard. "Exact computations with quasiseparable matrices." Proceedings of the 2023 International Symposium on Symbolic and Algebraic Computation. 2023.\
[D] Bertsch, Amanda, et al. "Unlimiformer: Long-range transformers with unlimited length input." Advances in Neural Information Processing Systems 36. 2023.\
[E] Ding, Jiayu, et al. "Longnet: Scaling transformers to 1,000,000,000 tokens." arXiv preprint arXiv:2307.02486. 2023.\
[F] Liu, Hao, Matei Zaharia, and Pieter Abbeel. "Ring attention with blockwise transformers for near-infinite context." International Conference on Learning Representations. 2024.