This repository contains the code for the paper "Solving Inverse Physics Problems with Score Matching" by Benjamin Holzschuh, Simona Vegetti, and Nils Thuerey. The paper can be found here.
Our works proposes a novel approach to solve inverse problems involving the temporal evolution of physics systems by leveraging the idea of score matching. The system’s current state is moved backward in time step by step by combining an approximate inverse physics simulator and a learned correction function. A central insight of our work is that training the learned correction with a single-step loss is equivalent to a score matching objective, while recursively predicting longer parts of the trajectory during training relates to maximum likelihood training of a corresponding probability flow. In the paper, we highlight the advantages of our algorithm compared to standard denoising score matching and implicit score matching, as well as fully learned baselines for a wide range of inverse physics problems. The resulting inverse solver has excellent accuracy and temporal stability and, in contrast to other learned inverse solvers, allows for sampling the posterior of the solutions.
Feel free to contact us if you have questions or suggestions regarding our work.
The code is written in Python 3.8 and tested with CUDA 11.4:
bash conda create -n smdp python=3.8
The majority of the code is based on JAX, which we install first with
conda activate smdp
# CUDA 11 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlFor each experiment, additional packages can be installed with
pip install -r requirements.txtwithin the experiment's directory.
Our code covers several experiments. Each experiment is located in a separate folder.
We learn score fields for 1D-processes with simple SDEs. The experiments are located in the folder toy-example.
The simple setup allows us to compare the score learned by our method with the analytical score and
evaluate how well the posterior distribution obtained from our method matches the true posterior distribution.
In a slightly more involved example, we learn the score field for Burgers' equation. The experiments are located in the folder burgers-equation.
A difficulty here is that the physics is very sensitive to small perturbations. Therefore, the 1-step training of our method is not sufficient to produce
stable trajectories over longer time horizons. This is why our proposed multi-step training is crucial for this example. Coming soon.
In this example, we learn the score field for the stochastic heat diffusion equation. The experiments are located in the folder heat-diffusion.
As the diffusive nature of the equation destroys information over time, small-scale structures need to be created during inference.
This highlights the advanatages of the SDE version of our method, as noise added to the trajectories can be used to create missing details.
This example is located in the folder buoyancy-flow. We learn the score field for a buoyancy-driven flow with obstacles.
What makes this experiment challenging is that it involves non-linear physics and randomly placed obstacles for each simulation.
This means that the learned score field needs to be able to generalize very well to unseen scenarios.
Finally, we learn the score field for isotropic forced turbulence. The experiments are located in the folder navier-stokes.
In this example, we do not have a numerical solver for the forward problem. Instead, we train a (time-independent) neural
network for the physics and a (time-dependent) neural network for the score field. Coming soon.
If you use our approach or code, please cite our paper:
@article{holzschuh2023smdp,
title="{Solving Inverse Physics Problems with Score Matching}",
author={Holzschuh, Benjamin and Vegetti, Simona and Thuerey, Nils},
journal={Advances in Neural Information Processing Systems},
year={2023}
}