Let SSMs be ConvNets: State-space Modeling with Optimal Tensor Contractions
Authors: Yan Ru Pei
Venue: ICLR 2025 Spotlight (2025)
Keywords: state-space models; convolution; tensor networks; audio processing; speech recognition
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: using deep SSMs like ConvNets to do audio processing

Abstract: We introduce Centaurus, a class of networks composed of generalized state-space model (SSM) blocks, where the SSM operations can be treated as tensor contractions during training. The optimal order of tensor contractions can then be systematically determined for every SSM block to maximize training efficiency. This allows more flexibility in designing SSM blocks beyond the depthwise-separable configuration commonly implemented. The new design choices will take inspiration from classical convolutional blocks including group convolutions, full convolutions, and bottleneck blocks. We architect the Centaurus network with a mixture of these blocks, to balance between network size and performance, as well as memory and computational efficiency during both training and inference. We show that this heterogeneous network design outperforms its homogeneous counterparts in raw audio processing tasks including keyword spotting, speech denoising, and automatic speech recognition (ASR). For ASR, Centaurus is the first network with competitive performance that can be made fully state-space based, without using any nonlinear recurrence (LSTMs), explicit convolutions (CNNs), or (surrogate) attention mechanism.

Published as a conference paper at ICLR 2025
LET SSMS BE CONVNETS: STATE-SPACE MODELING WITH OPTIMAL TENSOR CONTRACTIONS
Yan Ru Pei Brainchip Inc.
Laguna Hills, CA 92653, USA yanrpel @mm il 1 com
ABSTRACT
We introduce Centaurus, class of networks composed of generalized state-space model (SSM) blocks, where the SSM operations can be treated as tensor contrac- tions during training. The optimal order tensor contractions can then be sys- tematically determined for every SSM block to maximize training efficiency. This allows more ficxibility in designing SSM blocks beyond the depthwise-separabee configuration commonly implemented. The new design choices will take inspira- tion from classical convolutional blocks including group convolutions, full convo- lutions, and bottleneck blocks. We architect the Centaurus network with mix- ture of these blocks, to balance between network size and performance, as well memory and computational efficiency during both training and inference. We show that this heterogeneous network design outperforms its homogeneous coun- terparts in raw audio processing tasks including keyword spotting, speech denois- ing, and automatic speech recognition (ASR). For ASR, Centaurus is the first network with competitive performance that can be made fully state-space based, without using any nonlinear recurrence (LSTTMs), explicit convolutions (CNNs), or (surrogate) attention mechanism.
INTRODUCTION
Sequence or temporal modeling encompasses wide range of tasks from audio processing to lan- guage modeling. Traditionally, there have been many (related) statistical methods employed (Box
al. 2015). In the age of deep learning, neural networks have been predominantly used (LeCun et al. 2015), including recurrent neural networks (RNNs), convolutional neural networks (CNNs), transformers /Vaswani 2017), and neural ODEs (Chen al., 2018). In many cases, the model will inevitably suffer from one of two drawbacks: 1) cannot be efficiently trained (or fitted) in parallel due to the sequential nature of the model, 2) cannot be efficiently configured for online inference due to its large memory and computational requirement. To address this, deep state-space models (SSMs) were adapted for sequence modeling, and have shown incredible potential across a wide range of tasks (Gu et al.) 2021 Goel et al., 22022 Gu & Dao 2023). Due the linearity of the SSM layers, they can not only be configured for efficient online inference with small memory and com- putational resources, but also configured for efficient training using parallel hardware with unrolling strategies (Gu et al. 2022 Smith et al. 2022 Dao & Gu 20024]
Heinsen 2023).
Currently, most deep SSM networks (along with most neural networks in general) follow the chitectural recipe of transformers, where they are composed of uniform "SSM blocks" throughout the network, containing little to no variations in the shapes of the intermediate features or weights. This simplifies the designs of deep SSM networks, but may sacrifice performance and efficiency in practice. To explore the opposite direction, we go "back to the future" classical CNN designs instead, where much more heterogeneous design principle is followed. More specifically, we go deeper in the network, we will gradually downsample the temporal dimension, increase the channel dimension, and increase connective sparsity (Tan & Le 2021). to balance between size and effi- ciency (Pei & Coenen 2024). In order to allow for the possibility of such a heterogeneous design, make two novel contributions in this work. First, express SSM blocks tensor networks (or with e i innur um expressions), where we can easily observe the connective structure of existing SSM blocks as being mostly septhwise-separable. This motivates to design SSM blocks using this

---

Published as a conference paper at ICLR 2025
tensor network formalism, giving us connective structures such full and bottleneck SSM blocks, as inspired by classical CNN blocks. Then, we optimize the contraction order of each SSM block dynamically based the type of the SSM block, the shape of the input features, and the shapes of the SSM system matrices. This enables significant speedups during training for all SSM blocks (both old and new ones). We name this class of efficient CNN-like networks Centauru The model source code is at
RELATED WORK
2.1 DEEP STATE-SPACE MODELING
The seminal work proposing memory encoding using orthogonal Legendre polynomials in recur- rent state-space model is the Legendre Memory Unit (LMU) (Voelker et al. (019)) where Legendre
polynomials (a special case of Jacobi polynomials) are used. The HiPPO formalism (Gu et al 2020)
then generalized this to other common orthogonal functions. Later, this sparked cornucopia of works interfacing with deep state-space models including $4 {Gu al. 2021), H3 (Fu et al. 2022).
and Mamba (Gu & Dao 2023), achieving impressive results on a wide range of tasks from audio generation Goel et al. (2022) to language modeling. Besides a few recent exceptions (Smith et al.) 2022 Dao & Gu. 2024), These networks mostly use an underlying depthwise structure, which may limit the network capacity, albeit reducing the compute requirement of the network. An important focal point of this work is to generalize such SSM models to enable more connective flexibility, such that design choices of classical convolutional blocks can be carried over to the Centaurus model.
2.2 CLASSICAL CONVOLUTIONAL BLOCKS
Here, we look at the variants besides the standard full convolutional layer (where every pair of input and output channels connected). First, the simplest variant is the depthwise convolutional layer, made popular by the MobileNetVI architecture (Howard et al., 2017 This variant connects the input and output channels one-to-one, which drastically reduces the connectivity of the architecture, hence also its parameters and computational complexity. This architecture is prevalent for both computer vision and speech domains (Kriman et al. 2020; Hannun et al., 2019). Next, a variant with moderate connectivity is the grouped convolutional layer, appearing as early as the AlexNet work (Krizhevsky et al. 2012). In this structure, the input/output channels are divided into groups, and each input-channel group is associated with output-channel group one-to-one. The input and output channels are then only intra-connected within each group, but there are no inter-connections between groups Finally, there is a class of convolutional blocks known as bottlenecks, typically containing sequence of three convolutional layers. This structure is used prominently in the ResNet model (He et al. (016)) and also MobileNet (Sandler al.) 2018) onwards. Typically, the order of convolutional layers can be pointiise-depthiiiepppnneee resulting in output feature of the same tensor shape as the input feature.
2.3 TENSOR NETWORKS
Tensor networks were originally developed as a tool for approximating the wavefunctions of many- body quantum systems (Orús) 2019). It can be considered as an incredibly generalized form of low-rank decomposition, where high-dimensional features can be compressed as the contraction of a few low-dimensional tensors (Kolda & Bader 2009). It has a direct relationship to the ein- sum expression, where each einsum operand is considered node of the tensor network, and each contraction index is considered a (hyper)edge. The memory- and compute-optimal contraction or evaluation of tensor network can be formulated as a congestion optimization problem the base graph, allowing tensor networks to contest with "quantum supremacy" (Gray & Kourtis 202T). The application of tensor networks to machine learning began with the seminal work of Novikov et al. 2215), where is shown that the fully connected layers of a neural network can be exponentially
As a quick explanation of the name, we are exploring Tensorized Temporal Networks, or TenTeNs, from which we get a hundred, or CENTaurus. Admittedly, Centaurus is not actually composed of the word root "cent", so this wordplay is not too sensible. However, Centaurus is constellation, usually visualized as stars heing linked; this is visually similar what tensor network looks like, so some naming sense is recovered.
2

---

