<!DOCTYPE html><html lang="en">
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8">
<title>[2407.09941] Untitled Document</title>
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta name="twitter:card" content="summary">
<meta name="twitter:title" content="Untitled Document">
<meta name="twitter:image:src" content="https://ar5iv.labs.arxiv.org/assets/ar5iv_card.png">
<meta name="twitter:image:alt" content="ar5iv logo">
<meta property="og:title" content="Untitled Document">
<meta property="og:site_name" content="ar5iv">
<meta property="og:image" content="https://ar5iv.labs.arxiv.org/assets/ar5iv_card.png">
<meta property="og:type" content="article">
<meta property="og:url" content="https://ar5iv.labs.arxiv.org/html/2407.09941">

<!--Generated on Mon Aug  5 16:37:09 2024 by LaTeXML (version 0.8.8) http://dlmf.nist.gov/LaTeXML/.-->
<meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no">

<script>
  function detectColorScheme(){
    var theme="light";
    var current_theme = localStorage.getItem("ar5iv_theme");
    if(current_theme){
      if(current_theme == "dark"){
        theme = "dark";
      } }
    else if(!window.matchMedia) { return false; }
    else if(window.matchMedia("(prefers-color-scheme: dark)").matches) {
      theme = "dark"; }
    if (theme=="dark") {
      document.documentElement.setAttribute("data-theme", "dark");
    } else {
      document.documentElement.setAttribute("data-theme", "light"); } }

  detectColorScheme();

  function toggleColorScheme(){
    var current_theme = localStorage.getItem("ar5iv_theme");
    if (current_theme) {
      if (current_theme == "light") {
        localStorage.setItem("ar5iv_theme", "dark"); }
      else {
        localStorage.setItem("ar5iv_theme", "light"); } }
    else {
        localStorage.setItem("ar5iv_theme", "dark"); }
    detectColorScheme(); }
</script>
<link media="all" rel="stylesheet" href="/assets/ar5iv-fonts.0.7.9.min.css"><link media="all" rel="stylesheet" href="/assets/ar5iv.0.7.9.min.css"><link media="all" rel="stylesheet" href="/assets/ar5iv-site.0.2.2.css">
</head>
<body>
<div class="ltx_page_main">
<div class="ltx_page_content">
<article class="ltx_document">
<div id="p1" class="ltx_para">
<span id="p1.1" class="ltx_ERROR undefined">\newtoggle</span>
<p id="p1.2" class="ltx_p">arxiv
<span id="p1.2.1" class="ltx_ERROR undefined">\toggletrue</span>arxiv
</p>
</div>
</article>
</div>
<div class="ar5iv-footer"><a href="/html/2407.09940" class="ar5iv-nav-button ar5iv-nav-button-prev">◄</a>
    <a class="ar5iv-home-button" href="/"><img height="40" alt="ar5iv homepage" src="/assets/ar5iv.png"></a>
    <a href="/feeling_lucky" class="ar5iv-text-button">Feeling<br>lucky?</a>
    <a href="/log/2407.09941" class="ar5iv-text-button ar5iv-severity-error">Conversion<br>report</a>
    <a class="ar5iv-text-button" target="_blank" href="https://github.com/dginev/ar5iv/issues/new?template=improve-article--arxiv-id-.md&title=Improve+article+2407.09941">Report<br>an issue</a>
    <a href="https://arxiv.org/abs/2407.09941" class="ar5iv-text-button arxiv-ui-theme">View&nbsp;original<br>on&nbsp;arXiv</a><a href="/html/2407.09942" class="ar5iv-nav-button ar5iv-nav-button-next">►</a>
</div><footer class="ltx_page_footer">
<a class="ar5iv-toggle-color-scheme" href="javascript:toggleColorScheme()" title="Toggle ar5iv color scheme"><span class="color-scheme-icon"></span></a>
<a class="ar5iv-footer-button" href="https://arxiv.org/help/license" target="_blank">Copyright</a>
<a class="ar5iv-footer-button" href="https://arxiv.org/help/policies/privacy_policy" target="_blank">Privacy Policy</a>

