Source-linked AI summary

Diffusion-based Time Series Imputation and Forecasting with Structured State Space Models

Juan Miguel Lopez Alcaraz, Nils Strodthoff

arXiv:2208.09399v3cs.LGstat.ML

TL;DR

Missing values pose a major problem for real-world time-series analysis, especially when missingness is structured or forecasting requires an entirely absent future window. The paper proposes SSSD, combining conditional diffusion models with structured state-space models, and reports strong imputation and forecasting performance across data sets and missingness scenarios, including blackout missingness.

  • Problem

    Missing time-series values can impair downstream analyses, and existing methods may be limited to random missingness or fail qualitatively on blackout missingness.

  • Method

    SSSD combines conditional diffusion models with structured state-space models and applies diffusion noise to regions targeted for imputation.

  • Results

    SSSDS4 outperforms existing state-of-the-art imputers across data sets and missingness scenarios, with particularly strong performance in blackout missingness and forecasting.

  • Takeaways & Limitations

    The results support SSSD as a promising generative-model technology for time-series imputation and forecasting across varied missingness scenarios.

  • Takeaways & Limitations

    The strongest results and main evaluation focus on data sets with fewer than about 100 input channels because the channel dimension is handled implicitly during diffusion.

Abstract

from arXiv · show

The imputation of missing values represents a significant obstacle for many real-world data analysis pipelines. Here, we focus on time series data and put forward SSSD, an imputation model that relies on two emerging technologies, (conditional) diffusion models as state-of-the-art generative models and structured state space models as internal model architecture, which are particularly suited to capture long-term dependencies in time series data. We demonstrate that SSSD matches or even exceeds state-of-the-art probabilistic imputation and forecasting performance on a broad range of data sets and different missingness scenarios, including the challenging blackout-missing scenarios, where prior approaches failed to provide meaningful results.

1 Introduction

The paper addresses the challenge of missing time-series values, emphasizing probabilistic imputations across realistic missingness scenarios and forecasting. It proposes combining diffusion models with structured state-space models to capture long-term dependencies and improve imputation quality.

  • Missing input data can arise from entry errors, equipment failures, or file losses, and poor imputations can bias downstream analyses.
  • Time-series missingness is prevalent, including potentially non-at-random patterns caused by faulty sensors, while forecasting is treated as blackout imputation at the sequence end.
  • Probabilistic imputation is presented as appropriate for underspecified missingness because it produces samples representing different plausible imputations.
  • Existing approaches may be restricted to random missingness, unstable during training, or qualitatively meaningless for blackout missingness on some data sets.
  • SSSD combines conditional diffusion models with structured state-space models suited to capturing long-term dependencies in time series.
  • The work modifies DiffWave, applies diffusion noise only to regions being imputed, and evaluates the approach across data sets and missingness scenarios, especially blackout and forecasting settings.

2 Structured state space diffusion (SSSD) models for time series imputation

SSSD conditions diffusion-based generation on observed time-series information and missingness masks, while using structured state-space layers to model long-term dependencies efficiently. The paper develops several architectures, including SSSDS4, by adapting diffusion and state-space components for time-series imputation.

  • Time series imputation: Imputation targets are represented by binary masks matching the input shape, with ones marking conditioned values and zeros marking values to impute.
  • Time series imputation: Missing values already present in the input require a second mask that distinguishes observed entries from unavailable ones.
  • Time series imputation: Random missingness samples individual time steps uniformly across channels, whereas random block missingness samples consecutive segments for each channel.
  • Diffusion models: Diffusion models learn to remove sequentially added noise through a reverse process, thereby mapping noisy latent representations back toward the original signal space.
  • Diffusion models: Conditional diffusion uses the masked input and imputation masks as conditioning information during the backward process, with setups applying diffusion to the full signal or imputed regions only.
  • Structured state-space models: Structured state-space models connect input and output sequences through hidden states and can be evaluated efficiently after discretization as convolution operations.
  • Proposed approaches: SSSDS4 replaces DiffWave’s bidirectional dilated convolutions with S4 layers and adds another S4 layer after incorporating conditional information.
  • Proposed approaches: The paper also investigates SaShiMi and CSDIS4 variants, including replacing CSDI’s time-direction transformer with an S4 model.

3 Related Work

Related work spans recurrent, generative, graph-based, and modern transformer architectures for time-series imputation and forecasting. SSSD is positioned against diffusion-based methods, especially CSDI, while adapting ideas from image inpainting and speech generation.

  • Deep-learning imputation methods include RNN-, GAN-, VAE-, diffusion-, and graph-based approaches, with several earlier methods facing training instability or suboptimal performance.CSDI is identified as a diffusion-based method with strong results.
  • Diffusion models originated in related inpainting and speech-generation settings, and have been adapted to general time-series imputation.DiffWave provides a time-series precedent through speech-generation tasks.
  • CSDI and SSSD are DiffWave variants that differ in architecture, diffusion direction, and denoising objective, with experiments comparing them directly.SSSD uses SSMs rather than transformers, diffuses along time rather than feature and time directions, and denoises only imputed segments.
  • Long-sequence forecasting includes recurrent architectures and transformer encoder-decoder models such as LSTNet, Autoformer, and Informer.