Published as a conference paper at ICLR 2025
compressed with minimal degradation in performance. This spawned series of works applying tensor networks for compressing convolutional and transformer blocks. Our work extends this line of "tensorization" technique to deep state-space models.
BACKGROUND
3.1
STATE SPACE MODELS
State-space models (SSMs) are general representations of linear time-invariant (LTT) systems (Hamilton 1994,, and they can be uniquely specified by four matrices: A RNNN, B R CERB XN and D R #H'HH, The first-order ODE describing the LTI system given
i= Ar + Bu,
Cz + Du,
(1)
where E R is the input signal, E RN is the internal state, and E R H is the output. The parameters {H, H' N denote the number of features for the input, output, and internal state respectively, Setting H = H' 1 yields single-input, single-output (SISO) SSMs (Gu et al. 2021 2022). and letting > 1, > yields multiple-input, multiple-output (MIMO) SSMs (Smith et al. 2022). We discretize the system using zero-order hold (ZOH), which gives us the discrete-time SSM matrices and B as follows (Gu et al. 2022):
A exp(AA),
B (AA)-1 (exp(AA) 1) AB.
(2)
The discrete SSM is then given by slt + 1] Ax(1]
Bu(l)
v(t C C+(1
It is then straightforward check that the discrete-ting time impulse response is given as k[r] CA' B, where 7 denotes the kernel timestep. During training, & can be considered the "full" long 1D con- volutional kernel with shape (output channels, input channels, length), in the sense that the output can be computed via the long convolution y,
Σ, (n, [t).
Similar to previous works (Gu al. 2022 assume A to be complex diagonal, but restrict B and to be real projection matrices to reduce memory and computational loads. In addition, we ignore the term Du it can be absorbed into the SSM system. Justification of these restrictions is given in Appendix A Like previous works also, we allow the parameters {A, B, C, A} to be directly learnable, which indirectly trains the kernel k. Unlike previous works in deep SSMs, do not try to keep the sizes H, HI' N consistent, to allow for more flexibility in feature extraction at each layer, mirroring the flexibility in selecting channel and kernel sizes in CNN: The flexibility of the tensor shapes requires careful choice of the optimal order of operations during training to minimize memory and computation, which will be the focal point of this work.
3.2 EINSTEIN SUMMATION NOTATION
The Einstein summation notation (Linstein 1922), or einium,, is a concise representation of general tensor contractions. We do not give a formal description of this notation here, but instead introduce it in the context of describing MIMO SSM. Recall that the output of MIMO SSM is computed by convolving the input with the impulse response, which we normally would write as
Σ (
kg))tt
Σ ui (EBmi [t).
where we defined the basis kernels K[r] R(A'). We get rather messy expression involving three summation indices. i denotes the input channel, denotes the internal state index, and j denotes the output channel.
To lessen the notation burden, we observe that the summation expression redundant and can be inferred from the tensor indices. In particular, if an index appears on the RHS of the equation but
2The size of the internal state h can be interpreted the degree of parametrization of basis temporal kernel, implicit (dilated) "kemel size" in the frequency domain.
3

---

Published as a conference paper at ICLR 2025
not the LHS, then it must have been summed over (or contracted). This allows to reduce the expression down to
1 u, * (Brilnn [t],
(5)
which is much better than before, and we can always assume the indices i and to be summed over without the explicit hinting of summation symbol. We can further simplify equation S by discarding convolution operator x. Naturally, we can leverage the convolution theorem, which states that the convolution operator is mapped to pointwise multiplication in the frequency domain. If we index the Fourier modes using f. then equation 5
can be expressed in the frequency domain as
ر (,. or Poイ = u. B
(6)
where the hat symbol denotes the Fourier transformed features (see Appendix [A for details).
3.3 A GENERALIZATION
Looking Kn, realize that there is only one oscillation mode per internal coefficient n, which in certain cases may limit the expressiveness of the network. A natural generalization is to expand the basis kernels as Knm, and arrive at the following system that is more expressive:
jis á Bn. f Enm Cjn.
Here, Fu serves as additional weighting factors for the basis kernels (which again we restrict to be real), representing the importance of each oscillation mode. The recurrent form of this system is then
2n == = Bwi u,[t), Zumlt (t + 1] Arm Fam [t] + u/y 3 = Cjn Eam R(#am [t), (8) which be considered parallel cascade of SISO state-space systems, intra-connected by the Eum factors, and inter-connected by the B.d and Cjn projection matrices. Alternatively, we can consider to index a "state block", and m index "sub-states" within the state block, which is an extension beyond the pure SISO and MIMO configurations. As a sidenote, this configuration is similar the Mamba block architecture (Gu & Dao 2023 but without the data-gating mechanism and the gated MLP structure.
Like Mamba, is possible configure the SSM matrices be data-dependent: A(u), B(u). C D(u), in which case the system becomes time-varian (Katsch 2023 Gu & Dao 2023 Dao & Gu 2024). We will not place major focus on this configuration in our work for two (temporary) rea- sons. Under the present algorithmic understanding, such dynamic systems require materialization of the internal states for sean operations (in Mamba 1) or require the introduction of a new sequence dimension into the tensor operands (in Mamba 2), meaning that it can restrict the flexibility of ten- sor contraction orders. On the more practical end, such models generally require custom kernels (in Triton or CUDA) for specialized support currently, making it difficult to build heterogeneous net- works with different connective configurations. We believe however that there are no fundamental restrictions to combining our framework with the data- mata-gating mechanism in theory, and leave it as future direction of study to achieve this efficiently in practice.
4 SSMS AND CNNS ARE TENSOR NETWORKS
Within SSM layer, the temporal kernels are constructed via the basis kernels, which are further generated (or parameterized) by the recurrent coefficients in A. In other words, a temporal kernel can be generated by weighted sum of selected Fourier modes, where each mode n is associated with a complex frequency of A, and a (real) weighting factor Ex- A natural viewpoint is that the parameters of A and E are analogous to the parameters of convolutional filters. but not directly parameterized. For instance, if we consider A having 3 complex parameters and E having 3 real parameters, then they together may contain the same "expressivity" as a representable with 9 real numbers. One may note that standard CNN spatial filters are local in space, while our temporal II filters are global in time. Interestingly however, by themselves, the difference between local convolutions and global convolutions is not too fundamental. This is because, by the convolution theorem, a temporal convolution is just simply a pointwise product in the frequency domain, and vice versa. Therefore, dual viewpoint is that in the frequency domain, deep SSM models are simply pointwise operations with "non-local" activation functions in between playing the "frequency mixing" role. This idea is also explored in the Hyena work (Poli et al. 2023).
spatial filter, both
4

---

Published as a conference paper at ICLR 2025
4.1
GENERAL SSM OPERATIONS WITH EINSUM EXPRESSIONS
Eey Rag
95r Duy Roor
Input iz
depthwise
tull
weight
besis kernets
Joll Batn Rannt
95r Frnn
grouped
battleneck
Figure 1: A tensor network representation of the tensorial connection structure of the differ- ent Centaurus block configurations, inspired by classical convolutional blocks. Here, the indices {i,j, n. G f} index the input channel, output channel, internal state, group/head, and Fourier mode respectively. Each (hyper)edge is denoted with one of the indices, representing the dimension to contract over. Note that the hyperedge associated with f always connects to three tensors, as rep- resents a convolution in the temporal domain, involving two input tensors and one output tensor.
In standard CNN layers, besides the canonical convolution operation, there usually an additional interaction between the input and output channels. For instance, we have configurations such as "depthwise convolution", "group convolution", and "full convolution", in increasing degree of chan- nel connectivity. In this section, will detail how structures. Naturally, matrices that are purely channel-mixin
the einsum expression to endow the A
tensor with these Channel-interacee like B and C are akin to "pointwise convolution" layers simply projection operations, which will place less emphasis on.
From here on, will use (i,j,n, f) to index the input channel, output channel, internal state, group/head, and Fourier mode respectively. The simplest example is a "depthwise" SSM block, or Su = ú. Ein Kin noting the appearance of the "input channel" index i in every operand and the total absence of the "output channel" index j. Therefore, the layer by itself will see any interactions among the "channel" dimension, hence why an additional mixing matrix M is needed the end 87, h. My:- In short, we have a pure gqquence-mixing layer followed by a pure channel-mixing layer, forming depthwise-separable structure, also discussed Gu al. (022).
On the other extreme, we have "full" SSM block, ūز = رية Ejin Kinn where both indices and j appear for the basis kernels, as each input-output pairing needs to be separately parameter- ized. In this case, both the sequence-mise mixing and channel-mixing structures are "baked" into the full SSM block, so not additional mixing layers are necessary. Figure. [ depicts tensor network representation of the different connectivity structures corresponding to classical CNN blocks.
4.2 OPTIMAL CONTRACTION ORDERS TRAINING: THE DEPTHWISE EXAMPLE
In online inference mode, we have no choice but to explicitly materialize and evolve the internal states, leaving little room for optimization of the order of operations. This is also true for other parallelization strategies where the internal states materialized in some form (Smith et al. 2022 Heinsen 2023). Fortunately, we have the flexibility of choosing the order of contraction we perform the SSM operations in the frequency domain (c.g. via FFTs), and considerable speedups (at times, orders of magnitudes) can be achieved by choosing the optimal contraction order. A slight drawback of using FFT is that the time complexity is ( log T.) instead of the desired O
5

---

Published as a conference paper at ICLR 2025
However, this is typically not practical issue, FFT is rarely compute-bound operation unless the sequence length is exceedingly large or the operation is already kernel-uueed (Fu et al.) 2022). At this point of discussion, for clarity of exposition. In practice, the batch size will factor into the determination of the optimal contraction order.
will start accounting for the batch dimension as b, omitted previously
Taking again the depthwise SSM layer, or Dbey = uus Eng Kniy. intuitively we would opt to contract the F. and K tensors first to "collapse" down the basis kernels along the state dimension (indexed n) as such, k.s- Knif. It then becomes much more manageable to compute the output by simply taking the product Dbsf = fbackiis Note that
F(kit) F(Ems (ntr)= = F
due to the linearity of the Fourier operator F, hence justifying the freedom of the contraction order. To optimize further, we would ideally compute F(Fair knrr) = F(kir) instead of Emi (Kirr), ) as the former requires performing the Fourier transform much smaller tensor (vector). Even in this simple example, we observe two important points: 1) the contraction order matters, 2) the placement of the Fourier operators matters.
To make things clearer, it is useful to think of the Fourier operator Fft itself as an operand to be contracted. This will allow us to write the SSM operations
Ybar FF (Frruba) (Fr=EnsKnir))
(10)
Here, we represented the Fourier operator as matrix (i.e. the DFT matrix), but unlike standard matrix multiplication, we know that multiplying/contracte
with DFT matrix can be done via FFT, so that it will not incur the same computational complexity as the standard tensor contraction operation. Correctly accounting for the FFT complexity is important in determining the optimal order contractions, which in this picture also includes the placement of the FFT operators.
This is simple example where the optimal contraction is somewhat obvious and static, but there are cases where we have more terms to contract and the differences between the contraction paths are more subtle. And in these cases, the optimal contraction order is dynamically linked with the shapes of the tensor operands. Fortunately, there is a systematic way to evaluate the memory and compute requirement of an einsum contraction path (used prominently in packages such as opt-e nsum) (Daniel et al., 2018), and we make a simple augmentation to this prescription to handle "contrac- tions" involving FFT operators, while being somewhat mindful of the software and hardware back- ends (PyTorch and CUDA).
4.3 PRACTICAL WALKTHROUGH: THE BOTTLENECK BLOCK
The main example that we will walk through here is the bottleneck layer example, or 9osf ub.y Emmms Cnjn where 5 tensor operands are involved. If we include the Fourier operators as 3 additional operands, then we have total of 8 operands, which yields sufficiently complex design where the systematic optimization of contraction orders will show its power. See Listing 1 in Appendix C for minimal pseudocode of the bottleneck block operations with the order of operations optimized, along with benchmarking an A100 GPU.
In theory, it possible to have pseudocode would appear much simpler (i.e. we only need one einsum expression in line 43 of Listing 1). However, there are certain SW and HW idiosyncrasies that make more efficient ex- plicitly "force" " certain contraction patterns. For instance: 1) complex tensors are not yet "natively" supported by CUDA, meaning that is more efficient to operate on real tensors as much as possible. In other words, we perform computations in the complex frequency domain only if necessary: 2) terch. comp l ee may not yet be able to identify kernel fusion opportunities within sufficiently complex torch. einsum um expression. Therefore, it may be more efficient to explicitly modularize and kernel-fuse certain contraction steps, particularly those that are memory-bound bound.
handle optimizing all the contraction orders, and the
We try to achieve a happy medium between full automation with einsum expressions and full spe- cialization with custom CUDA kernels. other words, we "semi- "semi-manually" inspect all possible contraction paths and discard the clearly non-optimal ones. Then, we perform some amount of practical optimization for the "feasible" " contraction paths, as we will explore in the following sec- tions. Initially, for ease of exposition, will continue our discussion treating the temporal domain

---

Published as a conference paper at ICLR 2025
and frequency domain almost equivalent, as the two be easily traversed via FFTs (which is fairly light in compute). Only at the very end will we make the effort to separate the two domains as we begin to consider the optimal "insertion points" of the FFT operations, the second-order optimization.
First, based on visual inspection of the bottleneck tensor network representation in Fig. D can clearly see the first step should be to always "contract away" the inner edge m, or equiv- alently compute the kernel knr Eum knrr first. However, K Kumr themselves are generated by An and Anm- Therefore, under kernel fusion, it is possi- ble to not even materialize the basis kernels Kymr in the GPU VRAM at all, which shaves off memory and computational requirements considerably during training. This is encapsulated in the get kernel 1 function in Listing I, which is meant to (jit-) )compiled during training, for exam- ple with torch. torch.compi compile or custom triton kernels. In the frequency domain, we are now left with the expression ibi übي ; Buskinn Cjnn which still allows for many possible contraction paths in theory. However, heuristically speaking, we should try not to produce any intermediate tensor of dimension greater than 3, as it will be unfavorable materialize and operate it. Lemma will heavily restrict the "feasible" contraction paths based on this criteria, a full proof which is given in Appendix B
note that the basis kernels
Given the einsum expression Goif = - Buskkn arranged in this order, all inter- mediate tensors will have most 3 dimensions, if and only if â (and intermediate tensors resulting from it) is contracted with only its neighboring operands.
Lemma 1.
Remark. An interpretation of Lemma 7 is that the contraction order should roughly follow the "natural order of operations" of the underlying state-space system, in some form of "associative" fashion. The proof sketch is to try contracting u with its non-adjacent operands such as k or C, and observe the appearance of high dimensional intermediate tensors, for instance, ûbs ūorkko = (uk)
propected
projected
A
9
from "left" "right"
b) using the "full" kemma
Figure 2: The two feasible patterns for performing the SSM bottleneck operations. a) Follow the "natural" order induced by the underlying SSM system: first project the inputs, then convolve them with the kernels, and finally project the outputs. b) Generate first the "full" SSM kernels, then convolve them with the inputs in one stage get the output.
Under the restriction of Lemma I there are really only two "feasible" contraction patterns as showcased in Fig. First, can perform the operations from left right, corresponding to projecting the input, performing the FFT convolution, and then projecting the output. This is just the "natural" order of operations induced by the underlying state-space system. Altema- tively, we can build the "full kernel" first, then perform the full FFT convolution with the input as (Byskn Cūj) kaickji ūbjf in one sterf) We see that the optimal contraction order intimately linked with the dimensions of the tensor operands. More formally, the first "natural"
order of contraction is only optimal when B + + 17 + 1. H as shown in Appendix B For the first contraction pattern, additionally the condition < H is met, then is clear that we should
Thhis is how convolutional networks typically handle full convolutions, except there is no "kernel huilding" stage the kemels explicitly parameterized.
7

---

Published as a conference paper at ICLR 2025
perform the input projection Tbnt Ubat Bas first before performing the FFT on the projected input F(æbns) = fo as the projected input x is smaller than the imput u. Similarly, for the second contraction patter, if additionally the condition HH' I N is met, then we can produce the full kernel Kijr = Brakin Cjn first before performing the FFT on the full kernel F(k,r)) kjs,, as the full kernel k is actually smaller than k. These two contraction patterns are explicitly forced via the if statements in the opt-fft..co function in Listing 1.
5 EXPERIMENTS
Parameter Scaling
Inference Compute Scaling
3x10-1
3x 3x
DWS (S4D)
DWS (S4D)
DWS (56)
DWS (56)
2x100
PW bottleneck (55)
2x10
PW bottleneck (S5)
full
Classification Error
bottleneck
bottleneck
hybrid
hybrid
6x100-
6x 10-2
4 10
40
10*
10a
10'
10*
Parameters
FLOPS
Figure 3: The scaling of the classification error versus the number of parameters and FLOPs per second during inference. The performance is evaluated on the SC35 testset. For each network variant in the ablation studies, we performed 10 training trials and took the median metric. DWS is short of seephhise-sepprabbe (with the S6 variant using selective scan), and PW is short of pointwise.
In our experiments, we study various configurations of Centaurus model architectures (e.g. classi- fier, hourglass, multi-streams) and block configurations (e.g. depthwise, bottleneck, full). This is to showcase the structural flexibility of our model and adaptation to different tasks, reminiscent of clas- sical CNNs. At the same time, we try to keep consistent the training pipeline (e.g. optimizers and schedulers) and the auxiliary layers (e.g. normalization layers and activations), to isolate the effects of only the core architectural changes in the SSM blocks. See Appendix and the basic design of the Centaurus block. We will mainly be comparing the following variants: 1) networks with all SSM layers being depthwise-separable (S4D-like and S6-like with data gating), 2) networks with all SSM layers being pointwise bottlenecks (S5-like), 3) networks with all SSM layers being bottlenecks, 4) networks with all SSM layers being full convolutions, and 5) hybrid net- works with mixtures of SSM layer types (inspired by classical CNN designs). We will put a strong emphasis on the computational requirements of the network during online inference, in addition to the typical parameter count. This is complimentary to the theoretical and algorithmic components of work, where optimal contractions leveraged speed up training. See Appendix D for a dis- cussion of how the parameters and FLOPs are estimated for these blocks during online recurrence. We show that our hybrid Centaurus network can achieve competitive performance as a stand-alone SSM network. In addition, we show that it can be well augmented by inserting (explicit) convolu- tional layers or modules into it, following the modern trend of designing bybrid SSM models (Lieber al. 2024; Ren et al.) 22244] Patro & Agneeswaran] 2024] Glorioso et al. 2024). However, since we aim for solutions that can be configured for efficient online inference, we will avoid the use of any attention mechanism though we expect it to also integrate well with Centaurus.
E for the training details
VVariants of the attention mechanism, such as sliding-window attention, can be used to improve the online inference efficiency. Any surrogate (linear) attention mechanism admitting recurrent form (c.g. Linear trans- former, RWKV, Mamhaj can also be used. We think the Centaurus model can interface well with these blocks, but do consider them here limit the scope of this work.
8

---

Published as a conference paper at ICLR 2025
5.1 KEYWORD SPOTTING
We begin with the simple task of keyword spotting (KWS) with raw waveforms, using the Google Speech Commands 35-class subset (Warden, 2018), or SC35. For each SSM configuration, we train networks of various sizes, to estimate the scalability of the performance with respect to the number of parameters and FLOPs per second during online inference. See Fig. for the scalability plots, where our hybrid architecture outperformed all the homogeneous counterparts. All the network variants we tested have six layers. Our Centaurus hybrid model contains two layers of full SSM blocks, followed by two layers of bottleneck blocks, followed by two final layers of pointwise bottleneck blocks (S5). This is in accordance with the classical CNN design philosophy of using sparser connective structures at deeper layers where the number of channels becomes large. The training and network details are given in Appendix EE.1] along with comparison against other deep SSMs, where hybrid network achieved similar performance with roughly 100 times fewer FLOPs. Additionally in Appendix E. compare against hybrid variant where the projection matrices (B, C, and E) complex (similar to the original S4D and S5 layers), but it did not outperform the real counterpart; furthermore, we compare against variants with S6 layers of different state dimensions.
5.2
SPEECH ENHANCEMENT
Table I: Comparing real-time audio denoising networks. Note that only the Centaurus network
requires additional signal processing (e.g. STFT). To boost the denoising performance, we prepend every SSM block with a causal depthwise Conv ld layer with kernel size of 4 (last row).
Model
PESQ (VB-DMD)
Parameters
FLOPs sec
Real-time Baselines (with processing)
FRCRN (Zhao et al. 2022)
3.21
6.9M
76.2G
DeepFilterNet' (Schröter et al.
2023
3.16
2.13M
0.688G
SEMamba (causal) (Chao et al.
2024
2.76
3.60M
0.76G
PercepNet (Valin et al. 2020)
2.73
8.00M
1,60G
DEMUCS (Detossez et 220
2.65
33.53M
15.44G
Centaurus (DWS)
0.71M
0.39G
Centaurus (DWS data-gated, or S6)
0.66M
0.50G
Centaurus (bottleneck)
0.87M
1.15G
Centaurus (PW bottleneck)
3.06
0.83M
0.65G
Centaurus (full)
3.04
2.53M
1.12G
Centaurus (hybrid)
3.12
0.51M
0.29G
Centaurus (hybrid, with Causal Conv)
3.25
0.51M
0.29G
Here, we focus on the task of real-time speech denoising directly on raw audio waveforms. The macro network architecture used in this experiment is adapted from deep SSM hourglass autoen- coder proposed in Pei et al. (2024), where the SSM blocks were originally configured as DWS blocks (S4D). Here, we modify the network include multiple variants of SSM blocks enhance its con- nective flexibility, and in addition perform real-time denoising on raw waveforms directly in the +1]+ range without one-hot encoding (Goel et al., 2022) or spectral processing (Schröter al., 2023). We will test variants of this network by simply swapping out the pointwise- -orriiseebbtteeeeckkk SSM blocks with other variants. The encoder architecture is based on the 6-layer KWS network backbone (see Section 5.1). and the decoder is a reflection of the encoder. The performance of the network is evaluated on the VoiceBank DEMAND (VB-DMD) testset and given Table. I Details of the network architectures and the training pipeline are given in Appendix E.2
Surprisingly, the DWS and bottleneck homogeneous variants suffered from severe training plateaus, and were not able to generate any meaningful audio samples. In addition, using the data-gating mechanism (S6 layers) did not solve the issue. The variant with homogeneous full SSM blocks did manage train, but despite its size, it did not achieve the best performance. Out of all variants, the hybrid Centaurus model achieved competitive PESQ scores, while requiring much fewer parameters and FLOPss There are concurrent works applying deep SSMs to speech enhancement (in cases

---

Published as a conference paper at ICLR 2025
outperforming Centaurus), but they configure the SSM layers bidirectional, which sacrifices the
online capability of the network (Zhang et al. 2024 Chao et al. 2024
5.3 TOWARDS END-TO-END SPEECH RECOGNITION SSMs
The reported metrics are word error rates (WERs) for the clean/other splits, on the Lib- rispeech test/dev-sets. The FLOPS estimated based 20.5-second audio segment.
Table 2:
Model
test |
dev +
Parameters (M) FLOPs (G)
Offline (full-context)
Full Conv Zeghidour et al.
2018
3.3 10.5
3.1/9.9
Transformer Zhang ct al. 2020)
3.1/7.3
29
ContextNet (Han et al. 2020)
2,4/5.4
31.4
Wav2Vec2 Baevski et
2020
2.1/4.8
1.7
94.4
285
Conformer Gulati ct al.
2020)
2.0/4.3
30.7
45.2
Online (streaming)
Transformer
5.0/11.6
18.9
ContextNet
4.5 10.0
31.4
Conformer
4.6
30.7
Centaurus (online, no attention)
Base (full SSM)
6.0 13.1
1.9
12.4
20.6
with FFN with causal conv-block with Mamba macro-block bloock
5,4/
5.2/113
18.0
32.1
4 10.6
4
23.6
43.7
10.2 4.3
29.9
46.9
Since this is not a work focused on training methodology. we opt for a rather simple supervised/dis- tillation pipeline h At a high level, we use hybrid SSM backbone similar our KWS hybrid network (see Section 5.I. but wider and deeper (see Appendix E.3). We do not use any additional lexicon or language model decoding (Zhang et al. 2020 Gulati al., 2020). And unlike previous works, we are not integrating SSM layers as auxiliary layers in off-the- shelf transformer models (Miyazaki et al., 2023) Shan al., (023). This means that the Centaurus network fully SSM based end-to-end, without any non-linear recurrent mechanism (self-conditioned decoding) at- tention mechanism. In Table. 2] we report the performance and size of the base Centaurus model on Librispeech (Panayotov et al. 2015), along with two augmented variants: 1) one with feed-forwart (FF) module appended to every SSM block, 2) one with a convolutional (Conv) module appended,
3) and one with the Mamba gated macro-architecture (Gu & Dao 2023 wrapped around
SSM layers. We describe the architecture these modules in Appendix E.3 Being a fully online-
inference ASR model with minimal optimization, the Centaurus network still remains competitive with "benchmark" ASR networks like Conformer and Wav2Vec2 in streaming mode.
CONCLUSION
We introduced Centaurus, composed of SSM blocks with flexible connectivity structures, trained us- ing optimized tensor contractions. We designed task-speeific realizations of the Centaurus network balancing between size and performance. As inspired by classical CNN designs, we used blocks with dense connections in the shallower layers and opted for sparser connections in the deeper lay- ers. The hybrid networks achieved competitive results on multiple audio processing tasks, while being fully SSM based. In the future, may be fruitful to explore how the lata-gating mecha- nism can be applied to the Centaurus network, and explore its scalability to tasks that of larger scales, such as language modeling. In addition, it may also be of interest to explore SSM convo- lutional structures in higher dimensions, incorporating additional structures such stride, dilation, and transposition, to see whether Centaurus can be effective for computer vision tasks as well.
'To train ASR model from scratch, the prevailing method is typically to perform self self-supervised learning large amount of unlabeled speech, then fine-tune relatively small amount of labeled speech.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
We would like to thank Nolan Ardolino, Olivier Coenen, M. Anthony Lewis, Sidharth (in alphabeti- cal order) for many inspiring discussions. In addition, we thank Dhvani Kothari and Kurt Manninen for producing spectacular demos using the Centaurus network. We also thank Daniel Keller for pro- viding a critical proofreading the manuscript. Finally, we thank the ICLR reviewers for offering much helpful feedback for making our work more complete.
ETHICS STATEMENT
We presented a novel state-space network for several important audio processing tasks. We made sure that our networks are lightweight for both training and inference, so that they can be easily experimented with by both researchers and hobbyists alike. This can done without requiring too much computational resource, part of our commitment democratizing Al and reducing carbon emissions. All of the datasets used in our work are publicly available, with no private or sensitive data used in our experiments.
REPRODUCIBILITY STATEMENT
We gave the full details of the training pipeline of our networks in Section 5 of the main text and Appendix E Listing 1 in Appendix C provides pseudocode that be easily adapted into Py- Torch code, and we additionally made available the model source code at github.com//BB corbbbanc/////ppp..... Inc/Centaurus
REFERENCES
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A frame- work for self-supervised learning of speech representations. Advances in neural information processing systems, 33: 13:1444--22460, 2020.
George EP Box, Gwilym M Jenkins, Gregory C Reinsel, and Greta M Ljung. Time series analysis: forecasting and control. John Wiley & Sons, 2015.
Rong Chao, Wen-Huang Cheng, Moreno La Quatra, Sabato Marco Siniscalchi, Chao-Han Huck Yang, Wei Fu, and Yu Tsao. An investigation of incorporating mamba for speech enhance- ment. arXiv preprint arXiv:2405.06577 2024.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of computation, 29990))277-31, 1965.
Daniel, Johnnie Gray, et Opt' -einsuum python package for optimizing contraction order for einsum-like expressions. Journal of Open Source Software, 3(26):753, 2018.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21066 2024.
Alexandre Defossez, Gabriel Synnaeve, and Yossi Adi. Real time speech enhancement in the wave- form domain. arXiv preprint arXiv:2006. 12867.12277.... 2020.
Albert Einstein. The general theory of relativity. In The meaning of relativity, pp. 54-75. Springer, 1922.
Daniel Y Fu, Tri Dao. Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 14052, 2022.
11

