Lost in Latent Space

This repository contains the official implementation of the paper Lost in Latent Space: An Empirical Study of Latent Diffusion Models for Physics Emulation by François Rozet, Ruben Ohana, Michael McCabe, Gilles Louppe, François Lanusse, and Shirley Ho.

Abstract

The steep computational cost of diffusion models at inference hinders their use as fast physics emulators. In the context of image and video generation, this computational drawback has been addressed by generating in the latent space of an autoencoder instead of the pixel space. In this work, we investigate whether a similar strategy can be effectively applied to the emulation of dynamical systems and at what cost. We find that the accuracy of latent-space emulation is surprisingly robust to a wide range of compression rates (up to 1000x). We also show that diffusion-based emulators are consistently more accurate than non-generative counterparts and compensate for uncertainty in their predictions with greater diversity. Finally, we cover practical design choices, spanning from architectures to optimizers, that we found critical to train latent-space emulators.

Code

The majority of the code is written in Python. Neural networks are implemented and trained using the PyTorch automatic differentiation framework. To run the experiments, it is necessary to have access to a Slurm cluster, to login to a Weights & Biases account and to install the lola module as a package.

First, create a new Python environment, for example with venv.

python -m venv ~/.venvs/lola
source ~/.venvs/lola/bin/activate

Then, install the lola module as an editable package with its dependencies.

pip install --editable .[all] --extra-index-url https://download.pytorch.org/whl/cu121

Optionally, we provide pre-commit hooks to automatically detect code issues.

pre-commit install --config pre-commit.yaml

Organization

The lola directory contains the implementations of the neural networks, the autoencoders, the diffusion models, the emulation routines, and others.

The experiments directory contains the training scripts, the evaluation scripts and their configurations. The euler, rayleigh_benard and gravity_cooling directories contain the notebooks that produced the figures of the paper.

Data

We rely on a Ceph File System partition to store the data. If your cluster uses a different file system, we recommend to create a symbolic link in your home folder.

ln -s /mnt/filesystem/users/you ~/ceph

The datasets (Euler, Rayleigh-Bénard and Turbulence Gravity Cooling) are downloaded from The Well.

the-well-download --base-path ~/ceph/the_well --dataset euler_multi_quadrants_openBC
the-well-download --base-path ~/ceph/the_well --dataset euler_multi_quadrants_periodicBC
the-well-download --base-path ~/ceph/the_well --dataset rayleigh_benard
the-well-download --base-path ~/ceph/the_well --dataset turbulence_gravity_cooling

This could take a while!

Experiments

We start with training the autoencoders. For clarity, we provide the commands for a single compression rate. To replicate the other experiments, modify the number of latent channels.

python train_ae.py dataset=euler_all optim.learning_rate=1e-5 ae.lat_channels=64
python train_ae.py dataset=rayleigh_benard optim.learning_rate=1e-5 ae.lat_channels=64
python train_ae.py dataset=gravity_cooling optim.learning_rate=1e-5 ae=dcae_3d_f8c64_large ae.lat_channels=64

Each train_*.py script schedules a Slurm job to train a model, log the training statistics with wandb, and store the weights in the ~/ceph/lola/runs directory. You will likely have to adapt the requested resources, either in the config files or in the command line. You can inspect the training logs with the dawgz command.

dawgz       # list all submitted workflows
dawgz 7     # list all jobs in the 7th workflow
dawgz -1 0  # show the logs of the first job in the last workflow

Once the above jobs are completed (1-4 days), we encode the entire dataset with each trained autoencoder and cache the resulting latent trajectories permanently on disk. For instance, for the autoencoder run named 1e3z5x2c_rayleigh_benard_dcae_f32c64_large,

python cache_latents.py dataset=rayleigh_benard split=train repeat=4 run=~/ceph/lola/runs/ae/1e3z5x2c_rayleigh_benard_dcae_f32c64_large
python cache_latents.py dataset=rayleigh_benard split=valid run=~/ceph/lola/runs/ae/1e3z5x2c_rayleigh_benard_dcae_f32c64_large

