gabrieleaccarino/wavelet-flow-matching

Wavelet Flow Matching for Multi-Scale Physics Emulation — pre-trained weights

NeurIPS 2026 · Gabriele Accarino, Juan Nathaniel, Carla Roesch, Pierre Gentine, Sara Shamekh, Duncan Watson-Parris, Viviana Acquaviva

Code · Project page · arXiv

Checkpoints of every model evaluated in the paper: Wavelet Flow Matching (WFM) for each number of scales and mother wavelet, and the flow-matching and neural-operator baselines, on three systems from The Well. These are the exact checkpoints behind the paper's tables and figures (40 files, 3.1 GB).

Quick start

git clone https://github.com/gabrieleaccarino/wavelet_flow_matching.git && cd wavelet_flow_matching
conda env create -f env/environment.yml && conda activate wfm && pip install -e .

from wavelet_flow_matching.hub import list_models, load_pretrained

list_models("rayleigh_benard")                        # available names
model = load_pretrained("rayleigh_benard", "WFM_3scale_db4", map_location="cuda")

load_pretrained downloads the checkpoint (cached after the first call) and returns the model in eval mode, ready for model.rollout(...); see the example notebook. To only download files:

hf download gabrieleaccarino/wavelet-flow-matching --include "rayleigh_benard/*" --local-dir ./weights      # one system
hf download gabrieleaccarino/wavelet-flow-matching --local-dir ./weights                                    # everything

To reproduce the paper's metrics, set ckpt_path: in an inference config (configs/<pde>/inference_<pde>_seed.yaml) to the downloaded file.

Models

★ = mother wavelet selected in the paper (lowest average spectral coherence RMSE). Val loss is the training objective on the validation split and is not comparable across model families.

Turbulent Radiative Layer 2D (turbulent_radiative_layer_2D)

Model	Description	Params	Epoch	Val loss	File
FM_fourier	Flow matching on Fourier coefficients (32 modes)	30.3 M	198	0.0513	FM_fourier_turbulent_radiative_layer_2D.ckpt
FM_fourier16	Flow matching on Fourier coefficients (16 modes)	12.7 M	198	0.0756	FM_fourier16_turbulent_radiative_layer_2D.ckpt
FM_pixel	Flow matching in pixel space (same U-Net, no transform)	12.7 M	193	0.0078	FM_pixel_turbulent_radiative_layer_2D.ckpt
WFM_1scale_db2	Wavelet Flow Matching, 1 wavelet scale, db2 wavelet	12.7 M	188	0.0154	WFM_1scale_db2_turbulent_radiative_layer_2D.ckpt
WFM_1scale_db4	Wavelet Flow Matching, 1 wavelet scale, db4 wavelet	12.7 M	183	0.0141	WFM_1scale_db4_turbulent_radiative_layer_2D.ckpt
WFM_1scale_db6	Wavelet Flow Matching, 1 wavelet scale, db6 wavelet	12.7 M	189	0.0153	WFM_1scale_db6_turbulent_radiative_layer_2D.ckpt
WFM_1scale_haar ★	Wavelet Flow Matching, 1 wavelet scale, haar wavelet	12.7 M	193	0.0153	WFM_1scale_haar_turbulent_radiative_layer_2D.ckpt
WFM_3scale_db2	Wavelet Flow Matching, 3 wavelet scales, db2 wavelet	30.6 M	176	0.0220	WFM_3scale_db2_turbulent_radiative_layer_2D.ckpt
WFM_3scale_db4	Wavelet Flow Matching, 3 wavelet scales, db4 wavelet	30.6 M	193	0.0276	WFM_3scale_db4_turbulent_radiative_layer_2D.ckpt
WFM_3scale_db6	Wavelet Flow Matching, 3 wavelet scales, db6 wavelet	30.6 M	196	0.0381	WFM_3scale_db6_turbulent_radiative_layer_2D.ckpt
WFM_3scale_haar ★	Wavelet Flow Matching, 3 wavelet scales, haar wavelet	30.6 M	181	0.0228	WFM_3scale_haar_turbulent_radiative_layer_2D.ckpt
FNO	Fourier Neural Operator (deterministic)	19.0 M	25	0.0909	FNO_turbulent_radiative_layer_2D.pt
TFNO	Tensorized Fourier Neural Operator (deterministic)	7.9 M	70	0.0762	TFNO_turbulent_radiative_layer_2D.pt