---

Published as a conference paper at ICLR 2025
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: compact 7b ssm hybrid model. arXiv preprint arXiv:2405.16712, 2024.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state- space models. In International Conference on Machine Learning, pp. 7616-7633, PMLR, 2022.
Johnnie Gray and Stefanos Kourtis. Hyper-optimized tensor network contraction, Quantum, 5:410, 2021.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:23/2.00752, 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00330 2021.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022.
Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: Convolution-augmented transformer for speech recognition. arXiv preprint arXiv:2005.08100, 2020.
James D Hamilton. State-space models. Handbook of econometrics, 4:3039-3080, 1994.
Wei Han, Zhengdong Zhang, Yu Zhang, Jiahui Yu, Chung-Chene Soung-Cheng Chiu, James Qin, Anmol Gulati, Ruoming Pang, and Yonghui Wu. Contextnet: Improving convolutional neural networks for automatic speech recognition with global context. arXiv preprint arXiv:2005.03//9, 2020.
Awni Hannun, Ann Lee, Qiantong Xu, and Ronan Collobert. Sequence-to--sqquenee speech recog- nition with time-depth separable convolutions. arXiv preprint arXiv:1904.02611 2019.
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint M.XXx:2229.129555 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and patlern recognition, pp- 770-778, 2016.
Franz A Heinsen. Parallelization an ubiquitous sequential computation. arXiv preprint arXiv:2311.0628., 2023.
Andrew Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. arXiv preprint arXiv: 174.00074 2017.
Tobias Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv
preprint arXiv:2311.01927 2023.
Tamara G Kolda and Brett W Bader. Tensor decompositions and applications. SIAM review, 51(3): 455-500, 2009.
Samuel Kriman, Stanislav Beliaev. Boris Ginsburg, Jocelyn Huang. Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, and Yang Zhang Quartznet: Deep automatic speech recog- nition with 1d time-channel separable convolutions. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 6124-6128. IEEE, 2020.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
12