<div class="ltx_page_logo">Generated  on Mon Aug  5 16:37:09 2024 by <a target="_blank" href="http://dlmf.nist.gov/LaTeXML/" class="ltx_LaTeXML_logo"><span style="letter-spacing:-0.2em; margin-right:0.1em;">L<span class="ltx_font_smallcaps" style="position:relative; bottom:2.2pt;">a</span>T<span class="ltx_font_smallcaps" style="font-size:120%;position:relative; bottom:-0.2ex;">e</span></span><span style="font-size:90%; position:relative; bottom:-0.2ex;">XML</span><img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAsAAAAOCAYAAAD5YeaVAAAAAXNSR0IArs4c6QAAAAZiS0dEAP8A/wD/oL2nkwAAAAlwSFlzAAALEwAACxMBAJqcGAAAAAd0SU1FB9wKExQZLWTEaOUAAAAddEVYdENvbW1lbnQAQ3JlYXRlZCB3aXRoIFRoZSBHSU1Q72QlbgAAAdpJREFUKM9tkL+L2nAARz9fPZNCKFapUn8kyI0e4iRHSR1Kb8ng0lJw6FYHFwv2LwhOpcWxTjeUunYqOmqd6hEoRDhtDWdA8ApRYsSUCDHNt5ul13vz4w0vWCgUnnEc975arX6ORqN3VqtVZbfbTQC4uEHANM3jSqXymFI6yWazP2KxWAXAL9zCUa1Wy2tXVxheKA9YNoR8Pt+aTqe4FVVVvz05O6MBhqUIBGk8Hn8HAOVy+T+XLJfLS4ZhTiRJgqIoVBRFIoric47jPnmeB1mW/9rr9ZpSSn3Lsmir1fJZlqWlUonKsvwWwD8ymc/nXwVBeLjf7xEKhdBut9Hr9WgmkyGEkJwsy5eHG5vN5g0AKIoCAEgkEkin0wQAfN9/cXPdheu6P33fBwB4ngcAcByHJpPJl+fn54mD3Gg0NrquXxeLRQAAwzAYj8cwTZPwPH9/sVg8PXweDAauqqr2cDjEer1GJBLBZDJBs9mE4zjwfZ85lAGg2+06hmGgXq+j3+/DsixYlgVN03a9Xu8jgCNCyIegIAgx13Vfd7vdu+FweG8YRkjXdWy329+dTgeSJD3ieZ7RNO0VAXAPwDEAO5VKndi2fWrb9jWl9Esul6PZbDY9Go1OZ7PZ9z/lyuD3OozU2wAAAABJRU5ErkJggg==" alt="Mascot Sammy"></a>
</div></footer>
</div>

    <script>
      var canMathML = typeof(MathMLElement) == "function";
      if (!canMathML) {
        var body = document.querySelector("body");
        body.firstElementChild.setAttribute('style', 'opacity: 0;');
        var loading = document.createElement("div");
        loading.setAttribute("id", "mathjax-loading-spinner");
        var message = document.createElement("div");
        message.setAttribute("id", "mathjax-loading-message");
        message.innerText = "Typesetting Equations...";
        body.prepend(loading);
        body.prepend(message);

        var el = document.createElement("script");
        el.src = "https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js";
        document.querySelector("head").appendChild(el);

        window.MathJax = {
          startup: {
            pageReady: () => {
              return MathJax.startup.defaultPageReady().then(() => {
                body.removeChild(loading);
                body.removeChild(message);
                body.firstElementChild.removeAttribute('style');
              }); } } };
      }
    </script>
    <script>
    // Auxiliary function, building the preview feature when
    // an inline citation is clicked
    function clicked_cite(e) {
      e.preventDefault();
      let cite = this.closest('.ltx_cite');
      let next = cite.nextSibling;
      if (next && next.nodeType == Node.ELEMENT_NODE && next.getAttribute('class') == "ar5iv-bibitem-preview") {
        next.remove();
        return; }
      // Before adding a preview modal,
      // cleanup older previews, in case they're still open
      document.querySelectorAll('span.ar5iv-bibitem-preview').forEach(function(node) {
        node.remove();
      })

      // Create the preview
      preview = document.createElement('span');
      preview.setAttribute('class','ar5iv-bibitem-preview');
      let target = document.getElementById(this.getAttribute('href').slice(1));
      target.childNodes.forEach(function (child) {
        preview.append(child.cloneNode(true));
      });
      let close_x = document.createElement('button');
      close_x.setAttribute("aria-label","Close modal for bibliography item preview");
      close_x.textContent = "×";
      close_x.setAttribute('class', 'ar5iv-button-close-preview');
      close_x.setAttribute('onclick','this.parentNode.remove()');
      preview.append(close_x);
      preview.querySelectorAll('.ltx_tag_bibitem').forEach(function(node) {
        node.remove();
      });
      cite.parentNode.insertBefore(preview, cite.nextSibling);
      return;
    }
    // Global Document initialization:
    // - assign the preview feature to all inline citation links
    document.querySelectorAll(".ltx_cite .ltx_ref").forEach(function (link) {
      link.addEventListener("click", clicked_cite);
    });
    </script>
    </body>
</html>

2407.09941
 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, Hydra outperforms BERT by 0.8 points on the GLUE benchmark and ViT by 2%  Top-1 accuracy on ImageNet. 
abstract
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Abstract
2407.09941
 ``` Large-scale pretrained models such as GPT [34], 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-Attention1[41], 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 [39], 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 [32]. Enhancing computational                                                                                                                     efficiency has spurred the creation of low-rank structures in models like Linear Attention (LA) [22] and Linformer [44], as                                                                                                                     well as Monarch matrices [5] in models such as M2 [13].                                                         ∗Equal Contributions.                       1    In this paper, Attention [3] exclusively refers to Self-Attention.                                                      ```  ![1_image_1.png](1_image_1.png)  ![1_image_0.png](1_image_0.png)  Despite the achievements of these approaches, they often and 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 [47, 37] 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 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, 47, 37], the Hadamard product [43], or concatenation [43, 12, 37] - 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, 47, 37, 43, 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 [42] 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 and pretrained weights at https://github.com/goombalab/hydra.  ``` 
introduction
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Introduction
2407.09941
 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
The Matrix Mixer Framework: Bridging Sequence Mixers And Structured Matrices
2407.09941
 Definition 2.1 (The matrix mixer framework). Let X ∈ R × be the input sequence, consisting of  *elements, each with* channels. Let  : R × → R × *be the input preprocessing function that encapsulates common data transformations. Let* and  be the number of heads and head dimension respectively, such that  = *. Let* M ⊆ R ×*represent the underlying* class of mixer matrices. For each head ℎ ∈ []*, let* ℎ M : R × × Θ → M be the matrix construction function that maps input data to mixer matrices, where Θ *is the space of learnable parameters. We denote* Mℎ =  ℎ M (X, ), when M, , X are clear from the context. Then, we define matrix mixing as, Y ℎ = Mℎ( (X))ℎ, 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Formalizing The Matrix Mixer Framework
