# Deep Latent State Space Models for
Time-Series Generation

Linqi Zhou
Stanford University
&Michael Poli 11footnotemark: 1
Stanford University
&Winnie Xu
University of Toronto
Stefano Massaroli
MILA &Stefano Ermon11footnotemark: 1
Stanford University {linqizhou, poli,
ermon}@stanford.eduwinniexu<EMAIL_ADDRESS>
###### Abstract

Methods based on ordinary differential equations (ODEs) are widely used to
build generative models of time-series. In addition to high computational
overhead due to explicitly computing hidden states recurrence, existing
ODE–based models fall short in learning sequence data with sharp transitions –
common in many real-world systems – due to numerical challenges during
optimization. In this work, we propose LS4, a generative model for sequences
with latent variables evolving according to a state space ODE to increase
modeling capacity. Inspired by recent deep state space models (S4), we achieve
speedups by leveraging a convolutional representation of LS4 which bypasses
the explicit evaluation of hidden states. We show that LS4 significantly
outperforms previous continuous-time generative models in terms of marginal
distribution, classification, and prediction scores on real-world datasets in
the Monash Forecasting Repository, and is capable of modeling highly
stochastic data with sharp temporal transitions. LS4 sets state–of–the–art for
continuous–time latent generative models, with significant improvement of mean
squared error and tighter variational lower bounds on irregularly–sampled
datasets, while also being $\times 100$ faster than other baselines on long
sequences.

###  1 Introduction

Time series are a ubiquitous data modality, and find extensive application in
weather (Hersbach et al., 2020) engineering disciplines, biology (Peng et al.,
1995), and finance (Poli et al., 2019). The main existing approaches for deep
generative learning of temporal data can be broadly categorized into
autoregressive (Oord et al., 2016), latent autoencoder models (Chen et al.,
2018; Yildiz et al., 2019; Rubanova et al., 2019), normalizing flows (de
Bézenac et al., 2020), generative adversarial (Yoon et al., 2019; Yu et al.,
2022; Brooks et al., 2022), and diffusion (Rasul et al., 2021). Among these,
continuous-time methods (often based on underlying ODEs) are the preferred
approach for irregularly-sampled sequences as they can predict at arbitrary
time steps and can handle sequences of varying lengths. Unfortunately,
existing ODE–based methods (Rubanova et al., 2019; Yildiz et al., 2019) often
fall short in learning models for real-world data (e.g., with stiff dynamics)
due to their limited expressivity and numerical instabilities during backward
gradient computation (Hochreiter, 1998; Niesen & Hall, 2004; Zhuang et al.,
2020).

A natural way to increase the flexibility of ODE-based models is to increase
the dimensionality of their (deterministic) hidden states. State–of–the–art
methods explicitly compute hidden states by unrolling the underlying
recurrence over time (each time step parametrized by a neural network),
incurring in polynomial computational costs which prevent scaling.

An alternative approach to increasing modeling capacity is to incorporate
_stochastic_ latent variables into the model, a highly successful strategy in
generative modeling (Kingma & Welling, 2013; Chung et al., 2015; Song et al.,
2020; Ho et al., 2020). However, reference models like latent neural ODE
models (Rubanova et al., 2019) inject stochasticity only at the initial
condition of the system. In contrast, we introduce LS4, a latent generative
model where the sequence of latent variables is represented as the solution of
linear state space equations (Chen, 1984). Unrolling the recurrence equation
shows an autoregressive dependence in the sequence of latent variables, the
joint of which is highly expressive in representing time series distributions.
The high dimensional structure of the latent space, being equivalent to that
of the input sequence, allows LS4 to learn expressive latent representations
and fit the distribution of sequences produced by a family of dynamical
systems, a common setting resulting from non–stationarity. We further show how
LS4 can learn the dynamics of stiff (Shampine & Thompson, 2007) dynamical
systems where previous methods fail to do so. Inspired by recent works on deep
state space models, or stacks of linear state spaces and non-linearities (Gu
et al., 2020; 2021), we leverage a convolutional kernel representation to
solve the recurrence, bypassing any explicit computations of hidden states via
the recurrence equation, which ensures log–linear scaling in both the hidden
state space dimensionality as well as sequence length.

We validate our method across a variety of time series datasets, benchmarking
LS4 against an extensive set of baselines. We propose a set of $3$ metrics to
measure the quality of generated time series samples and show that LS4
performs significantly better than baselines on datasets with stiff
transitions and obtains on average $30\%$ lower MSE scores and ELBO. On
sequences with $\approx 20K$ lengths, our model trains $\times 100$ faster
than the baseline methods.

###  2 Related Work

Rapid progress on deep generative modeling of natural language and images has
consolidated diffusion (Ho et al., 2020; Song et al., 2020; Song & Ermon,
2019; Sohl-Dickstein et al., 2015) and autoregressive techniques (Brown et
al., 2020) as the state–of–the–art. Although various approaches have been
proposed for generative modeling of time series and dynamical systems,
consensus on the advantages and disadvantages of each method has yet to
emerge.

Deep generative modeling of sequences. All the major paradigms for deep
generative modeling have seen application to time series and sequences. Most
relevant to our work are latent continuous–time autoencoder models proposed by
Chen et al. (2018); Yildiz et al. (2019); Rubanova et al. (2019), where a
neural differential equation encoder is used to parametrize as distribution of
initial conditions for the decoder. Massaroli et al. (2021) proposes a variant
that parallelizes computation in time by casting solving the ODE as a root
finding problem. Beyond latent models, other continuous–time approaches are
given in Kidger et al. (2020), which develops a GAN formulation using SDEs.

State space models. State space models. (SSMs) are at the foundation of
dynamical system theory (Chen, 1984) and signal processing (Oppenheim, 1999),
and have also been adapted to deep generative modeling. Chung et al. (2015);
Bayer & Osendorfer (2014) propose VAE variants of discrete–time RNNs,
generalized later by (Franceschi et al., 2020), among others. These models all
unroll the recurrence and are thus challenging to scale to longer sequences.

Our work is inspired by recent advances in deep architectures built as stacks
of linear SSMs, notably S4 (Gu et al., 2021). Similar to S4, our generative
model leverages the convolutional representation of SSMs during training, thus
bypassing the need to materialize the hidden state of each recurrence.

###  3 Preliminaries

We briefly introduce relevant details of continuous-time SSMs and their
different representations. Then we introduce preliminaries of generative
models for sequences.

####  3.1 State Space Models (SSM)

A single-input single-output (SISO) linear state space model is defined by the
following differential equation

$\displaystyle\begin{split}\frac{d}{dt}\bm{h}_{t}&={\bm{A}}\bm{h}_{t}+{\bm{B}}x_{t}\\\
y_{t}&={\bm{C}}\bm{h}_{t}+{\bm{D}}x_{t}\end{split}$ (1)

with scalar input $x_{t}\in\mathbb{R}$, state $\bm{h}_{t}\in\mathbb{R}^{N}$
and scalar output $y_{t}\in\mathbb{R}$. The system is fully characterized by
the matrices ${\bm{A}}\in\mathbb{R}^{N\times N},{\bm{B}}\in\mathbb{R}^{N\times
1},{\bm{C}}\in\mathbb{R}^{1\times N},{\bm{D}}\in\mathbb{R}^{1\times 1}$. Let
$x,y\in\mathcal{C}([a,b],\mathbb{R})$ be absolutely continuous real signals on
time interval $[a,b]$. Given an initial condition
$\bm{h}_{0}\in\mathbb{R}^{N}$ the SSM (1) realizes a mapping $x\mapsto y$.

SSMs are a common tool for processing continuous input signals. We consider
single input single output (SISO) SSMs, noting that input sequences with more
than a single channel can be processed by applying multiple SISO SSMs in
parallel, similarly to regular convolutional layers. We use such SSMs as
building blocks to map each input dimension to each output dimension in our
generative model.

Discrete recurrent representation. In practice, continuous input signals are
often sampled at time interval $\Delta$ and the sampled sequence is
represented by $x=(x_{t_{0}},x_{t_{1}},\dots,x_{t_{L}})$ where
$t_{k+1}=t_{k}+\Delta$. The discretized SSM follows the recurrence

$\displaystyle\begin{split}\bm{h}_{t_{k+1}}&=\bar{{\bm{A}}}\bm{h}_{t_{k}}+\bar{{\bm{B}}}x_{t_{k}}\\\
y_{t_{k}}&={\bm{C}}\bm{h}_{t_{k}}+{\bm{D}}x_{t_{k}}\end{split}$ (2)

where $\bar{{\bm{A}}}=e^{{\bm{A}}\Delta}$,
$\bar{{\bm{B}}}={\bm{A}}^{-1}(e^{{\bm{A}}\Delta}-I){\bm{B}}$ with the
assumption that signals are constant during the sampling interval.

