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