2407.09941
 ``` This definition encapsulates the view that existing sequence mixers can be mathematically represented as  ×  mixer matrices that act on the sequence length (see Figure 1, left). The framework not only incorporates typical data preprocessing                                                                                                                      steps like projections [22, 38], and short convolutions [16, 6], but it also accommodates data dependent mixer matrices [38,                                                                                                                      33]. Furthermore, it is also powerful enough to capture the concept of head structure [41] 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 ,                                                                                                                      which is generally true in practice, then the computational bottleneck lies in the M (X) operation. For a general matrix M, this multiplication will incur a (                                      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.                                        ```  | Table 1: Categorization of existing methods as matrix mixers. 𝐿 denotes input sequence length.   |                       |                         |                                      |                          |                          | |---------------------------------------------------------------------------------------------------|-----------------------|-------------------------|--------------------------------------|--------------------------|--------------------------| | Matrix Structure M                                                                                | Formulation (𝑚𝑖𝑗)  | Complexity              | Sequence Aligned                     | Method Instantiations    |                          | | Dense                                                                                             | 𝑚𝑖𝑗                | 𝑂(𝐿 2 )               | MLP-Mixer [39]                       |                          |                          | | Dense                                                                                             | softmax𝑗 (q          | )                       | ✓                                    | Transformer [41]         |                          | | (Softmax Attention)                                                                               | 𝑇                    | 2                       |                                      |                          |                          | | k𝑗)                                                                                              | 𝑂(𝐿                 |                         |                                      |                          |                          | | 𝑖                                                                                                |                       |                         |                                      |                          |                          | | Low-Rank                                                                                          | k𝑗                   | 𝑂(𝐿)                  | ✓                                    | Linear Attention [22],   |                          | | 𝑇                                                                                                |                       |                         |                                      |                          |                          | |                                                                                                   | q                     |                         |                                      |                          |                          | | (Linear Attention)                                                                                | Linformer [44]        |                         |                                      |                          |                          | | 𝑖                                                                                                |                       |                         |                                      |                          |                          | | Butterfly                                                                                         | See [7, 5]            | 𝑂(𝐿 log 𝐿)           | Kaleidoscope [8, 7], Monarch [5, 13] |                          |                          | | Toeplitz                                                                                          | 𝑚𝑗−𝑖               | 𝑂(𝐿 log 𝐿)           | S4 [17, 18], H3 [14]                 |                          |                          | | (Convolution)                                                                                     | TNN [32], CKConv [36] |                         |                                      |                          |                          | | 𝑖𝑗                                                                                              | 𝑂(𝐿 log2 𝐿)        | FNet [23]               |                                      |                          |                          | | Discrete Fourier Transform                                                                        | 𝑤                    |                         |                                      |                          |                          | | Vandermonde                                                                                       | (𝑚𝑖) 𝑗             | 𝑂(𝐿 log2 𝐿)          | ✓                                    | Ours (Section 2.4)       |                          | | Cauchy                                                                                            | Í 𝑑                  | −1                      |                                      |                          |                          | | (𝑞𝑖𝑑 − 𝑘𝑗𝑑 )                                                                                |                       |                         |                                      |                          |                          | | Semiseparable                                                                                     | c 𝑇                  | × b𝑗1{𝑖≥𝑗 } (1), (2) | 𝑂(𝐿)                               | ✓                        | Mamba (S6 [16], SSD [6]) | | 𝑖 A 𝑖:𝑗                                                                                        |                       |                         |                                      |                          |                          | | Quasiseparable                                                                                    | Equation (3)          | 𝑂(𝐿)                  | ✓                                    | Ours (Hydra) (Section 3) |                          | 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Where ( (X))ℎ, Y ℎ Denote The Preprocessed Input And Output Slice Corresponding To Head ℎ.
2407.09941
 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 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  *be the sequence length and let* M ∈ R × *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* E : [] → E*, and, for each*  ∈ [], the sub-matrix M[:  + 1, :  + 1] *is composed* solely from the parameters in the subset ∪E,≤ E () ⊆ 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  is constructed using only the parameters corresponding sequence segment upto and including index  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  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 [30] implements it via a data-dependent 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. Using our definition of data dependency that aligns with the third interpretation, 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.              ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Sequence Aligned Matrices
2407.09941
 ``` Using the formalization of the Matrix Mixer framework, we categorize a wide array of previous works - MLP-Mixer [39],                                                                                                                       Transformer [41], Linear Attention [22], Linformer [44], S4 [18, 17], H3 [14], TNN [32], CKConv [36], 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  as an identity function, and its mixer matrix M has an unstructured parameterization with a single head ( = 1).  ```  Proposition 2.6. (LA is a structured matrix mixer with sequence alignment). LA employs  *as a linear projection*  (X, ) = X , with its low-rank mixer matrix being SAM. The mixer matrix M has  *heads with a structured parameterization, where* each (, )-element  *is defined as:*  =  ( ) ()  ,  (·) = elu(·) + 1, , ∈ R ×.  ``` Additionally, it is important to note that not all structured matrix mixers inherently exhibit sequence alignment. For                                                                                                                        example, consider the matrix mixer formulation M =                                                               where , ∈ R                                                                                   × 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  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 ( = 1) with a structured parameterization, where each (, )-element  is defined as  =                                                                                                               , where  =       −        2                   . 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.                                     ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Prior Sequence Models As (Structured) Matrix Mixers
2407.09941
 ``` 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                         × is given, and utilize the prevalent query-key concept that we employ Q = X , K = X where , ∈ R                   × . For simplicity, we show the single-headed ( = 1) case, where  = . 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 data-dependent Vandermonde-based matrix                                                                                                   mixer. Then, we can construct a mixer matrix from a combination of two separate Vandermonde matrices M, M, each generated from Q and K. Specifically, each (, )-element of the resulting mixer matrix M is formulated as  = Í      (cos(2                      ) − cos(2 )). 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 (, )-element of the mixer matrix M is constructed by  =Í  1/( −  +  ), where  is a trainable  ![5_image_0.png](5_image_0.png)  | "�!                 | �$ "�&                                          | �$ "�'                                                                 | �$ "�+                            |                                                                                      |               |                                   |                                   |                  | |---------------------|-------------------------------------------------|------------------------------------------------------------------------|-----------------------------------|--------------------------------------------------------------------------------------|---------------|-----------------------------------|-----------------------------------|------------------| | δ$                  | �$                                              | �! "�! " + �! "�! " �! "�!:$ ×" �$ " �& "�&:$ ×" �$ " �' "�':$ ×" �$ " | δ$                                | �! "�!:! ×" �! " �& "�&:! ×" �! " �' "�':! ×" �! "                                   |               |                                   |                                   |                  | | �! "�$              | δ!                                              | �! "�&                                                                 | �! "�'                            | �! "�+                                                                               | "�!:$ ×" �$ " | �& "�&:! ×" �! " �' "�':! ×" �! " |                                   |                  | | �!                  | "�$:$ ×" �$ "                                   | δ!                                                                     | �& "�&:& ×" �& " �' "�':& ×" �& " |                                                                                      |               |                                   |                                   |                  | | �$                  |                                                 |                                                                        |                                   |                                                                                      |               |                                   |                                   |                  | | �# "�# " + �# "�# " |                                                 |                                                                        |                                   |                                                                                      |               |                                   |                                   |                  | | "�!                 | �& "�&                                          | �& "�$                                                                 | �& "�!                            | δ&                                                                                   | �& "�'        | �& "�+                            | �& "�&:$ ×" �$ " �& "�&:! ×" �! " | �' "�':& ×" �& " | | �$ "�$ " + �$ "�$ " | �! "�!:$ ×" �$ " �! "�!:! ×" �! "               | δ&                                                                     | �' "�':' ×" �' "                  |                                                                                      |               |                                   |                                   |                  | | "�!                 | �' "�&                                          | �' "�'                                                                 | �' "�$                            | �' "�!                                                                               | �' "�&        | δ'                                | �' "�+                            |                  | | "�$                 | �+ "�!                                          | �+ "�&                                                                 | �+ "�'                            | δ+                                                                                   |               |                                   |                                   |                  | | �+                  | (b) Quasiseparable                              |                                                                        |                                   |                                                                                      |               |                                   |                                   |                  | | "�!                 | �+ "�&                                          | �+ "�'                                                                 | �+ "�+                            | "�':$ ×" �$ " �' "�':! ×" �! " �' "�':& ×" �& " �' (c) Addition of Two Separate SSMs | (d) Hydra     |                                   |                                   |                  | | "�% " + �% "�% " �% | "�&:$ ×" �$ " �& "�&:! ×" �! " �& "�&:& ×" �& " | δ'                                                                     |                                   |                                                                                      |               |                                   |                                   |                  | | �&                  |                                                 |                                                                        |                                   |                                                                                      |               |                                   |                                   |                  | | (a) Semiseparable   |                                                 |                                                                        |                                   |                                                                                      |               |                                   |                                   |                  |  ```                  +                   "$ +                      "! +                         "& +                            "' δ+                     (b) Quasiseparable Figure 2: (a) A semiseparable (SS) matrix. (b) A quasiseparable (QS) matrix. (c) A mixer matrix of addition-based bidirectional                                                                     SSMs. (d) A QS mixer matrix for Hydra. SS and QS matrices are characterized by rank conditions (Definition 3.1,                                                                     Definition 3.2). The rank characterization of SS matrices include the diagonals (e.g., green submatrices), whereas that of QS matrices hold for off-diagonal submatrices (e.g., yellow submatrices). Because of the similar rank properties, 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.                                                                  constant that prevents elements leading to excessive values when the denominator  −  approaches zero. To our knowledge, this is the first introduction of a Cauchy-based sequence mixer that achieves performance comparable to                                                                     Attention matrix mixers.                ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
The Matrix Mixer Framework As A Creative Toolbox
2407.09941
 ``` In 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).               ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Hydra: The Double-Headed Mamba
2407.09941
 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  × ↦→ Y ∈ R × - can be succinctly represented within the matrix mixer framework, as detailed below:  $\mathbf{y}_{t}=\sum_{s=0}^{t}\mathbf{C}_{t}^{T}\mathbf{A}_{t.s}^{\times}\mathbf{B}_{s}\mathbf{x}_{s}$  $\mathbf{y}=\text{SSM}(\mathbf{A},\mathbf{B},\mathbf{C})(\mathbf{X})=\text{MX}$  $$m_{i j}=\mathbf{c}_{i}^{T}\mathbf{A}_{i}\cdot\cdot\cdot\mathbf{A}_{j+1}\mathbf{b}_{j}$$ 6 (1)  $\frac{1}{2}$ .   $$\mathbf{A}_{i:j}^{\times}=\begin{cases}\prod_{k=j+1}^{i}\mathbf{A}_{k}&\text{for}i>j\\ 1&\text{for}i=j\end{cases}.\tag{2}$$ Here,  represents an (, )-element of the mixer matrix M ∈ R ×, with each matrix A ∈ R  × and vector c  , b ∈ R  ×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 *-semiseparable iff* any submatrix from the lower triangle (on or below the diagonal) has a rank of at most *. 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 [43, 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, 47, 37] can be conceptualized within our matrix mixer framework, as illustrated in Figure 2 (c). 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Background: State Space Models Are Semiseparable Matrix Mixers
2407.09941
 ``` 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: a matrix                                                                                                                         M is -quasiseparable if each element  satisfies  ```  where each is a scalar, b , c ∈ R  ×1, and A ∈ R  × [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 [29]). A matrix M is -quasiseparable iff any submatrix from either the strictly upper or lower triangle (off from the diagonal) has a rank of at most *. 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 , b ∈ R  ×1and matrix A ∈ R  × has a rank of at most . 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, 47, 37] 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 in their diagonal elements {                                                                   −→                                                                   c                                                                                                                                         −→b +                                                                          ←−                                                                          c                                                                                                                                                        ←−                                                                             b }, which are directly influenced by the shared non-diagonal construction vectors {                                          −→c                                            ,                                                                                          −→b                                               , ←−c                                                  ,                                                                                                      ←−                                                   b }. 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 Section 4.1.2).                                                                                                               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. -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 by P = P˜ = {A                                                                                , b                                                                                   , c                                                                                      ,  }. We consider the partition Π = {{A }, {b }, {c }, { }}, and for each element of the partition set, we choose the bijection that maps token  to A                                                                                                                      ,                                                                                                                        b    , c          , and  respectively. Finally, it is easy to see that the sub-matrix  [:  + 1, :  + 1] indeed only contains parameters in the set {A           , b              , c                 ,  }                                            , thus satisfying the last requirement of SAM matrices. □                                                                                                                         ```  $$\mathbf{\Phi}:\begin{cases}{\overrightarrow{\mathbf{c}_{i}^{\dagger}}}{\overrightarrow{\mathbf{A}_{i:j}^{\times}}}\mathbf{\Phi}_{j},&{\mathrm{if}}\;i>j\\ {\delta_{i},}&{\mathrm{if}}\;i=j\\ {\overleftarrow{\mathbf{c}_{j}^{\dagger}}}{\overleftarrow{\mathbf{A}_{j:i}^{\times}}}{\overleftarrow{\mathbf{b}_{i}}},&{\mathrm{if}}\;i<j\end{cases}$$ $$m_{i j}=\left\{\begin{array}{l}{{}}\\ {{}}\\ {{}}\\ {{}}\end{array}\right.$$  ![6_image_0.png](6_image_0.png)  $${\mathrm{(3)}}$$  , (3)  7 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Quasiseparable Matrices: A Principled Bidirectional Matrix Mixer
2407.09941
 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, subquadratic 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 × *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 , b , c }, and define D = *diag*(1, · · · , ), where *'s are the diagonal parameters of the quasiseparable matrix. Then,* 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Taming The Hydra
2407.09941
 ``` 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 [46,                                                                                                          28].     Furthermore, Hydra significantly improves parameter over the heuristic approaches to bidirectional modeling using SSMs                                                                                                          [15, 47, 43, 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                                                                                                          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.                                     ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Qs(X) = Shift(Ss(X)) + Flip(Shift(Ss(Flip(X)))) + Dx,
2407.09941
 ``` 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 (Section 4.1.2), 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 secondhighest scores are marked with an underline. Each number is the average of five runs.                                                                                    ```  Structure DD #Params C4 Pretrain GLUE Tasks GLUE  L Avg  Acc MNLI QNLI QQP RTE SST2 MRPC COLA STS  Dense 71M 2.05 59.6 73.3 76.2 85.3 64.4 90.8 84.7 45.7 76.8 74.7  Toeplitz 71M 1.97 60.8 74.6 79.6 86.6 66.1 90.9 84.2 45.7 79.1 75.8  Toeplitz ✓ 72M 1.91 61.9 77.3 81.8 88.1 67.1 90.7 87.3 45.3 84.1 77.7  DFT 71M 2.46 53.1 70.4 70.8 84.5 59.9 89.8 83.6 44.4 69.8 71.7  Vandermonde 71M 2.46 53.0 55.2 61.3 82.5 66.3 87.4 84.2 45.8 84.2 70.8  Vandermonde ✓ 70M 2.04 59.7 74.1 80.0 86.2 67.9 89.3 84.3 46.0 80.1 76.0  Cauchy 71M 2.25 56.2 75.3 81.3 86.7 66.6 88.8 84.5 27.4 83.2 74.2  Cauchy ✓ 70M 1.94 61.6 77.5 84.4 84.2 68.0 91.0 86.7 48.1 85.2 78.2  Low-Rank 71M 2.06 59.4 73.7 76.5 85.1 61.5 90.6 85.8 49.2 76.6 74.9  Low-Rank ✓ 70M 1.90 62.2 77.6 84.1 88.2 69.1 91.0 85.9 47.6 83.9 78.4  Attention 71M 2.08 59.1 71.0 70.4 83.5 62.3 89.9 83.3 **49.6** 65.2 71.9  Attention ✓ 70M 1.87 62.9 78.5 85.4 88.4 67.9 91.2 86.4 47.8 84.3 78.8  Quasiseparable 72M 2.03 59.8 73.8 78.1 87.1 64.3 90.2 84.4 45.5 77.2 75.1  Quasiseparable ✓ 71M 1.84 63.3 79.5 85.5 88.6 69.8 91.9 **88.4** 48.4 85.6 **79.7** 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Experiments
2407.09941
 ``` Our controlled experimental setting distinctly separates the mixer matrices from other architectural components, enabling                                                                                                                           a rigorous and focused comparison between different types of mixer matrices. Specifically, utilizing the recent Mamba [6]                                                                                                                           block, we only replace SSD with different mixer matrices M. In Table 2 we provide experimental results that showcase the expressivity of various matrix mixers that support bidirectional sequence processing, primarily by utilizing off-the-shelf                                                                                                                           structured matrix families for the mixer matrix. Further details of the experimental setup are provided in Appendix E.                                                                                                                    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, 44] was attributed to their low-rank based mixer                                                                                                                           matrices, particularly emphasized by the query-key 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.                                                                                        ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Analysis Of The Matrix Mixer Framework 9 4.1.1 Effects Of Different Structured Matrix Families