Among many approaches to efficiently computing $e^{{\bm{A}}\Delta}$, Gu et al.
(2021) use a bilinear transform to estimate
$e^{{\bm{A}}\Delta}\approx(I-\frac{1}{2}{\bm{A}}\Delta)^{-1}(I+\frac{1}{2}{\bm{A}}\Delta)$.

This recurrence equation can be used to iteratively solve for the next hidden
state $h_{t_{k+1}}$, allowing the states to be calculated like an RNN or a
Neural ODE (Chen et al., 2018; Massaroli et al., 2020).

Convolutional representation.

Recurrent representations of SSM are not practical in training because
explicit calculation of hidden states for every time step requires
$\mathcal{O}(N^{2}L)$ in time and $\mathcal{O}(NL)$ in space for a sequence of
length $L$111Further explanations in Appendix A.1. This materialization of
hidden states significantly restricts RNN-based methods in scaling to long
sequences. To efficiently train an SSM, the recurrence equation can be fully
unrolled, assuming zero initial hidden states, as

$\displaystyle\begin{array}[]{llll}\bm{h}_{t_{0}}=\bar{{\bm{B}}}x_{t_{0}}&\bm{h}_{t_{1}}=\bar{{\bm{A}}}\bar{{\bm{B}}}x_{t_{1}}+\bar{{\bm{B}}}x_{t_{0}}&\bm{h}_{t_{2}}=\bar{{\bm{A}}}^{2}\bar{{\bm{B}}}x_{t_{2}}+\bar{{\bm{A}}}\bar{{\bm{B}}}x_{t_{1}}+\bar{{\bm{B}}}x_{t_{0}}&\dots\\\
y_{t_{0}}={\bm{C}}\bar{{\bm{B}}}x_{t_{0}}&y_{t_{1}}={\bm{C}}\bar{{\bm{A}}}\bar{{\bm{B}}}x_{t_{1}}+{\bm{C}}\bar{{\bm{B}}}x_{t_{0}}&y_{t_{2}}={\bm{C}}\bar{{\bm{A}}}^{2}\bar{{\bm{B}}}x_{t_{2}}+{\bm{C}}\bar{{\bm{A}}}\bar{{\bm{B}}}x_{t_{1}}+{\bm{C}}\bar{{\bm{B}}}x_{t_{0}}&\dots\end{array}$

and more generally as,

$\displaystyle y_{t_{k}}$
$\displaystyle={\bm{C}}\bar{{\bm{A}}}^{k}\bar{{\bm{B}}}x_{t_{k}}+{\bm{C}}\bar{{\bm{A}}}^{k-1}\bar{{\bm{B}}}x_{k-1}+\dots+{\bm{C}}\bar{{\bm{B}}}x_{t_{0}}$

For an input sequence $x=(x_{t_{0}},x_{t_{1}},\dots,x_{t_{L}})$, one can
observe that the output sequence $y=(y_{t_{0}},y_{t_{1}},\dots,y_{t_{L}})$ can
be computed using a convolution with a skip connection

$\displaystyle y={\bm{C}}{\bm{K}}*x+{\bm{D}}x,\text{ where
}{\bm{K}}=(\bar{{\bm{B}}},\bar{{\bm{A}}}\bar{{\bm{B}}},\dots,\bar{{\bm{A}}}^{L-1}\bar{{\bm{B}}},\bar{{\bm{A}}}^{L}\bar{{\bm{B}}})$
(3)

This is the well-known connection between SSM and convolution (Oppenheim &
Schafer, 1975; Chen, 1984; Chilkuri & Eliasmith, 2021; Romero et al., 2021; Gu
et al., 2020; 2021; 2022)

and it can be computed very efficiently with a Fast Fourier Transform (FFT),
which scales better than explicit matrix multiplication at each step.

####  3.2 Variational Autoencoder (VAE)

VAEs (Kingma & Welling, 2013; Burda et al., 2015) are a highly successful
paradigm in learning latent representations of high dimensional data and is
remarkably capable at modeling complex distributions. A VAE introduces a joint
probability distribution between a latent variable ${\bm{z}}$ and an observed
random variable ${\bm{x}}$ of the form

$p_{\theta}({\bm{x}},{\bm{z}})=p_{\theta}({\bm{x}}\mid{\bm{z}})p({\bm{z}})$

where $\theta$ represents learnable parameters.

The prior $p({\bm{z}})$ over the latent is usually chosen as a standard
Gaussian distribution, and the conditional distribution
$p_{\theta}({\bm{x}}\mid{\bm{z}})$ is defined through a flexible non-linear
mapping (such as a neural network) taking ${\bm{z}}$ as input.

Such highly flexible non-linear mappings often lead to an intractable
posterior $p_{\theta}({\bm{z}}\mid{\bm{x}})$. Therefore, an inference model
with parameters $\phi$ parametrizing $q_{\phi}({\bm{z}}\mid{\bm{x}})$ is
introduced as an approximation which allows learning through a variational
lower bound of the marginal likelihood:

$\displaystyle\log p_{\theta}({\bm{x}})\geq-
D_{\mathrm{KL}}(q_{\phi}({\bm{z}}\mid{\bm{x}})\|p({\bm{z}}))+\mathbb{E}_{q_{\phi}({\bm{z}}\mid{\bm{x}})}\left[\log
p_{\theta}({\bm{x}}\mid{\bm{z}})\right]$ (4)

where $D_{\mathrm{KL}}(\cdot\|\cdot)$ is the Kullback-Leibler divergence
between two distributions.

VAE for sequences.

Sequence data can be modeled in many different ways since the latent space can
be chosen to encode information at different levels of granularity, _i.e_.
${\bm{z}}$ can be a single variable encoding entire trajectories or a sequence
of variables of the same length as the trajectories. We focus on the latter.

Given observed sequence variables ${\bm{x}}_{\leq T}$ up to time $T$
discretized into sequence $({\bm{x}}_{t_{0}},\dots,{\bm{x}}_{t_{L-1}})$ of
length $L$ where $t_{L-1}=T$, a sequence VAE model with parameters
$\theta,\lambda,\phi$ learns a generative and inference distribution

$\displaystyle p_{\theta,\lambda}({\bm{x}}_{\leq t_{L-1}},{\bm{z}}_{\leq
t_{L-1}})=\prod_{i=0}^{L-1}p_{\theta}({\bm{x}}_{t_{i}}\mid{\bm{x}}_{<t_{i}},{\bm{z}}_{\leq
t_{i}})p_{\lambda}({\bm{z}}_{t}\mid{\bm{z}}_{<t_{i}})$ $\displaystyle
q_{\phi}({\bm{z}}_{\leq t_{L-1}}\mid{\bm{x}}_{\leq
t_{L-1}})=\prod_{i=0}^{L-1}q_{\phi}({\bm{z}}_{t_{i}}\mid{\bm{x}}_{\leq
t_{t_{i}}})$

where ${\bm{z}}_{\leq t_{L-1}}=({\bm{z}}_{t_{0}},\dots,{\bm{z}}_{t_{L-1}})$ is
the corresponding latent variable sequence. The approximate posterior
$q_{\phi}$ is explicitly factorized as a product of marginals due to
efficiency reasons we shall discuss in the next section. Given this form of
factorization, the variational lowerbound has been considered for discrete
sequence data (Chung et al., 2015) with the objective

$\displaystyle\begin{split}\mathbb{E}_{q_{\phi}({\bm{z}}_{\leq
t_{L-1}}\mid{\bm{x}}_{\leq
t_{L-1}})}\left[-\sum_{i=0}^{L-1}D_{\mathrm{KL}}(q_{\phi}({\bm{z}}_{t_{i}}\mid{\bm{x}}_{\leq
t_{i}})\|p_{\lambda}({\bm{z}}_{i}\mid{\bm{z}}_{<t_{i}}))+\log
p_{\theta}({\bm{x}}_{t_{i}}\mid{\bm{x}}_{<{t_{i}}}{\bm{z}}_{\leq{t_{i}}})\right]\end{split}$
(5)

###  4 Method

In this section, we introduce Latent S4 (LS4), a latent variable generative
model parameterized using SSMs. We show how SSMs can parametrize the
generative distribution $p_{\theta}({\bm{x}}_{\leq T}|{\bm{z}}_{\leq
T})p_{\lambda}({\bm{z}}_{\leq T})$, the prior distribution
$p_{\lambda}({\bm{z}}_{\leq T})$ and the inference distribution
$q_{\phi}({\bm{z}}_{\leq T}\mid{\bm{x}}_{\leq T})$ effectively. For the
purpose of exposition, we can assume $z_{t},x_{t}$ are scalars at any time
step $t$. Their generalization to arbitrary dimensions is discussed in Section
4.4.