Rayleigh–Bénard convection (rayleigh_benard)

Model	Description	Params	Epoch	Val loss	File
FM_fourier	Flow matching on Fourier coefficients (32 modes)	30.3 M	109	0.0571	FM_fourier_rayleigh_benard.ckpt
FM_fourier16	Flow matching on Fourier coefficients (16 modes)	12.7 M	158	0.0795	FM_fourier16_rayleigh_benard.ckpt
FM_pixel	Flow matching in pixel space (same U-Net, no transform)	12.7 M	194	0.0100	FM_pixel_rayleigh_benard.ckpt
LatentFM	Flow matching in the latent space of a frozen DC-AE (f32c64)	8.4 M	195	0.2147	LatentFM_rayleigh_benard.ckpt
WFM_1scale_db2	Wavelet Flow Matching, 1 wavelet scale, db2 wavelet	12.7 M	188	0.0194	WFM_1scale_db2_rayleigh_benard.ckpt
WFM_1scale_db4 ★	Wavelet Flow Matching, 1 wavelet scale, db4 wavelet	12.7 M	178	0.0193	WFM_1scale_db4_rayleigh_benard.ckpt
WFM_1scale_db6	Wavelet Flow Matching, 1 wavelet scale, db6 wavelet	12.7 M	185	0.0189	WFM_1scale_db6_rayleigh_benard.ckpt
WFM_1scale_haar	Wavelet Flow Matching, 1 wavelet scale, haar wavelet	12.7 M	190	0.0197	WFM_1scale_haar_rayleigh_benard.ckpt
WFM_3scale_db2	Wavelet Flow Matching, 3 wavelet scales, db2 wavelet	30.6 M	188	0.0264	WFM_3scale_db2_rayleigh_benard.ckpt
WFM_3scale_db4 ★	Wavelet Flow Matching, 3 wavelet scales, db4 wavelet	30.6 M	185	0.0314	WFM_3scale_db4_rayleigh_benard.ckpt
WFM_3scale_db6	Wavelet Flow Matching, 3 wavelet scales, db6 wavelet	30.6 M	198	0.0356	WFM_3scale_db6_rayleigh_benard.ckpt
WFM_3scale_haar	Wavelet Flow Matching, 3 wavelet scales, haar wavelet	30.6 M	188	0.0258	WFM_3scale_haar_rayleigh_benard.ckpt
FNO	Fourier Neural Operator (deterministic)	19.0 M	19	0.0272	FNO_rayleigh_benard.pt
TFNO	Tensorized Fourier Neural Operator (deterministic)	7.9 M	289	0.0207	TFNO_rayleigh_benard.pt

Active Matter (active_matter)

