Spectral Mean Flow (PyTorch)

Sequence Modeling with Spectral Mean Flows
Jinwoo Kim, Max Beier, Petar Bevanda, Nayun Kim, Seunghoon Hong
NeurIPS 2025

This repository contains the official PyTorch implementation of spectral mean flow. The codebase has been tested with NVIDIA H100 GPUs.

Setup

We recommend using the official PyTorch Docker image with CUDA support.

docker pull pytorch/pytorch:2.8.0-cuda12.6-cudnn9-devel
docker run -it --gpus all --ipc host --name spectral-mean-flow -v /home:/home pytorch/pytorch:2.8.0-cuda12.6-cudnn9-devel bash

Inside the container, install the required packages:

pip3 install tensorboard ema-pytorch einops dm-control
pip3 install tensorflow==2.15.0 tensorflow-probability==0.23.0
pip3 install seaborn scikit-learn scipy
pip3 install omegaconf
pip3 install tf-slim
pip3 install opt-einsum flow_matching

Experiments

Synthetic experiments (tractability, 2D checkerboard):

cd src_synthetic

# Figure 2
python3 tractability.py

# Figure 3
# run checkerboard.ipynb

cd ..

Regular time-series experiments (Sines, Stocks, ETTh, Mujoco, Energy, fMRI):

# Table 1
cd src_base
bash scripts/spectral_flow.sh
cd ..

# Table 2
cd src_medium
bash scripts/spectral_flow.sh
cd ..

Longer time-series experiments (FRED-MD, NN5 Daily):

# Table 3
cd src_large
bash scripts/spectral_flow.sh
cd ..

Irregular time-series experiments (Stocks with 0%, 30%, 50%, 70% missing data):

# Table 4
cd src_small
bash scripts/spectral_flow_irregular.sh
cd ..

Physics-informed modeling experiments (nonlinear pendulum):

# Table 5
cd src_small
bash scripts/spectral_flow_pendulum.sh
cd ..

References

Our implementation is based on the code from the following repositories:

• LRU-pytorch for complex value parameterization and initialization
• Diffusion-TS and SDFormer for regular time-series experiments
• ImagenTime for longer time-series experiments
• KoVAE for irregular time-series and nonlinear pendulum experiments

Citation

If you find our work useful, please consider citing it:

@article{kim2025sequence,
  author    = {Jinwoo Kim and Max Beier and Petar Bevanda and Nayun Kim and Seunghoon Hong},
  title     = {Sequence Modeling with Spectral Mean Flows},
  journal   = {arXiv},
  volume    = {abs/2510.15366},
  year      = {2025},
  url       = {https://arxiv.org/abs/2510.15366}
}