The stored latent trajectories are then used to train latent-space emulators (deterministic and diffusion-based), without needing to load and encode high-dimensional samples on the fly.

python train_surrogate.py dataset=rayleigh_benard ae_run=~/ceph/lola/runs/ae/1e3z5x2c_rayleigh_benard_dcae_f32c64_large  # neural solver
python train_diffusion.py dataset=rayleigh_benard ae_run=~/ceph/lola/runs/ae/1e3z5x2c_rayleigh_benard_dcae_f32c64_large  # diffusion model

We also train pixel-space deterministic emulators, which require more compute resources.

python train_surrogate.py dataset=euler_all surrogate=vit_pixel compute.nodes=2
python train_surrogate.py dataset=rayleigh_benard surrogate=vit_pixel compute.nodes=2

Finally, we evaluate each trained emulator on the test set.

python eval.py start=16 seed=0 run=~/ceph/lola/runs/sm/2k83f6km_rayleigh_benard_f32c64_vit_large  # neural solver
python eval.py start=16 seed=0 run=~/ceph/lola/runs/dm/ny04m1tl_rayleigh_benard_f32c64_vit_large  # diffusion model

The results will be compiled in CSV files at ~/ceph/lola/results.

Pre-trained models

We provide the weights of all models evaluated in the paper. Note that latent emulators rely on an autoencoder whose weights are provided separately. Please refer to the evaluation notebooks in the experiments directory for examples of loading and using trained models.

List of models

Dataset	Model	Compression	Size	Links
Euler	Neural solver	--	3.3GB	weights
Euler	Autoencoder	80	907MB	weights
Euler	Diffusion model	80	851MB	weights
Euler	Neural solver	80	850MB	weights
Euler	Autoencoder	320	904MB	weights
Euler	Diffusion model	320	850MB	weights
Euler	Neural solver	320	850MB	weights
Euler	Autoencoder	1280	903MB	weights
Euler	Diffusion model	1280	850MB	weights
Euler	Neural solver	1280	850MB	weights
RB	Neural solver	--	3.3GB	weights
RB	Autoencoder	64	1.2GB	weights
RB	Diffusion model	64	851MB	weights
RB	Neural solver	64	850MB	weights
RB	Autoencoder	256	1.2GB	weights
RB	Diffusion model	256	850MB	weights
RB	Neural solver	256	850MB	weights
RB	Autoencoder	1024	1.2GB	weights
RB	Diffusion model	1024	850MB	weights
RB	Neural solver	1024	850MB	weights
TGC	Autoencoder	48	2.8GB	weights
TGC	Diffusion model	48	855MB	weights
TGC	Neural solver	48	854MB	weights
TGC	Autoencoder	192	2.8GB	weights
TGC	Diffusion model	192	854MB	weights
TGC	Neural solver	192	854MB	weights
TGC	Autoencoder	768	2.8GB	weights
TGC	Diffusion model	768	854MB	weights
TGC	Neural solver	768	854MB	weights

Citation

If you find this project useful for your research, please consider citing

@unpublished{rozet2025lost,
  title = {Lost in Latent Space: An Empirical Study of Latent Diffusion Models for Physics Emulation},
  author = {Rozet, François and Ohana, Ruben and McCabe, Michael and Louppe, Gilles and Lanusse, François and Ho, Shirley},
  year = {2025},
  url = {https://arxiv.org/abs/2507.02608}
}

Acknowledgements

We thank Géraud Krawezik and the Scientific Computing Core at the Flatiron Institute, a division of the Simons Foundation, for the compute facilities and support. We gratefully acknowledge use of the research computing resources of the Empire AI Consortium, Inc., with support from the State of New York, the Simons Foundation, and the Secunda Family Foundation. Polymathic AI acknowledges funding from the Simons Foundation and Schmidt Sciences, LLC.