Original: E FFICIENTLY M ODELING L ONG S EQUENCES WITH S TRUCTURED S TATE S PACES
Revised: E FFICIENTLY M ODELING L ONG S EQUENCES WITH S TRUCTURED S TATE S PACES

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies.
Revised: A central goal of sequence modeling is designing a single principled model that can address sequence data across a range of modalities and tasks, particularly on long-range dependencies.

Original: Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of 10000 or more steps.
Revised: Although conventional models including RNNs, CNNs, and Transformers have specialized variants for capturing long dependencies, they still struggle to scale to very long sequences of 10000 or more steps.

Original: A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) x (cid:48)
Revised: A promising recent approach proposed modeling sequences by simulating the fundamental state space model (SSM) x (cid:48)

Original: ( t ) =
Revised: ( t ) =

Original: Ax ( t ) + Bu ( t ) , y ( t ) =
Revised: Ax ( t ) + Bu ( t ) , y ( t ) =

Original: Cx ( t ) + Du ( t ) , and showed that for appropriate choices of the state matrix A , this system could handle longrange dependencies mathematically and empirically.
Revised: Cx ( t ) + Du ( t ) , and showed that for appropriate choices of the state matrix A , this system could handle longrange dependencies mathematically and empirically.

Original: However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution.
Revised: However, this method has prohibitive computation and memory requirements, rendering it infeasible as a general sequence modeling solution.

Original: We propose the Structured State Space (S4) sequence model based on a new parameterization for the SSM, and show that it can be computed much more efﬁciently than prior approaches while preserving their theoretical strengths.
Revised: We propose the Structured State Space (S4) sequence model based on a new parameterization for the SSM, and show that it can be computed much more efﬁciently than prior approaches while preserving their theoretical strengths.

Original: Our technique involves conditioning A with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel.
Revised: Our technique involves conditioning A with a low-rank correction, allowing it to be diagonalized stably and reducing the SSM to the well-studied computation of a Cauchy kernel.

Original: S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation 60 × faster (iii)
Revised: S4 achieves strong empirical results across a diverse range of established benchmarks, including (i) 91% accuracy on sequential CIFAR-10 with no data augmentation or auxiliary losses, on par with a larger 2-D ResNet, (ii) substantially closing the gap to Transformers on image and language modeling tasks, while performing generation 60 × faster (iii)

Original: SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efﬁcient as all competitors.
Revised: SoTA on every task from the Long Range Arena benchmark, including solving the challenging Path-X task of length 16k that all prior work fails on, while being as efﬁcient as all competitors.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: A central problem in sequence modeling is efﬁciently handling data that contains long-range dependencies (LRDs).
Revised: A central problem in sequence modeling is efﬁciently handling data that contains long-range dependencies (LRDs).

Original: Real-world time-series data often requires reasoning over tens of thousands of time steps, while few sequence models address even thousands of time steps.
Revised: Real-world time-series data often requires reasoning over tens of thousands of time steps, while few sequence models address even thousands of time steps.

Original: For instance, results from the long-range arena (LRA) benchmark (Tay et al., 2021) highlight that sequence models today perform poorly on LRD tasks, including one (Path-X) where no model performs better than random guessing.
Revised: For instance, results from the long-range arena (LRA) benchmark (Tay et al., 2021) highlight that sequence models today perform poorly on LRD tasks, including one (Path-X) where no model performs better than random guessing.

Original: Since LRDs are perhaps the foremost challenge for sequence models, all standard model families such as continuous-time models (CTMs), RNNs, CNNs, and Transformers include many specialized variants designed to address them.
Revised: Since LRDs are perhaps the foremost challenge for sequence models, all standard model families such as continuous-time models (CTMs), RNNs, CNNs, and Transformers include many specialized variants designed to address them.

Original: Modern examples include orthogonal and Lipschitz RNNs (Arjovsky et al., 2016; Erichson et al., 2021) to combat vanishing gradients, dilated convolutions to increase context size (Bai et al., 2018;
Revised: Modern examples include orthogonal and Lipschitz RNNs (Arjovsky et al., 2016; Erichson et al., 2021) to combat vanishing gradients, dilated convolutions to increase context size (Bai et al., 2018;

Original: Oord et al., 2016), and an increasingly vast family of efﬁcient
Revised: Oord et al., 2016), and an increasingly vast family of efﬁcient

Original: Transformers that reduce the quadratic dependence on sequence length (Katharopoulos et al., 2020;
Revised: Transformers that reduce the quadratic dependence on sequence length (Katharopoulos et al., 2020;

Original: Choromanski et al., 2020).
Revised: Choromanski et al., 2020).

Original: Despite being designed for LRDs, these solutions still perform poorly on challenging benchmarks such as LRA (Tay et al., 2021) or raw audio classiﬁcation (Gu et al., 2021).
Revised: Despite being designed for LRDs, these solutions still perform poorly on challenging benchmarks such as LRA (Tay et al., 2021) or raw audio classiﬁcation (Gu et al., 2021).

Original: An alternative approach to LRDs was recently introduced based on the state space model (SSM) (Fig. 1).
Revised: An alternative approach to LRDs was recently introduced based on the state space model (SSM) (Fig. 1).

Original: SSMs are a foundational scientiﬁc model used in ﬁelds such as control theory, computational neuroscience, and many more, but have not been applicable to deep learning for concrete theoretical reasons.
Revised: SSMs are a foundational scientiﬁc model used in ﬁelds such as control theory, computational neuroscience, and many more, but have not been applicable to deep learning for concrete theoretical reasons.

Original: In particular, Gu et al.
Revised: In particular, Gu et al.

Original: (2021) showed that deep SSMs actually struggle even on simple tasks, but can perform exceptionally well when equipped with special state matrices
Revised: (2021) showed that deep SSMs actually struggle even on simple tasks, but can perform exceptionally well when equipped with special state matrices

Original: A recently derivedto solve a problem of continuous-time memorization (Voelker et al., 2019; Gu et al., 2020a). Their
Revised: A recently derived

Original: [Equation]
Revised: [Equation]

Original: Continuous
Revised: [Figure]

Original: State Space
Revised: State Space

Original: Fast Discrete Representations
Revised: Fast Discrete Representations

Original: Long-Range
Revised: Long [Figure]

Original: Dependencies
Revised: Dependencies

Original: [Figure]
Revised: [Figure]

Revised: to solve a problem of continuous-time memorization (Voelker et al., 2019; Gu et al., 2020a).

Original:  Linear State Space Layer (LSSL) conceptually uniﬁes the strengths of CTM, RNN and CNN models, and provides a proof of concept that deep SSMs can address LRDs in principle.
Revised: Their Linear State Space Layer (LSSL) conceptually uniﬁes the strengths of CTM, RNN and CNN models, and provides a proof of concept that deep SSMs can address LRDs in principle.

Original: Unfortunately, the LSSL is infeasible to use in practice because of prohibitive computation and memory requirements induced by the state representation.
Revised: Unfortunately, the LSSL is infeasible to use in practice because of prohibitive computation and memory requirements induced by the state representation.

Original: For state dimension N and sequence length L , computing the latent state requires O ( N 2 L ) operations and O ( NL ) space – compared to a Ω( L + N ) lower bound for both.
Revised: For state dimension N and sequence length L , computing the latent state requires O ( N 2 L ) operations and O ( NL ) space – compared to a Ω( L + N ) lower bound for both.

Original: Thus for reasonably sized models (e.g. N = 256 in Gu et al.
Revised: Thus for reasonably sized models (e.g. N = 256 in Gu et al.

Original: (2021)), the LSSL uses orders of magnitude more memory than comparably-sized RNNs or CNNs.
Revised: (2021)), the LSSL uses orders of magnitude more memory than comparably-sized RNNs or CNNs.

Original: Although theoretically efﬁcient algorithms for the LSSL were proposed, we show that these are numerically unstable.
Revised: Although theoretically efﬁcient algorithms for the LSSL were proposed, we show that these are numerically unstable.

Original: In particular, the special A matrix is highly non-normal in the linear algebraic sense, which prevents the application of conventional algorithmic techniques.
Revised: In particular, the special A matrix is highly non-normal in the linear algebraic sense, which prevents the application of conventional algorithmic techniques.

Original: Consequently, although the LSSL showed that SSMs have strong performance, they are currently computationally impractical as a general sequence modeling solution.
Revised: Consequently, although the LSSL showed that SSMs have strong performance, they are currently computationally impractical as a general sequence modeling solution.

Original: In this work, we introduce the Structured State Space (S4) sequence model based on the SSM that solves the critical computational bottleneck in previous work.
Revised: In this work, we introduce the Structured State Space (S4) sequence model based on the SSM that solves the critical computational bottleneck in previous work.

Original: Technically, S4 reparameterizes the structured state matrices A appearing in Voelker et al.
Revised: Technically, S4 reparameterizes the structured state matrices A appearing in Voelker et al.

Original: ; Gu et al.
Revised: ; Gu et al.

Original: (2020a)
Revised: (2020a)

Original: by decomposing them as the sum of a low-rank and skew-symmetric term.
Revised: by decomposing them as the sum of a low-rank and skew-symmetric term.

Original: Additionally, instead of expanding the standard SSM in coefﬁcient space, we compute its truncated generating function in frequency space, which can be simpliﬁed into a multipole-like evaluation.
Revised: Additionally, instead of expanding the standard SSM in coefﬁcient space, we compute its truncated generating function in frequency space, which can be simpliﬁed into a multipole-like evaluation.

Original: Combining these two ideas, we show that the low-rank term can be corrected by the Woodbury identity while the skew-symmetric term can be diagonalized stably, ultimately reducing to a well-studied and theoretically stable Cauchy kernel (Pan, 2001; 2017).
Revised: Combining these two ideas, we show that the low-rank term can be corrected by the Woodbury identity while the skew-symmetric term can be diagonalized stably, ultimately reducing to a well-studied and theoretically stable Cauchy kernel (Pan, 2001; 2017).

Original: This results in ˜ O
Revised: This results in ˜ O

Original: ( N + L ) computation and O ( N + L ) memory usage, which is essentially tight for sequence models.
Revised: ( N + L ) computation and O ( N + L ) memory usage, which is essentially tight for sequence models.

Original: Compared to the LSSL, S4 is up to 30 × faster with 400 × less memory usage, while exceeding the LSSL’s performance empirically.
Revised: Compared to the LSSL, S4 is up to 30 × faster with 400 × less memory usage, while exceeding the LSSL’s performance empirically.

Original: Empirically, S4 signiﬁcantly advances the state-of-the-art for LRD.
Revised: Empirically, S4 signiﬁcantly advances the state-of-the-art for LRD.

Original: On the LRA benchmark for efﬁcient sequence models, S4 is as fast as all baselines while outperforming them by 20+ points on average.
Revised: On the LRA benchmark for efﬁcient sequence models, S4 is as fast as all baselines while outperforming them by 20+ points on average.

Original: S4 is the ﬁrst model to solve the difﬁcult LRA Path-X task (length- 16384 ), achieving 88% accuracy compared to 50% random guessing for all prior work.
Revised: S4 is the ﬁrst model to solve the difﬁcult LRA Path-X task (length- 16384 ), achieving 88% accuracy compared to 50% random guessing for all prior work.

Original: On speech classiﬁcation with length- 16000 sequences, S4 halves the test error ( 1 . 7% ) of specialized Speech CNNs – by contrast, all RNN and Transformer baselines fail to learn ( ≥ 70% error).
Revised: On speech classiﬁcation with length- 16000 sequences, S4 halves the test error ( 1 . 7% ) of specialized Speech CNNs – by contrast, all RNN and Transformer baselines fail to learn ( ≥ 70% error).

Original: Towards a general-purpose sequence model.
Revised: Towards a general-purpose sequence model.

Original: Beyond LRD, a broad goal of machine learning is to develop a single model that can be used across a wide range of problems.
Revised: Beyond LRD, a broad goal of machine learning is to develop a single model that can be used across a wide range of problems.

Original: Models today are typically specialized to solve problems from a particular domain (e.g. images, audio, text, time-series), and enable a narrow range of capabilities (e.g. efﬁcient training, fast generation, handling irregularly sampled data).
Revised: Models today are typically specialized to solve problems from a particular domain (e.g. images, audio, text, time-series), and enable a narrow range of capabilities (e.g. efﬁcient training, fast generation, handling irregularly sampled data).

Original: This specialization is typically expressed via domain-speciﬁc preprocessing, inductive biases, and architectures.
Revised: This specialization is typically expressed via domain-speciﬁc preprocessing, inductive biases, and architectures.

Original: Sequence models provide a general framework for solving many of these problems with reduced specialization – e.g. Vision Transformers for image classiﬁcation with less
Revised: Sequence models provide a general framework for solving many of these problems with reduced specialization – e.g. Vision Transformers for image classiﬁcation with less

Original: 2D information (Dosovitskiy et al., 2020).
Revised: 2D information (Dosovitskiy et al., 2020).

Original: However, most models such as Transformers generally still require substantial specialization per task to achieve high performance.
Revised: However, most models such as Transformers generally still require substantial specialization per task to achieve high performance.

Original: Deep SSMs in particular have conceptual strengths that suggest they may be promising as a general sequence modeling solution.
Revised: Deep SSMs in particular have conceptual strengths that suggest they may be promising as a general sequence modeling solution.

Original: These strengths include a principled approach to handling LRDs, as well as the ability to move between continuous-time, convolutional, and recurrent model representations, each with distinct capabilities (Fig. 1).
Revised: These strengths include a principled approach to handling LRDs, as well as the ability to move between continuous-time, convolutional, and recurrent model representations, each with distinct capabilities (Fig. 1).

Original: Our technical contributions enable SSMs to be applied successfully to a varied set of benchmarks with minimal modiﬁcation:
Revised: Our technical contributions enable SSMs to be applied successfully to a varied set of benchmarks with minimal modiﬁcation:

Original: • Large-scale generative modeling.
Revised: • Large-scale generative modeling.

Original: On CIFAR-10 density estimation, S4 is competitive with the best autoregressive models ( 2 . 85 bits per dim).
Revised: On CIFAR-10 density estimation, S4 is competitive with the best autoregressive models ( 2 . 85 bits per dim).

Original: On WikiText-103 language modeling, S4 substantially closes the gap to Transformers (within 0 .
Revised: On WikiText-103 language modeling, S4 substantially closes the gap to Transformers (within 0 .

Original: 8 perplexity), setting SoTA for attention-free models.
Revised: 8 perplexity), setting SoTA for attention-free models.

Original: • Fast autoregressive generation.
Revised: • Fast autoregressive generation.

Original: Like RNNs, S4 can use its latent state to perform 60 × faster pixel/token generation than standard autoregressive models on CIFAR-10 and WikiText-103.
Revised: Like RNNs, S4 can use its latent state to perform 60 × faster pixel/token generation than standard autoregressive models on CIFAR-10 and WikiText-103.

Original: • Irregularly sampled time-series.
Revised: • Sampling resolution change.

Original: Like specialized CTMs, S4 can adapt to changes in time series sampling frequency without retraining, e
Revised: Like specialized CTMs, S4 can adapt to changes in time-series sampling frequency without retraining, e

Original: g. at 0 .
Revised: g. at 0 .

Original: 5 × frequency on speech classiﬁcation.
Revised: 5 × frequency on speech classiﬁcation.

Original: • Learning with weaker inductive biases.
Revised: • Learning with weaker inductive biases.

Original: With no architectural changes, S4 surpasses Speech CNNs on speech classiﬁcation, outperforms the specialized Informer model on40 / 50 time-series forecasting problems, and matches a 2-D ResNet on sequential CIFAR with over 90% accuracy.
Revised: With no architectural changes, S4 surpasses Speech CNNs on speech classiﬁcation, outperforms the specialized Informer model on time-series forecasting problems, and matches a 2-D ResNet on sequential CIFAR with over 90% accuracy.

Original:  B ACKGROUND : S TATE S PACES
Revised:  B ACKGROUND : S TATE S PACES