4 Experiments

Experiments evaluate SSSD across imputation and forecasting tasks, missingness scenarios, datasets, and model variants. Results show strong performance, especially for blackout missingness and long-horizon forecasting, while revealing limitations on high-dimensional inputs and some specialized forecasting settings.

  • Experimental setup: SSSDS4 uses a channel-splitting design whose broader improvement is evaluated primarily on datasets with fewer than 100 input channels.The authors identify more elaborate channel-dimension handling as future work.
  • Experimental setup: Training and evaluation use identical missingness scenarios and ratios, with MSE as the loss function across experiments.Additional results, baselines, datasets, and preprocessing details are provided in supplementary material.
  • Imputation results: Applying diffusion only to imputation regions (D1) consistently outperforms applying it to the full signal (D0) across PTB-XL model types and missingness scenarios.Unconditional diffusion training performs clearly worse in the reported comparison.
  • Imputation results: SSSDS4 achieves more than 50% MAE reduction versus CSDI for PTB-XL blackout missingness and captures essential ECG signal features qualitatively.CSDIS4 improves qualitative imputations but still misses essential features; DiffWave, SSSDSA, and SSSDS4 capture them.
  • Imputation results: At 90% random missingness on MuJoCo, SSSDS4 achieves an MSE of 1.90e-3 with more than 50% error reduction against baselines.CSDI remains best suited to some lower-ratio random-missing settings.
  • Forecasting results: For forecasting, SSSDS4 achieves MSE 5.03e2 and a 27% error reduction versus TLAE on solar data, while ETTm1 results are strong but sometimes below Autoformer.On ETTm1, SSSDS4 generally outperforms baselines, except that Autoformer remains stronger in several settings and CSDI wins at forecast length 288.

5 Conclusion

SSSD combines structured state-space models with diffusion models for time-series generation. It outperforms existing state-of-the-art imputers across data sets and missingness scenarios, especially blackout-missing and forecasting settings, when input channels are not too numerous.

  • SSSD combines structured state-space models for long-term dependencies with diffusion models for generative modeling.
  • SSSD outperforms existing state-of-the-art imputers across diverse data sets and missingness scenarios.
  • SSSD performs particularly strongly in blackout-missing and forecasting scenarios.
  • The reported performance advantage holds provided the number of input channels does not grow too large.

A Technical and implementation details

The implementation maps input channels into the diffusion dimension and performs diffusion only along time, simplifying processing relative to CSDI. This design supports long-term time-series modeling but creates convergence and channel-reconstruction challenges as the number of input channels increases.

  • CSDI uses a separate diffusion dimension, creating a four-dimensional representation that requires alternating processing of time and feature dimensions.
  • SSSD maps input channels into the diffusion dimension and performs diffusion only along the time dimension.
  • This single time-direction diffusion avoids separate diffusion processes along time and feature axes.
  • The authors attribute SSSD’s stronger performance in the small-channel regime partly to this simpler time-diffusion design.
  • Training convergence issues arise above about 100 input channels, motivating channel splitting but potentially causing suboptimal predictions when channel identities differ across batches.
  • Future channel-splitting improvements include grouping channels by correlation and providing batch-specific information to learn channel-specific representations.

A.3 Implementation details on missingness scenarios

The experiments evaluate random, contiguous-block, and blackout masking, with training and sampling driven by conditioned diffusion over masked time-series inputs.

  • RM samples missing targets independently per channel, while RBM and BM partition sequences into contiguous segments sized by the missingness ratio.
  • The conditional input concatenates the masked signal with M, while the network receives the noisy signal, conditioning input, and diffusion steps.
  • Training linearly interpolates beta between β0 and β1, samples Gaussian noise and diffusion steps, and forms M by multiplying imputation and missing-value masks.
  • Sampling starts from Gaussian noise, reapplies the mask, predicts updates conditioned on the observed series, and adds variance except at the final diffusion step.

B.1 Model architecture

The ablation compares convolutional and S4-based architectures for PTB-XL imputation. Conditional training is favored, and the two-S4-layer variant becomes the default model.

  • Architecture variants: Replacing dilated convolutions with an S4 layer and adding a second S4 layer after conditional data addition defines SSSDS4, setting C.
  • Ablation setup: The ablation evaluates four architecture variants across RM, RBM, and BM at 20% and 50% missingness using averaged MAE and RMSE.Each result averages three trials with 10 generated samples per test example after 150,000 training iterations.
  • Architecture choice: Conditional training outperforms the compared alternative, leading the authors to select SSSDS4, setting C, as the default variant.Figure 5 provides a qualitative comparison of generated PTB-XL ECG imputations under unconditional training.

B.2 Additional missingness ratio on PTB-XL

