ACTION MATCHING: A VARIATIONAL METHOD FOR LEARNING STOCHASTIC DYNAMICS FROM SAMPLES

Abstract

Stochastic dynamics are ubiquitous in many fields of science, from the evolution of quantum systems in physics to diffusion-based models in machine learning. Existing methods such as score matching can be used to simulate these physical processes by assuming that the dynamics is a diffusion, which is not always the case. In this work, we propose a method called "Action Matching" that enables us to learn a much broader family of stochastic dynamics. Our method requires access only to samples from different time-steps, makes no explicit assumptions about the underlying dynamics, and can be applied even when samples are uncorrelated (i.e., are not part of a trajectory). Action Matching directly learns an underlying mechanism to move samples in time without modeling the distributions at each time-step. In this work, we showcase how Action Matching can be used for several computer vision tasks such as generative modeling, super-resolution, colorization, and inpainting; and further discuss potential applications in other areas of science.

1. INTRODUCTION

The problem of learning stochastic dynamics is one of the most fundamental problems in many different fields of science. In physics, porous medium equations (Vázquez, 2007) describe many natural phenomena from this perspective, such as Fokker Planck equation in statistical mechanics, Vlasov equation for plasma, and Nonlinear heat equation. Another prominent example is from Quantum Mechanics where the state of physical systems is a distribution whose evolution is described by the Schrödinger equation. Recently, stochastic dynamics have achieved very promising results in machine learning applications. The most promising examples of this approach are the diffusion-based generative models (Song et al., 2020b; Ho et al., 2020) . Informal Problem Setup In this paper we approach the problem of Learning Stochastic Dynamics from their samples. Suppose we observe the time evolution of some random variable X t with the density q t , from t 0 to t 1 . Having access to samples from the density q t at different points in time t ∈ [t 0 , t 1 ], we want to build a model of the dynamics by learning how to move samples in time such that they respect the marginals q t . In this work, we propose a method called "Action Matching" as a solution to this problem.

Learning Stochastic Dynamics vs. Time-Series

There is an important distinction between the problem of learning stochastic dynamics and time-series modeling (e.g., language, speech or video modeling). In time-series, the samples come in trajectories, where the samples in each trajectory are usually highly correlated. However, in learning stochastic dynamics, we only have access to independent samples at any given time-step (i.e., uncorrelated samples through time). This degree of freedom allows us to solve different types of problems that can not be approached by time-series modeling. We provide several examples in our experiment section, but also point out that sometimes it is even physically impossible to obtain samples along trajectories. For example, in Quantum Mechanics, the act of measurement at a given point collapses the wave function which prevents us from obtaining further samples along that trajectory. Generative Modeling with Action Matching From the Machine Learning perspective, the problem of learning stochastic dynamics is a generalization of generative modeling. One way so solve generative modeling is to first construct a distributional path (stochastic dynamics) from the data

