Rubato: Signature Attention for Irregular Multivariate Time Series Forecasting

NeurIPS 2026 · Spotlight

Official PyTorch implementation of Rubato.

Overview · Quickstart

Download datasets & checkpoints · Explore the model

---

Overview

Rubato forecasts irregular multivariate time series by combining per-channel patch representations with joint signature features. Signature attention learns relationships between patches from the multivariate trajectory, while a decoder uses target timestamps and channel embeddings to produce forecasts.

• Channel-specific content. Lightweight MLP tokenizers encode each variable's patches, with empty patches compacted before attention.
• Shared signature attention. Joint random Fourier signature features define a multivariate attention kernel, paired with per-channel value gates and channel mixing.
• Experiments ready to reproduce. Six datasets, five runs per dataset, saved configurations, checkpoints, and evaluation metrics.

Quickstart

1. Set up the environment

Use Python 3.12+, PyTorch 2.10, and an NVIDIA GPU with a CUDA toolkit for the bundled fast-hadamard-transform extension. Dependencies are declared in pyproject.toml and resolved in uv.lock.

From the repository root, with uv installed:

uv sync
source .venv/bin/activate

Keep this environment active when running the commands below. The bundled Hadamard extension includes the SORF kernels used by Rubato.

2. Prepare datasets and checkpoints

Download rubato_assets.zip and extract it into the repository root:

unzip rubato_assets.zip -d .

The archive contains the four public datasets USHCN, FNSPID, EPA_AIR, and ClusterTrace, together with the saved checkpoints. Data and experiment artifacts live under:

storage/
├── datasets/
│   ├── USHCN/ushcn_standard.csv
│   ├── FNSPID/FNSPID_standard.csv
│   ├── EPA_AIR/EPA-Air_standard.csv
│   └── ClusterTrace/ClusterTrace_standard.csv
└── results/<dataset>/<dataset_id>/Rubato/<model_id>/...

MIMIC and P12 are prepared separately. MIMIC requires credentialed access; P12 / PhysioNet Challenge 2012 is publicly available.

Prepare MIMIC and P12

The preprocessing scripts additionally require pandas; P12 also requires a torchvision version compatible with the installed PyTorch. These are not included in the base dependency list.

MIMIC. After obtaining access, place the prepared full_dataset.csv at storage/datasets/MIMIC/raw/full_dataset.csv. The preprocessor expects record_id, Time, and paired Value.* / Mask.* fields.

python -m data.preprocessing.mimic --root storage/datasets/MIMIC
mv storage/datasets/MIMIC/raw/mimic_standard.csv storage/datasets/MIMIC/

P12. The preprocessor downloads and processes sets A, B, and C automatically.

python -m data.preprocessing.physionet --root storage/datasets/P12
mv storage/datasets/P12/raw/physionet_standard.csv storage/datasets/P12/

3. Evaluate the saved checkpoints

Once all six datasets are available, run:

bash scripts/reproduce_main_table.sh

The script evaluates six datasets × five saved training runs (training seeds 2022–2026). It loads the latest matching training directory for each experiment; use the downloaded result directories to evaluate the released checkpoints. New metrics are written to:

storage/results/<dataset>/<dataset_id>/Rubato/<model_id>/<seq_len>_<pred_len>/<run>/iter*/eval_*/metric.json

4. Train a model

Dataset-specific configurations are provided in scripts/Rubato. For example:

# Train and evaluate on USHCN
bash scripts/Rubato/USHCN.sh

# Train and evaluate on FNSPID
bash scripts/Rubato/FNSPID.sh

Scripts are also available for MIMIC, P12, EPA_AIR, and ClusterTrace. Model-specific options use the --rubato_ prefix; see utils/configs.py for all arguments.

Acknowledgements

Our experiment framework builds on PyOmniTS. The bundled CUDA implementation extends fast-hadamard-transform. We thank their authors for sharing these tools.

License

This repository is released under the MIT License. Datasets remain subject to their original terms of use.

Citation

If you find this work useful, please cite our paper:

@inproceedings{yang2026rubato,
  title={Rubato: Signature Attention for Irregular Multivariate Time Series Forecasting},
  author={Jintao Yang and Meng Wang and Chenjuan Guo and Bin Yang},
  booktitle={NeurIPS},
  year={2026}
}