At 50% missingness on PTB-XL, SSSDS4 maintains stronger imputation performance than CSDI and DiffWave across the evaluated missingness scenarios and metrics.

  • At 50% missingness, DiffWave outperforms CSDI for RBM and BM but not RM, where CSDI is better suited.
  • SSSDS4 outperforms both CSDI and DiffWave across all three missingness scenarios and both reported metrics.
  • Diffusion-step ablation: Reducing diffusion steps from 200 to 50 cuts sampling time to 16% of the original but worsens target metrics to 166% and 131%.Training time decreases to 86% of the original.

B.4 Robustness analysis

The robustness analysis tests mismatches in missingness ratios and scenarios between training and evaluation. Performance is more robust when testing at higher missingness than training exposure in the reported experiment.

  • Missingness-ratio mismatch: Testing at a higher missingness ratio after training at a lower ratio incurs an 11% MAE penalty relative to matched high-ratio training.
  • Missingness-ratio mismatch: Training at a higher ratio and testing at a lower one produces metrics almost three times larger than matched low-ratio training.
  • Missingness-ratio mismatch: The experiment indicates greater robustness to testing at higher missingness ratios than to testing below the ratio seen during training.The authors stress that the tested mismatch is already large.
  • Missingness-scenario mismatch: The scenario-mismatch experiment compares RM, RBM, and BM on PTB-XL to assess whether matching training and test missingness characteristics is optimal.

C Training hyperparameters and architectures

The experiments compare diffusion and structured-state-space architectures using specified hyperparameters, with SSSDS4 receiving focused tuning and architectural choices aimed at convergence and computational complexity.

  • LAMC is implemented as a matrix-completion baseline, providing a qualitatively different methodology from probabilistic diffusion models.
  • CSDIS4 modifies official CSDI by replacing its temporal transformer layer with an S4 layer.The S4 implementation is bidirectional, layer-normalized, uses no dropout, and has internal state dimension N = 64.
  • SSSDSA is a 128-dimensional U-Net variant with six residual layers and S4 plus feed-forward blocks at each pooling level.Its pooling factors are 2 and 2, with corresponding feature expansion of 2.
  • SSSDS4 builds on DiffWave with 36 stacked residual layers and 256 residual and skip channels.It uses a three-layer diffusion embedding and an input convolution that doubles the channel dimension before the first S4 diffusion computation.
  • 200 diffusion steps were selected as a compromise between faster convergence with fewer steps and less accurate results.The linear noise schedule uses minimum noise level B0 = 0.001 and maximum noise level B1 = 0.5.
  • SSSDS4 uses 36 residual layers after ablation, reduced from 48 to improve convergence speed and reduce computational complexity when variant C requires two S4 layers.

D Data set descriptions

The study evaluates time-series imputation and forecasting across clinical, energy, physical-simulation, traffic, and solar datasets, while describing S4 as the sequence-modeling component for long dependencies.

  • PTB-XL: PTB-XL contains 21,837 clinical 12-lead ECGs lasting 10 seconds from 18,885 patients.Experiments use 100 Hz sampling, 12 input channels, and 20% target values across three imputation and forecasting scenarios.
  • Electricity, MuJoCo, and traffic: The benchmark includes Electricity usage from 370 clients, MuJoCo trajectories with 10,000 sequences of 100 points in 14 dimensions, and traffic datasets from PEMS-Bay and METR-LA.PEMS-Bay has 325 sensors, while METR-LA represents 207 loop detectors.
  • Solar: Solar forecasting conditions on 168 time steps to forecast the following 24 using hourly data from 137 photovoltaic power plants.
  • S4 modeling: S4 models long sequences through continuous state-space representations, HiPPO-based online signal compression, and recurrent discretization.The discretization yields a recurrent state-space mapping that can be computed like an RNN and handled through a transition matrix A.
  • Evaluation: Evaluation computes errors between targets and imputed values under conditional masking, using MAE, MSE, RMSE, and MRE.

E.3 Training times

Training information is organized by dataset and model in tables, while supplementary figures illustrate imputations and forecasting settings across several datasets and missingness patterns.

  • Training records: Tables 27–29 report the training iterations or epochs used for the datasets and models.
  • Imputation evaluations: Table 24 covers PTB-XL imputations under RM, RBM, and BM missingness scenarios.
  • Reported results: Tables 25 and 26 report traffic-dataset imputation metrics and PTB-XL forecasting results, respectively.
  • PTB-XL visualizations: PTB-XL figures show RM and RBM imputations with 100 samples at 50% missingness, and BM imputations at 30%.The RM caption states that diffusion models reconstruct time series under high missingness, while the BM figures illustrate qualitative imputation differences.
  • PTB-XL visualizations: The PTB-XL BM example notes that CSDI can learn the series trend but may generate outliers outside the 0.05 and 0.95 quantile range, while CSDIS4 improves generation quality.
  • Forecasting visualizations: Supplementary figures display complete ETTm1 forecasting samples for five horizons and Solar samples across six channels, including imputed areas.
Loading 2208.09399v3…