---

Published as a conference paper at ICLR 2025
Yann LeCun, Yoshua Bengio, and Geoffrey Hinton. Deep learning. nature, 521(7553):436 444, 2015.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403. 19887, 2024.
Koichi Miyazaki, Masato Murata, and Tomoki Koriyama. Structured state space decoder for speech recognition and synthesis. In ICASSP 2023-2023 IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Alexander Novikov, Dmitrii Podoprikhin, Anton Osokin, and Dmitry P Vetrov. Tensorizing neural networks. Advances in neural information processing systems, 28, 2015.
Román Orús. Tensor networks for complex quantum systems. Nature Reviews Physics, 1(9):538- 550, 2019.
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur, Librispeech: an asr corpus based public domain audio books. In 2015 IEEE international conference acoustics, speech and signal processing (ICASSP), pp. 5206-520222 5210. IEEE, 2015.
Badri Patro and Vijay S Agneeswaran. Simba: Simplified mamba-based architecture for vision and multivariate time series. arXiv preprint arXiv:2403. /5360, 2024.
Yan Ru Pei and Olivier Coenen. Building temporal kernels with orthogonal polynomials. arXiv preprint 1rXiv22455.21179 2024.
Yan Ru Pei, Ritik Shrivastava, and FNU Sidharth. Raw speech enhancement with deep state space modeling. arXiv preprint arXiv:2409.03377 033777 2024.
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 International Conference Machine Learning, pp. 28043-28078. PMLR, 2023.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang and Weizhu Chen. Samba: Sim- ple hybrid state space models for efficient unlimited context language modeling. arXiv preprint 2rXiv:2406.07522, 2024.
Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang- Chich Chen. Mo- bilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 4510-4520, 2018.
Hendrik Schröter, Tobias Rosenkranz, Andreas Maier, et al. Deepfilterne:: Perceptually motivated real-time speech enhancement. arXiv preprint urXiv:2305.08227 2023.
Haozhe Shan, Albert Gu, Zhong Meng, Weiran Wang, Krzysztof Choromanski, and Tara Sainath. Augmenting conformers with structured state space models for online speech recognition. arXiv preprint arXiv:2309. 08551,, 2023.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.0449) 00433,, 2022.
Mingxing Tan and Quoc Le. Efficientnet22: Smaller models and faster training. In International conference on machine learning, pp. 10096-10106. PMLR, 2021.
Jean-Mare Valin, Umut Isik, Neerad Phansalkar, Ritwik Giri, Karim Helwani, and Arvindh Kr- ishnaswamy. A perceptually-motie mottvaated approach for low-complexity, real-time enhancement of fullband speech. arXiv preprint arXiv:2008. 04259, 2020.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Aaron Voelker, Ivana Kajié, and Chris Eliasmith. Legendre memory units: Continuous-time repre- sentation in recurent neural networks. Advances in neural information processing systems, 32, 2019.
13

---