2407.09941
 ``` 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 [43], and concatenation [43, 12].                                          Experimental Setup. 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.  Table 2: Matrix mixer ablations. A systematic empirical study of matrix mixers on the GLUE benchmark by controlling for the architecture and varying only the matrix parameterization. Sequence-aligned matrices dynamically parameterize                                                                                                                               via input projections, becoming data-dependent (DD) that significantly increases performance. Most DD variants achieve                                                                                                                               competitive GLUE scores.                             ```  | Method   | #Params   | C4 Pretrain   | GLUE   |      | |----------|-----------|---------------|--------|------| | L        | Avg       |               |        |      | | 𝑐𝑒     | Acc (%)   |               |        |      | | Mamba    | 68M       | 3.45          | 39.5   | 77.7 | | Add      | 70M       | 1.68          | 65.6   | 80.6 | | Mult     | 70M       | 1.72          | 64.9   | 79.6 | | Concat   | 69M       | 1.71          | 65.4   | 81.1 | | Quasi    | 70M       | 1.66          | 65.9   | 81.7 |  Table 3: Performance of various approaches that extend Mamba to a bidirectional model. We compare our quasiseparable matrix mixer to element-wise addition (Add), the Hadamard product  (Mult), and concatenation (Concat) variants.  ![9_image_0.png](9_image_0.png)   Figure 3: Cross-entropy loss of various bidirectional variants, measured on the C4 validation set.  Table 4: **GLUE Results.** Evaluation of various sequence models that can be formulated as matrix mixers. For maximum performance, all models are trained using established recipes [31, 13].  | Method    | #Params   | C4 Pretrain   | GLUE Tasks   | GLUE   |      |      |      |      |      |       |      |      | |-----------|-----------|---------------|--------------|--------|------|------|------|------|------|-------|------|------| | L         | Avg       |               |              |        |      |      |      |      |      |       |      |      | | 𝑐𝑒      | Acc (%)   | MNLI          | QNLI         | QQP    | RTE  | SST2 | MRPC | COLA | STS  |       |      |      | | BERT      | 110M      | 1.59          | 67.3         | 84.4   | 90.3 | 89.7 | 77.1 | 92.3 | 90.7 | 54.2  | 89.1 | 83.5 | | MLP-Mixer | 112M      | 1.77          | 63.5         | 77.2   | 82.4 | 87.6 | 67.3 | 90.5 | 86.5 | 43.0  | 85.2 | 77.5 | | FNet      | 112M      | 1.94          | 61.3         | 74.9   | 82.1 | 85.7 | 63.6 | 87.6 | 86.4 | 42.72 | 83.1 | 75.8 | | M2        | 116M      | 1.65          | 65.9         | 80.5   | 86.0 | 87.0 | 69.3 | 92.3 | 89.2 | 56.0  | 86.9 | 80.9 | | Hydra     | 112M      | 1.46          | 69.1         | 84.5   | 90.0 | 91.3 | 77.5 | 93.5 | 91.2 | 57.2  | 88.9 | 84.3 |  Results. The results 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