Original: Sections 2.1 to 2.4 describe the four properties of SSMs in Fig.
Revised: Sections 2.1 to 2.4 describe the four properties of SSMs in Fig.

Original: 1: the classic continuous-time representation, addressing LRDs with the HiPPO framework, the discrete-time recurrent representation, and the parallelizable convolution representation.
Revised: 1: the classic continuous-time representation, addressing LRDs with the HiPPO framework, the discrete-time recurrent representation, and the parallelizable convolution representation.

Original: In particular, Section 2.4 introduces the SSM convolution kernel K , which is the focus of our theoretical contributions in Section 3.
Revised: In particular, Section 2.4 introduces the SSM convolution kernel K , which is the focus of our theoretical contributions in Section 3.

Original: 2. S TATE S PACE M ODELS : A C ONTINUOUS - TIME L ATENT S TATE M ODEL
Revised: 2. S TATE S PACE M ODELS : A C ONTINUOUS - TIME L ATENT S TATE M ODEL

Original: The state space model is deﬁned by the simple equation (1).
Revised: The state space model is deﬁned by the simple equation (1).

Original: It maps a 1-D input signal u ( t ) to an
Revised: It maps a 1-D input signal u ( t ) to an

Original: N -D latent state x ( t ) before projecting to a 1-D output signal y ( t ) .
Revised: N -D latent state x ( t ) before projecting to a 1-D output signal y ( t ) .

Original: [Equation]
Revised: [Equation]

Original: SSMs are broadly used in many scientiﬁc disciplines and related to latent state models such as Hidden
Revised: SSMs are broadly used in many scientiﬁc disciplines and related to latent state models such as Hidden

Original: Markov Models (HMM).
Revised: Markov Models (HMM).

Original: Our goal is to simply use the SSM as a black-box representation in a deep sequence model, where A , B , C , D are parameters learned by gradient descent.
Revised: Our goal is to simply use the SSM as a black-box representation in a deep sequence model, where A , B , C , D are parameters learned by gradient descent.

Original: For the remainder of this paper, we will omit the parameter D for exposition (or equivalently, assume D = 0 ) because the term D u can be viewed as a skip connection and is easy to compute.
Revised: For the remainder of this paper, we will omit the parameter D for exposition (or equivalently, assume D = 0 ) because the term D u can be viewed as a skip connection and is easy to compute.

Original: 2. A DDRESSING L ONG -R ANGE D EPENDENCIES WITH H I PPO
Revised: 2. A DDRESSING L ONG -R ANGE D EPENDENCIES WITH H I PPO

Original: Prior work found that the basic SSM (1) actually performs very poorly in practice.
Revised: Prior work found that the basic SSM (1) actually performs very poorly in practice.

Original: Intuitively, one explanation is that linear ﬁrst-order ODEs solve to an exponential function, and thus may suffer from gradients scaling exponentially in the sequence length (i.e., the vanishing/exploding gradients problem (Pascanu et al., 2013)).
Revised: Intuitively, one explanation is that linear ﬁrst-order ODEs solve to an exponential function, and thus may suffer from gradients scaling exponentially in the sequence length (i.e., the vanishing/exploding gradients problem (Pascanu et al., 2013)).

Original: To address this problem, the LSSL leveraged the HiPPO theory of continuous-time memorization (Gu et al., 2020a).
Revised: To address this problem, the LSSL leveraged the HiPPO theory of continuous-time memorization (Gu et al., 2020a).

Original: HiPPO speciﬁes a class of certain matrices
Revised: HiPPO speciﬁes a class of certain matrices

Original: A ∈ R N × N that when incorporated into (1), allows the state x ( t ) to memorize the history of the input u ( t ) .
Revised: A ∈ R N × N that when incorporated into (1), allows the state x ( t ) to memorize the history of the input u ( t ) .

Original: The most important matrix in this class is deﬁned by equation (2), which we will call the
Revised: The most important matrix in this class is deﬁned by equation (2), which we will call the

Original: HiPPO matrix.
Revised: HiPPO matrix.

Original: For example, the LSSL found that simply modifying an SSM from a random matrix
Revised: For example, the LSSL found that simply modifying an SSM from a random matrix

Original: A to equation (2) improved its performance on the sequential MNIST benchmark from 50% to 98% .
Revised: A to equation (2) improved its performance on the sequential MNIST benchmark from 60% to 98% .

Original: [Equation]
Revised: [Equation]

Original: 2. D ISCRETE - TIME SSM: T HE R ECURRENT R EPRESENTATION
Revised: 2. D ISCRETE - TIME SSM: T HE R ECURRENT R EPRESENTATION

