Forecasting time-series image data with generative models

This repository implements several approaches for learned forecasting of small video datasets.

The current datasets are:

• Moving MNIST, introduced in this paper.
• Elastic disk dynamics, a generated dataset of grayscale movies showing circular particles moving in a 2D box with reflecting or periodic boundaries. The simulation uses simple equal-mass elastic collisions and wall/periodic boundary handling. ElasticDisksDataset supports configurable image_size; Moving MNIST uses its native 64x64 resolution.

Moving MNIST	Elastic disk dynamics

_{Autoregressive SimVP rollout for a sample from the elastic disks dataset.
Left: model-generated frames. Center: ground truth frames. Right: error map.}

_{Autoregressive stochastic interpolant rollout for a sample from the elastic disks dataset.
Left: model-generated frames. Center: ground truth frames. Right: error map}

Notebooks

notebooks/
  moving_mnist/
    visualize_moving_mnist_data.ipynb
    train_moving_mnist_pixel_flow_matching.ipynb
    train_moving_mnist_pixel_stochastic_interpolant.ipynb
    train_moving_mnist_pixel_diffusion.ipynb
    train_moving_mnist_simvp.ipynb
    train_moving_mnist_latent_flow_matching.ipynb
    train_moving_mnist_latent_stochastic_interpolant.ipynb
    train_moving_mnist_latent_flow_matching_1dlatent.ipynb
    train_moving_mnist_latent_diffusion_1dlatent.ipynb
    train_moving_mnist_latent_transformer.ipynb
    train_moving_mnist_mdn_rnn.ipynb
    train_moving_mnist_mdn_rnn_1dvaelatent.ipynb
  elastic_disks/
    visualize_elastic_disks_data.ipynb
    train_elastic_disks_pixel_flow_matching.ipynb
    train_elastic_disks_pixel_stochastic_interpolant.ipynb
    train_elastic_disks_pixel_diffusion.ipynb
    train_elastic_disks_simvp.ipynb
    train_elastic_disks_latent_flow_matching.ipynb
    train_elastic_disks_latent_stochastic_interpolant.ipynb
    train_elastic_disks_latent_flow_matching_1dlatent.ipynb
    train_elastic_disks_latent_diffusion_1dlatent.ipynb
    train_elastic_disks_latent_transformer.ipynb
    train_elastic_disks_mdn_rnn.ipynb
    train_elastic_disks_mdn_rnn_1dvaelatent.ipynb

The forecasting notebooks cover pixel- and latent-space flow matching, pixel- and latent-space stochastic interpolants, pixel- and latent-space diffusion, SimVP, causal latent transformers, MDN-RNNs over spatial latents, and MDN-RNNs over 1D VAE latents.

Package Layout

src/video_forecasting/
  datasets/
    moving_mnist.py
    elastic_disks.py
  models/
    vae.py
    flow_matching.py
    stochastic_interpolants.py
    diffusion.py
    transformer.py
    mdn_rnn.py
    simvp.py
  runtime.py
  training.py
  visualization.py