We first define a structured state space model with two input streams and use
this as a building block for our generative model. It is an SSM of the form

$\displaystyle\begin{split}\frac{d}{dt}\bm{h}_{t}&={\bm{A}}\bm{h}_{t}+{\bm{B}}x_{t}+{\bm{E}}z_{t}\\\
y_{t}&={\bm{C}}\bm{h}_{t}+{\bm{D}}x_{t}+{\bm{F}}z_{t}\end{split}$ (6)

where $x,y,z\in\mathcal{C}([0,T],\mathbb{R})$ are continuous real signals on
time interval $[0,T]$. We denote
$H(x,z,{\bm{A}},{\bm{B}},{\bm{E}},\bm{h}_{0},t)=H_{\beta}(x,z,\bm{h}_{0},t)$,
where $\beta$ denotes trainable parameters ${\bm{A}},{\bm{B}},{\bm{E}}$, as
the deterministic function mapping from signals $x,z$ to $\bm{h}_{t}$, the
state at time $t$, given initial state $\bm{h}_{0}$ at time $0$. The above SSM
can be compactly represented by

$\displaystyle
y_{t}={\bm{C}}H_{\beta}(x,z,\bm{h}_{0},t)+{\bm{D}}x_{t}+{\bm{F}}z_{t}$ (7)

When the continuous-time input signals are discretized into discrete-time
sequences $(x_{t_{0}},\dots,x_{t_{L-1}})$ and $(z_{t_{0}},\dots,z_{t_{L}})$,
the corresponding hidden state at time $t_{k}$ has a convolutional view
(assuming ${\bm{D}}={\bm{F}}=\bm{0}$ for simplicity)

$\displaystyle\begin{split}&y_{t_{k}}={\bm{C}}{\bm{K}}_{t_{k}}*x_{t_{k}}+{\bm{C}}\hat{\bm{K}}_{t_{k}}*z_{t_{k}},\text{
where
}\;{\bm{K}}_{t_{k}}=\bar{{\bm{A}}}^{k}\bar{{\bm{B}}},~{}~{}\hat{\bm{K}}_{t_{k}}=\bar{{\bm{A}}}^{k}\bar{{\bm{E}}}\end{split}$
(8)

which can be evaluated efficiently using FFT. Additionally, ${\bm{A}}$ is
HiPPO-initialized (Gu et al., 2021) for all such SSM blocks.

####  4.1 Latent Space as Structured State Space

The goal of the prior model is to realize a sequence of random variables
$(z_{t_{0}},z_{t_{1}},\dots,z_{t_{L}})$, which the prior distribution
$p_{\lambda}(z_{\leq t_{L}})$ models autoregressively. Suppose
$(z_{t_{0}},z_{t_{1}},\dots,z_{t_{n}})$ is a sequence of latent variables up
to time $t_{n}$, we define the prior distribution of $z_{t_{n}}$
autoregressively as

$\displaystyle p_{\lambda}(z_{t_{n}}\mid
z_{<t_{n}})={\mathcal{N}}(\mu_{z,n}(z_{<t_{n}},\lambda),\text{$\sigma_{z,n}^{2}(z_{<t_{n}},\lambda)$})$
(9)

where the mean $\mu_{z,n}$ and standard deviation $\sigma_{z,n}$ are
deterministic functions of previously generated $z_{<t_{n}}$ parameterized by
$\lambda$. To parameterize the above distribution, we first define an
intermediate building block, a stack of which will produce the wanted
distribution.

LS4 prior block. The forward pass through our SSM is a two–step procedure:
first, we consider the latent dynamics of $z$ on $[t_{0},t_{n-1}]$ where we
simply leverage Equation 6 to define the hidden states to follow
$H_{\beta_{1}}(0,z,0,t)$. Second, on $(t_{n-1},t_{n}]$, since no additional
$z$ is available in this interval, we ignore additional input signals in the
ODE and only include the last given latent, _i.e_. $z_{t_{n-1}}$, as an
auxiliary signal for the outputs, which can be compactly denoted, with a final
GELU non-linearity, as

$\displaystyle\begin{split}y_{z,n}=\text{GELU}({\bm{C}}_{y_{z}}H_{\beta_{1}}(0,0,\underbrace{H_{\beta_{2}}(0,z_{[t_{0},t_{n-1}]},h_{t_{n-1}},\mathbf{0},t_{n-1})},t_{n})+{\bm{F}}_{y_{z}}z_{t_{n-1}})\end{split}$
(10)

We call this LS4 prior layer, which we use to build our LS4 prior block that
is built upon a ResNet structure with a skip connection, denoted as

$\displaystyle\begin{split}\text{LS4{}}_{\text{prior}}(z_{[t_{0},t_{n-1}]},\psi)=\text{LayerNorm}({\bm{G}}_{y_{z}}y_{z,n}+b_{y_{z}})+z_{t_{n-1}}\end{split}$
(11)

where $\psi$ denotes the union of parameters
$\beta_{i},{\bm{C}}_{y_{z}},{\bm{F}}_{y_{z}},{\bm{G}}_{y_{z}},b_{y_{z}}$. We
define the final parameters $\mu_{z,n}$ and $\sigma_{z,n}$ for the conditional
distribution in the autoregressive model as the result of a stack of LS4 prior
blocks. During generation, as an initial condition,
$z_{t_{0}}\sim{\mathcal{N}}(\mu_{z,0},\sigma_{z,0}^{2})$ where
$\mu_{z,0},\sigma_{z,0}$ are learnable parameters, and subsequent latent
variables are generated autoregressively. We specify our architecture in
Appendix C and use $\lambda$ to denote the union of all trainable parameters.

####  4.2 Generative Model

Given the latent variables, we now specify a decoder that represents the
distribution $p_{\theta}(x_{\leq t_{L}}|z_{\leq t_{L}})$. Suppose $z_{\leq
t_{L}}$ is a latent path generated via the latent state space model, the
output path $x_{\leq t_{L}}$ also follows the state space formulation.
Assuming we have generated $(x_{t_{0}},\dots,x_{t_{n-1}})$ and
$(z_{t_{0}},\dots,z_{t_{n}})$, the conditional distribution of $x_{t_{n}}$ is
parametrized as

$\displaystyle p_{\theta}(x_{t_{n}}|x_{<t_{n}},z_{\leq
t_{n}})={\mathcal{N}}(\mu_{x,n}(x_{<t_{n}},z_{\leq
t_{n}},\theta),\sigma_{x}^{2})$ (12)

where $\sigma_{x}$ is a pre-defined observation standard deviation and
$\mu_{x,n}$ is a deterministic function of $z_{\leq t_{n}}$ and $x_{<x_{n}}$.

LS4 generative block. Different from the prior block, both observation and
latent sequences are input into our model, and we define intermediate outputs
$g_{x,n}$ and $g_{z,n}$ as

$\displaystyle\begin{split}&h_{t_{n}}=H_{\beta_{3}}(0,z_{t_{n-1}},H_{\beta_{4}}(x_{[t_{0},t_{n-1}]},z_{[t_{0},t_{n-1}]},\mathbf{0},t_{n-1}),t_{n})\\\
&g_{x,n}=\text{GELU}({\bm{C}}_{g_{x}}h_{t_{n}}+{\bm{D}}_{g_{x}}x_{t_{n-1}}+{\bm{F}}_{g_{x}}z_{t_{n}})\\\
&g_{z,n}=\text{GELU}({\bm{C}}_{g_{z}}h_{t_{n}}+{\bm{D}}_{g_{z}}x_{t_{n-1}}+{\bm{F}}_{g_{z}}z_{t_{n}})\end{split}$
(13)

which are used to build a LS4 generative block defined as

$\displaystyle\begin{split}&\hat{g}_{x,n}=\text{LayerNorm}({\bm{G}}_{g_{x}}g_{x,n}+b_{g_{x}})+x_{t_{n-1}}\\\
&\hat{g}_{z,n}=\text{LayerNorm}({\bm{G}}_{g_{z}}g_{z,n}+b_{g_{z}})+z_{t_{n}}\\\
&\text{LS4{}}_{\text{gen}}(x_{[t_{0},t_{n-1}]},z_{[t_{0},t_{n}]},\psi)=(\hat{g}_{x,n},\hat{g}_{z,n})\end{split}$
(14)

where $\psi$ denotes all parameters inside the block. Note that the generative
block gives two streams of outputs, which can be used as inputs for the next
stack. We then define the final mean value $\mu_{x,n}$ as the result of a
stack of LS4 generative blocks. The initial condition for generation is given
as $x_{t_{0}}\sim{\mathcal{N}}(\mu_{x,0}(z_{0},\theta),\sigma_{x})$ where
$\mu_{x,0}$ exactly follows our formulation while taking only $z_{t_{0}}$ as
input. The subsequent $x_{t_{n}}$’s are generated autoregressively. We specify
our architecture in Appendix C and use $\theta$ to denote the union of all
trainable parameters.

