SCORE MATCHING VIA DIFFERENTIABLE PHYSICS

Abstract

Diffusion models based on stochastic differential equations (SDEs) gradually perturb a data distribution p(x) over time by adding noise to it. A neural network is trained to approximate the score ∇ x log p t (x) at time t, which can be used to reverse the corruption process. In this paper, we focus on learning the score field that is associated with the time evolution according to a physics operator in the presence of natural non-deterministic physical processes like diffusion. A decisive difference to previous methods is that the SDE underlying our approach transforms the state of a physical system to another state at a later time. For that purpose, we replace the drift of the underlying SDE formulation with a differentiable simulator or a neural network approximation of the physics. At the core of our method, we optimize the so-called probability flow ODE to fit a training set of simulation trajectories inside an ODE solver and solve the reverse-time SDE for inference to sample plausible trajectories that evolve towards a given end state. We demonstrate the competitiveness of our approach for different challenging inverse problems.

1. INTRODUCTION

Many physical systems are time-reversible on a microscopic scale. For example, a continuous material can be represented by a collection of interacting particles (Gurtin, 1982; Blanc et al., 2002) based on which we can predict future states of the material. We can also compute earlier states, meaning we can evolve the simulation backwards in time (Martyna et al., 1996) . When taking a macroscopic perspective, we only know the average quantities within specific regions (Farlow, 1993) , which constitutes a loss of information. It is only then that time-reversibility is no longer possible, since many macroscopic and microscopic initial states exist that evolve to yield the same macroscopic state. In the following, we target inverse problems to reconstruct the distribution of initial macroscopic states for a given end state. This genuinely tough problem has applications in many areas of scientific machine learning (Zhou et al., 1996; Gómez-Bombarelli et al., 2018; Delaquis et al., 2018; Lim & Psaltis, 2022) , and existing methods lack tractable approaches to represent and sample the distribution of states. We address this issue by leveraging continuous approaches for diffusion models in the context of physical simulations. In particular, our work builds on the reversediffusion theorem (Anderson, 1982) . Given the functions f (•, t) : R d → R d , called drift, and g(•) : R → R, called diffusion, it can be shown that under mild conditions, for the forward stochastic differential equation (SDE ) dx = f (x, t)dt + g(t)dw there is a corresponding reverse-time SDE dx = [f (x, t) -g(t) 2 ∇ x log p t (x)]dt + g(t)d w. In particular, this means that given a marginal distribution of states p 0 (x) at time t = 0 and p T (x) at t = T such that the forward SDE transforms p 0 (x) to p T (x), then the reverse-time SDE runs backward in time and transforms p T (x) into p 0 (x). The term ∇ x log p t (x) is called the score. This theorem is a central building block for SDE-based diffusion models and denoising score matching (Song et al., 2021c; Jolicoeur-Martineau et al., 2021) , which parameterize the drift and diffusion in such a way that the forward SDE corrupts the data and transforms it into random noise. By training a neural network to represent the score, the reverse-time SDE can be deployed as a generative model, which transforms samples from random noise p T (x) to the data distribution p 0 (x). In this paper, we show that a similar methodology can likewise be employed to model physical processes. We replace the drift f (x, t) by a physics model P(x) : R d → R d , which is implemented by a differentiable solver or a neural network that represent the dynamics of a physical system, thus deeply integrating physical knowledge into our method. The end state at t = T on which the forward SDE acts is not fully destroyed by the diffusion g(t), but instead, the noise acts as a perturbation of the system state over time. An overview of our method is shown in figure 1 . To the best of our knowledge, our work is the first to leverage the reverse-diffusion theorem as a method for solving inverse problems of physical systems. As such, our primary aim is to demonstrate how existing algorithms from this field can be used in the context of physics simulations. We showcase the efficacy of the score matching viewpoint on physics problems with a range of challenging inverse problems. Specifically, our contributions are: 1. We develop a framework in which we incorporate the reverse-diffusion theorem and score matching into a method for solving inverse problems that involve the time evolution of physical systems. We demonstrate its competitiveness against common baseline approaches using the heat equation as an example. 2. We highlight the effectiveness of our method with a more challenging inverse problem where we simulate a fluid-based transport process backwards in time in the presence of randomized obstacles. Here, we compare our method to different strategies for learned solvers. 3. Finally, we show that this approach can even be used when the underlying SDE is unknown. Our approach can be combined with operator learning methods and we demonstrate its effectiveness for learning the Navier-Stokes equation in the turbulent regime.

2. BACKGROUND AND RELATED WORK

Learned solvers: Numerical simulations benefit greatly from machine learning models (Tompson et al., 2017; Morton et al., 2018; Pfaff et al., 2020; Li et al., 2020) . By integrating a neural network inside differential equation solvers, it is possible learn to reduce numerical errors (Tompson et al., 2017; Kochkov et al., 2021; Brandstetter et al., 2022) or guide the simulation towards a desired target state (Holl et al., 2020b; Li et al., 2022) . As errors may accumulate quickly over time, trained networks benefit from gradients that are backpropagated over multiple time steps (Um et al., 2020) . Diffusion models: Diffusion models (Ho et al., 2020; Song et al., 2021c) have been considered for a wide range of applications. Most notably, diffusion models have been proposed for image (Dhariwal & Nichol, 2021 ), video (Ho et al., 2022; Höppe et al., 2022; Yang et al., 2022) and audio synthesis (Chen et al., 2021) . Recently, Bansal et al. (2022) have proposed to train generalized diffusion models for arbitrary transformations and suggest that fully deterministic models without any noise are sufficient for generative behaviour.



Figure 1: Overview: we employ a physics simulator P to learn the score field ∇x log pt(x) with a neural network s θ in the presence of noise or uncertainties. The trained model allows for sampling the posterior of p0, i.e. different states that explain an observation pT , via probability flow or by solving the reverse-time SDE.