Model	Description	Params	Epoch	Val loss	File
FM_fourier	Flow matching on Fourier coefficients (32 modes)	30.3 M	199	0.1335	FM_fourier_active_matter.ckpt
FM_fourier16	Flow matching on Fourier coefficients (16 modes)	12.7 M	187	0.2510	FM_fourier16_active_matter.ckpt
FM_pixel	Flow matching in pixel space (same U-Net, no transform)	12.7 M	179	0.0026	FM_pixel_active_matter.ckpt
WFM_1scale_db2	Wavelet Flow Matching, 1 wavelet scale, db2 wavelet	12.8 M	195	0.0586	WFM_1scale_db2_active_matter.ckpt
WFM_1scale_db4	Wavelet Flow Matching, 1 wavelet scale, db4 wavelet	12.8 M	192	0.0605	WFM_1scale_db4_active_matter.ckpt
WFM_1scale_db6	Wavelet Flow Matching, 1 wavelet scale, db6 wavelet	12.8 M	194	0.0639	WFM_1scale_db6_active_matter.ckpt
WFM_1scale_haar ★	Wavelet Flow Matching, 1 wavelet scale, haar wavelet	12.8 M	198	0.0612	WFM_1scale_haar_active_matter.ckpt
WFM_3scale_db2 ★	Wavelet Flow Matching, 3 wavelet scales, db2 wavelet	31.0 M	199	0.0163	WFM_3scale_db2_active_matter.ckpt
WFM_3scale_db4	Wavelet Flow Matching, 3 wavelet scales, db4 wavelet	31.0 M	187	0.0234	WFM_3scale_db4_active_matter.ckpt
WFM_3scale_db6	Wavelet Flow Matching, 3 wavelet scales, db6 wavelet	31.0 M	185	0.0306	WFM_3scale_db6_active_matter.ckpt
WFM_3scale_haar	Wavelet Flow Matching, 3 wavelet scales, haar wavelet	31.0 M	190	0.0140	WFM_3scale_haar_active_matter.ckpt
FNO	Fourier Neural Operator (deterministic)	19.0 M	57	0.0156	FNO_active_matter.pt
TFNO	Tensorized Fourier Neural Operator (deterministic)	7.9 M	73	0.0068	TFNO_active_matter.pt

Repository layout

README.md                          this model card
manifest.csv                       per file: model, size, SHA-256, epoch, val loss, original file name
<pde>/<model>/<model>_<pde>.ckpt   flow-matching models (WFM, FM_pixel, FM_fourier*, LatentFM)
<pde>/<model>/<model>_<pde>.pt     neural operators (FNO, TFNO)
<pde>/<model>/config.json          hyper-parameters stored in the checkpoint + provenance

Model names match the paper and the metric files in the code repository (experiments/<pde>/metrics/<model>_<pde>_seed<S>_metrics.csv).

Checkpoint format and loading

Each file is a self-contained FP32 PyTorch checkpoint (independent of Lightning/DeepSpeed) with the weights, the model hyper-parameters and the normalization statistics. Load them with the classes in the code repository:

Models	Loader
WFM_*, FM_pixel	WaveletFlowModule.from_checkpoint(path)
FM_fourier, FM_fourier16	FourierFlowModule.from_checkpoint(path)
LatentFM	LatentFlowModule.from_checkpoint(path, ae_state_path=...)
FNO, TFNO	NeuralOperatorModule.from_pretrained(path)

The files are Python pickles (torch.load(..., weights_only=False)), so only load files from this repository or other sources you trust.

Training

Data: The Well (turbulent_radiative_layer_2D, rayleigh_benard, active_matter), standard train/validation/test splits, 4 input frames → 1 output frame, autoregressive rollout at test time. Flow-matching models: 200 epochs, AdamW with 20 warm-up epochs and cosine decay. FNO/TFNO: Adam with cosine decay and early stopping on the validation loss. All models were trained on NCAR Derecho (NVIDIA A100 40GB GPUs). Full settings: config.json next to each checkpoint and configs/<pde>/ in the code repository.

Not included

• WNO and WDNO baselines: trained and evaluated with their authors' official code; see the paper for references.
• The LatentFM autoencoder (rayleigh_benard/LatentFM/autoencoder_dcae_f32c64.pth) is the pre-trained DC-AE of Rozet et al., Lost in Latent Space (https://github.com/polymathicai/lola), redistributed under its original license for convenience.

License

MIT, same as the code.

Citation

@inproceedings{accarino2026wfm,
  title     = {Wavelet Flow Matching for Multi-Scale Physics Emulation},
  author    = {Accarino, Gabriele and Nathaniel, Juan and Roesch, Carla and
               Gentine, Pierre and Shamekh, Sara and Watson-Parris, Duncan and
               Acquaviva, Viviana},
  booktitle = {Advances in Neural Information Processing Systems},
  year      = {2026}
}