####  4.3 Inference model

The latent variable model up to time $t_{n}$ has intractable posterior
$p_{\theta}(z_{\leq t_{n}}\mid x_{\leq t_{n}})$. Therefore, we approximate
this distribution with $q_{\phi}(z_{\leq t_{n}}\mid x_{\leq t_{n}})$ using
variational inference.

We parameterize the inference distribution at time $t_{n}$ to depend only on
the observed path $x_{\leq t_{n}}$:

$\displaystyle q_{\phi}(z_{t}\mid x_{\leq
t_{n}})={\mathcal{N}}(\hat{\mu}_{z,t_{n}}(x_{\leq
t_{n}},\phi),\text{$\hat{\sigma}_{z,t_{n}}^{2}(x_{\leq t_{n}},\phi)$})$ (15)

LS4 inference block. The inference block is defined as

$\displaystyle\begin{split}&\hat{y}_{z,n}=\text{GELU}({\bm{C}}_{\hat{y}_{z}}H_{\beta_{5}}(x_{[t_{0},t_{n}]},0,\mathbf{0},t_{n-1})+{\bm{D}}_{\hat{y}_{z}}x_{t_{n}})\\\
&\text{LS4{}}_{\text{inf}}(x_{[t_{0},t_{n}]},\psi)=\text{LayerNorm}({\bm{G}}_{\hat{y}_{z}}\hat{y}_{z,n}+b_{\hat{y}_{z}})+x_{t_{n}}\end{split}$
(16)

Notice that input $x$ is fully present in $[t_{0},t_{n}]$ unlike in the
generative model. $\hat{\mu}_{z,t}$ and $\hat{\sigma}_{z,t}$ are each a deep
stack of our inference blocks similarly defined as before. Here we also
justify our choice of factorizing posterior as a product of marginals as in
Equation 5. By having each $z_{t}$ explicitly depending on $x_{\leq t_{n}}$
only, we can leverage the fast convolution operation to obtain all $z_{t}$ in
parallel, thus achieving fast inference time, unlike the autoregressive nature
of the prior and generative model.

####  4.4 LS4: Properties and Practice

We highlight some properties of LS4. In particular, we compare in the
following proposition the expressive power of our generative model against
structured state space models.

######  Proposition 4.1.

(LS4 subsumes S4.) Given any autoregressive model $r(x)$ with conditionals
$r(x_{n}|x_{<n})$ parameterized via deep S4 models, there exists a choice of
$\theta,\lambda,\phi$ such that $p_{\theta,\lambda}(x)=r(x)$ and
$p_{\theta,\lambda}(z|x)=q_{\phi}(z|x)$, i.e. the variational lower bound
(ELBO) is tight.

A proof sketch is provided in Appendix B. This result shows that LS4 subsumes
autoregressive generative models based on vanilla S4 (Gu et al., 2021), given
that the architecture between SSM layers is the same. Crucially, with the
assumption that we are able to globally optimize the ELBO training objective,
LS4 will fit the data at least as well as vanilla S4.

Scaling to arbitrary feature dimensions. So far we have assumed the input and
latent signals are real numbers. The approach can be scaled to arbitrary
dimensions of inputs and latents by constructing LS4 layers for each dimension
which are input into a mixing linear layer. We call such parallelized SSMs
heads and provide a pseudo-code in Appendix C.

######  Proposition 4.2.

(Efficiency.) For a SSM with $H$ heads, an observation sequence of length $L$
and hidden dimension $N$ can be calculated in $\mathcal{O}(H(L+N)\log(L+N))$
time and $\mathcal{O}(H(L+N))$ space.

We provide proof in Appendix B. Note that our model is much more efficient in
both time and space than RNN/ODE-based methods (which requires
$\mathcal{O}(N^{2}L)$ in time and $\mathcal{O}(NL)$ in space as discussed in
Section 3.1). To demonstrate the computation efficiency, we additionally
provide below pseudo-code for a single LS4 prior layer 10. The other blocks
can be similarly constructed.

⬇

def LS4_prior_layer(z, A, B, C, F, h_0): # z: (B, L, 1)

K = C @ materialize_kernel(z, A, B, h_0) # O((L+N)log(L+N)) time

CH = fft_conv(K, z) # O(LlogL) time and O(L) space

y = gelu(CH + F * z)

return y

Note that in practice, ${\bm{A}}$ is HiPPO initialized (Gu et al., 2020) and
the materialized kernel includes ${\bm{C}}$ so that the convolution is
computed directly in the projected space, bypassing materializing the high-
dimensional hidden states.

###  5 Experiments

In this section, we verify the modeling capability of LS4 empirically. There
are three main questions we seek to answer: (1) How effective is LS4 in
modeling stiff sequence data? (2) How expressive is LS4 in scaling to real
time-series with a variety of temporal dynamics? (3) How efficient in training
and inference is LS4 in terms of wall-clock time?

####  5.1 Learning to generate data from stiff systems

Modeling data generated by dynamics with widely separated time scales has been
proven to be particularly challenging for vanilla ODE-based approaches which
make use of standard explicit solvers for inference and gradient calculation.
Kim et al. (2021) showed that as the learned dynamics stiffen up to track data
paths, the ODE numerics start to catastrophically fail; the inference cost
raises drastically and the gradient estimation process becomes
ill–conditioned. These issues can be mitigated by employing implicit ODE
solvers or ad-hoc rescalings of the learned vector field (see (Kim et al.,
2021) for further details).

In turn, state–space models do not suffer from stiffness of dynamics as the
numerical methods are sidestepped in favor of an exact evaluation of the
convolution operator. We hereby show that LS4 is able to model data generated
by a prototype stiff system.

FLAME problem We consider a simple model of flame growth (FLAME) (Wanner &
Hairer, 1996), which has been extensively studied as a representative of
highly stiff systems:

$\frac{\differential}{\differential t}x_{t}=x_{t}^{2}-x_{t}^{p}$

where $p\in\\{3,4,\dots,10\\}$. For each $p$, $1000$ trajectories are
generated for $t\in[0,1000]$ with unit increment.

Generation. In Figure 1(a), we show the mean trajectories and the distribution
at each time step and that our samples closely match the ground-truth data.
The Latent ODE (Rubanova et al., 2019) instead fails to do so and produces
non–stiff samples drastically different from the target.

(a) Generation of the stiff system.

(b) Marginal histograms at steps equally spaced between the $0.5\%$ and $10\%$
steps in log scale.

Marginal Distribution. We plot the marginal distribution of the real data and
the generated data from both our model and Latent ODE. Since the stiff
transitions are mostly distributed before $10\%$ of total steps, we visualize
the marginal histograms at $4$ time steps equally spaced between the $0.5\%$
and $10\%$ steps in log scale (see Figure 1(b)). We observe that the empirical
histogram matches the ground truth distribution significantly better than what
is produced by Latent ODEs, as also qualitatively visible from the samples in
(a).

####  5.2 Generation with Real Time-Series Datasets

We investigate the generative capability of LS4 on real time-series data. We
show that our model is capable of fitting a wide variety of time-series data
with significantly different temporal dynamics. We leave implementation
details to Appendix D.1.

Data. We use Monash Time Series Repository (Godahewa et al., 2021), a
comprehensive benchmark containing $30$ time-series datasets collected in the
real world, and we choose FRED-MD, NN5 Daily, Temperature Rain, and Solar
Weekly as our target datasets. Each dataset exhibits unique temporal dynamics,
which makes generative learning a challenging task. A sample from each dataset
can be visualized in Figure 2.

Figure 2: Monash data. The selected datasets exhibit a variety of temporal
dynamics ranging from relatively smooth to stiff transitions.

Metrics. We propose 3 different metrics for measuring generation performance,
namely Marginal, Classification, and Prediction scores. Marginal scores
calculate the absolute difference between empirical probability density
functions of two distributions – the lower the better (Ni et al., 2020).
Following Kidger et al. (2021), we define Classification scores as using a
sequence model to classify whether a sample is real or generated and use its
cross-entropy loss as a proxy for generation quality – the higher the less
distinguishable the samples, thus better the generation. Similarly, Prediction
scores use a train-on-synthetic-test-on-real seq2seq model to predict $k$
steps into the future – the lower the more predictable, thus better the
generation. We use a 1-layer S4 (Gu et al., 2021) for both Classification and
Prediction scores (see Appendix D.1 for more details). We will discuss the
necessity of all 3 metrics in the following discussion section.