Published as a conference paper at ICLR 2025
Pete Warden. Speech commands: A dataset for limited-vocabulary speech recognition. arXiv preprint Stti::8804032009 2018.
Neil Zeghidour, Qiantong Xu, Vitaliy Liptchinsky, Nicolas Usunier, Gabriel Synnaeve, and Ronan Collobert. Fully convolutional speech recognition. arXiv preprint arXiv:1812.06864, 2018.
Qian Zhang, Han Lu, Hasim Sak, Anshuman Tripathi, Erik McDermott, Stephen Koo, and Shankar Kumar. Transformer transducer: A streamable speech recognition model with transformer en- coders and mn-t loss. In ICASSP 2020-2020 JEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 7829-7833. IEFE, 2020.
Xiangyu Zhang, Qiquan Zhang, Hexin Liu, Tianyi Xiao, Xinyuan Qian, Beena Ahmed, Eliathamby Ambikairajah, Haizhou Li, and Julien Epps. Mamba in speech: Towards an alternative to self- attention. arXiv preprint arXiv:2405. 10:405//26000 /2609, 2024.
Shengkui Zhao, Bin Ma, Karn N Watcharasupat, and Woon-Seng Gan. From:: Boosting feature representation using frequency recurrence for monaural speech enhancement. In ICASSP 2022- 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 9281-9285, IEEE, 2022.
A CANONICAL FORM OF A STATE-SPACE MODEL
is generic property (though always true) that diagonal form exists for the state-space model, meaning that we can almost always assume A to be diagonal, at the expense of potentially requiring B and C' to be complex matrices. Since the original system a real system, the diagonal A matrix can only contain real elements and/or complex elements in conjugate pairs. In this work, we sacrifice a slight loss in expressivity by continuing to restrict 7 and C be real matrices and letting A be diagonal matrix with all complex elements, but not restricting them to come conjugate pairs. During online inference, we then need maintain the internal states æ as complex values but only need propagate their real parts to the next layer.
Recall that the discretized state-space system is given by
xlt + 1] = Ax + Bu[t],
y[t + 1] = Cxlt + +] Duft + 1],
(1I)
To see that the term Du can be absorbed into the state if we double the internal states, we simply need to make the following substitution
B
A
ONNN
-
c c+[CI
(12)
1NXH
ONXN
ONXN
In other words, we trivially set aside set of "memoryless" internal states whose sole purpose is to duplicate the input u, effectively routing directly the inputs to the outputs.
If allow B and C be complex matrices, then can further assume À to be diagonal without any loss of generality. To see why, we simply let A = p-1 AAA (where A is the diagonalized A matrix, and P is the similarity matrix) and observe the following:
k(4] C(P AA
CCP- (MPPP AA(PP PB
repeat
times
(13)
= (CCP-))A--
=C'AA-
where B' and C" are complex matrices that have "absorbed" the similarity matrix P. but WLOG we just redefine them to be B and C. Since A is real matrix, the complex cigenvalues in A must come conjugate pairs. And WLOG we can again redefine A as Note that in this section, we
TThis in deviation of prior works, such S4D and $5, where B and C complex.
14

---

Published as a conference paper at ICLR 2025
only need to ensure that A is square matrix, and all the other state-space matrices and state vectors can be of arbitrary shape, as long as they are conformable.
Since we still want to work with real feature ]
we then only take the real part of the impulse response
kernel as such:
k|r] RCCA B) - CB(AA)B,
(14)
which equivalently in the state-space equation can be achieved by simply letting yft] = C R(=ft)) This means that during online inference, need to maintain the internal states complex values, but only need to propagate their real parts to the next layer. This restriction of propagating real features is not required, and is merely for better CUDA support, as discussed in Section 4.3
We can then define the basis kernels be Kn [T] R(A"). Furthermore, we use the hat symbol to denote the Fourier transform of the features, or ã = F(x) where F is the Fourier operator. To be more explicitly, if r E {0, (0,1,2,...
1}. then we have
T-1
Σ KnlTTeexp ( 2xiTf ==
= =(kn))
(15)
T
noting that the index n is not of importance, and can be replaced with any (multi)index when needed. For now, we ignore the artifacts of cyclic convolution (needing padding to ensure causality) and algorithmic optimization except for the obvious fact that the Fourier transform can be computed efficiently via the Fast Fourier Transform (FFT) algorithm (Cooley & Tukey 1965). At this stage, then have a condensed e einsum expression in the frequency domain that allows to easily see the interactions between the input signals, basis kernels, and system matrices.
B BOTTLENECK SSM CONTRACTION ORDERS
Recall that Lemma. states that: Given the einsum expression Obiy = Bask Cjjn arranged in this order, all intermediate tensors will have most 3 dimensions, if and only â (and intermediate tensors resulting from it) is contracted with only its neighboring operands.
Proof. If we contract usiy with kaf. this immediately results in a 4D tensor (ūk) Similarly, we contract uiss with Cjn. this immediately results in a 5D tensor (ūC))jjny. This means that we can only contract fuy with Bas first (its only neighboring operand), resulting in the 3D tensor bf.)
At the second stage, if we contract (ūB) This means that we can only contract (.B) with kn (again its only neighboring operand), resulting in the 3D tensor (ūBk)
Jun with Cjn, this will result in a 4D tensor (iBC))jnf-
For the remaining contraction paths, we have to verify that any contraction paths within B BaskinsC (not including i) will only result in intermediate tensors of dimension at most 3: Buiknf = (Bk)nif, kay Cin = (kC)jng, and Cin (BC))nii.
An interpretation of Lemma is that the contraction order should roughly follow the "natural order of operations" of the underlying state-space system, in some form of "associative" fashion. For example, it clearly makes little sense contract the input u directly with the output project matrix C' before even passing the input through the internal states first, and this is formally reflected as the production of 5D intermediate tensor. Note that this not say that the contractions of an einsum expression should be restricted to neighboring operands. In fact, einsum expressions are agnostic to the ordering of operands, meaning that contractions can be performed on any two operands at any stage. The discussion here is only specific to the state-space system ÿbif = ueef Br,ck Cjn, for
7 This prior required, as technically we can configure our network complex-valued handle com- plex features. However, we do not explore this configuration this work (complex-valued neural networks). As example, for real signal, only need to extract half the Fourier modes, up the Nyquist frequency.
15

---