Original: To be applied on a discrete input sequence ( u 0 , u 1 , . .
Revised: To be applied on a discrete input sequence ( u 0 , u 1 , . .

Original: instead of continuous function u ( t ) , (1) must be discretized by a step size ∆ that represents the resolution of the input.
Revised: instead of continuous function u ( t ) , (1) must be discretized by a step size ∆ that represents the resolution of the input.

Original: Conceptually, the inputs u k can be viewed as sampling an implicit underlying continuous signal u ( t ) , where u k = u ( k ∆) .
Revised: Conceptually, the inputs u k can be viewed as sampling an implicit underlying continuous signal u ( t ) , where u k = u ( k ∆) .

Original: To discretize the continuous-time SSM, we follow prior work in using the bilinear method (Tustin, 1947), which converts the state matrix A into an approximation A .
Revised: To discretize the continuous-time SSM, we follow prior work in using the bilinear method (Tustin, 1947), which converts the state matrix A into an approximation A .

Original: The discrete SSM is
Revised: The discrete SSM is

Original: [Equation]
Revised: [Equation]

Original: Equation (3) is now a sequence-to-sequence map u k
Revised: Equation (3) is now a sequence-to-sequence map u k

Original: (cid:55)→ y k instead of function-to-function.
Revised: (cid:55)→ y k instead of function-to-function.

Original: Moreover the state equation is now a recurrence in x k , allowing the discrete SSM to be computed like an RNN.
Revised: Moreover the state equation is now a recurrence in x k , allowing the discrete SSM to be computed like an RNN.

Original: Concretely, x k ∈ R N can be viewed as a hidden state with transition matrix A .
Revised: Concretely, x k ∈ R N can be viewed as a hidden state with transition matrix A .

Original: Notationally, throughout this paper we use A , B , . . .
Revised: Notationally, throughout this paper we use A , B , . . .

Original: to denote discretized SSM matrices deﬁned by (3).
Revised: to denote discretized SSM matrices deﬁned by (3).

Original: Note that these matrices are a function of both A as well as a step size ∆ ; we suppress this dependence for notational convenience when it is clear.
Revised: Note that these matrices are a function of both A as well as a step size ∆ ; we suppress this dependence for notational convenience when it is clear.

Original: 2. T RAINING SSM S : T HE C ONVOLUTIONAL R EPRESENTATION
Revised: 2. T RAINING SSM S : T HE C ONVOLUTIONAL R EPRESENTATION

Original: The recurrent SSM (3) is not practical for training on modern hardware due to its sequentiality.
Revised: The recurrent SSM (3) is not practical for training on modern hardware due to its sequentiality.

Original: Instead, there is a well-known connection between linear time-invariant (LTI) SSMs such as (1) and continuous convolutions.
Revised: Instead, there is a well-known connection between linear time-invariant (LTI) SSMs such as (1) and continuous convolutions.

Original: Correspondingly, (3) can actually be written as a discrete convolution.
Revised: Correspondingly, (3) can actually be written as a discrete convolution.

Original: For simplicity let the initial state be x − 1 = 0 .
Revised: For simplicity let the initial state be x − 1 = 0 .

Original: Then unrolling (3) explicitly yields x 0 = B ux
Revised: Then unrolling (3) explicitly yields x 0 = B ux

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: This can be vectorized into a convolution (4) with an explicit formula for the convolution kernel (5).
Revised: This can be vectorized into a convolution (4) with an explicit formula for the convolution kernel (5).

Original: [Equation]
Revised: [Equation]

Original:  In other words, equation (4) is a single (non-circular) convolution and can be computed very efﬁciently with FFTs, provided that K is known.
Revised: [Equation] In other words, equation (4) is a single (non-circular) convolution and can be computed very efﬁciently with FFTs, provided that K is known.

Original: However, computing K in (5) is non-trivial and is the focus of our technical contributions in Section 3.
Revised: However, computing K in (5) is non-trivial and is the focus of our technical contributions in Section 3.

Original: We call K the SSM convolution kernel or ﬁlter.
Revised: We call K the SSM convolution kernel or ﬁlter.

Original:  M ETHOD : S TRUCTURED S TATE S PACES (S4)
Revised:  M ETHOD : S TRUCTURED S TATE S PACES (S4)

Original: Our technical results focus on developing the S4 parameterization and showing how to efﬁciently compute all views of the SSM (Section 2): the continuous representation ( A , B , C ) (1), the recurrent representation ( A , B , C ) (3), and the convolutional representation K (4).
Revised: Our technical results focus on developing the S4 parameterization and showing how to efﬁciently compute all views of the SSM (Section 2): the continuous representation ( A , B , C ) (1), the recurrent representation ( A , B , C ) (3), and the convolutional representation K (4).

Original: Section 3.1 motivates our approach, which is based on the linear algebraic concepts of conjugation and diagonalization, and discusses why the naive application of this approach does not work.
Revised: Section 3.1 motivates our approach, which is based on the linear algebraic concepts of conjugation and diagonalization, and discusses why the naive application of this approach does not work.

Original: Section 3. gives an overview of the key technical components of our approach and formally deﬁnes the S parameterization.
Revised: Section 3. gives an overview of the key technical components of our approach and formally deﬁnes the S parameterization.

Original: Section 3.3 sketches the main results, showing that S4 is asymptotically efﬁcient (up to log factors) for sequence models.
Revised: Section 3.3 sketches the main results, showing that S4 is asymptotically efﬁcient (up to log factors) for sequence models.

Original: Proofs are in Appendices B and C.
Revised: Proofs are in Appendices B and C.

Original: 3. M OTIVATION : D IAGONALIZATION
Revised: 3. M OTIVATION : D IAGONALIZATION

Original: The fundamental bottleneck in computing the discrete-time SSM (3) is that it involves repeated matrix multiplication by A .
Revised: The fundamental bottleneck in computing the discrete-time SSM (3) is that it involves repeated matrix multiplication by A .

Original: For example, computing (5) naively as in the LSSL involves L successive multiplications by A , requiring O ( N 2 L ) operations and O ( NL ) space.
Revised: For example, computing (5) naively as in the LSSL involves L successive multiplications by A , requiring O ( N 2 L ) operations and O ( NL ) space.

Original: To overcome this bottleneck, we use a structural result that allows us to simplify SSMs.
Revised: To overcome this bottleneck, we use a structural result that allows us to simplify SSMs.

Original: Input: S4 parametersΛ , p , q , B , C ∈ C N and step size ∆
Revised: [Table] [Table] parameters

Original:  Output: SSM convolution kernel K= K L ( A , B , C ) for A = Λ − pq ∗ (equation (5))
Revised: [Table] Output: SSM convolution kernel K

Original: [Table]
Revised: [Table]

Revised: A

Original: [Table]
Revised: [Table]

Original: (cid:46)

Original: Truncate SSM generating function (SSMGF) to length

Original: L
Revised: P

Original: [Table]
Revised: [Table]

Original: Black-box Cauchy kernel

Original: [Table]
Revised: [Table]

Original: SSMGF

Original: [Table]
Revised: [Table]

Original: (cid:46)
Revised: Truncate SSM generating function (SSMGF) to length L

Original: [Table]
Revised: [Table]

Original: Transform

Original: Lemma 3.1.
Revised: Lemma 3.1.

Original: Conjugation is an equivalence relation on SSMs ( A , B , C ) ∼ ( V − 1 AV , V − 1 B , CV ) .
Revised: Conjugation is an equivalence relation on SSMs ( A , B , C ) ∼ ( V − 1 AV , V − 1 B , CV ) .

Original: Proof.
Revised: Proof.

Original: Write out the two SSMs with state denoted by x and ˜ x respectively:
Revised: Write out the two SSMs with state denoted by x and ˜ x respectively:

Original: [Equation]
Revised: [Equation]

Original: After multiplying the right side SSM by V , the two SSMs become identical with x = V ˜ x .
Revised: After multiplying the right side SSM by V , the two SSMs become identical with x = V ˜ x .

Original: Therefore these compute the exact same operator u (cid:55)→ y , but with a change of basis by V in the state x .
Revised: Therefore these compute the exact same operator u (cid:55)→ y , but with a change of basis by V in the state x .

Original: Lemma 3.1 motivates putting A into a canonical form by conjugation 1 , which is ideally more structured and allows faster computation.
Revised: Lemma 3.1 motivates putting A into a canonical form by conjugation 2 , which is ideally more structured and allows faster computation.

Original: For example, if A were diagonal, the resulting computations become much more tractable.
Revised: For example, if A were diagonal, the resulting computations become much more tractable.

Original: In particular, the desired K (equation (4)) would be a Vandermonde product which theoretically only needs O (( N + L ) log 2 ( N + L )) arithmetic operations
Revised: In particular, the desired K (equation (4)) would be a Vandermonde product which theoretically only needs O (( N + L ) log 2 ( N + L )) arithmetic operations

Original: (Pan, 2001).
Revised: (Pan, 2001).

Original: Unfortunately, the naive application of diagonalization does not work due to numerical issues.
Revised: Unfortunately, the naive application of diagonalization does not work due to numerical issues.

Original: First,
Revised: First,

Original: Vandermonde multiplication is itself a famously ill-conditioned problem (Pan, 2016).
Revised: Vandermonde multiplication is itself a famously ill-conditioned problem (Pan, 2016).

Original: Furthermore, we derive the explicit diagonalization for the HiPPO matrix (2) and show it has entries exponentially large in the state size N , rendering the diagonalization numerically infeasible (e.g. CV in Lemma 3.
Revised: Furthermore, we derive the explicit diagonalization for the HiPPO matrix (2) and show it has entries exponentially large in the state size N , rendering the diagonalization numerically infeasible (e.g. CV in Lemma 3.

Original: would not be computable).
Revised: would not be computable).

Original: We note that Gu et al.
Revised: We note that Gu et al.

Original: (2021) proposed a different (unimplemented) algorithm to compute K faster than the naive algorithm.
Revised: (2021) proposed a different (unimplemented) algorithm to compute K faster than the naive algorithm.

Original: In Appendix B, we prove that it is also numerically unstable for related reasons.
Revised: In Appendix B, we prove that it is also numerically unstable for related reasons.

Original: Lemma 3.2.
Revised: Lemma 3.2.

Original: The HiPPO matrix A in equation (2) is diagonalized by the matrix V ij = (cid:0)
Revised: The HiPPO matrix A in equation (2) is diagonalized by the matrix V ij = (cid:0)

Original: i + j i − j (cid:1) .
Revised: i + j i − j (cid:1) .

Original: In particular, V 3 i,i = (cid:0) 

Revised: In particular, V 3 i,i = (cid:0)

Original: i 2 i (cid:1)
Revised: i 2 i (cid:1)

Original: ≈ 2 

Revised: ≈ 2

Original: i .
Revised: i .

Original: Therefore V has entries of magnitude up to 2 4 N/ 3 .
Revised: Therefore V has entries of magnitude up to 2 4 N/ 3 .

Original: 3. T HE S4 P ARAMETERIZATION : N ORMAL P LUS L OW -R ANK
Revised: 3. T HE S4 P ARAMETERIZATION : N ORMAL P LUS L OW -R ANK

Original: The previous discussion implies that we should only conjugate by well-conditioned matrices V .
Revised: The previous discussion implies that we should only conjugate by well-conditioned matrices V .

Original: The ideal scenario is when the matrix A is diagonalizable by a perfectly conditioned (i.e., unitary) matrix.
Revised: The ideal scenario is when the matrix A is diagonalizable by a perfectly conditioned (i.e., unitary) matrix.

Original: By the Spectral Theorem of linear algebra, this is exactly the class of normal matrices .
Revised: By the Spectral Theorem of linear algebra, this is exactly the class of normal matrices .

Original: However, this class of matrices is restrictive; in particular, it does not contain the HiPPO matrix (2).
Revised: However, this class of matrices is restrictive; in particular, it does not contain the HiPPO matrix (2).

Original: We make the observation that although the HiPPO matrix is not normal, it can be decomposed as the sum of a normal and low-rank matrix .
Revised: We make the observation that although the HiPPO matrix is not normal, it can be decomposed as the sum of a normal and low-rank matrix .

Original: However, this is still not useful by itself: unlike a diagonal matrix, powering up this sum (in (5)) is still slow and not easily optimized
Revised: However, this is still not useful by itself: unlike a diagonal matrix, powering up this sum (in ( )) is still slow and not easily optimized.

Original: We overcome this bottleneck by simultaneously applying three new techniques.
Revised: We overcome this bottleneck by simultaneously applying three new techniques.

Original: • Instead of computing K directly, we compute its spectrum by evaluating its truncated generating function (cid:80) L − 1 j =

Revised: • Instead of computing K directly, we compute its spectrum by evaluating its truncated generating function (cid:80) L − 1 j =

Original: K j ζ j at the roots of unity ζ .
Revised: K j ζ j at the roots of unity ζ .

Original: K can then be found by applying an inverse FFT.
Revised: K can then be found by applying an inverse FFT.

Original: • This generating function is closely related to the matrix resolvent, and now involves a matrix inverse instead of power .
Revised: • This generating function is closely related to the matrix resolvent, and now involves a matrix inverse instead of power .

Original: The low-rank term can now be corrected by applying the Woodbury identity which reduces ( A + pq ∗ ) − 1 in terms of A − 1 , truly reducing to the diagonal case.
Revised: The low-rank term can now be corrected by applying the Woodbury identity which reduces ( A + P Q ∗ ) − 1 in terms of A − 1 , truly reducing to the diagonal case.

Original: Finally, we show that the diagonal matrix case is equivalent to the computation of a Cauchy kernelω j − ζ k , a well-studied problem with stable near-linear algorithms (Pan, 2015; 2017).
Revised: Finally, we show that the diagonal matrix case is equivalent to the computation of a Cauchy kernelω j − ζ k , a well-studied problem with stable near-linear algorithms (Pan, 2015; 2017).

Original: Our techniques apply to any matrix that can be decomposed as Normal Plus Low-Rank (NPLR) .
Revised: Our techniques apply to any matrix that can be decomposed as Normal Plus Low-Rank (NPLR) .

Original: Theorem 1.
Revised: Theorem 1.

Original: All HiPPO matrices from (Gu et al., 2020a) have a NPLR representation
Revised: All HiPPO matrices from (Gu et al., 2020a) have a NPLR representation

Original: [Equation]
Revised: [Equation]

Original: for unitary V ∈ C N × N , diagonal Λ , and low-rank factorization p , q ∈ R N × r .
Revised: for unitary V ∈ C N × N , diagonal Λ , and low-rank factorization P , Q ∈ R N × r .

Original: These matrices
Revised: These matrices

Original: HiPPO- LegS, LegT, LagT all satisfy r = 1 or r = 2 .
Revised: HiPPO- LegS, LegT, LagT all satisfy r = 1 or r = 2 .

Original: In particular, equation (2) is NPLR with r = 1 .
Revised: In particular, equation (2) is NPLR with r = 1 .

Original: 3. S4 A LGORITHMS AND C OMPUTATIONAL C OMPLEXITY
Revised: 3. S4 A LGORITHMS AND C OMPUTATIONAL C OMPLEXITY

Original: By equation (6), note that NPLR matrices can be conjugated into diagonal plus low-rank (DPLR) form (now over C instead of R ).
Revised: By equation (6), note that NPLR matrices can be conjugated into diagonal plus low-rank (DPLR) form (now over C instead of R ).

Original: Theorems 2 and 3 describe the complexities of SSMs where A is in
Revised: Theorems 2 and 3 describe the complexities of SSMs where A is in

Original: DPLR form.
Revised: DPLR form.

Original: S4 is optimal or near-optimal for both recurrent and convolutional representations.
Revised: S4 is optimal or near-optimal for both recurrent and convolutional representations.

Original: Theorem 2 (S4 Recurrence) .
Revised: Theorem 2 (S4 Recurrence) .

Original: Given any step size ∆ , computing one step of the recurrence (3) can be done in O ( N ) operations where N is the state size.
Revised: Given any step size ∆ , computing one step of the recurrence (3) can be done in O ( N ) operations where N is the state size.

Original: Theorem 2 follows from the fact that the inverse of a DPLR matrix is also DPLR (e.g. also by the
Revised: Theorem 2 follows from the fact that the inverse of a DPLR matrix is also DPLR (e.g. also by the

Original: Woodbury identity).
Revised: Woodbury identity).

Original: This implies that the discretized matrix A is the product of two DPLR matrices and thus has O ( N ) matrix-vector multiplication.
Revised: This implies that the discretized matrix A is the product of two DPLR matrices and thus has O ( N ) matrix-vector multiplication.

Original: Appendix C.2 computes A in closed DPLR form.
Revised: Appendix C.2 computes A in closed DPLR form.

Original: Theorem 3 (S4 Convolution) .
Revised: Theorem 3 (S4 Convolution) .

Original: Given any step size ∆ , computing the SSM convolution ﬁlter K can be reduced to 4 Cauchy multiplies, requiring only (cid:101) O
Revised: Given any step size ∆ , computing the SSM convolution ﬁlter K can be reduced to 4 Cauchy multiplies, requiring only (cid:101) O

Original: ( N + L ) operations and O ( N + L ) space.
Revised: ( N + L ) operations and O ( N + L ) space.

Original: Appendix C, Deﬁnition 3 formally deﬁnes Cauchy matrices, which are related to rational interpolation problems.
Revised: Appendix C, Deﬁnition 3 formally deﬁnes Cauchy matrices, which are related to rational interpolation problems.

Original: Computing with Cauchy matrices is an extremely well-studied problem in numerical analysis, with both fast arithmetic and numerical algorithms based on the famous Fast Multipole
Revised: Computing with Cauchy matrices is an extremely well-studied problem in numerical analysis, with both fast arithmetic and numerical algorithms based on the famous Fast Multipole

Original: Method (FMM) (Pan, 2001; 2015; 2017).
Revised: Method (FMM) (Pan, 2001; 2015; 2017).

Original: The computational complexities of these algorithms under various settings are described in Appendix C, Proposition 5.
Revised: The computational complexities of these algorithms under various settings are described in Appendix C, Proposition 5.

Original: We reiterate that Theorem 3 is our core technical contribution, and its algorithm is the very motivation of the  S4 parameterization.
Revised: We reiterate that Theorem 3 is our core technical contribution, and its algorithm is the very motivation of the NPLR S4 parameterization.

Original: This algorithm is formally sketched in Algorithm 1.
Revised: This algorithm is formally sketched in Algorithm 1.

Original: A RCHITECTURE D ETAILS OF THE D EEP S4 L AYER
Revised: A RCHITECTURE D ETAILS OF THE D EEP S4 L AYER

Original: Concretely, an S4 layer is parameterized as follows.
Revised: Concretely, an S4 layer is parameterized as follows.

Original: First initialize a SSM with A set to the HiPPO matrix (2).
Revised: First initialize a SSM with A set to the HiPPO matrix (2).

Original: By Lemma 3.1 and Theorem 1, this SSM is unitarily equivalent to some ( Λ − pq ∗ , B , C ) for some diagonal Λ and vectors p , q , B , C ∈ C N × 1 .
Revised: By Lemma 3.1 and Theorem 1, this SSM is unitarily equivalent to some ( Λ − P Q ∗ , B , C ) for some diagonal Λ and vectors P , Q , B , C ∈ C N × 1 .

Original: These comprise S4’s 5 N trainable parameters.
Revised: These comprise S4’s 5 N trainable parameters.

Original: The overall deep neural network (DNN) architecture of S4 is similar to prior work.
Revised: The overall deep neural network (DNN) architecture of S4 is similar to prior work.

Original: As deﬁned above,
Revised: As deﬁned above,

Original: S4 deﬁnes a map from R L → R L , i.e. a 1-D sequence map.
Revised: S4 deﬁnes a map from R L → R L , i.e. a 1-D sequence map.

Original: Typically, DNNs operate on feature maps of size H instead of 1 .
Revised: Typically, DNNs operate on feature maps of size H instead of 1 .

Original: S4 handles multiple features by simply deﬁning H independent copies of itself, and then mixing the H features with a position-wise linear layer for a total of O ( H 2 )
Revised: S4 handles multiple features by simply deﬁning H independent copies of itself, and then mixing the H features with a position-wise linear layer for a total of O ( H 2 )

Original: + O ( HN ) parameters per layer.
Revised: + O ( HN ) parameters per layer.

Original: Nonlinear activation functions are also inserted between these layers.
Revised: Nonlinear activation functions are also inserted between these layers.

Original: Overall,
Revised: Overall,

Original: S4 deﬁnes a sequence-to-sequence map of shape (batch size, sequence length, hidden dimension), exactly the same as related sequence models such as Transformers, RNNs, and CNNs.
Revised: S4 deﬁnes a sequence-to-sequence map of shape (batch size, sequence length, hidden dimension), exactly the same as related sequence models such as Transformers, RNNs, and CNNs.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Section 4.1 benchmarks S4 against the LSSL and efﬁcient Transformer models.
Revised: Section 4.1 benchmarks S4 against the LSSL and efﬁcient Transformer models.

Original: Section 4.2 validates
Revised: Section 4.2 validates

Original: S4 on LRDs: the LRA benchmark and raw speech classiﬁcation.
Revised: S4 on LRDs: the LRA benchmark and raw speech classiﬁcation.

Original: Section 4.3 investigates whether S can be used as a general sequence model to perform effectively and efﬁciently in a wide variety of settings including image classiﬁcation, image and text generation, and time series forecasting
Revised: Section 4.3 investigates whether S can be used as a general sequence model to perform effectively and efﬁciently in a wide variety of settings including image classiﬁcation, image and text generation, and time series forecasting.

Original: . 4. S4 E FFICIENCY B ENCHMARKS
Revised: 4. S4 E FFICIENCY B ENCHMARKS

Original: We benchmark that S4 can be trained quickly and efﬁciently, both compared to the LSSL, as well as efﬁcient Transformer variants designed for long-range sequence modeling.
Revised: We benchmark that S4 can be trained quickly and efﬁciently, both compared to the LSSL, as well as efﬁcient Transformer variants designed for long-range sequence modeling.

Original: As outlined in Section 3,
Revised: As outlined in Section 3,

Original: S4 is theoretically much more efﬁcient than the LSSL, and Table 1 conﬁrms that the S4 is orders of magnitude more speed- and memory-efﬁcient for practical layer sizes.
Revised: S4 is theoretically much more efﬁcient than the LSSL, and Table 1 conﬁrms that the S4 is orders of magnitude more speed- and memory-efﬁcient for practical layer sizes.

Original: In fact, S4’s speed and memory use is competitive with the most efﬁcient Transformer variants benchmarked by Tay et al.
Revised: In fact, S4’s speed and memory use is competitive with the most efﬁcient Transformer variants benchmarked by Tay et al.

Original: Linear Transformer (Katharopoulos et al., 2020) and Performer (Choromanski et al., 2020)—in a parameter-matched setting (Table 2, following the protocol of Tay et al. (2021)).
Revised: Linear Transformer (Katharopoulos et al., 2020) and Performer (Choromanski et al., 2020)—in a parameter-matched setting (Table 2, following the protocol of Tay et al. (2021)).

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: 4. L EARNING L ONG R ANGE D EPENDENCIES
Revised: 4. L EARNING L ONG R ANGE D EPENDENCIES

Original: As described in Sections 2.2 and 3.
Revised: As described in Sections 2.2 and 3.

Original: [Table]
Revised: [Table]

Original: S4 uses a principled approach to address LRDs based on the HiPPO theory of continuous-time memorization.
Revised: S4 uses a principled approach to address LRDs based on the HiPPO theory of continuous-time memorization.

Original: Our goal in this section is to validate that S achieves high performance on difﬁcult tasks that require long-range reasoning.
Revised: Our goal in this section is to validate that S achieves high performance on difﬁcult tasks that require long-range reasoning.

Original: We focus here on two problems: (i) the Long-Range Arena, a well-known benchmark designed to test efﬁcient sequence models on LRDs
Revised: We focus here on two problems: (i) the Long-Range Arena, a well-known benchmark designed to test efﬁcient sequence models on LRDs

Original: [Table]
Revised: [Table]

Original: and (ii) a speech classiﬁcation problem as a real-world test of LRDs.
Revised: and (ii) a speech classiﬁcation problem as a real-world test of LRDs.

Original: Long Range Arena (LRA).
Revised: Long Range Arena (LRA).

Original: LRA (Tay et al., 2021) contains 6 tasks with lengths 1K-16K steps, encompassing modalities and objectives that require similarity, structural, and visuospatial reasoning.
Revised: LRA (Tay et al., 2021) contains 6 tasks with lengths 1K-16K steps, encompassing modalities and objectives that require similarity, structural, and visuospatial reasoning.

Original: Table 3 compares S4 against the 11 Transformer variants from Tay et al.
Revised: Table 3 compares S4 against the 11 Transformer variants from Tay et al.

Original: (2021) as well as follow-up work.
Revised: (2021) as well as follow-up work.

Original: S4 substantially advances the SoTA, outperforming all baselines on all tasks and averaging 80 .
Revised: S4 substantially advances the SoTA, outperforming all baselines on all tasks and averaging 80 .

Original: 88% compared to less than 60% for every baseline.
Revised: 48% compared to less than 60% for every baseline.

Original: Notably, S4 solves the Path-X task, an extremely challenging task that involves reasoning about LRDswith sparse signal over sequences of length 128 × 128 = 16384 .
Revised: Notably, S4 solves the Path-X task, an extremely challenging task that involves reasoning about LRDs over sequences of length 128 × 128 = 16384 .

Original: All previous models have failed (i.e. random guessing) due to memory or computation bottlenecks, or simply being unable to learn such long dependencies.
Revised: All previous models have failed (i [Table] e. random guessing) due to memory or computation bottlenecks, or simply being unable to learn such long dependencies.

Original: We analyze S4’s performance on Path-X by visualizing its learned representations, in particular 1-D convolution kernels K which are the focus of our technical results in Section 3.
Revised: We analyze S4’s performance on Path-X by visualizing its learned representations, in particular 1-D convolution kernels K which are the focus of our technical results in Section 3.

Original: Fig.
Revised: Fig.

Original: 2 shows that S  learns a variety of ﬁlters that display spatially consistent structure and demonstrate awareness of the 2-D nature of the data.
Revised: 2 shows that S  learns a variety of ﬁlters that display spatially consistent structure and demonstrate awareness of the 2-D nature of the data.

Original: In particular, the lower layers learn simple kernels that extract features from just a few rows of local context while ignoring the rest of the image.
Revised: In particular, the lower layers learn simple kernels that extract features from just a few rows of local context while ignoring the rest of the image.

Original: On the other hand, higher layers aggregate information globally across full columns of the image at varying spatial frequencies.
Revised: On the other hand, higher layers aggregate information globally across full columns of the image at varying spatial frequencies.

Original: Filters in these higher layers span the entire context ( 16384 pixels), conﬁrming S4’s ability to learn LRDs.
Revised: Filters in these higher layers span the entire context ( 16384 pixels), conﬁrming S4’s ability to learn LRDs.

Original: Raw Speech Classiﬁcation.
Revised: Raw Speech Classiﬁcation.

Original: Speech is a typical real-world time series domain, involving signals sampled from an underlying physical process at high frequency.
Revised: Speech is a typical real-world time series domain, involving signals sampled from an underlying physical process at high frequency.

Original: We perform speech classiﬁcation using the Speech Commands dataset (Warden, 2018).
Revised: We perform speech classiﬁcation using the Speech Commands dataset (Warden, 2018).

Original: While most sequence models for speech rely on extensive preprocessing (e.g. to MFCC features), we classify raw speech (length- 16000 ) following Romero et al.
Revised: While most sequence models for speech rely on extensive preprocessing (e.g. to MFCC features), we classify raw speech (length- 16000 ) following Romero et al.

Original: S4 achieves 98 .
Revised: S4 achieves 98 .

Original: 3% accuracy, higher than all baselines that use the 100 × shorter MFCC features, and validates that a powerful LRD model is able to extract more information from the raw data and outperform hand-crafted pre-processing.
Revised: 3% accuracy, higher than all baselines that use the 100 × shorter MFCC features, and validates that a powerful LRD model is able to extract more information from the raw data and outperform hand-crafted pre-processing.

Original: Additionally, we include a baseline CNN speciﬁcally designed for raw speech, the discriminator from the WaveGAN model (Donahue et al., 2019), which performs worse than S4 while having 90 × more parameters and incorporating many more architectural heuristics (Appendix D.2).
Revised: Additionally, we include a baseline CNN speciﬁcally designed for raw speech, the discriminator from the WaveGAN model (Donahue et al., 2019), which performs worse than S4 while having 90 × more parameters and incorporating many more architectural heuristics (Appendix D.2).

Original: 4. S4 AS A G ENERAL S EQUENCE M ODEL
Revised: 4. S4 AS A G ENERAL S EQUENCE M ODEL

Original: A key goal of sequence modeling research is to develop a single model that can be applied in many domains (e.g. images, audio, text, time-series) with a broad range of capabilities (e.g. efﬁcient training, fast generation, handling irregularly sampled data).
Revised: A key goal of sequence modeling research is to develop a single model that can be applied in many domains (e.g. images, audio, text, time-series) with a broad range of capabilities (e.g. efﬁcient training, fast generation, handling irregularly sampled data).

Original: As a fundamental scientiﬁc model,
Revised: As a fundamental scientiﬁc model,

Original: SSMs are a promising candidate that come with a range of capabilities, and S4’s strong results on LRD benchmarks spanning images, text, and speech are evidence of S4’s potential as a general sequence model.
Revised: SSMs are a promising candidate that come with a range of capabilities, and S4’s strong results on LRD benchmarks spanning images, text, and speech are evidence of S4’s potential as a general sequence model.

Original: In this section, we focus on understanding this question in more depth by highlighting key strengths of S4 in settings that usually require specialized models.
Revised: In this section, we focus on understanding this question in more depth by highlighting key strengths of S4 in settings that usually require specialized models.

Original: The tasks we focus on (generative modeling, image classiﬁcation, time-series forecasting) are considered as LRD tasks in the literature, and serve as additional validation that S4 handles LRDs efﬁciently.
Revised: The tasks we focus on (generative modeling, image classiﬁcation, time-series forecasting) are considered as LRD tasks in the literature, and serve as additional validation that S4 handles LRDs efﬁciently.

Original: Large-scale generative modeling.
Revised: Large-scale generative modeling.

Original: We investigate two well-studied image and text benchmarks to validate the scalability, ﬂexibility
Revised: We investigate two well-studied image and text benchmarks to validate the scalability, ﬂexibility

Original: [Table]
Revised: [Table]

Original: and efﬁciency of S4.
Revised: and efﬁciency of S4.

Original: These tasks require much larger models than our previous tasks – up to 250 M parameters.
Revised: These tasks require much larger models than our previous tasks – up to 250 M parameters.

Original: First
Revised: First

Original: [Table]
Revised: [Table]

Original: CIFAR density estimation is a popular benchmark for autoregressive models, where images are ﬂattened into a sequence of 3072 RGB subpixels that are predicted one by one.
Revised: CIFAR density estimation is a popular benchmark for autoregressive models, where images are ﬂattened into a sequence of 3072 RGB subpixels that are predicted one by one.

Original: Table 6 shows that with no 2D inductive bias , S4 is competitive with the best models designed for this task.
Revised: Table 6 shows that with no 2D inductive bias , S4 is competitive with the best models designed for this task.

Original: Second, WikiText-103 is an established benchmark for language modeling, an important task for large-scale sequence models where tokens are predicted sequentially based on past context.
Revised: Second, WikiText-103 is an established benchmark for language modeling, an important task for large-scale sequence models where tokens are predicted sequentially based on past context.

Original: Although
Revised: Although

Original: RNNs were the model of choice for many years, Transformers are now the dominant model in such applications that contain data that is inherently discrete.
Revised: RNNs were the model of choice for many years, Transformers are now the dominant model in such applications that contain data that is inherently discrete.

Original: We show that alternative models
Revised: We show that alternative models

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table] 
Revised: [Table] [Table] [Table]

Original: [Table] 
Revised: [Table] [Table] [Table] [Table]

Original: [Table]
Revised: [Table]

Original: to Transformers can still be competitive in these settings.
Revised: to Transformers can still be competitive in these settings.

Original: By simply taking a strong Transformer baseline (Baevski & Auli, 2018) and replacing the self-attention layers, S4 substantially closes the gap to Transformers (within 0 .
Revised: By simply taking a strong Transformer baseline (Baevski & Auli, 2018) and replacing the self-attention layers, S4 substantially closes the gap to Transformers (within 0 .

Original: 8 ppl), setting SoTA for attention-free models by over 2 ppl.
Revised: 8 ppl), setting SoTA for attention-free models by over 2 ppl.

Original: Fast autoregressive inference.
Revised: Fast autoregressive inference.

Original: A prominent limitation of autoregressive models is inference speed (e.g. generation), since they require a pass over the full context for every new sample.
Revised: A prominent limitation of autoregressive models is inference speed (e.g. generation), since they require a pass over the full context for every new sample.

Original: Several methods have been speciﬁcally crafted to overcome this limitation such as the Linear Transformer, a hybrid Transformer/RNN that switches to a stateful, recurrent view at inference time for speed.
Revised: Several methods have been speciﬁcally crafted to overcome this limitation such as the Linear Transformer, a hybrid Transformer/RNN that switches to a stateful, recurrent view at inference time for speed.

Original: As a stateful model, SSMs automatically have this ability (Fig. 1).
Revised: As a stateful model, SSMs automatically have this ability (Fig. 1).

Original: By switching to its recurrent representation (Section 2.3), S4 requires constant memory and computation per time step – in contrast to standard autoregressive models which scale in the context length.
Revised: By switching to its recurrent representation (Section 2.3), S4 requires constant memory and computation per time step – in contrast to standard autoregressive models which scale in the context length.

Original: On both CIFAR-10 and WikiText103, we report the throughput of various models at generation time, with S4 around 60 × faster than a vanilla Transformer on both tasks (details in Appendix D.3.3).
Revised: On both CIFAR-10 and WikiText103, we report the throughput of various models at generation time, with S4 around 60 × faster than a vanilla Transformer on both tasks (details in Appendix D.3.3).

Original: Sampling resolution change.
Revised: Sampling resolution change.

Original: As a continuous-time model, S4 automatically adapts to data sampled at different rates, a challenging setting for time series with a dedicated line of work (Rubanova et al., 2019; De Brouwer et al., 2019; Romero et al., 2021).
Revised: As a continuous-time model, S4 automatically adapts to data sampled at different rates, a challenging setting for time series with a dedicated line of work (Rubanova et al., 2019; De Brouwer et al., 2019; Romero et al., 2021).

Original: Without re-training, S4 achieves 96 .
Revised: Without re-training, S4 achieves 96 .

Original: 3% accuracy at 0 .
Revised: 3% accuracy at 0 .

Original: 5 × the frequency on Speech Commands (Table 4), simply by changing its internal step size ∆ (Section 2.3).
Revised: 5 × the frequency on Speech Commands (Table 4), simply by changing its internal step size ∆ (Section 2.3).

Original: Learning with weaker inductive bias.
Revised: Learning with weaker inductive bias.

Original: Beyond our results on speech (Section 4.2), we further validate that S4 can be applied with minimal modiﬁcations on two domains that typically require specialized domain-speciﬁc preprocessing and architectures.
Revised: Beyond our results on speech (Section 4.2), we further validate that S4 can be applied with minimal modiﬁcations on two domains that typically require specialized domain-speciﬁc preprocessing and architectures.

Original: First, we compare S4 to the Informer
Revised: First, we compare S4 to the Informer

Original: (Zhou et al., 2021), a new Transformer architecture that uses a complex encoder-decoder designed for time-series forecasting problems.
Revised: (Zhou et al., 2021), a new Transformer architecture that uses a complex encoder-decoder designed for time-series forecasting problems.

Original: A simple application of S4 that treats forecasting as a generic sequence-to-sequence transformation (Appendix D.3.5) outperforms the Informer and other baselines on 40 / 50 settings across 5 forecasting tasks.
Revised: A simple application of S4 that treats forecasting as a masked sequence-to-sequence transformation (Fig. 3) outperforms the Informer and other baselines on 40 /settings across 5 forecasting tasks.

Original: Notably, S4 is better on the longest setting in each task, e.g. reducing MSE by 32% when forecasting 30 days of weather data (Appendix D.3.5).
Revised: Notably, S4 is better on the longest setting in each task, e.g. reducing MSE by 37% when forecasting 30 days of weather data (Appendix D.3.5).

Original: Finally, we evaluate S4 on pixel-level sequential image classiﬁcation tasks (Table 5), popular benchmarks which were originally LRD tests for RNNs (Arjovsky et al., 2016).
Revised: Finally, we evaluate S4 on pixel-level sequential image classiﬁcation tasks (Table 5), popular benchmarks which were originally LRD tests for RNNs (Arjovsky et al., 2016).

Original: Beyond LRDs, these benchmarks point to a recent effort of the ML community to solve vision problems with reduced domain knowledge, in the spirit of models such as Vision Transformers (Dosovitskiy et al., 2020) and
Revised: Beyond LRDs, these benchmarks point to a recent effort of the ML community to solve vision problems with reduced domain knowledge, in the spirit of models such as Vision Transformers (Dosovitskiy et al., 2020) and

Original: MLP-Mixer (Tolstikhin et al., 2021).
Revised: MLP-Mixer (Tolstikhin et al., 2021) .

Original: Sequential CIFAR is a particularly challenging dataset where outside of SSMs, all sequence models have a gap of over 25% to a simple 2-D CNN.
Revised: Sequential CIFAR is a particularly challenging dataset where outside of SSMs, all sequence models have a gap of over 25% to a simple 2-D CNN.

Original: By contrast, S is competitive with a larger ResNet18 (7.9M vs. 11.0M parameters), both with ( 93.16% vs. 95.62%) or without ( 91.12% vs. 89.46%) data augmentation.
Revised: By contrast, S is competitive with a larger ResNet18 (7.9M vs. 11.0M parameters), both with ( 93.16% vs. 95.62%) or without ( 91.12% vs. 89.46%) data augmentation.

Original: Moreover, it is much more robust to other architectural choices (e.g. 90.46% vs. 79.52% when swapping BatchNorm for LayerNorm).
Revised: Moreover, it is much more robust to other architectural choices (e.g. 90.46% vs. 79.52% when swapping BatchNorm for LayerNorm).

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We introduce S4, a sequence model that uses a new parameterization for the state space model’s continuous-time, recurrent, and convolutional views to efﬁciently model LRDs in a principled manner.
Revised: We introduce S4, a sequence model that uses a new parameterization for the state space model’s continuous-time, recurrent, and convolutional views to efﬁciently model LRDs in a principled manner.

Original: Results across established benchmarks evaluating a diverse range of data modalities and model capabilities suggest that S4 has the potential to be an effective general sequence modeling solution.
Revised: Results across established benchmarks evaluating a diverse range of data modalities and model capabilities suggest that S4 has the potential to be an effective general sequence modeling solution.

Original: A D ISCUSSION
Revised: A D ISCUSSION

Original: Related Work.
Revised: Related Work.

Original: Our work is most closely related to a line of work originally motivated by a particular biologically-inspired SSM, which led to mathematical models for addressing LRDs.
Revised: Our work is most closely related to a line of work originally motivated by a particular biologically-inspired SSM, which led to mathematical models for addressing LRDs.

Original: Voelker (2019); Voelker et al.
Revised: Voelker (2019); Voelker et al.

Original: (2019) derived a non-trainable SSM motivated from approximating a neuromorphic spiking model, and Chilkuri & Eliasmith (2021) showed that it could be sped up at train time with a convolutional view.
Revised: (2019) derived a non-trainable SSM motivated from approximating a neuromorphic spiking model, and Chilkuri & Eliasmith (2021) showed that it could be sped up at train time with a convolutional view.

Original: Gu et al.
Revised: Gu et al.

Original: (2020a) extended this special case to a general continuous-time function approximation framework with several more special cases of A matrices designed for long-range dependencies.
Revised: (2020a) extended this special case to a general continuous-time function approximation framework with several more special cases of A matrices designed for long-range dependencies.

Original: However, instead of using a true SSM, all of these works ﬁxed a choice of A and built RNNs around it.
Revised: However, instead of using a true SSM, all of these works ﬁxed a choice of A and built RNNs around it.

Original: Most recently, Gu et al.
Revised: Most recently, Gu et al.

Original: (2021) used the full (1) explicitly as a deep SSM model, exploring new conceptual views of SSMs, as well as allowing A to be trained.
Revised: (2021) used the full (1) explicitly as a deep SSM model, exploring new conceptual views of SSMs, as well as allowing A to be trained.

Original: As mentioned in Section 1, their method used a naive instantiation of SSMs that suffered from an additional factor of N in memory and N 2 in computation.
Revised: As mentioned in Section 1, their method used a naive instantiation of SSMs that suffered from an additional factor of N in memory and N 2 in computation.

Original: Beyond this work, our technical contributions (Section 3) on the S4 parameterization and algorithms are applicable to a broader family of SSMs including these investigated in prior works, and our techniques for working with these models may be of independent interest.
Revised: Beyond this work, our technical contributions (Section 3) on the S4 parameterization and algorithms are applicable to a broader family of SSMs including these investigated in prior works, and our techniques for working with these models may be of independent interest.

Original: Implementation.
Revised: Implementation.

Original: The computational core of S4’s training algorithm is the Cauchy kernel discussed in Sections 3.2 and 3.3 and Appendix C.3.
Revised: The computational core of S4’s training algorithm is the Cauchy kernel discussed in Sections 3.2 and 3.3 and Appendix C.3.

Original: As described in Appendix C.3 Proposition 5, there are many algorithms for it with differing computational complexities and sophistication.
Revised: As described in Appendix C.3 Proposition 5, there are many algorithms for it with differing computational complexities and sophistication.

Original: Our current implementation of S4 actually uses the naive O ( NL ) algorithm which is easily parallelized on GPUs and has more easily accessible libraries allowing it to be implemented; we leverage the pykeops library for memory-efﬁcient kernel operations.
Revised: Our current implementation of S4 actually uses the naive O ( NL ) algorithm which is easily parallelized on GPUs and has more easily accessible libraries allowing it to be implemented; we leverage the pykeops library for memory-efﬁcient kernel operations.

Original: However, this library is a much more general library that may not be optimized for the Cauchy kernels used here, and we believe that a dedicated CUDA implementation can be more efﬁcient.
Revised: However, this library is a much more general library that may not be optimized for the Cauchy kernels used here, and we believe that a dedicated CUDA implementation can be more efﬁcient.

Original: Additionally, as discussed in this work, there are asymptotically faster and numerically stable algorithms for the Cauchy kernel (Proposition 5).
Revised: Additionally, as discussed in this work, there are asymptotically faster and numerically stable algorithms for the Cauchy kernel (Proposition 5).

Original: However, these algorithms are currently not implemented for GPUs due to a lack of previous applications that require them.
Revised: However, these algorithms are currently not implemented for GPUs due to a lack of previous applications that require them.

Original: We believe that more efﬁcientimplementations of theseself-contained computational kernels are possible, and that S4 (and SSMs at large) may have signiﬁcant room for further improvements in efﬁciency.
Revised: We believe that more efﬁcient these kernels are possible, and that S4 (and SSMs at large) may have signiﬁcant room for further improvements in efﬁciency.

Original: Limitations and Future Directions.
Revised: Limitations and Future Directions.

Original: In this work, we show that S4 can address a wide variety of data effectively.
Revised: In this work, we show that S4 can address a wide variety of data effectively.

Original: However, it may not necessarily be the most suitable model for all types of data.
Revised: However, it may not necessarily be the most suitable model for all types of data.

Original: For example, Table 7 still found a gap compared to Transformers for language modeling.
Revised: For example, Table 7 still found a gap compared to Transformers for language modeling.

Original: An interesting future direction is exploring combinations of S4 with other sequence models to complement their strengths.
Revised: An interesting future direction is exploring combinations of S4 with other sequence models to complement their strengths.

Original: We are excited about other directions, including continuing to explore the beneﬁts of
Revised: We are excited about other directions, including continuing to explore the beneﬁts of

Original: S4 on audio data (e.g. pre-training or generation settings), and generalizing HiPPO and S4 to higher-dimensional data for image and video applications.
Revised: S4 on audio data (e.g. pre-training or generation settings), and generalizing HiPPO and S4 to higher-dimensional data for image and video applications.

Original: B N UMERICAL I NSTABILITY OF LSSL
Revised: B N UMERICAL I NSTABILITY OF LSSL

Original: This section proves the claims made in Section 3.1 about prior work.
Revised: This section proves the claims made in Section 3.1 about prior work.

Original: We ﬁrst derive the explicit diagonalization of the HiPPO matrix, conﬁrming its instability because of exponentially large entries.
Revised: We ﬁrst derive the explicit diagonalization of the HiPPO matrix, conﬁrming its instability because of exponentially large entries.

Original: We then discuss the proposed theoretically fast algorithm from (Gu et al., 2021)
Revised: We then discuss the proposed theoretically fast algorithm from (Gu et al., 2021)

Original: (Theorem 2) and show that it also involves exponentially large terms and thus cannot be implemented.
Revised: (Theorem 2) and show that it also involves exponentially large terms and thus cannot be implemented.

Original: B. H I PPO D IAGONALIZATION
Revised: B. H I PPO D IAGONALIZATION

Original: Proof of Lemma 3.2.
Revised: Proof of Lemma 3.2.

Original: The HiPPO matrix (2) is equal, up to sign and conjugation by a diagonal matrix, to
Revised: The HiPPO matrix (2) is equal, up to sign and conjugation by a diagonal matrix, to

Original: A =   −  − − −  − − − − −  − − − − − − − 13...
Revised: A =   −  − − −  − − − − −  − − − − − − − 13...

Original:  
Revised:  

Original: [Equation]
Revised: [Equation]

Original: Our goal is to show that this A is diagonalized by the matrix
Revised: Our goal is to show that this A is diagonalized by the matrix

Revised: V

Original: [Equation]
Revised: [Equation]

Original:  [Equation] [Equation] [Equation]
Revised: =             1 10  1 15 35 28 9...

Original: 
Revised: 

Original: [Equation] or in other words that columns of this matrix are eigenvectors of A .
Revised:    , or in other words that columns of this matrix are eigenvectors of A .

Original: Concretely, we will show that the j -th column of this matrix v ( j ) with elements
Revised: Concretely, we will show that the j -th column of this matrix v ( j ) with elements

Original: [Equation]
Revised: [Equation]

Original: is an eigenvector with eigenvalue j + 1 .
Revised: is an eigenvector with eigenvalue j + 1 .

Original: In other words we must show that for all indices k ∈ [ N ] ,
Revised: In other words we must show that for all indices k ∈ [ N ] ,

Original: [Equation]
Revised: [Equation]

Original: If k < j , then for all i inside the sum, either k < i or i < j .
Revised: If k < j , then for all i inside the sum, either k < i or i < j .

Original: In the ﬁrst case A ki = 0 and in the second case v ( j ) i = 0 , so both sides of equation (7) are equal to 0 .
Revised: In the ﬁrst case A ki = 0 and in the second case v ( j ) i = 0 , so both sides of equation (7) are equal to 0 .

Original: It remains to show the case k ≥ j , which proceeds by induction on k .
Revised: It remains to show the case k ≥ j , which proceeds by induction on k .

Original: Expanding equation (7) using the formula for A yields
Revised: Expanding equation (7) using the formula for A yields

Original: [Equation]
Revised: [Equation]

Original: In the base case k = j , the sum disappears and we are left with ( Av ( j ) )
Revised: In the base case k = j , the sum disappears and we are left with ( Av ( j ) )

Original: j = ( j +1) (cid:0) 2 j 2 j (cid:1) =
Revised: j = ( j + [Equation] (cid:0) 2 j 2 j (cid:1) =

Original: ( j +1) v ( j ) j , as desired.
Revised: ( j +1) v ( j ) j , as desired.

Original: Otherwise, the sum for ( Av ) ( j ) k is the same as the sum for ( Av ) ( j ) k − 1 but with sign reversed and a few edge terms.
Revised: Otherwise, the sum for ( Av ) ( j ) k is the same as the sum for ( Av ) ( j ) k − 1 but with sign reversed and a few edge terms.

Original: The result follows from applying the inductive hypothesis and algebraic simpliﬁcation:
Revised: The result follows from applying the inductive hypothesis and algebraic simpliﬁcation:

Original: [Equation]
Revised: [Equation]

Original: (cid:18) k − 1 + j 2 j (cid:19)

Original: +

Original: ( k + 1) (cid:18) k
Revised: (cid:18) k

Original: + j 2 j (cid:19) = − ( j + k ) (cid:18) k − 1 + j 2 j (cid:19) + ( k + 1) (cid:18) k
Revised: [Equation]

Original: + j 2 j (cid:19) =
Revised: j 2 j (cid:19) =

Original: − ( j + k ) ( k − 1 + j )! ( k − 1 − j )! (2 j )!
Revised: [Equation]

Original: + ( k + 1) (cid:18) k
Revised: ( k [Equation] ( k

Original: + j 2 j (cid:19) = − ( k + j )! ( k − 1 − j )! (2 j )!
Revised: [Equation]

Original: + ( k + 1) (cid:18) k + j 2 j (cid:19) = − ( k − j ) ( k + j )! ( k − j )!
Revised: [Equation]

Original: (2 j )!
Revised: [Equation]

Original: + ( k + 1) (cid:18) k
Revised: (cid:18) [Equation]

Original: + j 2 j (cid:19) =
Revised: j 2 j (cid:19) [Equation]

Original: ( j − k )( k + 1) (cid:18) k
Revised: [Equation] (cid:18) [Equation]

Original: + j 2 j (cid:19) 
Revised: (cid:19) [Equation]

Original: + ( k + 1) (cid:18)k
Revised: [Equation] (cid:18)

Original: + j 2 j (cid:19)=
Revised: [Equation] (cid:19)

Original: ( j + 1) v ( j ) k .
Revised: [Equation] (cid:18) [Equation]

Original: B. F AST BUT U NSTABLE LSSL A LGORITHM
Revised: B. F AST BUT U NSTABLE LSSL A LGORITHM

Original: Instead of diagonalization, Gu et al.
Revised: Instead of diagonalization, Gu et al.

Original: (2021, Theorem 2) proposed a sophisticated fast algorithm to compute
Revised: (2021, Theorem 2) proposed a sophisticated fast algorithm to compute

Original: K L ( A , B , C ) = ( CB , CAB , . . . , CA L − 1 B ) .
Revised: [Equation]

Original: This algorithm runs in O ( N log 2 N + L log L ) operations
Revised: This algorithm runs in O ( N log 2 N + L log L ) operations

Original: and
Revised: and

Original: O ( N + L ) space.
Revised: O ( N + L ) space.

Original: However, we now show that this algorithm is also numerically unstable.
Revised: However, we now show that this algorithm is also numerically unstable.

Original: There are several reasons for the instability of this algorithm, but most directly we can pinpoint a particular intermediate quantity that they use.
Revised: There are several reasons for the instability of this algorithm, but most directly we can pinpoint a particular intermediate quantity that they use.

Original: Deﬁnition 1.
Revised: Deﬁnition 1.

Original: The fast LSSL algorithm computes coefﬁcients of p ( x ) , the characteristic polynomial of A , as an intermediate computation.
Revised: The fast LSSL algorithm computes coefﬁcients of p ( x ) , the characteristic polynomial of A , as an intermediate computation.

Original: Additionally, it computes the coefﬁcients of its inverse, p ( x ) −(mod x L ) .
Revised: Additionally, it computes the coefﬁcients of its inverse, p ( x ) −(mod x L ) .

Original: We now claim that this quantity is numerically unfeasible.
Revised: We now claim that this quantity is numerically unfeasible.

Original: We narrow down to the case
Revised: We narrow down to the case

Original: when A = I is the identity matrix.
Revised: when A = I is the identity matrix.

Original: Note that this case is actually in some sense the most typical case: when discretizing the continuous-time SSM to discrete-time by a step-size ∆ , the discretized transition matrix A is brought closer to the identity.
Revised: Note that this case is actually in some sense the most typical case: when discretizing the continuous-time SSM to discrete-time by a step-size ∆ , the discretized transition matrix A is brought closer to the identity.

Original: For example, with the Euler discretization A
Revised: For example, with the Euler discretization A

Original: = I + ∆ A , we have A → I as the step size ∆ → 0 .
Revised: = I + ∆ A , we have A → I as the step size ∆ → 0 .

Original: Lemma B.1.
Revised: Lemma B.1.

Original: When A =
Revised: When A =

Original: I , the fast LSSL algorithm requires computing terms exponentially large in
Revised: I , the fast LSSL algorithm requires computing terms exponentially large in

Original: N .
Revised: N .

Original: Proof.
Revised: Proof.

Original: The characteristic polynomial of I is
Revised: The characteristic polynomial of I is

Original: [Equation]
Revised: [Equation]

Original: These coefﬁcients have size up to (cid:0) 
Revised: These coefﬁcients have size up to (cid:0) N

Original: [Equation]
Revised: [Equation]

Revised: 2 (cid:1) ≈

Revised: [Equation]

Original: The inverse of p ( x ) has even larger coefﬁcients.
Revised: The inverse of p ( x ) has even larger coefﬁcients.

Original: It can be calculated in closed form by the generalized binomial formula:
Revised: It can be calculated in closed form by the generalized binomial formula:

Original: [Equation]
Revised: [Equation]

Original: Taking this (mod x L ) , the largest coefﬁcient is
Revised: Taking this (mod x L ) , the largest coefﬁcient is

Original: [Equation]
Revised: [Equation]

Original: When L = N − 1 this is
Revised: When L = N − 1 this is

Original: [Equation]
Revised: [Equation]

Original: already larger than the coefﬁcients of (1 − x ) N , and only increases as L grows.
Revised: already larger than the coefﬁcients of (1 − x ) N , and only increases as L grows.

Original: C S4 A LGORITHM D ETAILS
Revised: C S4 A LGORITHM D ETAILS

Original: This section proves the results of Section 3.3, providing complete details of our efﬁcient algorithms for S4.
Revised: This section proves the results of Section 3.3, providing complete details of our efﬁcient algorithms for S4.

Original: Appendices C.1 to C.3 prove Theorems 1 to 3 respectively.
Revised: Appendices C.1 to C.3 prove Theorems 1 to 3 respectively.

Original: C. NPLR R EPRESENTATIONS OF H I PPO M ATRICES
Revised: C. NPLR R EPRESENTATIONS OF H I PPO M ATRICES

Original: We ﬁrst prove Theorem 1, showing that all HiPPO matrices for continuous-time memory fall under the S4 normal plus low-rank (NPLR) representation.
Revised: We ﬁrst prove Theorem 1, showing that all HiPPO matrices for continuous-time memory fall under the S4 normal plus low-rank (NPLR) representation.

Original: Proof of Theorem 1.
Revised: Proof of Theorem 1.

Original: We consider each of the three cases HiPPO-LagT, HiPPO-LegT, and HiPPOLegS separately.
Revised: We consider each of the three cases HiPPO-LagT, HiPPO-LegT, and HiPPOLegS separately.

Original: Note that the primary HiPPO matrix deﬁned in this work (equation (2)) is the
Revised: Note that the primary HiPPO matrix deﬁned in this work (equation (2)) is the

Original: HiPPO-LegT matrix.
Revised: HiPPO-LegT matrix.

Original: HiPPO-LagT.
Revised: HiPPO-LagT.

Original: The HiPPO-LagT matrix is simply
Revised: The HiPPO-LagT matrix is simply

Original: [Equation]
Revised: [Equation]

Original:  [Equation]  [Equation]
Revised:    .

Original: Adding the matrix of all 12 , which is rank 1, yields −   −−− −−  −   
Revised: Adding the matrix of all 12 , which is rank 1, yields −   −−− −−  −   

Original:  .
Revised:  .

Original: This matrix is now skew-symmetric.
Revised: This matrix is now skew-symmetric.

Original: Skew-symmetric matrices are a particular case of normal matrices with pure-imaginary eigenvalues.
Revised: Skew-symmetric matrices are a particular case of normal matrices with pure-imaginary eigenvalues.

Original: Gu et al.
Revised: Gu et al.

Original: (2020a) also consider a case of HiPPO corresponding to the generalized Laguerre polynomials that generalizes the above HiPPO-LagT case.
Revised: (2020a) also consider a case of HiPPO corresponding to the generalized Laguerre polynomials that generalizes the above HiPPO-LagT case.

Original: In this case, the matrix A (up to conjugation by a diagonal matrix) ends up being close to the above matrix, but with a different element on the diagonal.
Revised: In this case, the matrix A (up to conjugation by a diagonal matrix) ends up being close to the above matrix, but with a different element on the diagonal.

Original: After adding the rank-1 correction, it becomes the above skew-symmetric matrix plus a multiple of the identity.
Revised: After adding the rank-1 correction, it becomes the above skew-symmetric matrix plus a multiple of the identity.

Original: Thus after diagonalization by the same matrix as in the LagT case, it is still reduced to diagonal plus low-rank (DPLR) form, where the diagonal is now pure imaginary plus a real constant.
Revised: Thus after diagonalization by the same matrix as in the LagT case, it is still reduced to diagonal plus low-rank (DPLR) form, where the diagonal is now pure imaginary plus a real constant.

Original: HiPPO-LegS.
Revised: HiPPO-LegS.

Original: We restate the formula from equation (2) for convenience.
Revised: We restate the formula from equation (2) for convenience.

Original: [Equation]
Revised: [Equation]

Original: if n > k n +if n = k if n < k .

Original: Adding 12 (2 n + 1) 1 / 2 (2 k + 1) 1 / 2 to the whole matrix gives− [Equation] if [Equation]
Revised: Adding 12 (2 n + 1) 1 / 2 (2 k + 1) 1 / 2 to the whole matrix gives

Original: if

Original: [Equation]
Revised: [Equation]

Original: k

Original: [Equation] if [Equation] Note that this matrix is not skew-symmetric, but is 12 I + S where S is a skew-symmetric matrix.
Revised: Note that this matrix is not skew-symmetric, but is 12 I + S where S is a skew-symmetric matrix.

Original: This is diagonalizable by the same unitary matrix that diagonalizes S .
Revised: This is diagonalizable by the same unitary matrix that diagonalizes S .

Original: HiPPO-LegT.
Revised: HiPPO-LegT.

Original: Up to the diagonal scaling, the LegT matrix is
Revised: Up to the diagonal scaling, the LegT matrix is

Original: A = −   1 − − 1 . . .
Revised: A = −   1 − − 1 . . .

Original: −    −    
Revised: −    −    ..

Original: [Equation]
Revised: [Equation]

Revised:   .

Original: By adding − 1 to this matrix and then the matrix [Equation]
Revised: By adding − 1 to this matrix and then the matrix     

Original:  the matrix becomes [Equation]
Revised:  the matrix becomes   −− −  

Original:  which is skew-symmetric.
Revised:  which is skew-symmetric.

Original: In fact, this matrix is the inverse of the Chebyshev Jacobi.
Revised: In fact, this matrix is the inverse of the Chebyshev Jacobi.

Original: An alternative way to see this is as follows.
Revised: An alternative way to see this is as follows.

Original: The LegT matrix is the inverse of the matrix [Equation]
Revised: The LegT matrix is the inverse of the matrix   −  − − − 1 − 

Original: This can obviously be converted to a skew-symmetric matrix by adding a rank 2 term.
Revised: This can obviously be converted to a skew-symmetric matrix by adding a rank 2 term.

Original: The inverses of these matrices are also rank-2 differences from each other by the Woodbury identity.
Revised: The inverses of these matrices are also rank-2 differences from each other by the Woodbury identity.

Original: A ﬁnal form is
Revised: A ﬁnal form is

Original: [Equation]
Revised: [Equation]

Original: 

Original:  +   1 0 10 1 01 0 10 1 0  =      −    −  − − 1

Original: 
Revised: 

Original: This has the advantage that the rank-2 correction is symmetric (like the others), but the normal skew-symmetric matrix is now 2 -quasiseparable instead of 1 -quasiseparable.
Revised: This has the advantage that the rank-2 correction is symmetric (like the others), but the normal skew-symmetric matrix is now 2 -quasiseparable instead of 1 -quasiseparable.

Original: C. C OMPUTING THE S4 R ECURRENT V IEW
Revised: C. C OMPUTING THE S4 R ECURRENT V IEW

Original: We prove Theorem 2 showing the efﬁciency of the S4 parameterization for computing one step of the recurrent representation (Section 2.3).
Revised: We prove Theorem 2 showing the efﬁciency of the S4 parameterization for computing one step of the recurrent representation (Section 2.3).

Original: Recall that without loss of generality, we can assume that the state matrix A = Λ − pq ∗ is diagonal plus low-rank (DPLR), potentially over C .
Revised: Recall that without loss of generality, we can assume that the state matrix A = Λ − P Q ∗ is diagonal plus low-rank (DPLR), potentially over C .

Original: Our goal in this section is to explicitly write out a closed form for the discretized matrix A .
Revised: Our goal in this section is to explicitly write out a closed form for the discretized matrix A .

Original: Recall from equation (3)
Revised: Recall from equation (3)

Original: that
Revised: that

Original: [Equation]
Revised: [Equation]

Original: We ﬁrst simplify both terms in the deﬁnition of A independently.
Revised: We ﬁrst simplify both terms in the deﬁnition of A independently.

Original: Forward discretization.
Revised: Forward discretization.

Original: The ﬁrst term is essentially the Euler discretization motivated in Section 2.3.
Revised: The ﬁrst term is essentially the Euler discretization motivated in Section 2.3.

Original: [Equation]
Revised: [Equation]

Original: where A 0 is deﬁned as the term in the ﬁnal brackets.
Revised: where A 0 is deﬁned as the term in the ﬁnal brackets.

Original: Backward discretization.
Revised: Backward discretization.

Original: The second term is known as the Backward Euler’s method.
Revised: The second term is known as the Backward Euler’s method.

Original: Although this inverse term is normally difﬁcult to deal with, in the DPLR case we can simplify it using Woodbury’s
Revised: Although this inverse term is normally difﬁcult to deal with, in the DPLR case we can simplify it using Woodbury’s

Original: Identity (Proposition 4).
Revised: Identity (Proposition 4).

Original: [Equation]
Revised: [Equation]

Original: where D = (cid:0) 2∆ − Λ (cid:1) − 1 and A 1 is deﬁned as the term in the ﬁnal brackets.
Revised: where D = (cid:0) 2∆ − Λ (cid:1) − 1 and A 1 is deﬁned as the term in the ﬁnal brackets.

Original: Note that (1 + q ∗ Dp ) is actually a scalar, and A 1 is actually itself a DPLR matrix.
Revised: Note that (1 + Q ∗ DP ) is actually a scalar in the case when the low-rank term has rank 1 .

Original: S4 Recurrence.
Revised: S4 Recurrence.

Original: Finally, the full bilinear discretization can be rewritten in terms of these matrices as
Revised: Finally, the full bilinear discretization can be rewritten in terms of these matrices as

Original: [Equation]
Revised: [Equation]

Original: The discrete-time SSM (3) becomes
Revised: The discrete-time SSM (3) becomes

Original: [Equation]
Revised: [Equation]

Original: Note that A 0 , A 1 are accessed only through matrix-vector multiplications.
Revised: Note that A 0 , A 1 are accessed only through matrix-vector multiplications.

Original: Since they are both DPLR, they have O ( N ) matrix-vector multiplication, showing Theorem 2.
Revised: Since they are both DPLR, they have O ( N ) matrix-vector multiplication, showing Theorem 2.

Original: C. C OMPUTING THE C ONVOLUTIONAL V IEW
Revised: C. C OMPUTING THE C ONVOLUTIONAL V IEW

Original: The most involved part of using SSMs efﬁciently is computing K .
Revised: The most involved part of using SSMs efﬁciently is computing K .

Original: This algorithm was sketched in
Revised: This algorithm was sketched in

Original: Section 3.2 and is the main motivation for the S4 parameterization.
Revised: Section 3.2 and is the main motivation for the S4 parameterization.

Original: In this section, we deﬁne the necessary intermediate quantities and prove the main technical result.
Revised: In this section, we deﬁne the necessary intermediate quantities and prove the main technical result.

Original: The algorithm for Theorem 3 falls in roughly three stages.
Revised: The algorithm for Theorem 3 falls in roughly three stages, leading to Algorithm 1.

Original: Assuming A has been conjugated into diagonal plus low-rank form, we successively simplify the problem of computing K by applying the techniques outlined in Section 3.2. 
Revised: Assuming A has been conjugated into diagonal plus low-rank form, we successively simplify the problem of computing K by applying the techniques outlined in Section 3.2. Remark C.1.

Revised: We note that for the remainder of this section , we transpose C to be a column vector of shape C N or C N × 1 instead of matrix or row vector C 1 × N as in (1) .

Revised: In other words the SSM is

Revised: [Equation]

Revised: This convention is made so that C has the same shape as B , P , Q and simpliﬁes the implementation of S4.

Original: Reduction 0: Diagonalization By Lemma 3.1, we can switch the representation by conjugating with any unitary matrix.
Revised: Reduction 0: Diagonalization By Lemma 3.1, we can switch the representation by conjugating with any unitary matrix.

Original: For the remainder of this section, we can assume that A is (complex) diagonal plus low-rank (DPLR).
Revised: For the remainder of this section, we can assume that A is (complex) diagonal plus low-rank (DPLR).

Original: Note that unlike diagonal matrices, a DPLR matrix does not lend itself to efﬁcient computation of
Revised: Note that unlike diagonal matrices, a DPLR matrix does not lend itself to efﬁcient computation of

Original: K .
Revised: K .

Original: The reason is that K computes terms CA i B which involve powers of the matrix A .
Revised: The reason is that K computes terms C ∗ A i B which involve powers of the matrix A .

Original: These are trivially computable when A is diagonal, but is no longer possible for even simple modiﬁcations to diagonal matrices such as DPLR.
Revised: These are trivially computable when A is diagonal, but is no longer possible for even simple modiﬁcations to diagonal matrices such as DPLR.

Original: Reduction 1: SSM Generating Function To address the problem of computing powers of A , we introduce another technique.
Revised: Reduction 1: SSM Generating Function To address the problem of computing powers of A , we introduce another technique.

Original: Instead of computing the SSM convolution ﬁlter K directly, we introduce a generating function on its coefﬁcients and compute evaluations of it.
Revised: Instead of computing the SSM convolution ﬁlter K directly, we introduce a generating function on its coefﬁcients and compute evaluations of it.

Original: Deﬁnition 2 (SSM Generating Function) .
Revised: Deﬁnition 2 (SSM Generating Function) .

Original: We deﬁne the following quantities:
Revised: We deﬁne the following quantities:

Original: • The SSM convolution function is K
Revised: • The SSM convolution function is K

Original: ( A , B , C ) =
Revised: ( A , B , C ) =

Original: ( CB , CAB , . . . ) and the (truncated) SSM ﬁlter of length LK L ( A , B , C ) = ( CB , CAB , . . .
Revised: ( C ∗ B , C ∗ AB , . . . ) and the (truncated) SSM ﬁlter of length L

Original: [Equation]
Revised: [Equation]

Original: • The SSM generating function at node z is
Revised: • The SSM generating function at node z is

Original: [Equation]
Revised: [Equation]

Original: and the truncated SSM generating function at node z is
Revised: and the truncated SSM generating function at node z is

Original: [Equation]
Revised: [Equation]

Original: • The truncated SSM generating function at nodes Ω ∈ C M is
Revised: • The truncated SSM generating function at nodes Ω ∈ C M is

Original: [Equation]
Revised: [Equation]

Original: Intuitively, the generating function essentially converts the SSM convolution ﬁlter from the time domain to frequency domain.
Revised: Intuitively, the generating function essentially converts the SSM convolution ﬁlter from the time domain to frequency domain.

Original: Importantly, it preserves the same information, and the desired SSM convolution ﬁlter can be recovered from evaluations of its generating function.
Revised: Importantly, it preserves the same information, and the desired SSM convolution ﬁlter can be recovered from evaluations of its generating function.

Original: Lemma C.1.
Revised: Lemma C.2.

Original: The SSM function K L ( A , B , C ) can be computed from the SSM generating function
Revised: The SSM function K L ( A , B , C ) can be computed from the SSM generating function

Original:  K L (Ω; A , B , C  at the roots of unity Ω = { exp(2 π k L : k ∈ [ L ] } stably in O
Revised: ˆ K L (Ω; A , B , C ) at the roots of unity Ω = { exp( − 2 πi kL : k ∈ [ L ] } stably in O

Original: ( L log L ) operations.
Revised: ( L log L ) operations.

Original: Proof.
Revised: Proof.

Original: For convenience deﬁne
Revised: For convenience deﬁne

Original: [Equation]
Revised: [Equation]

Original: Note that
Revised: Note that

Original: [Equation]
Revised: [Equation]

Original: Note that this is exactly the same as the Discrete Fourier Transform (DFT):
Revised: Note that this is exactly the same as the Discrete Fourier Transform (DFT):

Original: [Equation]
Revised: [Equation]

Original: Therefore K can be recovered from ˆ K with a single inverse DFT, which requires O ( L log L ) operations with the Fast Fourier Transform (FFT) algorithm.
Revised: Therefore K can be recovered from ˆ K with a single inverse DFT, which requires O ( L log L ) operations with the Fast Fourier Transform (FFT) algorithm.

Original: Reduction 2: Woodbury Correction
Revised: Reduction 2: Woodbury Correction

Original: The primary motivation of Deﬁnition 2 is that it turns powers of A into a single inverse of A (equation (9)).
Revised: The primary motivation of Deﬁnition 2 is that it turns powers of A into a single inverse of A (equation (10)).

Original: While DPLR  cannot be powered efﬁciently due to the low-rank term, they can be inverted efﬁciently by the well-known Woodbury identity.
Revised: While DPLR matrices cannot be powered efﬁciently due to the low-rank term, they can be inverted efﬁciently by the well-known Woodbury identity.

Original: Proposition 4 (Binomial Inverse Theorem or Woodbury matrix identity Woodbury (1950); Golub &
Revised: Proposition 4 (Binomial Inverse Theorem or Woodbury matrix identity Woodbury (1950); Golub &

Original: Van Loan (2013)) .
Revised: Van Loan (2013)) .

Original: Over a commutative ring R , let A ∈ R N × N and U , V ∈ R N × p .
Revised: Over a commutative ring R , let A ∈ R N × N and U , V ∈ R N × p .

Original: Suppose A and A + UV ∗ are invertible.
Revised: Suppose A and A + UV ∗ are invertible.

Original: Then I p + V ∗
Revised: Then I p + V ∗

Original: A − 1 U ∈ R
Revised: A − 1 U ∈ R

Original: p × p is invertible and
Revised: p × p is invertible and

Original: [Equation]
Revised: [Equation]

Original: With this identity, we can convert the SSM generating function on a DPLR matrix A into one on just its diagonal component.
Revised: With this identity, we can convert the SSM generating function on a DPLR matrix A into one on just its diagonal component.

Original: Lemma C.2.
Revised: Lemma C.3.

Original: Let A = Λ − pq ∗ be a diagonal plus low-rank representation.
Revised: Let A = Λ − P Q ∗ be a diagonal plus low-rank representation.

Original: Then for any evaluation node z ∈ C , the truncated generating function satisﬁes
Revised: Then for any root of unity z ∈ Ω , the truncated generating function satisﬁes

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Directly expanding Deﬁnition 2 yields
Revised: Directly expanding Deﬁnition 2 yields

Original: [Equation]
Revised: [Equation]

Original: where ˜ C ∗ =
Revised: where ˜ C ∗ =

Original: C ∗ (cid:16) I − A L (cid:17) .
Revised: C ∗ (cid:16) I − A L (cid:17) .

Original: We can now explicitly expand the discretized SSM matrices A and B back in terms of the original
Revised: We can now explicitly expand the discretized SSM matrices A and B back in terms of the original

Original: SSM parameters A and B .
Revised: SSM parameters A and B .

Original: Lemma C.3 provides an explicit formula, which allows further simplifying˜
Revised: Lemma C.4 provides an explicit formula, which allows further simplifying

Original: [Equation]
Revised: [Equation]

Original: The last line applies the Woodbury Identity (Proposition 4) where R ( z ) = (cid:16) 2 ∆ 1 − z 1+ z − ˜ A (cid:17) − 1 .
Revised: The last line applies the Woodbury Identity (Proposition 4) where R ( z ) = (cid:16) 2∆ 1 − z 1+ z − Λ (cid:17) − 1 .

Original: The previous proof used the following self-contained result to back out the original SSM matrices from the discretization.
Revised: The previous proof used the following self-contained result to back out the original SSM matrices from the discretization.

Original: Lemma C.3.
Revised: Lemma C.4.

Original: Let A , B be the SSM matrices A , B discretized by the bilinear discretization with step size ∆ .
Revised: Let A , B be the SSM matrices A , B discretized by the bilinear discretization with step size ∆ .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: Recall that the bilinear discretization that we use (equation (3)) is
Revised: Recall that the bilinear discretization that we use (equation (3)) is

Original: [Equation]
Revised: [Equation]

Original: The result is proved algebraic manipulations.
Revised: The result is proved algebraic manipulations.

Original: [Equation]
Revised: [Equation]

Original: B Note that in the S4 parameterization, instead of constantly computing ˜ C
Revised: Note that in the S4 parameterization, instead of constantly computing ˜ C

Original: = C (cid:16) I − A L (cid:17) , we can simply reparameterize our parameters to learn ˜ C directly instead of C , saving a minor computation cost and simplifying the algorithm.
Revised: = C (cid:16) I − A L (cid:17) , we can simply reparameterize our parameters to learn ˜ C directly instead of C , saving a minor computation cost and simplifying the algorithm.

Original: Reduction 3:Core Cauchy Kernel We have reduced the original problem of computing K to the problem of computing the SSM generating function ˆ K  (Ω; A, B, C ) in the case that A is a diagonal matrix.
Revised: Reduction 3: Cauchy Kernel We have reduced the original problem of computing K to the problem of computing the SSM generating function ˆ K L (Ω; A , B , C ) in the case that A is a diagonal matrix.

Original: We show that this is exactly the same as a Cauchy kernel, which is a well-studied problem with fast and stable numerical algorithms.
Revised: We show that this is exactly the same as a Cauchy kernel, which is a well-studied problem with fast and stable numerical algorithms.

Original: Deﬁnition 3.
Revised: Deﬁnition 3.

Original: A Cauchy matrix or kernel on nodes Ω = ( ω i ) ∈
Revised: A Cauchy matrix or kernel on nodes Ω = ( ω i ) ∈

Original: C M and Λ = (
Revised: C M and Λ = (

Original: λ j ) ∈ C N is
Revised: λ j ) ∈ C N is

Original: [Equation]
Revised: [Equation]

Original: The computation time of a Cauchy matrix-vector product of size M × N is denoted by C ( M, N ) .
Revised: The computation time of a Cauchy matrix-vector product of size M × N is denoted by C ( M, N ) .

Original: Computing with Cauchy matrices is an extremely well-studied problem in numerical analysis, with both fast arithmetic algorithms and fast numerical algorithms based on the famous Fast Multiple
Revised: Computing with Cauchy matrices is an extremely well-studied problem in numerical analysis, with both fast arithmetic algorithms and fast numerical algorithms based on the famous Fast Multipole

Original: Method (FMM) (Pan, 2001; 2015; 2017).
Revised: Method (FMM) (Pan, 2001; 2015; 2017).

Original: Proposition 5 (Cauchy) .
Revised: Proposition 5 (Cauchy) .

Original: A Cauchy kernel requires O ( M + N ) space, and operation count
Revised: A Cauchy kernel requires O ( M + N ) space, and operation count

Original: C ( M, N )=
Revised: C ( M, N )

Revised: [Equation]

Original:  
Revised: 

Revised: [Equation]

Original: O ( MN ) naively
Revised: ( MN [Equation] naively

Original: O (cid:0) ( M + N ) log 2 ( M + N ) (cid:1) in exact arithmetic
Revised: O (cid:0) ( M + N ) log 2 ( M + N ) (cid:1) in exact arithmetic

Original: O (cid:0) ( M + N ) log( M + N ) log 1 ε (cid:1) numerically to precision ε.
Revised: O (cid:0) ( M + N ) log( M + N ) log 1 ε (cid:1) numerically to precision ε.

Original: Corollary C.4.
Revised: Corollary C.5.

Original: Evaluating q ∗ R (Ω; Λ) p (deﬁned in Lemma C.2) for any set of nodes Ω ∈ C L , diagonal matrix Λ , and vectors p , q can be computed in C ( L, N ) operations and O ( L + N ) space, where C ( L, N ) =
Revised: Evaluating Q ∗ R (Ω; Λ) P (deﬁned in Lemma C.3) for any set of nodes Ω ∈ C L , diagonal matrix Λ , and vectors P , Q can be computed in C ( L, N ) operations and O ( L + N ) space, where C ( L, N ) =

Original: ˜ O
Revised: ˜ O

Original: ( L + N ) is the cost of a Cauchy matrix-vector multiplication.
Revised: ( L + N ) is the cost of a Cauchy matrix-vector multiplication.

Original: Proof.
Revised: Proof.

Original: For any ﬁxed ω ∈ Ω , we want to compute (cid:80) j q ∗ j p
Revised: For any ﬁxed ω ∈ Ω , we want to compute (cid:80) j q ∗ j p

Original: j ω − λ j .
Revised: j ω − λ j .

Original: Computing this over all ω
Revised: Computing this over all ω

Original: i is therefore exactly a Cauchy matrix-vector multiplication.
Revised: i is therefore exactly a Cauchy matrix-vector multiplication.

Original: This completes the proof of Theorem 3.
Revised: This completes the proof of Theorem 3.

Original: In Algorithm 1, note that the work is dominated by Step 2, which has a constant number of calls to a black-box Cauchy kernel, with complexity given by
Revised: In Algorithm 1, note that the work is dominated by Step 2, which has a constant number of calls to a black-box Cauchy kernel, with complexity given by

Original: Proposition 5.
Revised: Proposition 5.

Original: D E XPERIMENT D ETAILS AND F ULL R ESULTS
Revised: D E XPERIMENT D ETAILS AND F ULL R ESULTS

Original: This section contains full experimental procedures and extended results and citations for our experimental evaluation in Section 4.
Revised: This section contains full experimental procedures and extended results and citations for our experimental evaluation in Section 4.

Original: Appendix D.1 corresponds to benchmarking results in Section 4.1,
Revised: Appendix D.1 corresponds to benchmarking results in Section 4.1,

Original: Appendix D.2 corresponds to LRD experiments (LRA and Speech Commands) in Section 4.2, and Appendix D.3 corresponds to the general sequence modeling experiments (generation, image classiﬁcation, forecasting) in Section 4.3.
Revised: Appendix D.2 corresponds to LRD experiments (LRA and Speech Commands) in Section 4.2, and Appendix D.3 corresponds to the general sequence modeling experiments (generation, image classiﬁcation, forecasting) in Section 4.3.

Original: D. B ENCHMARKING
Revised: D. B ENCHMARKING

Original: Benchmarking results from Table 1 and Table 2 were tested on a single A100 GPU.
Revised: Benchmarking results from Table 1 and Table 2 were tested on a single A100 GPU.

Original: Benchmarks against LSSL For a given dimension H , a single LSSL or S4 layer was constructed with H hidden features.
Revised: Benchmarks against LSSL For a given dimension H , a single LSSL or S4 layer was constructed with H hidden features.

Original: For LSSL, the state size N was set to H as done in (Gu et al., 2021).
Revised: For LSSL, the state size N was set to H as done in (Gu et al., 2021).

Original: For S4, the state size N was set to parameter-match the LSSL, which was a state size of N 4 due to differences in the parameterization.
Revised: For S4, the state size N was set to parameter-match the LSSL, which was a state size of N 4 due to differences in the parameterization.

Original: Table 1 benchmarks a single forward+backward pass of a single layer.
Revised: Table 1 benchmarks a single forward+backward pass of a single layer.

Original: Benchmarks against Efﬁcient Transformers Following (Tay et al., 2021), the Transformer models had 4 layers, hidden dimension 256 with 4 heads, query/key/value projection dimension 128 , and batch size 32 , for a total of roughly 600 k parameters.
Revised: Benchmarks against Efﬁcient Transformers Following (Tay et al., 2021), the Transformer models had 4 layers, hidden dimension 256 with 4 heads, query/key/value projection dimension 128 , and batch size 32 , for a total of roughly 600 k parameters.

Original: The S4 model was parameter tied while keeping the depth and hidden dimension constant (leading to a state size of N = 256 ).
Revised: The S4 model was parameter tied while keeping the depth and hidden dimension constant (leading to a state size of N = 256 ).

Original: We note that the relative orderings of these methods can vary depending on the exact hyperparameter settings.
Revised: We note that the relative orderings of these methods can vary depending on the exact hyperparameter settings.

Original: Other settings found S4 faster and more memory efﬁcient than all Transformer variants, for example when the query/key/value dimension was set to 256 instead of 128 , as is more standard, and when the batch size is increased.

Original: However, we chose to report a setting closer to the ones in (Tay et al., 2021).

Original: D. L ONG -R ANGE D EPENDENCIES
Revised: D. L ONG -R ANGE D EPENDENCIES

Original: This section includes information for reproducing our experiments on the Long-Range Arena and
Revised: This section includes information for reproducing our experiments on the Long-Range Arena and

Original: Speech Commands long-range dependency tasks.
Revised: Speech Commands long-range dependency tasks.

Original: Long Range Arena Table 8 contains extended results table with all 11 methods considered in (Tay et al., 2021).
Revised: Long Range Arena Table 8 contains extended results table with all 11 methods considered in (Tay et al., 2021).

Original: For the S4 model, hyperparameters for all datasets are reported in Table 9.
Revised: For the S4 model, hyperparameters for all datasets are reported in Table 9.

Original: For all datasets, we used the AdamW optimizer with a constant learning rate schedule.
Revised: For all datasets, we used the AdamW optimizer with a constant learning rate schedule with decay on validation plateau.

Revised: However, the learning rate on HiPPO parameters (in particular Λ , P , Q , B , C , ∆ ) were reduced to a maximum starting LR of 0 .

Revised: 001 , which improves stability since the HiPPO equation is crucial to performance.

Original: The S4 state size was always ﬁxed to N = 64 .
Revised: The S4 state size was always ﬁxed to N = 64 .

Original: As S4 is a sequence-to-sequence model with output shape (batch, length, dimension) and LRA tasks are classiﬁcation, mean pooling along the length dimension was applied after the last layer.
Revised: As S4 is a sequence-to-sequence model with output shape (batch, length, dimension) and LRA tasks are classiﬁcation, mean pooling along the length dimension was applied after the last layer.

Original: We note that most of these results were trained for far longer than what was necessary to achieve
Revised: We note that most of these results were trained for far longer than what was necessary to achieve

Original: SotA results (e.g., the Image task reaches SotA in 1 epoch).
Revised: SotA results (e.g., the Image task reaches SotA in 1 epoch).

Original: Results often keep improving with longer training times.
Revised: Results often keep improving with longer training times.

Original: Hardware.
Revised: Hardware.

Original: All models were run on single GPU.
Revised: All models were run on single GPU.

Original: Some tasks used an A100 GPU (notably, the Path-X experiments), which has a larger max memory of 40Gb.
Revised: Some tasks used an A100 GPU (notably, the Path-X experiments), which has a larger max memory of 40Gb.

Original: To reproduce these on smaller GPUs, the batch size can be reduced or gradients can be accumulated for two batches.
Revised: To reproduce these on smaller GPUs, the batch size can be reduced or gradients can be accumulated for two batches.

Original: [Table] [Table] Path-X. We remark that an earlier version of this paper reported a higher score for Path-X.
Revised: Path-X. We remark that an earlier version of this paper reported a higher score for Path-X.

Original: This earlier version used a different variant of the dataset, due to a misunderstanding of the properties of the dataset.
Revised: This earlier version used a different variant of the dataset, due to a misunderstanding of the properties of the dataset.

Original: More speciﬁcally [Table] we found that S4 scored 93 .
Revised: More speciﬁcally, we found that S4 scored 93 .

Original: 68% accuracy on a version that involved  taking the 256 × 256 resolution version of the  dataset and averaging every 2 × 2 square; we erroneously thought that this version of the dataset was equivalent to the original Path-X.
Revised: 68% accuracy on a version that involved [Table] [Table] taking the 256 × 256 resolution version of the Pathﬁnder dataset and averaging every 2 × 2 square; we erroneously thought that this version of the dataset was equivalent to the original Path-X.

Original: After discussions with the LRA authors, we discovered that this is not equivalent to the 128 ×resolution  dataset (the correct Path-X), which is in fact much harder.
Revised: After discussions with the LRA authors [Table] we discovered that this is not equivalent to the 128 ×resolution Pathﬁnder dataset (the correct Path-X), which is in fact much harder.

Original: In fact, Path-X is so difﬁcult that a 2D ResNet without global receptive ﬁeld (e.g. ResNet-18 or ResNet-34)  cannot achieve above chance.
Revised: In fact, Path-X is so difﬁcult that a 2D CNN without global receptive ﬁeld (e.g. ResNet-18 or ResNet-34) also cannot achieve above chance.

Original: This fact led to the original misunderstanding, as we could not solve this image classiﬁcation task even with a ResNet and thought the data might have errors.
Revised: This fact led to the original misunderstanding, as we could not solve this image classiﬁcation task even with a ResNet and thought the data might have errors.

Original: Speech Commands We provide details of sweeps run for baseline methods run by us—numbers for all others method are taken from Gu et al.
Revised: Speech Commands We provide details of sweeps run for baseline methods run by us—numbers for all others method are taken from Gu et al.

Original: The best hyperparameters used for S4 are included in Table 9.
Revised: The best hyperparameters used for S4 are included in Table 9.

Original: Transformer (Vaswani et al., 2017)
Revised: Transformer (Vaswani et al., 2017)

Original: For MFCC, we swept the number of model layers { 2 , 4 } , dropout { 0 , 0 . 1 } and learning rates { 0 . 001 , 0 . 0005 } .
Revised: For MFCC, we swept the number of model layers { 2 , 4 } , dropout { 0 , 0 . 1 } and learning rates { 0 . 001 , 0 . 0005 } .

Original: We used 8 attention heads, model dimension 128 , prenorm, positional encodings, and trained for 150 epochs with a batch size of 100 .
Revised: We used 8 attention heads, model dimension 128 , prenorm, positional encodings, and trained for 150 epochs with a batch size of 100 .

Original: For Raw, the
Revised: For Raw, the

Original: Transformer model’s memory usage made training impossible.
Revised: Transformer model’s memory usage made training impossible.

Original: Performer (Choromanski et al., 2020)
Revised: Performer (Choromanski et al., 2020)

Original: For MFCC, we swept the number of model layers { 2 , 4 } , dropout { 0 , 0 . 1 } and learning rates { 0 . 001 , 0 . 0005 } .
Revised: For MFCC, we swept the number of model layers { 2 , 4 } , dropout { 0 , 0 . 1 } and learning rates { 0 . 001 , 0 . 0005 } .

Original: We used 8 attention heads, model dimension 128 , prenorm, positional encodings, and trained for 150 epochs with a batch size of 100 .
Revised: We used 8 attention heads, model dimension 128 , prenorm, positional encodings, and trained for 150 epochs with a batch size of 100 .

Original: For Raw, we used a model dimension of 128 , 4 attention heads, prenorm, and a batch size of 16 .
Revised: For Raw, we used a model dimension of 128 , 4 attention heads, prenorm, and a batch size of 16 .

Original: We reduced the number of model layers to 4 , so the model would ﬁt on the single GPU.
Revised: We reduced the number of model layers to 4 , so the model would ﬁt on the single GPU.

Original: We trained for 100 epochs with a learning rate of 0 .
Revised: We trained for 100 epochs with a learning rate of 0 .

Original: 001 and no dropout.
Revised: 001 and no dropout.

Original: ExpRNN (Lezcano-Casado & Mart´ınez-Rubio, 2019)
Revised: ExpRNN (Lezcano-Casado & Mart´ınez-Rubio, 2019)

Original: For MFCC, we swept hidden sizes { 256 , 512 } and learning rates { 0 . 001 , 0 . 002 , 0 . 0005 } .
Revised: For MFCC, we swept hidden sizes { 256 , 512 } and learning rates { 0 . 001 , 0 . 002 , 0 . 0005 } .

Original: Training was run for 200 epochs, with a single layer model using a batch size of 100 .
Revised: Training was run for 200 epochs, with a single layer model using a batch size of 100 .

Original: For Raw, we swept hidden sizes { 32 , 64 } and learning rates { 0 . 001 , 0 . 0005 } (however, ExpRNN failed to learn).
Revised: For Raw, we swept hidden sizes { 32 , 64 } and learning rates { 0 . 001 , 0 . 0005 } (however, ExpRNN failed to learn).

Original: Lipschitz RNN (Erichson et al., 2021)
Revised: Lipschitz RNN (Erichson et al., 2021)

Original: For MFCC, we swept hidden sizes { 256 , 512 } and learning rates { 0 . 001 , 0 . 002 , 0 . 0005 } .
Revised: For MFCC, we swept hidden sizes { 256 , 512 } and learning rates { 0 . 001 , 0 . 002 , 0 . 0005 } .

Original: Training was run for 150 epochs, with a single layer model using a batch size of 100 .
Revised: Training was run for 150 epochs, with a single layer model using a batch size of 100 .

Original: For Raw, we found that LipschitzRNN was too slow to train on a single GPU (requiring a full day for 1 epoch of training alone).
Revised: For Raw, we found that LipschitzRNN was too slow to train on a single GPU (requiring a full day for 1 epoch of training alone).

Original: WaveGAN Discriminator (Donahue et al., 2019)
Revised: WaveGAN Discriminator (Donahue et al., 2019)

Original: The WaveGAN-D in Table 4 is actually our improved version of the discriminator network from the recent WaveGAN model for speech (Donahue et al., 2019).
Revised: The WaveGAN-D in Table 4 is actually our improved version of the discriminator network from the recent WaveGAN model for speech (Donahue et al., 2019).

Original: This CNN actually did not work well out-of-the-box, and we added several features to help it perform better.
Revised: This CNN actually did not work well out-of-the-box, and we added several features to help it perform better.

Original: The ﬁnal model is highly specialized compared to our model, and includes:
Revised: The ﬁnal model is highly specialized compared to our model, and includes:

Original: • Downsampling or pooling between layers, induced by strided convolutions, that decrease the sequence length between layers.
Revised: • Downsampling or pooling between layers, induced by strided convolutions, that decrease the sequence length between layers.

Original: • A global fully-connected output layer; thus the model only works for one input sequence length and does not work on MFCC features or the frequency-shift setting in Table 4.
Revised: • A global fully-connected output layer; thus the model only works for one input sequence length and does not work on MFCC features or the frequency-shift setting in Table 4.

Original: •It only works with Batch Normalization turned on, whereas S4 works equally well with either Batch Normalization or Layer Normalization.
Revised: • Batch Normalization is essential, whereas S4 works equally well with either Batch Normalization or Layer Normalization.

Original: • Almost 90 × as many parameters as the S4 model ( 26 . 3 M vs. 0 .
Revised: • Almost 90 × as many parameters as the S4 model ( 26 . 3 M vs. 0 .

Original: 3 M)
Revised: 3 M).

Original: D. G ENERAL S EQUENCE M ODELING
Revised: D. G ENERAL S EQUENCE M ODELING

Original: This subsection corresponds to the experiments in Section 4.3.
Revised: This subsection corresponds to the experiments in Section 4.3.

Original: Because of the number of experiments in this section, we use subsubsection dividers for different tasks to make it easier to follow: CIFARdensity estimation Appendix D.3.1, WikiText-103 language modeling Appendix D.3.2, autoregressive generation Appendix D.3.3, sequential image classiﬁcation Appendix D.3.4, and time-series forecasting Appendix D.3.5.
Revised: Because of the number of experiments in this section, we use subsubsection dividers for different tasks to make it easier to follow: CIFARdensity estimation Appendix D.3.1, WikiText-103 language modeling Appendix D.3.2, autoregressive generation Appendix D.3.3, sequential image classiﬁcation Appendix D.3.4, and time-series forecasting Appendix D.3.5.

Original: D.3. CIFAR D ENSITY E STIMATION
Revised: D.3. CIFAR D ENSITY E STIMATION

Original: This task used a different backbone than the rest of our experiments.
Revised: This task used a different backbone than the rest of our experiments.

Original: We used blocks of alternating
Revised: We used blocks of alternating

Original: S4 layers and position-wise feed-forward layers (in the style of Transformer blocks).
Revised: S4 layers and position-wise feed-forward layers (in the style of Transformer blocks).

Original: Each feedforward intermediate dimension was set to 2 × the hidden size of the incoming S4 layer.
Revised: Each feedforward intermediate dimension was set to 2 × the hidden size of the incoming S4 layer.

Original: Similar to
Revised: Similar to

Original: Salimans et al.
Revised: Salimans et al.

Original: (2017), we used a UNet-style backbone consisting of B identical blocks followed by a downsampling layer.
Revised: (2017), we used a UNet-style backbone consisting of B identical blocks followed by a downsampling layer.

Original: The downsampling rates were 3 , 4 , 4 (the 3 chosen because the sequence consists of RGB pixels).
Revised: The downsampling rates were 3 , 4 , 4 (the 3 chosen because the sequence consists of RGB pixels).

Original: The base model had B = 8 with starting hidden dimension 128, while the large model had B = 16 with starting hidden dimension 192.
Revised: The base model had B = 8 with starting hidden dimension 128, while the large model had B = 16 with starting hidden dimension 192.

Original: We experimented with both the mixture of logistics from (Salimans et al., 2017) as well as a simpler 256-way categorical loss.
Revised: We experimented with both the mixture of logistics from (Salimans et al., 2017) as well as a simpler 256-way categorical loss.

Original: We found they were pretty close and ended up using the simpler softmax loss along with using input embeddings.
Revised: We found they were pretty close and ended up using the simpler softmax loss along with using input embeddings.

Original: We used the LAMB optimizer with learning rate 0.005.
Revised: We used the LAMB optimizer with learning rate 0.005.

Original: The base model had no dropout, while the large model had dropout 0.1 before the linear layers inside the S4 and FF blocks.
Revised: The base model had no dropout, while the large model had dropout 0.1 before the linear layers inside the S4 and FF blocks.

Original: D.3. W IKI T EXT -103 L ANGUAGE M ODELING
Revised: D.3. W IKI T EXT -103 L ANGUAGE M ODELING

Original: The RNN baselines included in Table 7 are the AWD-QRNN (Merity et al., 2018), an efﬁcient linear gated RNN, and the LSTM + Cache + Hebbian + MbPA (Rae et al., 2018), the best performing pure RNN in the literature.
Revised: The RNN baselines included in Table 7 are the AWD-QRNN (Merity et al., 2018), an efﬁcient linear gated RNN, and the LSTM + Cache + Hebbian + MbPA (Rae et al., 2018), the best performing pure RNN in the literature.

Original: The CNN baselines are the CNN with GLU activations (Dauphin et al., 2017), the TrellisNet (Bai et al., 2019), Dynamic Convolutions (Wu et al., 2019), and TaLK
Revised: The CNN baselines are the CNN with GLU activations (Dauphin et al., 2017), the TrellisNet (Bai et al., 2019), Dynamic Convolutions (Wu et al., 2019), and TaLK

Original: Convolutions (Lioutas & Guo, 2020).
Revised: Convolutions (Lioutas & Guo, 2020).

Original: The Transformer baseline is (Baevski & Auli, 2018), which uses Adaptive Inputs with a tied Adaptive
Revised: The Transformer baseline is (Baevski & Auli, 2018), which uses Adaptive Inputs with a tied Adaptive

Original: Softmax.
Revised: Softmax.

Original: This model is a standard high-performing Transformer baseline on this benchmark, used for example by Lioutas & Guo (2020) and many more.
Revised: This model is a standard high-performing Transformer baseline on this benchmark, used for example by Lioutas & Guo (2020) and many more.

Original: Our S4 model uses the same Transformer backbone as in (Baevski & Auli, 2018).
Revised: Our S4 model uses the same Transformer backbone as in (Baevski & Auli, 2018).

Original: The model consists of 16 blocks of S4 layers alternated with position-wise feedforward layers, with a feature dimension of 1024.
Revised: The model consists of 16 blocks of S4 layers alternated with position-wise feedforward layers, with a feature dimension of 1024.

Original: Because our S4 layer has around 1/4 the number of parameters as a self-attention layer with the same dimension, we made two modiﬁcations to match the parameter count better: (i) we used a GLU activation after the S4 linear layer (Section 3.4) (ii) we used two S4 layers per block.
Revised: Because our S4 layer has around 1/4 the number of parameters as a self-attention layer with the same dimension, we made two modiﬁcations to match the parameter count better: (i) we used a GLU activation after the S4 linear layer (Section 3.4) (ii) we used two S4 layers per block.

Original: Blocks use Layer Normalization in the pre-norm position.
Revised: Blocks use Layer Normalization in the pre-norm position.

Original: The embedding and softmax layers were the Adaptive Embedding from (Baevski & Auli, 2018) with standard cutoffs 20000, 40000, 200000.
Revised: The embedding and softmax layers were the Adaptive Embedding from (Baevski & Auli, 2018) with standard cutoffs 20000, 40000, 200000.

Original: Evaluation was performed similarly to the basic setting in (Baevski & Auli, 2018), Table 5, which involves sliding non-overlapping windows of width 1024 tokens.
Revised: Evaluation was performed similarly to the basic setting in (Baevski & Auli, 2018), Table 5, which involves sliding non-overlapping windows of width 1024 tokens.

Original: Other settings are reported in (Baevski & Auli, 2018) that include more context at training and evaluation time and improves the score.
Revised: Other settings are reported in (Baevski & Auli, 2018) that include more context at training and evaluation time and improves the score.

Original: Because such evaluation protocols are orthogonal to the basic model, we do not consider them and report the base score from (Baevski & Auli, 2018)
Revised: Because such evaluation protocols are orthogonal to the basic model, we do not consider them and report the base score from (Baevski & Auli, 2018)

Original: Table 5.
Revised: Table 5.

Original: Instead of SGD+Momentum with multiple cosine learning rate annealing cycles, our S4 model was trained with the simpler AdamW optimizer with a single cosine learning rate cycle with a maximum of 800000 steps.
Revised: Instead of SGD+Momentum with multiple cosine learning rate annealing cycles, our S4 model was trained with the simpler AdamW optimizer with a single cosine learning rate cycle with a maximum of 800000 steps.

Original: The initial learning rate was set to 0.0005.
Revised: The initial learning rate was set to 0.0005.

Original: We used 8 A100 GPUs with a batch size of 8 per gpu and context size 1024.
Revised: We used 8 A100 GPUs with a batch size of 8 per gpu and context size 1024.

Original: We used no gradient clipping and a weight decay of 0.1.
Revised: We used no gradient clipping and a weight decay of 0.1.

Original: Unlike (Baevski & Auli, 2018) which speciﬁed different dropout rates for different parameters, we used a constant dropout rate of 0.25 throughout the network, including before every linear layer and on the residual branches.
Revised: Unlike (Baevski & Auli, 2018) which speciﬁed different dropout rates for different parameters, we used a constant dropout rate of 0.25 throughout the network, including before every linear layer and on the residual branches.

Original: D.3. A UTOREGRESSIVE G ENERATION S PEED
Revised: D.3. A UTOREGRESSIVE G ENERATION S PEED

Original: Protocol.
Revised: Protocol.

Original: To account for different model sizes and memory requirements for each method, we benchmark generation speed by throughput. measured in images per second (Table 6) or tokens per second (Table 7).
Revised: To account for different model sizes and memory requirements for each method, we benchmark generation speed by throughput, measured in images per second (Table 6) or tokens per second (Table 7).

Original: Each model generates images on a single A 100 GPU, maximizing batch size to ﬁt in memory.
Revised: Each model generates images on a single A 100 GPU, maximizing batch size to ﬁt in memory.

Original: (For CIFAR-10 generation we limited memory to 16Gb, to be more comparable to the
Revised: (For CIFAR-10 generation we limited memory to 16Gb, to be more comparable to the

Original: Transformer and Linear Transformer results reported from (Katharopoulos et al., 2020).)
Revised: Transformer and Linear Transformer results reported from (Katharopoulos et al., 2020).)

Original: Baselines.
Revised: Baselines.

Original: The Transformer and Linear Transformer baselines reported in Table 6 are the results reported directly from Katharopoulos et al.
Revised: The Transformer and Linear Transformer baselines reported in Table 6 are the results reported directly from Katharopoulos et al.

Original: Note that the Transformer number is the one in their Appendix, which implements the optimized cached implementation of self-attention.
Revised: Note that the Transformer number is the one in their Appendix, which implements the optimized cached implementation of self-attention.

Original: For all other baseline models, we used open source implementations of the models to benchmark generation speed.
Revised: For all other baseline models, we used open source implementations of the models to benchmark generation speed.

Original: For the PixelCNN++, we used the fast cached version by Ramachandran et al.
Revised: For the PixelCNN++, we used the fast cached version by Ramachandran et al.

Original: (2017), which sped up generation by orders of magnitude from the naive implementation.
Revised: (2017), which sped up generation by orders of magnitude from the naive implementation.

Original: This code was only available in TensorFlow, which may have slight differences compared to the rest of the baselines which were implemented in PyTorch.
Revised: This code was only available in TensorFlow, which may have slight differences compared to the rest of the baselines which were implemented in PyTorch.

Original: We were unable to run the Sparse Transformer (Child et al., 2019) model due to issues with their custom CUDA implementation of the sparse attention kernel, which we were unable to resolve.
Revised: We were unable to run the Sparse Transformer (Child et al., 2019) model due to issues with their custom CUDA implementation of the sparse attention kernel, which we were unable to resolve.

Original: The Transformer baseline from Table 7 was run using a modiﬁed GPT-2 baackbone from the HuggingFace repository, conﬁgured to recreate the architecture reported in (Baevski & Auli, 2018).
Revised: The Transformer baseline from Table 7 was run using a modiﬁed GPT-2 backbone from the HuggingFace repository, conﬁgured to recreate the architecture reported in (Baevski & Auli, 2018).

Original: These numbers are actually slightly favorable to the baseline, as we did not include the timing of the embedding or softmax layers, whereas the number reported for S4 is the full model.
Revised: These numbers are actually slightly favorable to the baseline, as we did not include the timing of the embedding or softmax layers, whereas the number reported for S4 is the full model.

Original: D.3. P IXEL -L EVEL S EQUENTIAL I MAGE C LASSIFICATION
Revised: D.3. P IXEL -L EVEL S EQUENTIAL I MAGE C LASSIFICATION

Original: Our models were trained with the AdamW optimizer for up to 200 epochs.
Revised: Our models were trained with the AdamW optimizer for up to 200 epochs.

Original: Hyperparameters for the
Revised: Hyperparameters for the

Original: CIFAR-10 model is reported in Table 9.
Revised: CIFAR-10 model is reported in Table 9.

Original: For our comparisons against ResNet-18, the main differences between the base models are that
Revised: For our comparisons against ResNet-18, the main differences between the base models are that

Original: S4 uses LayerNorm by default while ResNet uses BatchNorm.
Revised: S4 uses LayerNorm by default while ResNet uses BatchNorm.

Original: The last ablation in Section 4. swaps the normalization type, using BatchNorm for S4 and LayerNorm for ResNet, to ablate this architectural difference.
Revised: The last ablation in Section 4. swaps the normalization type, using BatchNorm for S4 and LayerNorm for ResNet, to ablate this architectural difference.

Original: The experiments with augmentation take the base model and train with mild data augmentation: horizontal ﬂips and random crops (with symmetric padding).
Revised: The experiments with augmentation take the base model and train with mild data augmentation: horizontal ﬂips and random crops (with symmetric padding).

Original: D.3. T IME S ERIES F ORECASTING COMPARED TO I NFORMER
Revised: D.3. T IME S ERIES F ORECASTING COMPARED TO I NFORMER

Original: We include a simple ﬁgure (Appendix D.3.5) contrasting the architecture of S4 against that of the Informer (Zhou et al., 2021).[Table]
Revised: We include a simple ﬁgure (Fig. 3) contrasting the architecture of S4 against that of the Informer (Zhou et al., 2021).

Original: [Figure]
Revised: In Fig.

Original: In Appendix D.3.5, the goal is to forecast a contiguous range of future predictions (Green, length F ) given a range of past context (Blue, length C ).
Revised: 3, the goal is to forecast a contiguous range of future predictions (Green, length F ) given a range of past context (Blue, length C ).

Original: We simply concatenate the entire context with a sequence of masks set to the length of the forecast window.
Revised: We simply concatenate the entire context with a sequence of masks set to the length of the forecast window.

Original: This input is a single sequence of length 
Revised: This input is a single sequence of length C + F that [Table]

Original: C + F that is run through the same simple deep S4 model used throughout this work, which maps to an output of length C + F .
Revised: [Figure] is run through the same simple deep S4 model used throughout this work, which maps to an output of length C + F .

Original: We then use just the last F features as the forecasted predictions.
Revised: We then use just the last F features as the forecasted predictions.

Original: Tables 11 and 12 contain full results on all 50 settings considered by Zhou et al.
Revised: Tables 11 and 12 contain full results on all 50 settings considered by Zhou et al.

Original: S4 sets the best results on 40 out of 50 of these settings. 
Revised: S4 sets the best results on 40 out of 50 of these settings. [Table] [Table]

Original: D. V ISUALIZATIONS
Revised: D. V ISUALIZATIONS

Original: We visualize the convolutional ﬁlter ¯ K learned by S4 for the Pathﬁnder and CIFAR-10 tasks in
Revised: We visualize the convolutional ﬁlter ¯ K learned by S4 for the Pathﬁnder and CIFAR-10 tasks in

Original: AppendixD.4 and Appendix D.4.[Table] [Table]
Revised: Appendix D.4.