We compare our model with several time-series generative models, namely the
VAE-based methods such as RNN-VAE (Rubanova et al., 2019), GP-VAE (Fortuin et
al., 2020), ODE2VAE (Yildiz et al., 2019), Latent ODE (Rubanova et al., 2019),
GAN-based methods such as TimeGAN (Yoon et al., 2019) and SDE GAN (Kidger et
al., 2021), and SaShiMi (Goel et al., 2022). We show quantitative results in
Table 1.

Data | Metric | RNN-VAE | GP-VAE | ODE2VAE | Latent ODE | TimeGAN | SDEGAN | SaShiMi | LS4 (Ours)
---|---|---|---|---|---|---|---|---|---
FRED-MD | Marginal $\downarrow$ | 0.132 | 0.152 | 0.122 | 0.0416 | 0.0813 | 0.0841 | 0.0482 | 0.0221
| Class. $\uparrow$ | 0.0362 | 0.0158 | 0.0282 | 0.327 | 0.0294 | 0.501 | 0.00119 | 0.544
| Prediction $\downarrow$ | 1.47 | 2.05 | 0.567 | 0.0132 | 0.0575 | 0.677 | 0.232 | 0.0373
NN5 Daily | Marginal $\downarrow$ | 0.137 | 0.117 | 0.211 | 0.107 | 0.0396 | 0.0852 | 0.0199 | 0.00671
| Class. $\uparrow$ | 0.000339 | 0.00246 | 0.00102 | 0.000381 | 0.00160 | 0.0852 | 0.0446 | 0.636
| Prediction $\downarrow$ | 0.967 | 1.169 | 1.19 | 1.04 | 1.34 | 1.01 | 0.849 | 0.241
Temp Rain | Marginal $\downarrow$ | 0.0174 | 0.183 | 1.831 | 0.0106 | 0.498 | 0.990 | 0.758 | 0.0834
| Class. $\uparrow$ | 0.00000212 | 0.0000123 | 0.0000319 | 0.0000419 | 0.00271 | 0.0169 | 0.0000167 | 0.976
| Prediction $\downarrow$ | 159 | 2.305 | 1.133 | 145 | 1.96 | 2.46 | 2.12 | 0.521
Solar Weekly | Marginal $\downarrow$ | 0.0903 | 0.308 | 0.153 | 0.0853 | 0.0496 | 0.147 | 0.173 | 0.0459
| Class. $\uparrow$ | 0.0524 | 0.000731 | 0.0998 | 0.0521 | 0.6489 | 0.591 | 0.00102 | 0.683
| Prediction $\downarrow$ | 1.25 | 1.47 | 0.761 | 0.973 | 0.237 | 0.976 | 0.578 | 0.141

Table 1: Generation results on FRED-MD, NN5 Daily, Temperature Rain, and Solar
Weekly.

Our model significantly outperforms the baselines on all datasets. We note
that baseline models have a hard time modeling NN5 Daily and Temperature Rain
where the transition dynamics are stiff. For Temperature Rain where most data
points lie around $x$-axis with sharp spikes throughout, latent ODE generates
mostly closely to the $x$-axis, thus achieving lower marginal scores, but its
generation is easily distinguishable from data, thus making it a less
favorable generative model. We demonstrate that Marginal scores alone are an
insufficient metric for generation quality. SaShiMi, an autoregressive model
based on S4, does not perform as well on time series generation in the tasks
considered. We further discuss the reason in Appendix D.1.

####  5.3 Interpolation & Extrapolation

We also show that our model is expressive enough to fit to irregularly-sampled
data and performs favorably in terms of interpolation and extrapolation.
Interpolation refers to the task of predicting missing data given a subset of
a sequence while extrapolation refers to the task that data is separated into
2 segments each with half the length of the full sequence, and one predicts
the latter segment given the former.

Data. Following Rubanova et al. (2019); Schirmer et al. (2022), we use USHCN
and Physionet as our datasets of choice. The United States Historical
Climatology Network (USHCN) (Menne et al., 2015) is a climate dataset
containing daily measurements form 1,218 weather stations across the US for
precipitation, snowfall, snow depth, minimum and maximum temperature.
Physionet (Silva et al., 2012) is a dataset containing health measurements of
41 signals from 8,000 ICU patients in their first 48 hours. We follow
preprocessing steps of Schirmer et al. (2022) for training and testing.

Metrics. We use mean squared error (MSE) to evaluate both interpolation and
extrapolation.

We compare our model with RNN (Rumelhart et al., 1985), RNN-VAE (Chung et al.,
2014; Rubanova et al., 2019), ODE-RNN (Rubanova et al., 2019), GRU-D (Rubanova
et al., 2019), Latent ODE (Chen et al., 2018; Rubanova et al., 2019), and CRU
Schirmer et al. (2022). Results are shown in Table 2.

Task | Data | RNN | RNN-VAE | ODE-RNN | GRU-D | Latent ODE | CRU 222Numbers are taken from the original paper. We keep the significant digits unchanged | LS4 (Ours)
---|---|---|---|---|---|---|---|---
Interp. | Physionet | 2.918 | 5.930 | 2.234 | 3.325 | 8.341 | 1.82 | 0.6287
USHCN | 4.322 | 7.561 | 2.473 | 3.395 | 6.859 | 0.16 | 0.0594
Extrap. | Physionet | 3.406 | 3.064 | 3.014 | 3.120 | 4.212 | 6.29 | 4.942
USHCN | 9.474 | 9.083 | 9.045 | 8.964 | 8.959 | 12.73 | 2.975

Table 2: Interpolation and extrapolation MSE $(\times 10^{-3})$ scores. Lower
scores are better.

We observe that our model outperforms previous continuous-time methods. Our
model performs less well on extrapolation for Physionet compared to ODE-RNN
and latent ODE. We postulate that this is due to the high variability granted
by our latent space. Since new latent variables are generated as we
extrapolate, our model generates paths that are more flexible (hence less
predictable) than those of Latent ODE, which instead uses a single latent
variable to encode an entire path. We additional observe that our model
achieves ELBO of $-669.0$ and $-250.2$ on Physionet interpolation and
extrapolation tasks respectively. These bounds are much tighter lower bounds
than other VAE-based methods, _i.e_. RNN-VAE, which reports $-412.8$ and
$-220.2$, and latent ODE, which reports $-410.3$ and $-168.5$. We leave
additional ELBO comparisons in Appendix D.

####  5.4 Runtime

We additionally verify the computational efficiency of our model for both
training and inference. We do so by training and inferring on synthetic data
with controlled lengths specified below.

Data. We create a set of synthetic datasets with lengths
$\\{80,320,1280,5120,20480\\}$ to investigate scaling of training/inference
time with respect to sequence length. Training is done with 100 iterations
through the synthetic data, and inference is performed given one batch of
synthetic data (see Appendix D.3).

Metrics. We use wall-clock runtime measured in milliseconds.

Figure 3: Runtime comparison. The $y$-axis shows run-time (ms) of each setting
in log scale. Our runtime stays consistently lower across all sequence lengths
investigated.

Figure 3 shows model runtime in log scale. ODE2VAE fails to finish training on
the last sequence length within a reasonable time frame, so we omit its plot
of the last data point. Our model performs consistently and significantly
lower than baselines, which are observed to scale linearly with input lengths,
and is $\times 100$ faster than baselines in both training and inference on
$20480$ length.

###  6 Conclusion

We introduce LS4, a powerful generative model with latent space evolution
following a state space ODE. Our model is built with a deep stack of LS4
prior/generative/inference blocks, which are trained via standard sequence VAE
objectives. We also show that under specific choices of model parameters, LS4
subsumes autoregressive S4 models. Experimentally, we demonstrate the modeling
power of LS4 on datasets with a wide variety of temporal dynamics and show
significant improvement in generation/interpolation/extrapolation quality. In
addition, our model shows $\times 100$ speed-up in training and inference time
on long sequences. LS4 demonstrates improved expressivity and computational
efficiency, and we believe that it has a further role to play in modeling
general time-series sequences.