Published as a conference paper at ICLR 2025
which this artificial ordering restriction is only needed for efficiency concerns, and not functional correctness (which will be guaranteed regardless of the contraction path taken).
Under the restriction of Lemma [1] there are really only two "feasible" contraction patterns. First, we can perform the operations from left to right, corresponding to projecting the input, performing the FFT convolution, and then projecting the output. This is just the "natural" order of operations induced by the underlying state-space system. Alternatively, we can build the "full kernel" first,
I' (kr = Poss in one step If we only focus on the computational requirements of the forward pas- [ then the first contraction order will result in BN H BN B NF BNF units computation. The second contraction order will result in H'N + JNHFF- + BH F s H + N) units of computation. Therefore, we see that the optimal contraction order intimately linked with the dimensions of the tensor operands. More formally, the first "natural" order of contraction is only optimal when BN + H') < H'HF( N) or + N > H
then perform the full FFT convolution with the input as (Backni
C BOTTLENECK SSM BLOCK PSEUDOCODE AND BENCHMARKING
Listing 1: Contraction operations for the bottleneck SSM block
assumes opt_einssm backend is enabled
def fft_conv (equation, x, k, +args)
PFT conv with additional nen-sequence tensors as args
I. = x.shape[-1) [-1] x_f rfft (x, length-2+L)
k_ rfft (k, length=2+I)
y_f einsun(equation x_,, k_ *args)
return irfft (y_ [..., :L]
10
@compilc
12 def get_kerncl (delta,, A, E, length):
13
dtA einsum nm-2nm', nm- delta, A)
14
K einsum (m, t->mmt', dtA, (eange).exp) return einsum (f nmt, nm->nt', K.real, E)
17 def sft_ftt_cce opt_ift_comm cift_coniiu, k_reat, B discrete, C):
18
# all the input arguments are real tensors u: (batch, H_in, L)
19
k_real: (N, L)
B_discrete: (N, H_in)
22
C: out, N)
23
force certain contraction paths based on the tensor shapes batch, Hinn, shape Houut C.shape
26
27
28
# project inputs, perform FFT cenv, then project outputs
if (1 / H + H_out) (1 batch + N)
30
if H_in:
31
x einsum nit u, discrrete)
32
x = fft__cov ntontttt x, k r_rell)
33
return einsum(' ('bnt, jn->bjt', x, C)
evaluate the full kernel, then perform full FFT conv with inputs directly one stage
35
36
else:
37
if iin A Hoout <= N:
k_full einsum (' jn, ni, C, B B_discrete, k_real) fft_conv ('bit, jit->bjt', x, k_full)
Thiis how convolutional networks typically handle full convolutions, except there is "kemel huilding" stage as the kemnels are explicitly parameterized. Under reverse-mode automatic differentiation, it not hard to show that the compute units required for the backward pass are 2 times that of the forward pass, for any einsum expression.
16

---

Published as a conference paper at ICLR 2025
41
can simply just return this and ignore the above if some loss of efficiency can be accepted return fft_conv ni, jn->bjt', u, Kreaa,, __discrete, C) nt,
42
43
def sam sem_bottlennck (u, delta, A, 46
E):
length u. shape[-11
47
reall get ( (delta, A, E, length)
48 B gisscrete einsum (' ni delta, B) & zon discretization return opt_ fft conv (u, reall, p_discrete,
Here, acknowledge that this semi-manual method, though efficient, is not the most elegant or scalable approach, as requires some amount of "hard coding" for every SMM block variant. For- tunately, the bottleneck block walked through in this section is the most involved example, and the other SSM block variants explored in this work are much simpler in terms of feasible contraction paths. Nevertheless, we are working on extension of the opt_einn opt_einsum library specifically for general temporal sequence modeling using tensor networks that can: 1) identify any kernel-fusion opportunities in the underlying tensor network, identify the optimal "insertion points" for FFT and iFFT operations accounting for the preference of real tensors over complex ones.
We compare the training performance of the contraction-optimizze SSM block as given in Listing against the naive contraction order following the "natural" order of input projection, state evolution, and output projection (as typically done for other SSM networks). optimal contraction, we focus the total time it takes for the forward plus backward pass on single bottleneck SSM block. We perform SXM4 with PyTorch 2.5.1 under CUDA 12.4. The baseline dimensions are batch 256, H 16, H' 32, length = 2048, 256, and M 16. We perform three separate scaling studies along the batch, N, and length dimensions respectively, and scale them from 32 to 2048.
isolate just the effect of
benchmark in fp32 precision on A100 40GB
30
20
307
107
107
12
197h
Rake Domnactin Optimteed Contraction
Figure 4: The training time (forward plus backward pass time) scaling with respect to the batch, state, and length dimensions. Note that the x-axis is in log seale.
17

---

Published as a conference paper at ICLR 2025
D ESTIMATION OF PARAMETERS AND FLOPS
It is very tempting to train network variants that are feasible by optimal contractions, but eventually turn out to be incredibly memory or computationally expensive during inference. An example would be to aggressively use "full" SSM blocks with large channels and internal states, which will result in the internal states of size (out channels, input channels, states) needing to be maintained and updated during inference. Even though the memory and computational bottlenecks during training via optimal contractions, this is not possible during online inference time if none of the internal states can be "contracted away" Therefore, following classical designs of lightweight convolutional networks (Howard et al. 2017 2017:SSaddlet et al.) 2018), we use full SSM blocks sparingly only in the first few layers where the channel dimensions are still small, and we will mostly be using (pointwise) bottleneck blocks in the deeper layers where the channel dimensions become larger.
be mitigated
Note that the number of parameters and FLOPs differ between online inference and training. For inference parameters, certain learnable parameters can be absorbed into the system matrices, such as the A parameter. For inference computations, the number of FLOPs is always linear with respect to the sequence length; however, this is often outweighed by the suboptimal order of operations imposed by the explicit materialization of internal states. In Table. 3 we provide the parameter count and the FLOPs per recurrent step when various SSM block variants are configured for online inference. There are a couple of points to note:
During online inference, the internal states need to be explicitly maintained and updated, meaning that the order of operation is always first the input projection, internal states up- date, output projection, and an optional mixer layer (for S4D).
The internal states are maintained as complex tensors, but we only take the real parts for the output projection. Note that complex weight is doubled the parameter count of the real counterpart, and a complex multiplication is 6 times the FLOPs of the real counterpart (4 multiplications 2 additions). In addition, updating the internal state with the projected real input requires 1 FLOP, and 2 FLOPs if the project input is also complex.
In the case where the projection matrices (e.g. B and C') are also complex, complex dot product (used in matrix multiplications) is performed and will incur additionally 2 extra FLOPs per accumulation, resulting total of 8 FLOPs per matrix element. However, if we only need the real part of the projected outputs (or analogously only having real inputs), the projection operation will incur half the number of FLOPs compared the full complex projection, or 4 FLOPs per matrix element.
We do not count peripheral parameters and FLOPs for biases or affine transformations in normalization layers, as they are negligible compared to the SSM operations.
In Centaurus, we restrict all model parameters except for the state transition matrix A be real. If we make the projection matrices complex, we can then recover the original S4D and S5 implemen- tations as DWS and PW bottleneck blocks respectively. Following the original implementations, we still restrict the inputs and outputs of the SSM layers to be real. Besides using complex projection matrices, there are some additional idiosyncratic differences:
The original S4D layer uses a standard where a complex internal state is considered to be two states, whereas S5 and Centaurus do not and simply consider it as one full state.
The original S4D layer uses the GLU activation, meaning that the pre-activations hence the C' matrix will have double the size compared to Centaurus.
Note that if a direct residual connection is added to the SSM block, there will only be a trivial amount of FLOPs added equaling H'. In the case where the residual connection projection operation (necessary if H & H'), then there will be H' additional parameters added and roughly H' additional FLOPs. Residual projections will be used in our Centaurus networks for keyword spotting and automatic speech recognition.
Of course, this is a non non-problem if the network is to perform inference in an offline setting, where no internal states need to be explicitly materialized during inference, and the flexibility of contraction order is restored.
18

---

Published as a conference paper at ICLR 2025
Table 3: The number of parameters and FLOPs per recurrent step during online inference mode for the SSM block variants. We assume that all the optimizations for online inference have already been done, including the pre-computation of the SSM discretization and diagonalization, except for the S6 block whose SSM matrices need to be dynamically generated.
SSM Block
Parameters
FLOPs step
Depthwise
3HN
9HN
Depthwise Separable (S4D-like)
3HN + HH' HN + + H'N
9HN + 2HH1'
Pointwise Bottleneck (S5-like)
2HN + 7N + 2
Bottleneck
HN +3NM H
2HN + 99 M 2H'N
Full
3
9HH'N
Depthwise (complex)
4HN
1
Depthwise Separable (complex)
4HN + HH'
IHH + 2HH
Pointwise Bottleneck (complex)
2HN + 2N + 2
4HN + SN + +HHN
Bottleneck (complex)
2HN + 4NM + 2H'N
4HN + 16NM + N
Full (complex)
4HH'N
11HH'N
Depthwise Separable (S6-like) Mamba
4HN + /12rr++ HH'
14HN + +r 2H 16H + 28HN + 16H
8H + 8HN + 4HH//r
we only extract the core selective scan op- eration, and replace the depthwise SSM layer with it for our ablations, while still maintaining the output mixer layer. Note that every parameter and operation involved in the S6 layer is fully real space. The core of the S6 layer is the selective scan mechanism, where the state matrices are data- controlled. This mechanism only meaningful the channel (feature) dimension is greater than 1, so for mono-channel layers will always fallback to the standard depthwise layer. The estimation of parameters and FLOPs for this layer as follows:
For the S6 layer used Mamba (Gu Dao 2023
For the data gattaaggatng operation, the generation of A requires roughly H²/r parameters and 4H2 /r FLOPs, where is the low-rank factor and the low-rank dimension being [H/16], chosen to be 16. The generation of the B and C matrices requires a total of 2HN parame- ters and N FLOPs.
The discretization process has to occur dynamically due to the dynamicity of the A pa- rameter. It does not incur additional parameters but does require additional FLOPs. This involves applying softplus to A, performing element- wise multiplication with the input and the transition matrix A, and exponentiating the latter. We conservatively estimate the FL.PPs required to be 5HN, under the assumption that the expf operation takes at least FLOPss per element.
Finally, we have the actual depthwise SSM operations, which in real space, incur 2HN parameters and 5 FLOPs.
In the original Mamba block, = 16 is kept constant. The channel dimension H also gets doubled before passing into the core S6 layer. On top of this, the network has an additional non-linear "gating" path that also has 2H features. The input and output projection layers allowing for this macro structure will then result
total 8H²² parameters and 16/ FLOPs, in addition to the core S6 operations. There is also lightweight causal depthwise ConviD ID layer that has negligible parameters and FLOPs. The Mamba block like "bottleneck" block that does not alter the channel dimension H. hence can perform channel projection during downsampling project H H' (see next Section).
E EXPERIMENT DETAILS
Unless otherwise mentioned, all of our network variants are trained with:
AdamW W optimizer with the PyTorch default configs
19

---

Published as a conference paper at ICLR 2025
Projection
Down sampling
SSM Block
LayerNorm
Dropouttd
SILU
We use basic design for Centaurus block. where the lighter-shaded blocks optional. The skip path typically just an identity connection, but can be a pointwise projection layer to change the channel dimension. Similarly, the (optional) downsampling layer can be simple average pooling to preserve the channel dimension, or downsampling convolutional layer to project the channel dimension at the same time. The downsampling layer can be replaced by the upsampling layer if needed (e.g. in a decoder block of an auto-encoder).
Figure 5:
cosine decay scheduler with linear warmup period equal to 0.01 of the total training steps, updating after every optimizer step
gradient clip value of 1
layer normalization (over the feature dimension) with elementwise affine parameters
SiLU activation
no dropout layers except for the keyword-sponne spotting network
Additionally, trained with automatic mixed precision (AMP) along with torch. comp le, except for the FFT convolution operations which are performed in full fp32 precision and without compilation (due inability handle complex data types currently).
Training configs not mentioned in the above list will be mentioned separately in the following in- dividual subsections. We initialize the and C projection matrices with the standard Kaiming uniform distribution. For the initialization of the A and A parameters, we follow the S4D initializa- tion strategy due to its simplicity. In other words. we set Å be geometrically ranged from 0.001 to 0, and A, = 1/2 inn. We suspect any sensible initialization method should also work fine. Below, we explicitly mention the dimension for which the initialization range applied, with the indices denoting the dimensions of the A and A parameters (slight abuse of notation).
For S4D, initialize Acn over the dimension, and Acn the dimension.
For S5, we split the internal states into groups of 4. A, is initialized across the groups, and An is initialized within each group.
For the full SSM block, we initialize Ain over the dimension, and Ajin over the n dimen- sion.
For the bottleneck block, we initialize An over the dimension, and Anm over the m dimension.
All of our training runs and trials are done with Py Porch with torch. comp l enabled except for operations involving complex numbers (e.g. FFTs). In addition, enabled tensorf o.2.32 for matrix multiplications, and the opt_einsum backend for all torch. e einsum operations.
F.
KEYWORD SPOTTING
For all our trials in this experiment, we
train for 200 epochs
use learning rate of 0.01 with weight decay 0.05
use linear warmup period of 0.1 for the scheduler.
Dropout d with probability 0.1, only applied if the number of features is greater than 4
20

---

Published as a conference paper at ICLR 2025
Architecture
Channels
States
septhwise-separable (S4D/S6-like)
[2 4. 8, 16, 32, 64]
[4, 16, 64, 128)
64,
4
88 1,16,30 32, 64, 128, 256]
4
16, 64, 128, 256,512)
pointwise bottleneck (S5-like)
[2, 4, 8, 16,32, 64] [4 84 16, 32, 64, 128] [s, 1616, 32, 64, 128, 256]
[4, 8, 16, 32, 64, 128]
[8, 16, 32, 64, 128, 256
[16, 64, 128, 256, 512]
full
[2, 4. 8, 16, 32, 64] [4, 8, 16, 32, 64, 128)
4
8, 16, 32, 64, 128, 256]
bottleneck
[2,4, 32, 64)
[4,8. 16,32, 32, 64, 128)
[4, 8, 16, 32, 64, 128)
88 16, 32, 64, 128, 256)
8, 16, 32, 64, 128, 256]
16,32, 64, 128, 256, 512]
full x 2 bottleneck x 2 S5 x
2, 8, 32, 64]
[4,4, 16, 16, 32, 64, 128]
[4. 8, 16, 32, 64, 128)
4, 32, 64, 128, 256] [4, 4, 64, 128, 256, 512)
[8 16, 32, 64, 128, 256]
Table 4:
The different network variants tested. For the bottleneck blocks, it is always assumed that the number of sub-states is 4.
* perform training single NVIDIA A30 GPU with a batch size of 512
The network contains six SSM blocks, with output channels and internal states given in Table. F. Besides the first block of the network, every SSM block is a residual block where the skip path is a pointwise projection. We apply LayerNorm after the SSM operations but before the skip merge, and apply the SiLU activation after the skip merge, following the classical ResNet residual block design. After each SSM layer, we perform ID average pooling with window sizes 4, 2} respectively. We apply this network over one-second durations of raw audio waveforms, and a hooballavraggepoollin (GAP) on the features generated by the SSM backbone, followed by a 2- layer MLP classification head
In addition, we test against hybrid network variant where the projection matrices are made com- plex, shown in Fig. 6 Over the model sizes tested, it appears that real variants outperform the complex variants when accounting for the model parameters and inference FLOPs. Furthermore, compare against network variants with S6 layers having different state dimensions, the num- ber of states appears to be an important factor for data-controlled SSM layers (Gu & Dao 2023).
To compare with the top results achieved KWS by other state-space models, Centaurus hybrid network architectured as given Table. E. We report the performance and size of each model in Table.50on the Speech Command 10-class subset (SC10), as it is common benchmark for all the models. Centaurus achieves SOTA results with less parameters and orders of magnitude less FLOPs.
our largest
E.2 RAW SPEECH DENOISING
The high-level training pipeline for the raw audio denoising model to simply generate synthetic noisy audios by randomly mixing clean and noise audio sources. The noisy sample is then used as input to the model, and the clean sample is used as the target output. For the clean training samples, use the processed VCTK and LibriVos datasets that can be downloaded from the Microsoft DNS4 challenge. We also use the noise training samples from the DNS4 challenge as well, which contains
The application of the GAP layer is in accordance with previous works. In a practical real-time sal-time inference setting, the GAP layer will likely be removed, and the classification head will be likely attached directly to the streaming output features of the SSM hackbone. Standard post-processing techniques such as majority- and early prediction then he applied.
21