2 Ablating Approaches For Bidirectionality
2407.09941
 ``` We pretrain our models on the masked language modeling objective using the Colossal Cleaned Common Crawl (C4)                                                                                                              corpus [35], then finetune and evaluate them on the GLUE benchmark [42].                                                                       Experimental Setup. 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 [39], 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 [31, 13] for optimal performance of each method. We relegate further experimental details in Appendix D.1.                                                                                                ```  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 [45] - in both pretraining and GLUE  benchmark scores. BERT has shown a noticeable gap in the masked language modeling task compared to other previous methods [39, 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.  2We adjust the learning rate to 1 − 5 to address instabilities observed in the training of FNet on COLA. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Evaluation Results Of Hydra 4.2.1 Bidirectional Masked Language Modeling
2407.09941
 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 [30], 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 [30], 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.1.  Table 5: Top 1 & 5 image classification accuracies evaluated on Method \#Params Top-1 (%) Top-5 (%) Acc AccEMA Acc AccEMA  ViT-B 87M 78.8 80.6 94.2 95.2 S4-ViT-B 89M 79.4 80.4 94.2 95.1 Hyena-ViT-B 88M 78.4 76.4 94.0 93.0 Mamba-ViT-B 89M 79.1 80.0 94.2 94.9 Hydra-ViT-B 91M 81.0 81.6 95.3 **95.6** Results. The results, as presented in Table 5, compare the performance of Hydra with ViT [11] and other variants [27, 30] 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
2 Image Classification
2407.09941
 ``` 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.                          ``` 
conclusion
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Conclusion
2407.09941
 ``` We gratefully acknowledge Ratish Puduppully for his assistance during the initial stages of this project. This research was                                                                                                                             made possible by the generous support of computational resources provided by Cartesia AI.                                                                                                ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
Acknowledgements
2407.09941
 [1] Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. "Zoology: Measuring and improving recall in efficient language models". In: *arXiv preprint arXiv:2312.04927* (2023).  [2] Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. "Simple linear attention language models balance the recall-throughput tradeoff". In: arXiv preprint arXiv:2402.18668 (2024).  [3] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. "Neural machine translation by jointly learning to align and translate". In: *arXiv preprint arXiv:1409.0473* (2014).  [4] Tom Bella, Yuli Eidelman, Israel Gohberg, and Vadim Olshevsky. "Computations with quasiseparable polynomials and matrices". In: *Theoretical Computer Science* (2008).  [5] Tri Dao, Beidi Chen, Nimit S Sohoni, Arjun Desai, Michael Poli, Jessica Grogan, Alexander Liu, Aniruddh Rao, Atri Rudra, and Christopher Ré. "Monarch: Expressive structured matrices for efficient and accurate training". In:  International Conference on Machine Learning. PMLR. 2022, pp. 4690–4721.  [6] Tri Dao and Albert Gu. "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality". In: *International Conference on Machine Learning (ICML)*. 2024.  [7] Tri Dao, Albert Gu, Matthew Eichhorn, Atri Rudra, and Christopher Ré. "Learning fast algorithms for linear transforms using butterfly factorizations". In: *International conference on machine learning*. PMLR. 2019, pp. 1517–1527.  [8] Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszczynski, Atri Rudra, and Christopher Ré. "Kaleidoscope: An efficient, learnable representation for all structured linear maps". In: *arXiv preprint* arXiv:2012.14966 (2020).  [9] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. "Imagenet: A large-scale hierarchical image database". In: *CVPR*. 2009.  [10] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. "Bert: Pre-training of deep bidirectional transformers for language understanding". In: *NAACL* (2019).  [11] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. "An image is worth 16x16 words:  Transformers for image recognition at scale". In: *ICLR* (2021).  [12] Yassir Fathullah, Chunyang Wu, Yuan Shangguan, Junteng Jia, Wenhan Xiong, Jay Mahadeokar, Chunxi Liu, Yangyang Shi, Ozlem Kalinli, Mike Seltzer, and Mark J. F. Gales. "Multi-Head State Space Model for Speech Recognition". In:  Proc. INTERSPEECH 2023. 2023, pp. 241–245. doi: 10.21437/Interspeech.2023-1036.  [13] Daniel Y Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. "Monarch Mixer: A simple sub-quadratic GEMM-based architecture".  In: *NeurIPS* (2023).  [14] Daniel Y. Fu, Tri Dao, Khaled Kamal Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. "Hungry Hungry Hippos: Towards Language Modeling with State Space Models". In: *The Eleventh International Conference on Learning* Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. url: https://openreview.net/ pdf?id=COZDy0WYGg.  [15] Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. "It's raw! audio generation with state-space models". In:  ICML. 2022.  [16] Albert Gu and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces". In: arXiv preprint arXiv:2312.00752 (2023).  [17] Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. "On the parameterization and initialization of diagonal state space models". In: *NeurIPS* (2022).  [18] Albert Gu, Karan Goel, and Christopher Ré. "Efficiently modeling long sequences with structured state spaces". In:  ICLR (2022).  [19] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. "Deep residual learning for image recognition". In:  Proceedings of the IEEE conference on computer vision and pattern recognition. 2016, pp. 770–778.  [20] Peter Izsak, Moshe Berchansky, and Omer Levy. "How to train BERT with an academic budget". In: *EMNLP*. 2021.  [21] Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. "Repeat after me: Transformers are better than state space models at copying". In: *arXiv preprint arXiv:2402.01032* (2024).  [22] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. "Transformers are rnns: Fast autoregressive transformers with linear attention". In: *ICML*. 2020.  [23] James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. "Fnet: Mixing tokens with fourier transforms".  ```        In: NAACL (2022). [24] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. "Swin transformer:                                                                                                                                                  Hierarchical vision transformer using shifted windows". In: ICCV. 2021. [25] Zhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer, Trevor Darrell, and Saining Xie. "A convnet for the                                                                                                                                                  2020s". In: Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022, pp. 11976–11986. [26] Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. "Long Range Language Modeling via Gated                                                                                                                                                  State Spaces". In: The Eleventh International Conference on Learning Representations. 2023. url: https://openreview.        net/forum?id=5MkYIYCbva. [27] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré.                                                                                                                                                  "S4nd: Modeling images and videos as multidimensional signals with state spaces". In: NeurIPS (2022). [28] Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De.                                                                                                                                                  "Resurrecting recurrent neural networks for long sequences". In: International Conference on Machine Learning. PMLR.        2023, pp. 26670–26698.                                   [29] Clément Pernet, Hippolyte Signargout, and Gilles Villard. "Exact computations with quasiseparable matrices". In:                                                                                                                                                  Proceedings of the 2023 International Symposium on Symbolic and Algebraic Computation. 2023, pp. 480–489. [30] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon,                                                                                                                                                  and Christopher Ré. "Hyena Hierarchy: Towards Larger Convolutional Language Models". In: ICML. 2023. [31] Jacob Portes, Alex Trott, Sam Havens, Daniel King, Abhinav Venigalla, Moin Nadeem, Nikhil Sardana, Daya Khudia,                                                                                                                                                  and Jonathan Frankle. "MosaicBERT: A Bidirectional Encoder Optimized for Fast Pretraining". In: arXiv preprint        arXiv:2312.17482 (2023). [32] Zhen Qin, Xiaodong Han, Weixuan Sun, Bowen He, Dong Li, Dongxu Li, Yuchao Dai, Lingpeng Kong, and Yiran                                                                                                                                                  Zhong. "Toeplitz Neural Network for Sequence Modeling". In: ICLR. 2023. [33] Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. "The devil in                                                                                                                                                  linear transformer". In: EMNLP (2022). [34] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. "Language models are                                                                                                                                                  unsupervised multitask learners". In: OpenAI blog 1.8 (2019), p. 9. [35] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and                                                                                                                                                  Peter J Liu. "Exploring the limits of transfer learning with a unified text-to-text transformer". In: JMLR (2020). [36] David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomczak, and Mark Hoogendoorn. "Ckconv: Continuous                                                                                                                                                  kernel convolution for sequential data". In: arXiv preprint arXiv:2102.02611 (2021). [37] Yair Schiff, Chia-Hsiang Kao, Aaron Gokaslan, Tri Dao, Albert Gu, and Volodymyr Kuleshov. "Caduceus: Bi-directional                                                                                                                                                  equivariant long-range dna sequence modeling". In: arXiv preprint arXiv:2403.03234 (2024). [38] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive        Network: A Successor to Transformer for Large Language Models. 2023. arXiv: 2307.08621 [cs.CL]. [39] Ilya O Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica                                                                                                                                                  Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, et al. "Mlp-mixer: An all-mlp architecture for vision". In:                                                                                                                                                  NeurIPS (2021). [40] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. "Training                                                                                                                                                  data-efficient image transformers & distillation through attention". In: ICML. 2021. [41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and                                                                                                                                                  Illia Polosukhin. "Attention is all you need". In: NeurIPS (2017). [42] Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel Bowman. "GLUE: A Multi-Task                                                                                                                                                  Benchmark and Analysis Platform for Natural Language Understanding". In: ICLR. 2019. [43] Junxiong Wang, Jing Yan, Albert Gu, and Alexander Rush. "Pretraining Without Attention". In: EMNLP. 2023. [44] Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. "Linformer: Self-attention with linear complexity".                                                                                                                                                  In: arXiv preprint arXiv:2006.04768 (2020). [45] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac,                                                                                                                                                  Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite,                                                                                                                                                  Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush.                                                                                                                                                  "Transformers: State-of-the-Art Natural Language Processing". In: EMNLP. 2020. [46] Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. "Gated linear attention transformers                                                                                                                                                  with hardware-efficient training". In: arXiv preprint arXiv:2312.06635 (2023).  ```  [47] Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. "Vision mamba: Efficient visual representation learning with bidirectional state space model". In: *arXiv preprint arXiv:2401.09417* (2024). 
references
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
References
2407.09941
 In 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
A Discussions A.1 Limitations
2407.09941
 ``` 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.                                                                                                            ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
A.2 Broader Impact
2407.09941
 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 MLPMixer does not pre-process the input sequence, we set  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 B R ×, and set M (X; ) B  3 , for  ∈ Θ, where Θ B R ×. Since MLP-Mixer shares the matrix across all channels, the number of heads  = 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. □  Proposition 2.6. LA is a Structured Matrix Mixer Proof. Recall that for each head ℎ ∈ [], LA computes, Y ℎ = (Q ℎ )(K ℎ ) V ℎ ,  where Qℎ, K ℎ , V ℎ and projections of the input. Since LA preprocesses the input sequence via projection, we set  = W ∈ R  × , with W as a learned parameter. The class of structured matrices M is clearly the class of Low Rank matrices. We define the set Θ = (R × ) 2 to correspond to the learned parameters W, W for the {Qℎ, K ℎ }ℎ matrices, where  is the internal dimension for keys and queries. We set the matrix-generating function ℎ M to compute Mℎas Mℎ = (Qℎ)(K ℎ )  through linear projections, followed by applying element-wise nonlinearity , and a matrix multiplication. The output Y ℎ = Mℎ( (X))ℎ matches the expected functional form of LA's sequence mixing primitive. □ Proposition 3.2. Let X ∈ R × *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 , b , c }, and define D = *diag*(1, · · · , ), where '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. The proof follows by observing that the first term, shift(SS(X)), models the effect of the lower triangular part of the QS matrix. The next term, flip(shift(SS(flip(X)))), represents the upper triangular part of the QS matrix, and the last term, DX, accounts for the QS matrix's diagonal. □  ![15_image_0.png](15_image_0.png)  ![15_image_1.png](15_image_1.png)  Figure 5: Pseudo code for Hydra. *, , ,*  denote batch size, sequence length, number of heads, and head dimension respectively. The suffices _f and _b denote forward and backward.  shift: Right-shift. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
B Proofs
2407.09941
 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
C Architectural Details
2407.09941
 ``` 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 [31] which has been widely adopted in recent works - the models are trained for 70 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 [31].               ```  ImageNet-1K. For Table 5, we use 12 layers for the Transformer [41], S4ND [27], and Hyena [30] 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 1 −4 and a batch size of 1024. We also incorporate regular ImageNet-1K training techniques [40], 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. 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
D Additional Details Of Main Results In Section 4.2 D.1 Training Details Of Hydra
2407.09941
 ``` 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 inclusion of projection layers for data dependency results in a marginal increase in parameter count for                                                                                                                           ```  n_epochs 3 10 5 3 3 10 10 10  seq_len 256 256 256 256 256 256 256 256  | [20] [13] Ours   | |------------------|  n_epochs 3 10 10 6 3 10 10 10  seq_len 128 128 128 128 128 128 128 128  n_epochs 2 7 3 3 2 10 10 10  seq_len 256 256 256 256 256 256 256 256  LR 5e-5 1e-5 3e-5 1e-5 3e-5 8e-5 5e-5 3e-5  WD 5e-6 1e-6 3e-6 1e-6 3e-6 8e-6 5e-6 3e-6  LR 5e-5 5e-5 3e-5 1e-5 3e-5 5e-5 5e-5 7e-5  WD 5e-6 1e-6 1e-2 1e-2 3e-6 1e-2 5e-6 1e-2  LR 1e-4 5e-5 5e-5 1e-5 5e-5 8e-5 1e-4 3e-5  WD 5e-6 1e-6 3e-6 1e-6 3e-6 8e-6 5e-6 3e-6  Recipe #Params Pretrain GLUE Tasks GLUE  L Avg  Acc (%)  MNLI QNLI QQP RTE SST2 MRPC COLA STS  [13] 115M 1.45 69.3 83.7 89.7 89.7 77.4 92.8 91.5 54.7 90.1 83.7  Ours 112M 1.46 69.1 84.5 90.0 91.3 77.5 93.5 91.2 57.2 88.9 84.3  | Method    | Source   | #Params   | GLUE Tasks   | GLUE   |      |      |      |      |      |      |      | |-----------|----------|-----------|--------------|--------|------|------|------|------|------|------|------| | MNLI      | QNLI     | QQP       | RTE          | SST2   | MRPC | COLA | STS  | Avg  |      |      |      | | Ours      | 110M     | 84.4      | 90.3         | 89.7   | 77.1 | 92.3 | 90.7 | 54.2 | 89.1 | 83.5 |      | | [31]      | 110M     | 84.1      | 89.8         | 91.2   | 77.2 | 91.2 | 87.5 | 54.6 | 88.9 | 83.2 |      | | BERT      | [23]     | 110M      | 82.5         | 91.0   | 87.0 | 69.0 | 93.0 | 83.0 | 73.0 | 89.0 | 83.3 | | [13]      | 110M     | -         | -            | -      | -    | -    | -    | -    | -    | 79.6 |      | | MLP-Mixer | Ours     | 112M      | 77.2         | 82.4   | 87.6 | 67.3 | 90.5 | 86.5 | 43.0 | 85.2 | 77.5 | | Ours      | 112M     | 74.9      | 82.1         | 85.7   | 63.6 | 87.6 | 86.4 | 42.7 | 83.1 | 75.8 |      | | FNet      | [23]     | 83M       | 72.5         | 80.0   | 83.0 | 63.0 | 95.0 | 76.0 | 69.0 | 79.0 | 76.7 | | [23]      | 238M     | 77.0      | 85.0         | 85.0   | 69.0 | 94.0 | 88.0 | 78.0 | 84.0 | 81.9 |      | | M2        | [13]     | 116M      | 80.5         | 86.0   | 87.0 | 69.3 | 92.3 | 89.2 | 56.0 | 86.9 | 80.9 | | Hydra     | Ours     | 112M      | 84.5         | 90.0   | 91.3 | 77.5 | 93.5 | 91.2 | 57.2 | 88.9 | 84.3 |  Parameter Value  Image size 2242  Optimizer AdamW  Optimizer momentum 1, 2 = 0.9, 0.999  Weight init trunc. normal (std=0.02)   ViT base learning rate 1 − 3  ViT weight decay 0.05  Dropout None  Batch size 1024  Training epochs 300  Learning rate schedule cosine decay  Warmup epochs 10  Warmup schedule linear  Randaugment (9,0.5,layers=2)   Mixup 0.8  Cutmix 1.0  Random erasing 0.25  Label smoothing 0.1  Stochastic depth 0.5  Exp. mov. avg (EMA)  None  models with SAM. Consequently, we precisely configure qk_dim = 16, headdim = 128 for SAM variants, and qk_dim = 64, headdim = 64 for non-SAM variants. For dense, Toeplitz (SAM and non-SAM) , and Vandermonde DFT, we slightly adjust  the hyperparameters to match the number of parameters.  To facilitate understanding and reproducability of these methods, we provide PyTorch codes in Figure 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 [39], employing a mixer matrix of R                    ×.               Toeplitz. Toeplitz matrix mixers without SAM properties are equivalent to convolution operations [18, 17], formulated with R              2−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                                                                                                                    ∈ 0, . . .  − 1, each token  generates two convolution parameters − and                                                                           . Using all  parameters, a data-dependent                                                                                                                   dynamic convolutional weight is generated as {−+1,−+2, . . . ,−1,0,1, . . .−2,−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  , a matrix M is Cauchy if each (, ) -entry  *satisfies*  =1  −  ;  −  ≠ 0.  ``` class Dense ( nn . Module ) :     def __init__ (         self ,         d_model ,         max_seq_len , # max_seq_len is necessary for Dense .         expand =2 ,         headdim =128 ,         device =None ,         dtype =None ,     ) :         factory_kwargs = {" device ": device , " dtype ": dtype }         super ()  . __init__ ()          self . d_model = d_model         self . max_seq_len = max_seq_len         self . expand = expand         self . d_inner = self . expand * self . d_model         self . headdim = headdim         assert self . d_inner % self . headdim == 0         self . nheads = self . d_inner // headdim         self . std_dev = 1 / np . sqrt ( self . max_seq_len )          self . M = nn . Parameter (             torch . empty ( self . nheads , self . max_seq_len , self . max_seq_len , ** factory_kwargs )          )          nn . init . xavier_normal_ ( self . M )      def forward ( self , hidden_states ) :         residual = hidden_states         # Rearrange hidden states to shape [batch , n_heads , length , headdim ]         hidden_states = rearrange ( hidden_states , 'b l (n h)  -> b n l h', n= self . nheads )          output = torch . einsum ('b n t h, n l t -> b n l h', hidden_states , self . M )          output = self . std_dev * output         output = rearrange ( output , 'b n l h -> b l (n h) ')  + residual         return output                               Figure 6: PyTorch code for the Dense variant in Table 2.                                                                                       ``` 
body
Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers
E Additional Details Of Ablation Studies In Table 2