### References

  * Bayer & Osendorfer (2014) Justin Bayer and Christian Osendorfer.  Learning stochastic recurrent networks.  _arXiv preprint arXiv:1411.7610_ , 2014.
  * Brooks et al. (2022) Tim Brooks, Janne Hellsten, Miika Aittala, Ting-Chun Wang, Timo Aila, Jaakko Lehtinen, Ming-Yu Liu, Alexei A Efros, and Tero Karras.  Generating long videos of dynamic scenes.  _arXiv preprint arXiv:2206.03429_ , 2022.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Burda et al. (2015) Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov.  Importance weighted autoencoders.  _arXiv preprint arXiv:1509.00519_ , 2015.
  * Chen (1984) Chi-Tsong Chen.  _Linear system theory and design_.  Saunders college publishing, 1984.
  * Chen et al. (2018) Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud.  Neural ordinary differential equations.  _Advances in neural information processing systems_ , 31, 2018.
  * Chilkuri & Eliasmith (2021) Narsimha Reddy Chilkuri and Chris Eliasmith.  Parallelizing legendre memory unit training.  In _International Conference on Machine Learning_ , pp. 1898–1907. PMLR, 2021.
  * Chung et al. (2014) Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio.  Empirical evaluation of gated recurrent neural networks on sequence modeling.  _arXiv preprint arXiv:1412.3555_ , 2014.
  * Chung et al. (2015) Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Bengio.  A recurrent latent variable model for sequential data.  _Advances in neural information processing systems_ , 28, 2015.
  * de Bézenac et al. (2020) Emmanuel de Bézenac, Syama Sundar Rangapuram, Konstantinos Benidis, Michael Bohlke-Schneider, Richard Kurle, Lorenzo Stella, Hilaf Hasson, Patrick Gallinari, and Tim Januschowski.  Normalizing kalman filters for multivariate time series analysis.  _Advances in Neural Information Processing Systems_ , 33:2995–3007, 2020.
  * De Brouwer et al. (2019) Edward De Brouwer, Jaak Simm, Adam Arany, and Yves Moreau.  Gru-ode-bayes: Continuous modeling of sporadically-observed time series.  _Advances in neural information processing systems_ , 32, 2019.
  * Fabius & Van Amersfoort (2014) Otto Fabius and Joost R Van Amersfoort.  Variational recurrent auto-encoders.  _arXiv preprint arXiv:1412.6581_ , 2014.
  * Fortuin et al. (2020) Vincent Fortuin, Dmitry Baranchuk, Gunnar Rätsch, and Stephan Mandt.  Gp-vae: Deep probabilistic time series imputation.  In _International conference on artificial intelligence and statistics_ , pp. 1651–1661. PMLR, 2020.
  * Franceschi et al. (2020) Jean-Yves Franceschi, Edouard Delasalles, Mickaël Chen, Sylvain Lamprier, and Patrick Gallinari.  Stochastic latent residual video prediction.  In _International Conference on Machine Learning_ , pp. 3233–3246. PMLR, 2020.
  * Godahewa et al. (2021) Rakshitha Godahewa, Christoph Bergmeir, Geoffrey I. Webb, Rob J. Hyndman, and Pablo Montero-Manso.  Monash time series forecasting archive.  In _Neural Information Processing Systems Track on Datasets and Benchmarks_ , 2021.
  * Goel et al. (2022) Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré.  It’s raw! audio generation with state-space models.  _arXiv preprint arXiv:2202.09729_ , 2022.
  * Gu et al. (2020) Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré.  Hippo: Recurrent memory with optimal polynomial projections.  _Advances in Neural Information Processing Systems_ , 33:1474–1487, 2020.
  * Gu et al. (2021) Albert Gu, Karan Goel, and Christopher Ré.  Efficiently modeling long sequences with structured state spaces.  _arXiv preprint arXiv:2111.00396_ , 2021.
  * Gu et al. (2022) Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré.  On the parameterization and initialization of diagonal state space models.  _arXiv preprint arXiv:2206.11893_ , 2022.
  * Hersbach et al. (2020) Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz-Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers, et al.  The era5 global reanalysis.  _Quarterly Journal of the Royal Meteorological Society_ , 146(730):1999–2049, 2020.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Hochreiter (1998) Sepp Hochreiter.  The vanishing gradient problem during learning recurrent neural nets and problem solutions.  _International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems_ , 6(02):107–116, 1998.
  * Kidger et al. (2020) Patrick Kidger, James Morrill, James Foster, and Terry Lyons.  Neural controlled differential equations for irregular time series.  _Advances in Neural Information Processing Systems_ , 33:6696–6707, 2020.
  * Kidger et al. (2021) Patrick Kidger, James Foster, Xuechen Li, and Terry J Lyons.  Neural sdes as infinite-dimensional gans.  In _International Conference on Machine Learning_ , pp. 5453–5463. PMLR, 2021.
  * Kim et al. (2021) Suyong Kim, Weiqi Ji, Sili Deng, Yingbo Ma, and Christopher Rackauckas.  Stiff neural ordinary differential equations.  _Chaos: An Interdisciplinary Journal of Nonlinear Science_ , 31(9):093122, 2021.
  * Kingma & Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Li et al. (2020) Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud.  Scalable gradients for stochastic differential equations.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 3870–3882. PMLR, 2020.
  * Massaroli et al. (2020) Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama.  Dissecting neural odes.  _Advances in Neural Information Processing Systems_ , 33:3952–3963, 2020.
  * Massaroli et al. (2021) Stefano Massaroli, Michael Poli, Sho Sonoda, Taiji Suzuki, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama.  Differentiable multiple shooting layers.  _Advances in Neural Information Processing Systems_ , 34:16532–16544, 2021.
  * Menne et al. (2015) MJ Menne, CN Williams Jr, and RS Vose.  Long-term daily climate records from stations across the contiguous united states, 2015.
  * Ni et al. (2020) Hao Ni, Lukasz Szpruch, Magnus Wiese, Shujian Liao, and Baoren Xiao.  Conditional sig-wasserstein gans for time series generation.  _arXiv preprint arXiv:2006.05421_ , 2020.
  * Niesen & Hall (2004) Jitse Niesen and Trinity Hall.  _On the global error of discretization methods for ordinary differential equations_.  PhD thesis, Citeseer, 2004.
  * Oord et al. (2016) Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu.  Wavenet: A generative model for raw audio.  _arXiv preprint arXiv:1609.03499_ , 2016.
  * Oppenheim (1999) Alan V Oppenheim.  _Discrete-time signal processing_.  Pearson Education India, 1999.
  * Oppenheim & Schafer (1975) Alan V Oppenheim and Ronald W Schafer.  Digital signal processing(book).  _Research supported by the Massachusetts Institute of Technology, Bell Telephone Laboratories, and Guggenheim Foundation. Englewood Cliffs, N. J., Prentice-Hall, Inc., 1975. 598 p_ , 1975.
  * Peng et al. (1995) C-K Peng, Shlomo Havlin, H Eugene Stanley, and Ary L Goldberger.  Quantification of scaling exponents and crossover phenomena in nonstationary heartbeat time series.  _Chaos: an interdisciplinary journal of nonlinear science_ , 5(1):82–87, 1995.
  * Poli et al. (2019) Michael Poli, Jinkyoo Park, and Ilija Ilievski.  Wattnet: learning to trade fx via hierarchical spatio-temporal representation of highly multivariate time series.  _arXiv preprint arXiv:1909.10801_ , 2019.
  * Rasul et al. (2021) Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf.  Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting.  In _International Conference on Machine Learning_ , pp. 8857–8868. PMLR, 2021.
  * Romero et al. (2021) David W Romero, Anna Kuzina, Erik J Bekkers, Jakub M Tomczak, and Mark Hoogendoorn.  Ckconv: Continuous kernel convolution for sequential data.  _arXiv preprint arXiv:2102.02611_ , 2021.
  * Rubanova et al. (2019) Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud.  Latent ordinary differential equations for irregularly-sampled time series.  _Advances in neural information processing systems_ , 32, 2019.
  * Rumelhart et al. (1985) David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.  Learning internal representations by error propagation.  Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
  * Schirmer et al. (2022) Mona Schirmer, Mazin Eltayeb, Stefan Lessmann, and Maja Rudolph.  Modeling irregular time series with continuous recurrent units.  In _International Conference on Machine Learning_ , pp. 19388–19405. PMLR, 2022.
  * Shampine & Thompson (2007) Lawrence F Shampine and Skip Thompson.  Stiff systems.  _Scholarpedia_ , 2(3):2855, 2007.
  * Silva et al. (2012) Ikaro Silva, George Moody, Daniel J Scott, Leo A Celi, and Roger G Mark.  Predicting in-hospital mortality of icu patients: The physionet/computing in cardiology challenge 2012.  In _2012 Computing in Cardiology_ , pp. 245–248. IEEE, 2012.
  * Sohl-Dickstein et al. (2015) Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In _ICML_ , 2015.
  * Song & Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _NeurIPS_ , 32, 2019.
  * Song et al. (2020) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020.
  * Wanner & Hairer (1996) Gerhard Wanner and Ernst Hairer.  _Solving ordinary differential equations II_ , volume 375.  Springer Berlin Heidelberg New York, 1996.
  * Yildiz et al. (2019) Cagatay Yildiz, Markus Heinonen, and Harri Lahdesmaki.  Ode2vae: Deep generative second order odes with bayesian neural networks.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Yoon et al. (2019) Jinsung Yoon, Daniel Jarrett, and Mihaela Van der Schaar.  Time-series generative adversarial networks.  _Advances in neural information processing systems_ , 32, 2019.
  * Yu et al. (2022) Sihyun Yu, Jihoon Tack, Sangwoo Mo, Hyunsu Kim, Junho Kim, Jung-Woo Ha, and Jinwoo Shin.  Generating videos with dynamics-aware implicit generative adversarial networks.  _arXiv preprint arXiv:2202.10571_ , 2022.
  * Zhuang et al. (2020) Juntang Zhuang, Nicha Dvornek, Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Duncan.  Adaptive checkpoint adjoint method for gradient estimation in neural ode.  In _International Conference on Machine Learning_ , pp. 11639–11649. PMLR, 2020.

