DEEP LATENT STATE SPACE MODELS FOR TIME-SERIES GENERATION

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 realworld 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 ×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 because they can make predictions 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. However, that leads to quadratic scaling in the hidden dimensionality due to the need of explicitly computing hidden states by unrolling the underlying recurrence over time, thus preventing scaling to long sequences. 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, this leads to computational costs, and existing 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 ≈ 20K lengths, our model trains ×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. 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 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 d dt h t = Ah t + Bx t y t = Ch t + Dx t



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 continuoustime 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 continuoustime approaches are given in Kidger et al. (

1)with scalar input x t ∈ R, state h t ∈ R N and scalar output y t ∈ R. The system is fully characterized by the matricesA ∈ R N ×N , B ∈ R N ×1 , C ∈ R 1×N , D ∈ R 1×1 . Let x, y ∈ C([a, b], R) be absolutely continuous real signals on time interval [a, b]. Given an initial condition h 0 ∈ R N the SSM (1) realizes a mapping x → y.