---

Published as a conference paper at ICLR 2025
Parameter Scaling
Inference Compute Scaling
hybrid
hybrid
hybrid (complex proj.)
hybrid (complex proj.)
Classification Error
6
10
4x
10a Parameters
10
FLOPs
Figure 6: The scaling of the multiclass classification error with respect the number of parameters and the number of FLOPs per second during inference. The performance is evaluated on the SC35 testset. We compare the scaling of the hybrid network against its counterpart where the projection matrices are made complex.
Parameter Scaling
Inference Compute Scaling
3x 10-
10-1
S6 (state
S6
2x
S6 (state 16)
20 4.
S4D (state 4)
Classification Error
Centaurus (hybrid)
10-2
10-1
S6 (state 4)
10-2
10-2
(state 16)
S4D (state
Centaurus thybrid)
10'
Parameters
FLOPs
Figure 7:
We compare the scaling of the hybrid network compared homogeneous network variants where all SSM layers are S6 layers (or equivalently depthwise-separable SSM blocks where the SSM matrices are data-controlled).
Comparing Centaurus hybrid network against other deep state-space models the SC10 testset. All networks operate on raw input waveforms sampled at 16000kHz. Note that only our Centaurus network performs downsampling, hence the small computational footprint (despite having more parameters). The numbers with an asterisk are estimated and not reported in the original works.
Table 5:
Model
Accuracy 1
Parameters
FLOPs sec 1
S4 (Gu et al. 2021
98.32
0.306M
11.8G*
S4D-Lin Gu et al. 2022
96.25
0.306M
11.8G*
LiquidS4 (Hasani al. 2022
98.51
0.224M
S5 (Smith et al. 2022
0.280M
13.5G*
Centaurus (hybrid)
98.53
0.378M
0.134G
22

---

Published as a conference paper at ICLR 2025
Resampling factor and output channel of each block of the network, which consists of an encoder performing down-samplings, an intermediate bottleneck, decoder performing up- samplings, and finally output processor. The number of sub-states for the "neck" SSM block is
Table 6:
always
Layers
Resampling Factor
Channels
States
Encoder
Block (full)
16
16
Block 2 (full)
32
4
Block (neck)
64
128 x
Block 4 (neck)
96
128 x 4
Block (pw-neck)
128
256
Block 6 (pw-neck)
256
256
Bottleneck
Block (pw-neck)
256
256
Block 2 (pw-neck)
256
256
Decoder
Block (pw-neck)
128
256
Block (pw-neck)
96
256
Block 3 (neck)
64
128 X 4
Block (neck)
32
128 x
Block 5 (full)
16
4
Block 6 (full)
1
Output
Block (full)
1
16
Block 2 (full)
1
16
the Audioset, Freesound, and DEMAND datasets. For all audio samples, we use the librosa library to resample them to 16 kHz and load them numpy arrays.
For the LibriVox audio samples which form long continuous segments of human subjects reading from book, simply concatenate all the numpy arrays, and pad at the very end such that the array can reshaped into (segments, 277), or roughly 8.192 second segments. For all the other audio samples consisting of short disjoint segments,
perform intermediate paddings when necessary, to ensure single recording does not span two rows in the final array. For audio samples longer than length 27, we simply discard them. The input length to our network during training then also 2
For every epoch, we use the entirety of the VCTK dataset and 10 percent of randomly sampled subset of the LibriVox dataset. For each clean segment, we pair with a randomly sampled noise segment (with replacement). The clean and noise samples are added together with an SNR sampled from -5 to 15 dB, and the synthesized noisy sample is then rescaled to a random level from -35 dB to -15 dB. Furthermore, we perform random temporal and frequency masking (part of the SpecAugment transform) on only the input noisy samples.
For all our trials in this experiment, we
train for 500 epochs
use learning rate of 0.005 with weight decay 0.02
perform training on a single NVIDIA A30 GPU with a batch size of 1 192, resulting around 1573 seconds of audio data per batch.
For the loss function, we combine SmoothL with spectral loss on the ERB scale, with- out any equalization procedure on the waveforms spectrograms. The parameter of the LmoothL.lLoss is set at 0.5, and the spectral loss is weighted by a factor that grows from 0 to linearly during training.
See Table. for the Centaurus network architecture. The SSM blocks themselves do not change the channel dimensions of the features. The downsampling and upsampling layers trainable
23

---

Published as a conference paper at ICLR 2025
convolutional layers, which at the same time perform channel projections. For example, if the input features are monoaural with shape (1, 16000), then the first SSM block will retain this shape, but the following downsampling block will project it to the shape (16, 4000). The network follows an hourglass architecture, with a long-range skip connection branching after every SSM block in the encoder, and merging before every corresponding SSM block in the decoder (i.e. output features of encoder block is added to the input features of decoder block 6).
The homogeneous variants of the network is configured with the same architecture, with the same resampling factors and channels. To keep the parameters and FLOPs of the variants roughly at the same order, we use the following state configurations:
For the depthwise-separable variants (S4D-like and S6-like),
64 states.
For the pointwise bottleneck variant (S5-like), we use 256 states.
For the bottleneck variant, we use 256 states, with 4 sub-states per state.
For the full variant, we use 4 states per input-output channel pairing.
Surprisingly, using the Mamba blocks in their original form (including the highly complex nested macro-architecture) resulted in severe training plateaus, and resulted in many runs that did not out- perform random guessing, despite much effort in gested "no weight decay" training recipe for the A and D matrices.
tuning and following the sug-
E.3
SPEECH RECOGNITION
Table Resampling factor, input/output channels, and internal states of each block of the Centau- ASR network. The number of sub-states for the "neck" SSM block is always 4. The SSM block itself maintains the feature shape, followed by the downsampling layer projecting the temporal and channel dimension simultaneously. An optional feedforward convolutional module can be addi- tionally appended to the end of each block (again maintaining the feature shape). A block with a multiplier means that the block is stacked multiple times, with only the first occurrence of the block containing the downsampling layer (feature projection).
Layers
Resampling Factor
In Channels
Out Channels
States
Backbone (12 layers)
Full Block
5
1
16
64
Full
4
16
32
Neck x2
2
32
64
128 x 4
Neck x2
2
64
128
256 x 4
PW-Neck x2
2
128
256
512
PW-Neck x4
2
256
512
1024
Head (6 layers)
PW-Neck
512
512
1024
PW-Neck
2
512
512
1024
PW-Neck x4
1
512
512
1024
The backbone is 12-layer encoder, performing total downsampling factor of 320, effectively streaming output features period of 20 ms. The output features from the encoder are then branched into two heads: 1) smaller "character" head that is 2-layer MLP, producing charac- ter logits, 2) and a larger "language" head, consisting of six additional pointwise-bottlencck SSM blocks, producing BERT token logits[33 additionally downsampling the features by a factor of 8. The character head is supervised with the soft logits from the WAV2VEC2 LARGE LV50K_9600 teacher model, and the language head is supervised via CTC loss with the tokenized transcripts. During inference, we simply extract the most probable token produced by the final layer of the "language head".
Note that nowhere in our solution will actually use the BERT network. We only use the BERT tokenizer for the token encoding of transcriptions and the token decoding of the network predictions.
24

---

Published as a conference paper at ICLR 2025
The training pipeline involves supervised learning transcribed audio samples. The training dataset includes the LibriSpeech full 960h training set, and the Multilingual LibriSpeech (MLS) English training set. The training lasts 100 epochs, with each epoch containing the full LibriSpeech training set and 1% of the ML. English training set randomly sampled (so that a full training run will see all training data). To guarantee no data leakage, we check every sample in the training and remove it if it is sufficiently similar to any sample in the development and testing sets. All the audio data are sampled at 16kHz before being fed into the network. Similar the data pre-processing pipeline for speech denoising (see E.2), we pack the audio clips into segments of length 327680 (or roughly 20.5 seconds), introducing paddings when necessary to ensure a single recording does not span two segments.
For training the network, we
train for 100 epochs
use a learning rate of 0.02 with weight decay of 0.1
perform training with 8x 40GB A100 with batch size of 32 per GPU. This results in around 1.46 hours worth of audio data per effective batch.
The base Centaurus network can be augmented by appending additional module after every SSM block, to promote more local modeling capabilities. We experiment with three types of (residual) modules here.
Feed-forwart network (FFN): This consists of two pointwise layers in the main path, mak- ing it purely feature mixing. Both layers preserve the channel dimension, with no expansion factors.
Convolutional module (Conv): This is residual module adapted from the Conformer ar- chitecture (Gulati et al.] 2020). It consists of the pointwise-depthwin deethhwsse-pinntisse pattern with an expansion ratio of 2 in the main path; the depthwise layer is causal convolutional layer with kernel size of
This complex residual gated block containing three nested skip con- nections, adapted from the Mamba archt-architecetre (Gu & Dao 2023). It begins by projecting the input channels into 4 times the original dimension, and splitting the channel features into two paths, one path for the core SSM operations, and another path for gating. The SSM layer itself contains an inner skip connection via the matrix (which we ignored for our Centaurus block). Additionally, there is outer residual connection wrapping the entire Mamba block itself.
Mamba block:
For using the Mamba block, we only applied it for the last 4 "language head" blocks they are homogeneous. We also tested replacing the shallower blocks with Mamba blocks but it resulted in unstable training, which we believe was due to the highly complex Mamba blocks being too "close" to the less featurized raw audio signals. In addition, we also tested the Mamba blocks in their original form (with DWS gated layers), but it also resulted in unstable training.
25


Reviews:
[{"review_id": "yjUulr43QJ", "reviewer": "Reviewer_Td3n", "summary": "The paper presents a series of hybrid state-space models (SSMs), inspired by the design of convnets. It achieves this by generalizing 1D kernels to 2D and designing the depthwise, group, and full convolution counterparts for SSMs. Experiments on several sequence modeling tasks validate the effectiveness of the proposed method.", "questions": "My concerns are listed in the weaknesses section. \n\nI have one additional question: For the generalized SSM blocks with group/full convolution patterns, can I understand that these blocks integrate both token-mixing and channel-mixing within a single block, making other channel-mixing blocks, such as FFNs, unnecessary?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The concept of generalized SSM blocks, configurable with flexible connectivity structures, demonstrates good soundness and could complement existing SSM designs.\n\n2. The proposed method is supported by both theoretical analysis and empirical experiments.", "weaknesses": "1. The major concern is that the real-device training/inference efficiency of the proposed generalized SSM blocks is not provided. It is unclear whether the flexibility of these building blocks comes at the cost of reduced real-device efficiency.\n\n2. Another major concern is that the position of the proposed method among the latest SOTA SSMs is unclear. For example, Mamba has introduced input adaptivity in their selective state design and employed an advanced macro-structure with gating mechanisms and FFNs to enhance channel mixing in addition to token mixing. These advancements also address the limited expressiveness of vanilla SSMs, which are mostly depthwise-separable. The authors are expected to demonstrate whether the proposed method is a better choice or is orthogonal and can be combined with these advanced designs.\n\n3. The proposed method is evaluated on small-scale sequence modeling tasks. It would be highly desirable for it to be evaluated on language modeling and commonsense reasoning tasks.\n\n4. The writing clarity could be improved by better explaining certain terms. For example, the concept of \"depthwise-separable\" in the context of SSMs is important to the logic of this paper and should be better clarified.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "weK5LsaE2f", "reviewer": "Reviewer_j8MZ", "summary": "The authors present an approach to augment the number of features in the inputs, outputs and internal states of LTI systems, leveraging state space models, in order to facilitate architectural flexibility towards building abstractions that vary feature dimensions in neural models, such as those in CNNs. The authors present connections between tensor networks, which generalize low-rank decompositions in a way that can be applied to SSM einsums.\n\nThe authors introduce a weighting for basis kernels which can be adapted to a recurrent system. Convolutions, in this way, are shown as sums of various real and complex combinations of Fourier modes. The authors use this idea to present Centaurus, a configuration of tensor networks that allows retains SSM blocks but projects internal state to different sizes, making it possible to mimic various CNN abstractions such as depthwise, grouped, and bottleneck convolutions.\n\nThe authors construct each of these layers and present implementation details, i.e. a tradeoff between optimization of einsum expressions (for which we have few compiler tools at the moment) versus GPU kernel customization.\n The authors present experiments on three tasks within sequence modeling/speech: keyword spotting, raw speech denoising, and speech transcription, showing strong results in each area with small models with good computational scaling properties.", "questions": "- The authors should cite other work in speech that explores other convolutional architectures, such as ([Kirman et al.](https://ieeexplore.ieee.org/iel7/9040208/9052899/09053889.pdf), time-depthwise separable convolutions, [Hannun et al.](https://arxiv.org/pdf/1904.02619), or other extensive architectural ablations based on computational efficiency for training and inference: [Pratap et al.](https://arxiv.org/pdf/2001.0972))\n- ASR baselines can examine other convolutional architectures such as those presented in [Synnaeve et al.](https://arxiv.org/pdf/1911.08460); these baselines are slightly better comparisons given that they are also end-to-end and also present results language-model based decoding.\n- Experiments in Appendix E leverage samples from Libri-Vox — was they Libri-Light or samples of raw data?\n- A brief discussion of future work/applications to other domains (outside of speech) for CNNs (such as computer vision tasks) might strengthen the paper.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The presentation of methodology is convincing and strongly tied to first principle sin SSMs.\n- The application of tensor networks to solve the problem of projecting kernel matrices to and from frequency spaces of different dimensions is novel and elegant.\n- The presentation of implementation and computational considerations shows care was given to scaling tradeoffs, i.e. memory-boundedness of this regime of compute regime, opportunities (or lack thereof) for operator fusion, and operation/contraction ordering per kernel construction.\n- Results are compelling — in particular, efficiency of the proposed models is impressive given performance.\n- Attention to detail with the model name is nice.", "weaknesses": "- Baselines could be much stronger, in that there are no other SSM model baselines. Ablations are mostly within the architectural innovations present with Centaurus.\n- The constraint that some projection matrices from basis kernels must be real constrains the expressiveness of the model, although it’s likely that such generalizations in future work might enable this.\n- A more comprehensive architectural ablation would make the paper stronger.  While there are explanations of architectures in appendix E, there is no study or comparison of components to CNNs. Do architectural ideas useful in building CNNs apply to building Centaurus models?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "G2ZthZyrec", "reviewer": "Reviewer_x1nf", "summary": "The authors present Centaurus, an adaption of state-space models (SSMs) / linear RNNs (lRNNs) to audio modeling / processing, where they embed these layers of different dimenionalities into a deeper hybrid architecture, where they take inspiration from classical CNN networks that use different strides, kernel sizes and feature dimensions at different levels in the network. Their ablations show that the hybrid composition is superior to homogenous depth-wise, bottleneck and point-wise bottleneck architectures, and their models show strong performance while being superiorly FLOP-efficient compared to the current SOTA. Theoretically, they connect their hybrid architecture adaptions to tensor networks combined with Fourier operations and show the optimal tensor contraction strategy for their network.", "questions": "Since both RNNs and State Space models are well-known, a detailed introduction of these architectures can be omitted, except for the parallelization / transformations to/from the Fourier domain.\nEinsum notation is definitely superior for Tensor expressions, so also here the extensive introduction could be reduced.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Their hybrid architecture shows better results at lower FLOPs, while retaining a similar scaling behaviour compared to other homogeneous models across model sizes.", "weaknesses": "The proposed model is a combination of known primitives (SSMs and inhomogeneous scaling from CNN architectures).\nThe work does not cite and compare to other relevant, related work that potentially outperforms the trained models on the given datasets, e.g. Zhao et al. 2022: \"Monaural speech enhancement with complex convolutional block attention module and joint time frequency losses.\" for the VB-DMD dataset.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SFrhmZcqJg", "reviewer": "Reviewer_HmD1", "summary": "This paper proposes an efficient deep state-space model, Centaurus. Inspired by traditional ConvNets, the authors explore a heterogeneous design using block variants similar to convolutional blocks. The model is framed through the perspective of Tensor Networks and optimized by ordering tensor contractions for efficiency. Experiments demonstrate that the model achieves competitive performance across multiple tasks, including keyword spotting, raw speech denoising, and ASR.", "questions": "- The expression of the depthwise SSM operation in Section 4.1, Figure 1, differs slightly from the expression in Section 4.2 (above Equation 9), particularly in the subindex. The expression in Section 4.2 includes a subindex “b” that is not present in the earlier section. Could you clarify the difference and explain the meaning of “b”? If it refers to the bit representation in quantization, it would be helpful to explicitly cover this detail in the paper.\n\n- The authors provided an example of the most complex case—the bottleneck block—but skipped examples of other operations. Are there any general theories or principles that could be further distilled from these examples?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The authors proposed a novel perspective by viewing state-space models (SSMs) as tensor networks, enabling efficiency optimizations through contraction order analysis. This tensor network approach is advantageous for visualization, allowing the authors to address efficiency as a congestion optimization problem on the base graph. The approaches demonstrated in the paper contributes a valuable exploration into efficient SSM models.\n    - One of the optimizations involves applying FFT to transform the SSM into the frequency domain, allowing the optimal contraction order to be analyzed for improved efficiency. While FFT introduces a slight computational cost, which is not compute-bound, is minimal compared to the efficiency gains achieved through optimized contraction order.\n    - The authors explore multiple variants of SSM operations and illustrate them in a figure, broadening the possibilities for combining heterogeneous blocks to enhance model performance. Earlier SSM models only used identical operations throughout.\n    - The authors converted SSM into Frequency domain by FFT, which allows the optimization of determining the optimal “contraction order”. The paper also discussed the approach to evaluate the memory and computational requirements of an einsum contraction path involving FFT.\n    - As an example, the author explained how they find the optimal contraction order of Bottleneck SSM Operation. The proof involved Lemma 1 that reveals the intermediate tensors have at most 3 dims, which restricted the potential optimization paths. This largely reduced the manual work of comparing different contraction orders.\n\n- The proposed method demonstrates competitive performance with significantly enhanced efficiency. For keyword spotting, the model achieved the best accuracy with 100x fewer FLOPs. In raw speech denoising, the authors conducted an ablation study on different SSM variants and concluded that the hybrid Centaurus model achieved the highest PESQ score with the smallest parameter count and the fastest FLOPs/sec. For ASR, the model also showed competitive performance.", "weaknesses": "The paper includes a large amount of appendix material, likely due to page limitations. While the authors present key results in the main paper, leaving detailed proofs in the appendix, it would be beneficial to adjust the structure to highlight more of the original contributions and slightly condense other sections, such as the background. Although the background on SSM and einsum notation is important and well-explained, it occupies two full pages. For instance, a potential improvement could be briefly summarizing the proof of Lemma 1 and explaining how it restricts feasible patterns in the main paper. As it is an important detail of explaining how you explore the optimal order. It would reduce the need for readers to frequently switch between the appendix and the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)