Time-Series Generation via Deep Latent State Space Models

## Part I

###  Appendix A Derivations

####  A.1 Computational Complexity of Vanilla Recurrent Representation

Assuming recurrence of the simplest form in Equation 2, fulling computing
matrix multiplication $\bar{{\bm{A}}}\bar{h}_{t_{k}}$ requires
$\mathcal{O}(N^{2})$. Fully computing all hidden states sequentially requires
$\mathcal{O}(N^{2}L)$. In space, saving each hidden state requires
$\mathcal{O}(N)$ and in total requires $\mathcal{O}(NL)$.

###  Appendix B Proof

To prove this result in Proposition 4.1, we first prove the following
proposition.

######  Proposition B.1.

(Expressivity.) Given any deep autoregressive S4 model
$r:(x_{<t_{n}},t_{n})\mapsto y_{t_{n}}$ evaluated at time $t_{n}$ given input
sequence $x_{<t_{n}}$, there exists a choice of $\theta$ such that
$\mu_{x,n}(x_{<t_{n}},0,\theta)=r(x_{<t_{n}},t_{n})$.

###### Proof sketch..

Consider SSM of the form in Equation 6 as a building block to our generative
model with parameter $\theta$. We can choose ${\bm{E}}={\bm{F}}=0$ for all
layers, which exactly reduces it to the SSM of an S4 model. Keeping all other
hyperparameters (_e.g_. non-linearities, number of stacking layers) the same,
the final model is exactly the same as a deep autoregressive S4 model. ∎

Now we give a proof sketch to Proposition 4.1, See 4.1

###### Proof sketch..

From Proposition B.1 we know that we can choose $\theta$ so that
$p_{\theta}(x|z)=p_{\theta}(x)=r(x)$ for all $z$, i.e., choose a decoder that
ignores the latent variables $z$ and uses the same autoregressive structure
over the observed variables as $r(x)$. This implies the posterior
$p_{\theta,\lambda}(z|x)$ is equal to the prior $p_{\lambda}(z)$. We can then
choose $\lambda$ and $\phi$ so that $p_{\lambda}(z)=\mathcal{N}(0,I)$ and
$q_{\phi}(z|x)=\mathcal{N}(0,I)$ for all $x$. ∎

See 4.2

###### Proof.

Recall SISO SSM of the form

$\displaystyle\begin{split}\frac{d}{dt}\bm{h}_{t}&={\bm{A}}\bm{h}_{t}+{\bm{B}}x_{t}\\\
y_{t}&={\bm{C}}\bm{h}_{t}+{\bm{D}}x_{t}\end{split}$ (17)

The calculation of $(y_{t_{0}},\dots,y_{t_{L}})$ involves materializing the
convolution filter, which can be calculated in $\mathcal{O}((L+N)\log(L+N))$
time and $\mathcal{O}(L+N)$ space for diagonal-plus-low-rank matrices (Gu et
al., 2021). Since the convolution is constant time in frequency domain,
another computation cost comes from Fast Fourier Transform (FFT) and its
inverse, which is $\mathcal{O}(L\log L)$ in time. The computation scales
linearly with heads, Thus, a multi-input-multi-output (MIMO) SSM with $H$
heads can be processed in $\mathcal{O}(H(L+N)\log(L+N))$ time and
$\mathcal{O}(H(L+N))$ space. ∎

###  Appendix C Architecture

We parametrize our models using a similar architecture as in Goel et al.
(2022), but there is no pooling operation because for general time-series the
time length is hardly divisible by a reasonable factor. Before we present the
full structure, we present how a multi-channel inputs are parametrized (in the
case of LS4 prior layer (10):

⬇

def LS4_prior_layer_multi(z, psi):

# z: (B, L, C)

for c in range(C):

z[:,:,c] = LS4_prior_layer(z[:,:,c], *psi.LS4_params)

z = linear(z) # (B, L, C) channel-wise mixing

return z

The for loop is presented for demonstration purposes. In practice, the
channels can be processed in parallel.

####  C.1 Prior Model

We specify the parametrization of $\mu_{z,n}$ and $\sigma_{z,n}$ in pseudo-
code as the outputs of the following functions

⬇

def prior_model(z, lambda):

# z: (B, L, z_dim) this is for time [t_0, t_{n-1}]

z = linear(z) # (B,L,H) encoding to H

outputs = []

outputs.append(z)

for i in range(lambda.num_layers1):

z = linear(z) # (B, L, H) -> (B, L, 2H)

outputs.append(z)

for i in range(lambda.num_layers2):

z = LS4_prior_block_multi(z, *lambda.LS4_params)

# (B, L, H) -> (B, L, H) multi-channel SSMs

z = ResBlock(z) # this is a general 1 layer residual block

z = z + outputs.pop()

for i in range(lambda.num_layers1):

z = z + outputs.pop()

outputs.append(z)

z = linear(z) # (B, L, 2H) -> (B, L, H)

for i in range(lambda.num_layers2):

z = LS4_prior_block_multi(z, *lambda.LS4_params)

# (B, L, H) -> (B, L, H) multi-channel SSMs

z = ResBlock(z) # this is a general 1 layer residual block

z = z + outputs.pop()

z = layernorm(z)

z = linear(z) # (B,L,z_dim)

mu_z = LS4_prior_block_multi(z, *lambda.LS4_params) # (B,L,z_dim)

sigma_z = LS4_prior_block_multi(z, *lambda.LS4_params)

# (B,L,z_dim)

return mu_z, sigma_z

####  C.2 Generative Model

⬇

def prior_model(x, z, theta):

# z: (B, L, z_dim) this is for time [t_0, t_{n-1}]

z = linear(z) # (B,L,H) encoding to H

x = linear(x) # (B,L,H) encoding to H

outputs_x, outputs_z = [], []

outputs_z.append(z)

outputs_x.append(x)

for i in range(lambda.num_layers1):

z = linear(z) # (B, L, H) -> (B, L, 2H)

x = linear(x) # (B, L, H) -> (B, L, 2H)

outputs_z.append(z)

outputs_x.append(x)

for i in range(lambda.num_layers2):

z, x = LS4_gen_block_multi(z, *lambda.LS4_params)

# (B, L, H) -> (B, L, H) multi-channel SSMs

zx = ResBlock(concat(z, x))

# this is a general 1 layer residual block

z, x = split(z, x)

z = z + outputs_z.pop()

x = x + outputs_x.pop()

for i in range(lambda.num_layers1):

z = z + outputs_z.pop()

x = x + outputs_x.pop()

outputs_z.append(z)

outputs_x.append(z)

z = linear(z) # (B, L, 2H) -> (B, L, H)

x = linear(x) # (B, L, 2H) -> (B, L, H)

for i in range(lambda.num_layers2):

x, z = LS4_gen_block_multi(x, z, *lambda.LS4_params)

# (B, L, H) -> (B, L, H) multi-channel SSMs

zx = ResBlock(concat(z, x))

# this is a general 1 layer residual block

z, x = split(z, x)

z = z + outputs_z.pop()

x = x + outputs_x.pop()

x = layernorm(x)

z = layernorm(z)

x = linear(concat(x, z)) # (B,L,2H) -> (B,L,x_dim)

return mu_x

In practice, we find that only using $z$ input for the entire generative model
produces better generation better than including $x$. We hypothesize that $x$
presents too strong of a signal for the model to reconstruct, and so the model
learns to ignore signals from $z$ in that case.

####  C.3 Inference Model

⬇

def inference_model(x, phi):

# x: (B, L, x_dim) this is for time [t_0, t_{n-1}]

x = linear(x) # (B,L,H) encoding to H

outputs = []

outputs.append(x)

for i in range(phi.num_layers1):

x = linear(x) # (B, L, H) -> (B, L, 2H)

outputs.append(x)

for i in range(phi.num_layers2):

x = LS4_inf_block_multi(x, *phi.LS4_params)

# (B, L, H) -> (B, L, H) multi-channel SSMs

z = ResBlock(z) # this is a general 1 layer residual block

x = x + outputs.pop()

for i in range(phi.num_layers1):

x = x + outputs.pop()

outputs.append(x)

x = linear(x) # (B, L, 2H) -> (B, L, H)

for i in range(phi.num_layers2):

x = LS4_inf_block_multi(x, *phi.LS4_params)

# (B, L, H) -> (B, L, H) multi-channel SSMs

z = ResBlock(z) # this is a general 1 layer residual block

x = x + outputs.pop()

x = layernorm(x)

x = linear(x) # (B,L,x_dim)

mu_z = LS4_inf_block_multi(x, *phi.LS4_params) # (B,L,x_dim)

sigma_z = LS4_inf_block_multi(x, *phi.LS4_params)

# (B,L,x_dim)

return mu_z, sigma_z

###  Appendix D Experiments

For all experiments we use AdamW optimizer with learning rate $0.001$. We use
batch size $64$ and train for $7000$ epochs for FRED-MD, NN5 Daily, and Solar
Weekly, $1000$ epochs for Temperature Rain, and $500$ epochs for Physionet and
USHCN. The datasets are split into $80\%$ training data and $20\%$ testing
data.

####  D.1 MONASH Forecasting Repository

Data. For all selected MONASH data, FRED-MD, NN5 Daily, and Solar Weekly are
normalized per sequence such that each trajectory is centered at its own mean
and normally distributed. We make this choice from the observation that for
some datasets such as NN5 Daily the min and max can vary significantly across
different data points such that normalizing sequences with dataset-wise
statistics makes it difficult to learn the temporal dynamics, which would be
on a widely different range. For Temperature Rain we squash each sequence into
$[0,1]$. This is due to the fact that the dataset is always positive and lands
mostly around $x$-axis with sharp spikes in between. For the former 3
datasets, we do not use output activation while for the last, we use sigmoid
as our activation.

Hyperparameters. For all MONASH experiments, we use AdamW optimizer with
learning rate 0.001 and no weight decay. For each of
prior/generative/inference model, we use 4 stacks for each for loop in the
pseudocode. For each LS4 block, we use 64 as the dimension of
${\mathbf{h}}_{t}$ and 64 SSM channels in parallel, same as used in S4 and
SaShiMi. Each residual block consists of 2 linear layers with skip connection
at the output level where the first linear layer has 2 times output size as
the input size and the second layer squeezes it back to the input size of the
residual block. We generally find 5-dimensional latent space gives better
performance than 1, and so uses this setting throughout. We also employ EMA
for model weights and use 0.999 as the lambda value, but we do not find this
choice crucial. We also use 0.1 as the standard deviation for the observation
as this gives better ELBO than other choices we experimented with such as
$1,0.5,0.01$. For baselines, we reuse the code from official repo and follow
their suggestions for training. To keep representation power similar, we use
the same size for the latent space (for latent variable models) and the same
output standard deviation for ELBO evaluation.

Evaluation. For generation evaluation. The classification model and the
prediction model uses a linear encoder and linear decoder with a single S4
layer in between. The S4 layer uses 16 hidden state dimensions. For
classification model, encoder maps data dimension to 16 hidden state
dimension, and averages over the sequence output from S4 layer before feeding
into decoder that outputs logit for binary classification. We use cross
entropy loss. For prediction model, we use the same linear encoder and a
decoder that maps 16 hidden dimension to data dimension. We predict $k=10$
steps into the future. The evaluation models are trained using AdamW with
$0.01$ learning rate for $100$ epochs with batch size $128$. We generate
samples equal to the number of testing data, which together are used to train
the two models.

Additional discussion. We also briefly discuss the surprising result that
SaShiMi does not perform as well on general time-series generation. We
speculate that not using a quantization scheme to define discrete output
conditionals, as standard in autoregressive models for e.g., audio and images,
is the cause behind this drop in performance. LS4 does not requires
quantization and sets best performance with a simple Gaussian conditional on
the data space.

Additional comparisons. We additionally compare with two more relevant
baselines (Fabius & Van Amersfoort, 2014) and (Li et al., 2020) present result
in Table 3. We note that Latent SDE has an abnormally high classification
score for Temperature Rain data, and demonstrate that this is when the
classification score is not reliable. Upon visually examining generated
results for Latent SDE (Figure 4) compared to ground-truths (Figure 7), one
can observe that the variation is extremely noisy around the x-axis and that
the selected classifier is not powerful enough to capture the distinction from
real data due to the considerable noise that exists in both generated and real
data, resulting in high classification score. Marginal and predictive scores
are much worse in comparison and are more indicative of generation quality.

Data | Metric | VRNN | Latent SDE | LS4 (Ours)
---|---|---|---|---
FRED-MD | Marginal $\downarrow$ | 0.165 | 0.122 | 0.0221
| Class. $\uparrow$ | 0.000970 | 0.687 | 0.544
| Prediction $\downarrow$ | 0.371 | 1.62 | 0.0373
NN5 Daily | Marginal $\downarrow$ | 0.151 | 0.125 | 0.00671
| Class. $\uparrow$ | 0.00176 | 0.601 | 0.636
| Prediction $\downarrow$ | 1.22 | 0.957 | 0.241
Temp Rain | Marginal $\downarrow$ | 1.20 | 0.999 | 0.0834
| Class. $\uparrow$ | 0.479 | 14.534 | 0.976
| Prediction $\downarrow$ | 0.864 | 1.798 | 0.521
Solar Weekly | Marginal $\downarrow$ | 0.297 | 0.234 | 0.0459
| Class. $\uparrow$ | 0.00164 | 0.764 | 0.683
| Prediction $\downarrow$ | 0.964 | 1.01 | 0.141
Table 3: Additional generation results on FRED-MD, NN5 Daily, Temperature
Rain, and Solar Weekly. Figure 4: Latent SDE generation on Temperature Rain.

####  D.2 Physionet & USHCN

We follow the code provided by Rubanova et al. (2019) to process Physionet and
follow the code provided by De Brouwer et al. (2019) for USHCN. For Physionet
we do not use any activation to constrain the output space, and for USHCN, we
use sigmoid activation for output.

We present the variational lowerbound results in Table 4.

Task | Data | RNN-VAE | Latent ODE | LS4 (Ours) | LS4IWAE (Ours)
---|---|---|---|---|---
Interp. | Physionet | -412.8 | -410.3 | -669.0 | -684.3
USHCN | -244.9 | -251.0 | -312.2 | -315.6
Extrap. | Physionet | -220.2 | -168.5 | -250.2 | -288.7
USHCN | -113.3 | -110.3 | -194.4 | -211.8
Table 4: ELBO comparisons with VAE-based models.

Hyperparameters. In general we keep the hyperparameter choices the same as in
MONASH, and we describe a few differences for these 2 datasets. For USHCN, we
use 10 as the dimension for latent space, same as in Rubanova et al. (2019)
and we use Sigmoid as the output activation with output standard deviation
0.01. For Physionet, we use no output activation and 0.05 standard deviation,
and use 20-dimensional latent space, same as in baselines.

Additional Metric We also present CRPS (continuous ranked probability score)
as a more appropriate metric for time-series forecasting. With the same
baselines, we show CRPS result in Table 5.

Task | Data | RNN | RNN-VAE | ODE-RNN | GRU-D | Latent ODE | LS4 (Ours)
---|---|---|---|---|---|---|---
Interp. | Physionet | 2.09 | 5.59 | 2.40 | 2.71 | 6.16 | 1.25
USHCN | 3.33 | 4.68 | 3.18 | 4.69 | 4.68 | 0.438
Extrap. | Physionet | 3.30 | 2.17 | 2.16 | 13.9 | 2.43 | 2.36
USHCN | 72.1 | 5.01 | 5.09 | 5.01 | 5.04 | 2.76

Table 5: Interpolation and extrapolation CRPS $(\times 10^{-2})$ scores. Lower
scores are better.

####  D.3 Runtime

We test all models on a single RTX A5000 GPU. To set up the dataset, we need
to fully populate the GPU for each dataset during training for our
benchmarking. For sequence lengths $\\{80,320,1280,5120,20480\\}$, we build
dataset of length $\\{102400,25600,6400,1600,400\\}$ each with batch size
$\\{1024,256,64,16,4\\}$ so that for each dataset the models are trained with
$100$ iterations.

###  Appendix E Generation Results

We present ground-truths and generations on the two hardest selected datasets,
NN5 Daily and Temperature Rain because these are the hardest to model.

Figure 5: Normalized NN5 Daily data. Figure 6: NN5 Daily generation. Figure 7:
Normalized Temperature Rain data. Figure 8: Temperature Rain generation.