Probabilistic Forecasting with Stochastic Interpolants and Föllmer Processes

Introduction

Forecasting the future state of a dynamical system given complete or partial information about the current state is a ubiquitous problem across science and engineering, with diverse applications in climate modeling (Smagorinsky, 1963; Palmer et al., 1992; Gneiting et al., 2005; Pathak et al., 2022), fluid dynamics (Buaria & Sreenivasan, 2023), video prediction (Oprea et al., 2022; Finn et al., 2016; Lee et al., 2018), and extrapolation of time series data (Lim & Zohren, 2021; Masini et al., 2023). Forecasting can be performed deterministically, whereby the goal is to compute a single prediction for the future state (Giannakis et al., 2023), or probabilistically, whereby the goal is to predict a distribution over future states consistent with the current information (Gneiting & Katzfuss, 2014). Probabilistic forecasting is the natural formulation when the underlying system dynamics are stochastic, when the full system state cannot be measured, or when measurements are corrupted by noise, as is the case for most real-world systems of interest. Moreover, while deterministic forecasting appears to be a simpler problem, recent works have uncovered difficulties with deterministic forecasting methods when the underlying dynamics are chaotic (Jiang et al., 2023).

Motivated by the recent success of generative models built upon dynamical transport of measure, such as score-based diffusion (Ho et al., 2020; Song et al., 2020), flow matching (Lipman et al., 2022), and stochastic interpolants (Albergo & Vanden-Eijnden, 2022; Albergo et al., 2023), here we introduce a generative modeling approach for probabilistic forecasting. This approach maps the current state of the system onto the ensemble of possible future outcomes after a fixed time lag. From a transport perspective, this requires pushing a point mass measure onto a probability distribution with larger support. In what follows, we show that the interpolant framework enables us to design an artificial dynamics that performs this task using a stochastic differential equation. In practice, the drift fields entering the SDEs we introduce can be learned via square loss regression. We also show that the diffusion coefficient in these SDEs can be tuned a-posteriori (i.e. without having to retrain a drift). We show that a specific choice that minimizes the impact of the estimation error recovers a Föllmer process (Föllmer, 1986), a specific instantiation of the Schrödinger bridge problem (Schrödinger, 1932; Léonard, 2014) in which the base distribution is a point mass measure.

To demonstrate the utility and scalability of our approach, we consider several examples: an illustrative low-dimensional scenario based on a multi-modal jump diffusion process, a high-dimensional problem generated by a stochastic Navier-Stokes equation on the torus, and video generation on the KTH (Schuldt et al., 2004) and CLEVRER datasets (Yi et al., 2019). In the case of Navier-Stokes, we show that our probabilistic forecasting method is capable of reproducing quantitative metrics, such as the enstrophy spectrum of the dataset, using either high- or low-resolution measurement of the current system state. In each case, we highlight that we are able to obtain diverse samples consistent with the conditional distribution of interest, and demonstrate the need for probabilistic, as opposed to deterministic, forecasting. For the video generation tasks, we show that our learned models are more effective than standard conditional generative modeling. We also show that this forecasting procedure can be iterated autoregressively without retraining to compute a predicted trajectory. Our main contributions can be summarized as follows:

We design new generative models for probabilistic forecasting based on stochastic differential equations (SDEs) that map a point mass measure to a distribution with full support by incorporating stochasticity in a principled way, enabling us to initialize the SDE directly at the measured system state.

We prove that the drifts entering these SDEs can be learned via square loss regression over the data, and that the resulting loss has bounded variance.

We show that the drift and noise terms in these SDEs can be adjusted post-training, and that the specific choice of noise that minimizes the Kullback-Leibler (KL) divergence between the path measures of the exact forecasting process and the estimated one is realizable. We show that this drift and diffusion pair gives a Föllmer process.

We validate our theoretical results empirically on several challenging high-dimensional forecasting tasks, including the Navier-Stokes equation and video prediction.

Related Work

There is a vast body of literature on forecasting. Methodologies can be broadly classified into two main categories: forecasting a single output, typically achieved through regression and operator theoretic approaches (Kutz et al., 2016; Alexander & Giannakis, 2020; Li et al., 2021), and probabilistic forecasting (Gneiting & Katzfuss, 2014), based on stochastic and generative modeling.

For the class of methodologies focused on generating a single output, a widely used approach is regression or supervised learning, which entails directly learning a map in the state space. One can also work in the space of probability densities or functionals on the state space, under the setting of the Frobenius-Perron or Koopman operator approach (Dellnitz & Junge, 1999; Kaiser et al., 2021); this leads to linear dynamics in an infinite-dimensional space, and forecasting reduces to finding a tractable finite-dimensional approximation. In all formulations, the key is to identify and learn accurate representations of the dynamics, for example through nonparametric approaches such as diffusion maps (Berry et al., 2015) and kernel regression (Alexander & Giannakis, 2020), or parametric approaches such as linear regression, dynamical mode decomposition (Kutz et al., 2016), neural networks (Li et al., 2021; Gu et al., 2021) and operators (Lu et al., 2021; Jiang et al., 2023; Li et al., 2020). Ultimately, these approaches produce a function that maps the current state to a single output which is the deterministic forecast. Many methods under this category train models with the MSE (mean square error) or RMSE (root mean square error) as the objective, but these losses may be poor signals for training forecasters for chaotic systems (Jiang et al., 2023).

For dynamics that are inherently stochastic, or with incomplete information, a stochastic forecast is important to incorporate uncertainties. This stochasticity can be introduced by fitting a probabilistic model, such as a stochastic process or a graphical model, to data. They can also be approached from the perspective of stochastic Koopman operators (Wanner & Mezic, 2022; Zhao & Jiang, 2023). Although many approaches have been proposed, most of them target low-dimensional problems or simple conditional statistics such as Gaussians.

Probabilistic Forecasting with Generative Models.

Recently, generative modeling techniques in machine learning have received increasing attention for handling high-dimensional, complex distributions. Probabilistic forecasting can be seen as a form of the conditional generation problem. Various conditional generative models, including conditional GAN (Mirza & Osindero, 2014), VAE (Sohn et al., 2015), and normalizing flows (de Bézenac et al., 2020; Kidger et al., 2021) have been developed. More recently, diffusion generative models (Ho et al., 2020; Song et al., 2020) have gained popularity due to their state-of-the-art performance. The stochastic interpolant methodology, which is related to contemporary work like flow matching (Lipman et al., 2022), is a general framework that encompasses diffusion models. In the literature, there has been some work to develop conditional models for time series (Rasul et al., 2021; Lienen et al., 2023), including diffusion (Ho et al., 2022; Blattmann et al., 2023) and flow matching models for video prediction (Davtyan et al., 2023). These approaches learn ODEs or SDEs that map a Gaussian base to the conditional distribution of interest. There are also works that forecast stochastic dynamics by adding noise to the neural network layers (Cachay et al., 2023) using techniques such as dropout. In contrast, we incorporate stochasticity rigorously through interpolants and direct SDE modeling. Our construction of the stochastic generative model, which maps the current state to the distribution of the forecasted state, is new and can be seen as a direct stochastic extension of the deterministic map approach that is predominant in single output forecasting.

Föllmer processes.

The construction of SDEs that map a point mass to a target distribution dates back at least to the Föllmer process (Föllmer, 1986), which is a particular solution of the Schrödinger bridge problem (Schrödinger, 1932; Léonard, 2014; Chen et al., 2021) that minimizes the relative entropy with respect to the Wiener process. This approach has the desirable feature that it offers an entropy-regularized solution to the optimal transport problem. For this reason, the concepts of the Föllmer process and the Schrödinger bridge have found many applications in sampling densities with unknown normalization constants (Zhang & Chen, 2021; Huang et al., 2021; Jiao et al., 2021; Vargas et al., 2023), generative modeling (Tzen & Raginsky, 2019; Wang et al., 2021; De Bortoli et al., 2021; Liu et al., 2023; Peluchetti, 2023; Shi et al., 2024) as well as stochastic analysis and functional inequalities (Lehec, 2013; Eldan & Lee, 2018; Eldan et al., 2020). In this article, we show that the stochastic interpolant framework offers a simple way to construct new types of Föllmer processes. In addition, we give a new interpretation of Föllmer processes as minimizers of a KL divergence between the path measure of an SDE that forecasts exactly and the approximate, learned SDE. These results support the uses of Föllmer processes in probabilistic forecasting.

Setup and Main Results

Signal recovery from corrupted data.

Probabilistic forecasting from noisy observations.

2 Generation with Stochastic Interpolants

The generative models that we develop here are based on stochastic differential equations that map a fixed initial condition $X_{s=0}=x_{0}$ to samples from the conditional distribution $X_{s=1}\sim\rho_{c}(\cdot|x_{0})$. Towards the design of such SDEs, we first introduce the stochastic interpolant

where $(x_{0},x_{1})\sim\rho(x_{0},x_{1})$ and $W=(W_{s})_{s\in}$ is a Wiener process with $W\perp(x_{0},x_{1})$. In addition, we impose that $\alpha,\beta,\sigma\in C^{1}()$ satisfy the boundary conditions $\alpha_{0}=\beta_{1}=1$ and $\alpha_{1}=\beta_{0}=\sigma_{1}=0$. To facilitate some calculations, we assume that $\dot{\beta}_{s}>0$ for all $s\in(0,1]$ and $\dot{\sigma}_{s}<0$ for all $s\in$. Here we will use $\alpha_{s}=\sigma_{s}=1-s$, and $\beta_{s}=s$ or $\beta_{s}=s^{2}$ (see Appendix A.1). This second choice for $\beta_{s}$ has some advantages that we discuss below.

The boundary conditions on $\alpha,\beta,$ and $\sigma$ guarantee that $I_{s=0}=x_{0}$ and $I_{s=1}=x_{1}$, so that the probability distribution of $I_{s}|x_{0}$ bridges the point mass measure at $x_{0}$ to $\rho_{c}(\cdot|x_{0})$ as $s$ varies from 0 to 1. The following result shows that this probability distribution is also the law of the solution to a specific SDE that can be used as a generative model.

Let $b_{s}(x,x_{0})$ be the unique minimizer over all $\hat{b}_{s}(x,x_{0})$ of the objective

This theorem is proven in Appendix A.2. The result is formulated in a way that is tailored to practical approximation of the drift $b_{s}$, since the objective (3) can be estimated empirically by generating samples of $I_{s}$ and $R_{s}$ using sample pairs $(x_{0},x_{1})$ from $\rho$ and realizations of $W_{s}\stackrel{{\scriptstyle d}}{{=}}\sqrt{s}z$ with $z\sim{\sf N}(0,\text{\it Id})$. That is, $b_{s}$ may be learned over neural networks by minimizing the simulation-free loss (3) over the parameters. It is easy to see that the minimizer is given by

3 Generalizations with Tunable Diffusion

We now show that learning the drift coefficient (6) gives access to a broader set of SDEs to use as a generative model beyond just (5), and that selecting from them has appealing theoretical motivation.

Let $\rho_{s}(x|x_{0})$ be the PDF of $X_{s}\stackrel{{\scriptstyle d}}{{=}}I_{s}|x_{0}$. From (5), $\rho_{s}$ solves the Fokker-Planck equation

Given a candidate diffusion coefficient $g_{s}$, we can use the identity $\tfrac{1}{2}\sigma^{2}_{s}\Delta\rho_{s}=\tfrac{1}{2}g^{2}_{s}\Delta\rho_{s}-\tfrac{1}{2}(g_{s}^{2}-\sigma_{s}^{2})\nabla\cdot(\rho_{s}\nabla\log\rho_{s})$ to trade diffusion for transport in (7). This construction leads to a family of SDEs with tunable diffusion

Given any $g\in C^{0}()$ such that $\lim_{s\to 0^{+}}s^{-1}[g_{s}^{2}-\sigma^{2}_{s}]$ and $\lim_{s\to 1^{-}}g^{2}_{s}\sigma^{-1}_{s}$ exist, define

where $b_{s}(x,x_{0})$ is the minimizer of (3) given in (6) and $\rho_{s}(x|x_{0})$ is the PDF of $X_{s}\stackrel{{\scriptstyle d}}{{=}}I_{s}|x_{0}$. Then the solutions to the SDE

This theorem is proven in Appendix A.4, where we explain why the conditions on $g_{s}$ guarantee that the SDE (9) is well-posed. Working with this SDE requires the score $\nabla\log\rho_{s}$. Interestingly, this score can be expressed in terms of the drift $b_{s}$. A direct calculation reported in Appendix A.4 shows that

Using (10) in (8) shows that, to work with the SDE (9), we can estimate $b$ first and then adjust both the noise amplitude $g_{s}$ and the drift $b^{g}$ a-posteriori without having to retrain $b$.Using (10) in (8) requires some care at $s=0$ and $s=1$ due to the factor $[s\sigma_{s}]^{-1}$ in $A_{s}$, but this leads to no issue see Algorithm 2 and also Appendices A.1 and A.4. This offers flexibility at sampling time that can be leveraged to maximize performance, as shown in our numerical experiments below.

4 KL Optimization and Föllmer Processes

This expression is well-defined for all $s\in$, since $\lim_{s\to 0^{+}}2s\beta_{s}^{-1}\dot{\beta}_{s}<\infty$ because $\beta_{s}$ is differentiable at $s=0$ by assumption. The result in (13) is also amenable to an interesting interpretation:

If $\beta_{s}/[\sqrt{s}\sigma_{s}]$ is non-decreasing, then the process $X^{\text{\sc F}}\equiv X^{g^{\text{\sc F}}}$ that solves (9) with $g_{s}=g^{\text{\sc F}}_{s}$ is a Föllmer process.

This theorem is proven in Appendix A.6. To understand its significance, recall that the Föllmer process is the solution to the Schrödinger bridge problem when one of the endpoint measures is a point mass (in this case, at $x_{0}$). As such, it offers an entropy-regularized solution to the optimal transport problem. The Föllmer process is usually defined by minimizing its KL divergence with respect to the Wiener process subject to constraints on the endpoints. Theorem 3.3 offers a generalization and new interpretation of this process as the minimizer of the KL divergence of the exact forecasting process from the estimated one, which is more tailored to statistical inference. For more details about Föllmer processes and the Schrödinger bridge problem we refer the reader to Appendix A.5. We also test the performance of (13) in Appendix B.2.

5 Implementation

where $B_{K^{\prime}}\subset\{0:K\}$ is a subset of indices of cardinality $K^{\prime}$ and

with $z_{k}\sim{\sf N}(0,Id)$, $z_{k}\perp(x_{k\tau},x_{(k+1)\tau})$. To arrive at (14) we used that $W_{s}\stackrel{{\scriptstyle d}}{{=}}\sqrt{s}z$ with $z\sim{\sf N}(0,Id)$ at all $s\in$. In (14) and (16) below, the physical lag $\tau>0$ is fixed, while $s$ varies over $$, and the integral over$s$can be approximated via an empirical expectation over draws of$s\sim{\sf U}()$. By approximating$\hat{b}$ in an expressive parametric class such as a class of neural networks, we can optimize (14) over the parameters with standard gradient-based methods. This can be performed by batching over subsequences in the available time series, or via online learning if a stream of data is continuously observed.

Having learned an approximation $\hat{b}$, we can construct an approximation of $\hat{b}^{g}$ using (8), (10), and (11). We may then form our model given the new observation $x_{k\tau}$ by solving

with various realizations of the noise $W_{s}$ to generate a set of $\hat{X}^{k}_{s=1}$ that approximately samples $\rho_{c}(\cdot|x_{k\tau})$. This generates an ensemble of forecasts with statistics consistent with those of the time-series seen during training. This process can also be iterated autoregressively by setting $\hat{X}^{k+1}_{s=0}=\hat{X}^{k}_{s=1}$ and by solving the SDE (16) with $k$ replaced by $k+1$ to get an approximate sample $\hat{X}^{k+1}_{s=1}$ of $\rho_{c}(\cdot|x_{(k+1)\tau})$. This iteration does not require any additional training, since it uses the same $\hat{b}_{s}$. These procedures are summarized in Algorithms 1 and 2. The first step to get $X_{s=s_{1}}=\hat{X}_{1}$ in Algorithm 2 is consistent and designed so that it avoids computing $\hat{b}^{g}_{s=0}$, since using (8) and (10) can exhibit numerical singularities even though $b^{g}_{s=0}$ is well defined when the conditions of Theorem 3.2 are met.

Numerical illustrations

In what follows, we test our proposed method in several application domains. For all tests, an interpolant with coefficients $\alpha_{s}=1-s$, $\sigma_{s}=\varepsilon(1-s)$ for some $\varepsilon>0$, and $\beta_{s}=s^{2}$ is used. The condition that $\dot{\beta}_{0}=0$ empirically ensures that the norm of the parameter gradients used to train our neural networks are well behaved. We report results with the diffusion coefficient in (16) chosen to be $g_{s}=\sigma_{s}$, as we found that the impact of learning $\hat{b}_{s}$ well by choice of the right interpolant outweighed the effect of varying the SDE for the systems we study. For additional numerical experiments with $\alpha_{s}=1-s,\sigma_{s}=\varepsilon(1-s)$ and $\beta_{s}=s$, and with $g_{s}=g_{s}^{\text{\sc F}}$, we refer the reader to Appendix B.2. Investigation of the Föllmer SDE described in Theorem 3.3, in both theory and experiment, will be saved for future work.

Our first example is synthetic, and consists of forecasting a two-dimensional jump-diffusion process with invariant PDF given by a Gaussian mixture (Fig. 1). We study a particle governed by Langevin dynamics that is randomly kicked in the counterclockwise direction, where the times between kicks are specified by a Poisson process (see Appendix B.1 for details).

In this example, if the process starts at a point in one mode, its PDF spreads in the mode and leaks into the other nearby modes in the counterclockwise direction (see Fig. 2). Correspondingly, the conditional PDF of $x_{\tau}$ given $x_{0}$ is itself a Gaussian mixture PDF that is sharply peaked around $x_{0}$ when $\tau$ is small, and which slowly evolves into the 5-mode invariant measure associated with the potential as $\tau$ increases. We generate a long time series of this process and use it at different lags $\tau$ in the empirical loss (14) to learn the drift velocity $\hat{b}$, which we model as a fully connected neural network. We then use the estimated $\hat{b}^{g}$ in the SDE (16) to generate probabilistic forecasts. The results (Fig. 2) indicate that the law of these forecasts is in excellent agreement with the true $\rho_{c}(x_{\tau}|x_{0})$. We can also iterate using the procedure described in Sec. 3.5 to estimate $\rho_{c}(x_{k\tau}|x_{0})$ for $k>1$ without additional retraining. We find excellent agreement in doing so, including beyond the decorrelation time, when the conditional PDF relaxes into the equilibrium distribution independent of $x_{0}$.

We note that this is an example in which probabilistic forecasting is key, as regressing $x_{\tau}$ given $x_{0}$ would give results with little information. Probabilistic forecasting is needed to capture the non-Gaussian and multimodal nature of the forecasts, which would be hard to capture with deterministic methods.

2 Forecasting the 2d Navier-Stokes Equations

Here $v=\nabla^{\perp}\psi=(-\partial_{y}\psi,\partial_{x}\psi)$ is the velocity expressed in terms of the stream function $\psi$, which is a solution to $-\Delta\psi=\omega$, $d\eta$ is white-in-time random forcing acting on a few Fourier modes, and $\nu,\alpha,\varepsilon>0$ are parameters (see Appendix B.2 for details). We work in a setting where (17) is provably ergodic with a unique invariant measure (Hairer & Mattingly, 2006). Our objective is to forecast the solution to (17) at time $t+\tau$ given its solution at time $t$ after the process has reached a statistically steady state. We do so using both full- and low-resolution data at time $t$, but our goal is always to forecast at full resolution.

We employ a pseudo-spectral method to simulate (17) and hence to obtain a dataset of snapshots of the vorticity field. We set the timestep $\Delta t=10^{-4}$ and grid size to $256\times 256$. We store snapshots at regular intervals of $\Delta t=0.5$. We conduct simulations for 2000 trajectories within the time range of $t\in$; we then exclude the initial phase $t\in$ from our data. Ultimately, we collect a total of $2\times 10^{5}$ snapshots, which are treated as samples from the invariant measure of (17). To reduce memory requirements, we downsize the dataset to a resolution of $128\times 128$. In all our experiments, we used a UNet (Ho et al., 2020) as our network for approximating the velocity field. Detailed parameters for the training and dataset generation can be found in the Appendix B.2.

Forecasting at Full Resolution.

First, we consider predicting the distribution of the vorticity field that may evolve from a given realization. To this end, we learn the SDE that samples the conditional distribution of vorticity fields after lag $\tau=0.5$, and we iterate this SDE to get forecasted predictions after lag $2\tau$, $3\tau$, etc. In the top row of Fig. 3, the first panel shows a snapshot of a vorticity field, while the next three panels show different samples of vorticity fields generated after lag $2$ by iterating our forecasting procedure.

Note that these generated vorticity fields are different from one another, emphasizing the need for probabilistic forecasting. This is corroborated by the true and forecasted conditional means of the field shown in the first two panels on the bottom row of Fig. 3. While correctly captured by our approach, this conditional mean is clearly not informative on its own, as it averages over the spatial features present in the actual forecast. The spread of this ensemble of forecasts is also apparent from the standard deviation of the field, shown in the third and fourth panel on the bottom row of Fig. 3. Also shown in the right panel is the enstrophy spectrum of the true vorticity field and the ensemble of forecasts (see Section B.2), showing that our method captures this important physical quantity correctly.

Forecasting from Low-Resolution Data.

We now consider forecasting the vorticity field at a resolution of $128\times 128$ from a downsized version at resolution $32\times 32$ by learning the SDE as before with a drift velocity that is conditioned on the low-resolution field used as $x_{0}$. The outcome of this task is shown in Fig. 4, where the four panels on the right show the low-resolution field used as input surrounded by three full-resolution forecasts generated after lag $=1$. We also plot the enstrophy spectrum of the low-resolution field (which does not go past wavenumber 16) and the spectra of the true and the forecasted field at full resolution in the left panel. As can be seen, our approach recovers the true spectrum very accurately.

We present additional experiments such as superresolution, forecasting efficiency compared to simulation time ($100\times$ faster), comparisons between using $\sigma_{s}$ and $g^{\rm F}_{s}$ (Föllmer process) in terms of KL accuracy, and comparisons to flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo & Vanden-Eijnden, 2022) and deterministic forecasting in Appendix B.2.

3 Video Forecasting

We model the KTH and CLEVRER datasets. We follow RIVER (Davtyan et al., 2023) and model these videos in the latent space of a VQGAN (Esser et al., 2021) trained to auto-encode the datasets. Modeling in a latent space is a computationally efficient strategy for generative modeling of high resolution images (Vahdat et al., 2021; Rombach et al., 2022; Peebles & Xie, 2023; Ma et al., 2024) and video (Blattmann et al., 2023; Davtyan et al., 2023), because it reduces the dimensionality of the dataset.

Generation.

Because conditioning on the whole set of $C$ previous time slices is costly, we follow RIVER (Davtyan et al., 2023) and use a Monte-Carlo estimator that generates our estimate of the $t^{th}$ latent frame $\hat{y}^{t}$ conditional on the estimated latent frame $y^{t-1}$ and an additional estimated frame $y^{t-j}$ randomly chosen for $j>1$. To give the network context for the conditioned frame, we also condition on the time index $t-j$. This random conditioning set $(y^{t-1},y^{t-j},t-j)$ avoids the need to compute functions of the entire conditioning context, which consists of many frames. Samples are then decoded by the auto-encoder to produce images $x^{t}=\text{Decode}(y^{t})$. To sample a full video, we apply this forecasting strategy autoregressively. For further details, see Algorithm 3 in Appendix B.3.

FVD metric.

The Fréchet Video Distance (FVD) (Unterthiner et al., 2018) extends the Fréchet Inception Distance (FID) (Heusel et al., 2017). We select 256 test set videos and generate 100 completions for each one, thereby comparing 256 real videos to 25,600 generated videos. We also report several qualitative features of the generated videos. One consequence of using the VQGAN is that the performance is bounded by the FVD of the decoded-encoded data since the generative model targets the encoded distribution.

Baseline.

We compare with the setup from (Davtyan et al., 2023), which learns a deterministic flow using flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo & Vanden-Eijnden, 2022) to map a Gaussian sample to the next video frame, conditioned on the same information as in our setup. By contrast, we generate with an SDE sampler initialized at the previous video frame, which is more proximal to the next target frame than pure noise. The RIVER method requires training a new embedding model for the videos using a VQGAN. We use the VQGAN checkpoints from their work so that we can study our proposed generative modeling method in a controlled context.

Training.

We train models for 250k gradient steps using AdamW starting at a learning rate of 2e-4. We use the UNet architecture popularizedWe use the lucidrains repository. in (Ho et al., 2020). We modify the architecture to condition on past frames by concatenating them along the channel dimension of the input. Each model is trained on four A100 GPUs for approximately 1-2 days.

Datasets.

The KTH dataset (Schuldt et al., 2004) consists of black-and-white videos of 25 people completing one of 6 actions such as jogging and hand-waving. We use the last 5 people as the test set. The $1\times 64\times 64$-dimensional data space is mapped to a $4\times 8\times 8$-dimensional latent space by the VQGAN. During generation, we start with 10 given video frames, and we generate the next 30 frames. The CLEVRER datasethttp://clevrer.csail.mit.edu/ (Yi et al., 2019) contains videos created for studying reasoning and physics tasks. The videos feature cubes, spheres, and other shapes traveling across the screen and interacting through collisions while being subject to forces and rotations. One interpretation of this data is that a generative model needs to deduce physical phenomena to succeed at the generation. For example, objects should not go through each other, and instead should bounce off one another. The $3\times 128\times 128$-dimensional data space is mapped to a $4\times 16\times 16$ latent space. During generation, we condition on just 2 real frames and generate 14.

Results.

Table 1 shows the FVD performance of our model, probabilistic forecasting with interpolants (PFI), as compared to the RIVER baseline. We train both models under the same conditions for a controlled comparison. On both the KTH and CLEVRER datasets, PFI surpasses the standard flow matching approach. In addition to the numerical comparison, we demonstrate that our models produce diverse forecasts that capture physical rules inherent to the videos. In Figure 5, we illustrate two trajectories of the animation. For each, we supply the initial condition, the dataset trajectory, and a generated trajectory based off of the same initial condition to show that the continuation of frames is probabilistic. For example, there is variation in the dataset trajectory of a green cube colliding with a green cylinder as compared to the forecasted trajectory, while also preserving the animated physics. In Figure 6, we show that, from the same initial frame, the forecasts give varied realizations of the hand-waving video category.

Conclusion and Future Work

In this work, we introduced a principled approach to the use of generative modeling for probabilistic forecasting. By introducing stochastic processes that transport a point mass centered at a current observation of the system to a distribution over future states of that system, our proposed method uses dynamical measure transport in a way that naturally aligns with the framework of probabilistic forecasting. It also allows us to minimize the impact of the estimation error by tuning of the diffusion coefficient, which can be done after training and offers a new perspective on the Föllmer process. We have shown various uses of this approach, ranging from predicting the evolution of stochastic fluid dynamics to video completion tasks. Future work will consider the use of these models for empirical weather data and incorporation of physical structure into the generative model.

Impact Statement

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here. We note, however, that video generation, which is one of the possible applications of the forecasting framework presented, is a less explored domain that could promote harm through biases in the model. We surmise that it carries the same risks as image generation technologies.

Acknowledgments

We thank Joan Bruna, Jon Niles-Weed, Loucas Pillaud-Vivien, and Valentin De Bortoli for useful discussions.

References

Appendix A Details on stochastic interpolants

In this section, we present the formulas for some specific choices of $\alpha_{s},\beta_{s}$, and $\sigma_{s}$, the corresponding optimal (Föllmer) drift $g^{\rm F}_{s}$, and the expression for $b^{g}_{s}$ in terms of $b_{s}$.

For $\alpha_{s}=1-s,\sigma_{s}=\varepsilon(1-s)$ where $\varepsilon>0$ is a tunable parameter, and $\beta_{s}=s$, we have

Moreover, $b_{s}$ is the minimizer of the loss (with respect to $\hat{b}$)

so that, for any $g_{s}$ such that $\lim_{s\to 0^{+}}s^{-1}[g^{2}_{s}-\varepsilon^{2}(1-s)^{2}]$ and $\lim_{s\to 1^{-}}g^{2}_{s}/(\varepsilon(1-s))$ exists,

For $\alpha_{s}=1-s,\sigma_{s}=\varepsilon(1-s)$, and $\beta_{s}=s^{2}$, we have

Moreover, $b_{s}$ is the minimizer of the loss

so that, for any $g_{s}$ such that $\lim_{s\to 0^{+}}s^{-1}[g^{2}_{s}-\varepsilon^{2}(1-s)^{2}]$ and $\lim_{s\to 1^{-}}g^{2}_{s}/(\varepsilon(1-s))$ exist,

We summarize the above calculations in the following table:

A.2 Proof of Theorem 3.1

The stochastic interpolant $I_{s}$ is the stochastic process defined as

$\alpha,\beta,\sigma\in C^{1}()$ satisfy $\alpha_{s}^{2}+\beta^{2}_{s}+\sigma^{2}_{s}>0$ for all $s\in$, $\dot{\beta}_{s}>0$ for all $s\in(0,1]$, and $\dot{\sigma}_{s}<0$ for all $s\in$, as well as the boundary conditions $\alpha_{0}=\beta_{1}=1$, $\alpha_{1}=\beta_{0}=\sigma_{1}=0$.

$W=(W_{s})_{s\in}$ is a standard Wiener process with $W\perp(x_{0},x_{1})$.

In view of this definition, let us give a more precise formulation of Theorem 3.1:

Let $I_{s}$ be the stochastic interpolant introduced in Definition A.1 and let

Note that the objective (33) is the same as (3) because $x_{s}\stackrel{{\scriptstyle d}}{{=}}I_{s}$ and $\dot{\alpha}_{s}x_{0}+\dot{\beta}_{s}x_{1}+\dot{\sigma}_{s}\sqrt{s}\,z\stackrel{{\scriptstyle d}}{{=}}R_{s}$ at all $s\in$. Note also that, since $\alpha,\beta,\sigma\in C^{1}()$, the factors $\dot{\alpha}_{s},\dot{\beta}_{s},\dot{\sigma}_{s}\sqrt{s}$ in the loss (33) are all bounded.

where the first conditional expectation is over $(x_{0},x_{1})\sim\rho$ and $z\sim\sf{\sf N}(0,\text{\it Id})$ and the second over $(x_{0},x_{1})\sim\rho$ and $W$. By the Itô formula we have

Integrating this equation in time over $[0,s]$, taking the expectation conditional on $x_{0}$, and using both $W_{s}\perp(x_{0},x_{1})$ and the Itô isometry, we deduce that

where we used the fact that $x_{s}$ and $I_{s}$ share the same law at each $s$. Also, $W_{s}$ and $\sqrt{s}z$ share the same law at each $s$.

Using the tower property of the conditional expectation, (37) can also be written as

A.3 Regularity of the drift at s=0𝑠0s=0

In this section, we discuss the regularity of the drift $b_{s}(x,x_{0})$.

The conditional distribution $\rho_{c}(\cdot|x_{0})$ is exponential tailed. More precisely, there exist constants $C_{1},C_{2}>0$ (which can depend on $x_{0}$), such that

The aforementioned assumption is needed for technical reasons, and is used to ensure the validity of a step that involves the interchange of limits and integrations in the proof of Theorem A.4.

By linearity of $x_{s}$ in $x_{0}$, $x_{1}$, and $z$, we can also establish a few properties of the velocity $b$ in (31) which we state as:

Under Assumption A.3, if $\dot{\beta}_{s=0}=0$, then the velocity field $b_{s}(x,x_{0})$ can be decomposed as

In addition, these two functions satisfy the constraint

Note that (43) implies that we can get $\eta_{1}$ from $\eta_{z}$ at any time such that $\beta_{s}\not=0$, and $\eta_{z}$ from $\eta_{1}$ at any $s\in(0,1]$ such that $\sigma_{s}>0$: in particular

The proof of Theorem A.4 will rely the following result:

where $\delta(x)$ denotes the Dirac delta distribution. For any $s\in(0,1]$ such that $\sigma_{s}>0$, we can perform the integration over $z$ explicitly to get

Since the factors $e^{-\frac{1}{2}s^{-1}\sigma^{-2}_{s}|x-\alpha_{s}x_{0}|^{2}}$ at the numerator and the denominator do not depend on $x_{1}$, they can be taken out of the integrals and simplified, and we arrive at (45). ∎

The only statement that is not a direct consequence of Theorem A.2 is (42). To establish these limits, notice that, since $\beta\in C^{2}()$ and $\dot{\beta}_{0}=0$, we must have $\beta_{s}=O(s^{2})$ as $s\to 0$. As a result, since $\sigma\in C^{1}()$ and $\sigma_{0}>0$, we have

which establishes the limits in (42), where we used the formula in (44) to derive the limits of $\eta_{z}$ from that of $\eta_{1}$. Note that in the above derivation, we need to verify the interchange of limits and integrations; it is guaranteed by using Assumption A.3 and the Lebesgue dominated convergence theorem since, for a fixed $x,x_{0}$, when $s$ is sufficiently small, the factor $\rho_{c}(x_{1}|x_{0})e^{-\frac{1}{2}M_{s}|x_{1}|^{2}+m_{s}x_{1}\cdot(x-\alpha_{s}x_{0})}$ is dominated by $\rho_{c}(x_{1}|x_{0})e^{\frac{1}{2}C_{1}|x|}$, which is integrable as a function of $x_{1}$ due to our Assumption A.3.

To analyze the limit of $\nabla_{x}b_{s}(x,x_{0})$, using Lemma A.6, we have

To derive the above formula, we also need verify the interchange of limits and integrations; it is again guaranteed by using Assumption A.3 and the Lebesgue dominated convergence theorem, for sufficiently small $s$. We know that

and thus $\lim_{s\to 0}\nabla_{x}\eta_{1}(s,x,x_{0})=0$ as $\lim_{s\to 0}m_{s}=0$. Using the formula in (44), we get $\lim_{s\to 0}\nabla_{x}\eta_{z}(s,x,x_{0})=\frac{1}{\sigma_{0}}\text{\it Id}$. Therefore,

A.4 Changing the Diffusion Coefficient: Proof of Theorem 3.2

Theorem 3.2 and formula (10) are consequences of the following result:

Let $g\in C^{0}()$ be such that $\lim_{s\to 0^{+}}s^{-1}[g^{2}_{s}-\sigma^{2}_{s}]$ and $\lim_{s\to 1^{-}}g^{2}_{s}\sigma^{-1}_{s}$ exist. Let $b$ be given by (31) and define

where $\rho_{s}(x|x_{0})$ denotes the PDF of $X_{s}\stackrel{{\scriptstyle d}}{{=}}X^{g}_{s}\stackrel{{\scriptstyle d}}{{=}}I_{s}|x_{0}$. In addition the drift

Let us first establish (55). By a direct extension of Stein’s formula (Albergo et al., 2023), we have

where $x_{s}=\alpha_{s}x_{0}+\beta_{s}x_{1}+\sqrt{s}\sigma_{s}z$ and where we used $x_{s}|x_{0}\stackrel{{\scriptstyle d}}{{=}}I_{s}|x_{0}\stackrel{{\scriptstyle d}}{{=}}X_{s}$ for all $s\in$. Since

which can be inserted in (58) to show that (55) holds.

Second, note that the PDF $\rho_{s}(x|x_{0})$ of the solution to the SDE (57) is the same as the PDF of the solution to

since we can use the identity $\tfrac{1}{2}\sigma^{2}_{s}\Delta\rho_{s}=\tfrac{1}{2}g^{2}_{s}\Delta\rho_{s}-\tfrac{1}{2}(g_{s}^{2}-\sigma_{s}^{2})\nabla\cdot(\rho_{s}\nabla\log\rho_{s})$ to show that both densities satisfy the same Fokker-Planck equation (7).

It remains to show that the drift coefficient (56) is well-defined. To this end let us write it explicitly as

Our assumptions that $\alpha,\beta,\sigma\in C^{1}()$ and satisfy $\alpha_{0}=\beta_{1}=1$, $\alpha_{1}=\beta_{0}=\sigma_{1}=0$, and $\dot{\beta}_{s}>0$ for all $s\in(0,1]$ and $\dot{\sigma}_{s}<0$ for all $s\in$ guarantee that $\beta_{s}>0$ for all $s\in(0,1]$ and $\sigma_{s}>0$ for all $s\in[0,1)$ and, as a result, $[s\sigma_{s}(\dot{\beta}_{s}\sigma_{s}-\beta_{s}\dot{\sigma}_{s})]^{-1}$ is positive and finite for all $s\in(0,1)$. These assumptions also guarantee that $\beta_{s}[\sigma_{s}\dot{\beta}_{s}-\dot{\sigma}_{s}\beta_{s}]^{-1}$ and $\dot{\beta}_{s}[\sigma_{s}\dot{\beta}_{s}-\dot{\sigma}_{s}\beta_{s}]^{-1}$ have finite limits at $s=0,1$. Therefore the only factor in (62) that can be singular is $(g_{s}^{2}-\sigma_{s}^{2})/(2s\sigma_{s})$ at $s=0$, because of the factor $s^{-1}$, and at $s=1$ because of the factor $\sigma^{-1}_{1}$. These singularities disappear under our assumptions that $\lim_{s\to 0^{+}}s^{-1}[g^{2}_{s}-\sigma^{2}_{s}]$ and $\lim_{s\to 1^{-}}g^{2}_{s}\sigma^{-1}_{s}$ exist and are finite. Therefore the drift $b^{g}_{s}$ has the same regularity properties has $b_{s}$. ∎

A.5 Maximizing the likelihood with respect to the noise schedule

Let $X^{g}=(X^{g}_{s})_{s\in}$ be the process defined by the SDE (57), and let $\hat{X}^{g}=(\hat{X}^{g}_{s})_{s\in}$ denote an approximate, learned process governed by

With a slight abuse of notation, let us denote by $D_{\text{KL}}(X^{g}||\hat{X}^{g})$ the KL divergence of the path measure of $X^{g}$ from the path measure of $\hat{X}^{g}$. By Girsanov’s theorem, it is given by

Since $L_{s}$ is independent of $g_{s}$, minimizing (65) amounts to minimizing

for all $s\in$. Since $A_{s}>0$, (67) is minimized at

Together, these relations show that the minimizer of (67) is $g_{s}=g_{s}^{\text{\sc F}}$ with

where we used the expression for $A_{s}$ in (11) to get the second equality. This result is (13). Notice that we can also write

Since the sign of this expression is the same as the sign of $1-\tfrac{1}{2}\beta_{s}A_{s}\sigma^{2}_{s}$, it shows that (68) and (69) arise at values of $s\in$ where $\beta_{s}/[\sqrt{s}\sigma_{s}]$ is respectively decreasing or increasing in $s$. Since $\dot{\beta}_{s}>0$ and $\dot{\sigma}_{s}<0$, $\beta_{s}/[\sqrt{s}\sigma_{s}]$ cannot be decreasing for all $s\in$, i.e. the minimum of (67) cannot be zero for all $s\in$. As a result, the minimum of the KL divergence (65) must be positive if $L_{s}>0$.

Finally, let us investigate the conditions on $g_{s}$ in Theorem A.7 if we use the SDE (63) with $g_{s}=g_{s}^{\text{\sc F}}$. It is easy to see from (70) that the second condition is always satisfied since

since $\sigma_{1}=0$. Regarding the first condition, we have

since $\sigma_{0}>0$, $\dot{\sigma}_{0}<0$. Since $\beta_{0}=0$, if $\dot{\beta}_{0}>0$ and $\ddot{\beta}_{0}$ exists, we have

If, however $\dot{\beta}_{0}=0$, then $g_{0}^{\text{\sc F}}\not=\sigma_{0}$, and the first condition is not satisfied since $\lim_{s\to 0^{+}}s^{-1}[|g^{\text{\sc F}}_{s}|^{2}-\sigma^{2}_{s}]$ does not exist. In this case, we need to consider more carefully how to define the solution to the SDE (63). In the specific case when $\alpha_{s}=1-s$, $\sigma_{s}=\varepsilon(1-s)$ and $\beta_{s}=s^{2}$, if we set $g_{s}=g^{\text{\sc F}}_{s}=\varepsilon\sqrt{(1-s)(3-s)}$ in (63), this SDE reduces to (see the explicit formulas given in Appendix A.1 and denoting $X^{\text{\sc F}}_{s}=X^{g^{\text{\sc F}}}_{s}$):

The drift in this equation is singular at $s=0$ because of the term $2(X^{\text{\sc F}}_{s}-x_{0})/[s(2-s)]$. Nevertheless, the solution to this SDE is well-defined for the initial condition $X^{\text{\sc F}}_{s=0}=x_{0}$, and satisfies the integral equation

Since $b_{0}(x_{0},x_{0})=-x_{0}$ when $\dot{\beta}_{0}=0$, this equation implies that

which is also the law of $I_{s}|x_{0}$ as it should.

A.6 Connection with Föllmer Processes and Proof of Theorem 3.3

To begin, we give some background on the Föllmer process (Föllmer, 1986; Tzen & Raginsky, 2019). Originally, the Föllmer process was defined as the process $X=(X_{s})_{s\in}$ whose path measure has minimal KL divergence from the path measure of the Wiener process $W=(W_{s})_{s\in}$ under the constraint that $X_{s=1}$ be distributed according to some target distribution. This Föllmer process can be generalized to “reference processes” that differ from the standard Wiener process: as we will show next, in our context the natural reference process is the solution to the linear SDE

where $g^{\text{\sc F}}_{s}$ is given by (70) and where we have defined

We can then define the Föllmer process $X^{\text{\sc F}}=(X_{s}^{\text{\sc F}})_{s\in}$ by adjusting the drift $\check{b}_{s}(x,x_{0})$ in

in such a way that the KL divergence of the path measure of $\check{X}$ from the path measure of $Y$ (the solution to (79)) is minimized subject to the constraint that $\check{X}_{s=1}\sim\rho_{c}(\cdot|x_{0})$. That is, the Föllmer process is defined via

This traditional minimization problem for the Föllmer process is distinct from the minimization problem (66) considered in Appendix A.5. Nevertheless, our next result shows that the minimizers of (66) and (82) coincide:

Assume that $\beta_{s}/[\sqrt{s}\sigma_{s}]$ is non-decreasing on $s\in$. Then, the Föllmer process associated with the reference process $(Y_{s})_{s\in}$ that solves (79) is the process $X^{\text{\sc F}}\equiv(X^{g^{\rm F}}_{s})_{s\in}$ that solves (57) with $g_{s}=g^{\text{\sc F}}_{s}$ given by (70).

We closely follow the steps of Föllmer’s original construction involving time-reversal (Föllmer, 1986). To begin, notice that we can solve (55) to express $b_{s}$ in terms of the score $\nabla\log\rho_{s}$. We can then use the resulting expression in (56) to write the SDE (57) as

where $A_{s}$ and $c_{s}(x,x_{0})$ are defined in (11). Using the score, we can time-reverse (83) and derive the following SDE for $X_{s}^{\text{\sc R}}\overset{d}{=}X^{g}_{1-s}$

If we take $X^{\text{\sc R}}_{s=0}\sim\rho_{c}(\cdot|x_{0})$ in (84), then by construction we have that $X^{\text{\sc R}}_{s=1}=x_{0}$. Remarkably, if we use $g_{s}=g^{\text{\sc F}}_{s}$ with $g_{s}^{\text{\sc F}}$ given in (70), and if we choose $1-\frac{1}{2}\beta_{s}A_{s}\sigma^{2}_{s}\geq 0$ (i.e. choose $\beta_{s}/[\sqrt{s}\sigma_{s}]$ nondecreasing), the SDE (84) reduces to

This reverse-time SDE has the remarkable property that its drift is independent of the score $\nabla\log\rho_{s}$, meaning that the information about the target PDF $\rho_{c}(\cdot|x_{0})$ only enters through its initial condition $X^{\text{\sc R}}_{s=0}\sim\rho_{c}(\cdot|x_{0})$. This means that we can change the initial condition in (85) to construct (after reversing time back) a reference process. We can use any density for this purpose, but for simplicity it is convenient to choose a Gaussian, and therefore to consider

Using the explicit form of $c_{s}$ in (54), we can write this SDE explicitly as

Using that $d\alpha_{1-s}=-\dot{\alpha}_{1-s}ds$ and $d\beta_{1-s}=-\dot{\beta}_{1-s}ds$, we may rewrite this as

Using the explicit form of $g^{\rm F}_{s}$ given in (70) together with (71), we deduce that

where we used $\sigma_{1}=0$. This means that

where $z\sim{\sf N}(0,\text{\it Id})$ with $z\perp W$. Note that, unsurprisingly, the process on the right-hand side is simply the stochastic interpolant (2) conditioned on $x_{0}$ fixed, with $x_{1}$ replaced by a Gaussian $z$. Denoting by $\rho_{s}^{Y}(y|x_{0})$ the PDF of $Y^{\text{\sc R}}_{1-s}$, (91) implies that

We can use this result to time reverse (86) and obtain the following SDE for $Y_{s}\stackrel{{\scriptstyle d}}{{=}}Y^{\text{\sc R}}_{1-s}$

If we insert the explicit form of $|g^{\rm F}_{s}|^{2}=2s\sigma_{s}(\beta_{s}^{-1}\dot{\beta}_{s}\sigma_{s}-\dot{\sigma}_{s})-\sigma_{s}^{2}$ given in (70) into (93) we get

i.e. we recover the SDE (79). Since $a_{s}=\frac{d}{ds}\log\frac{\beta_{s}^{2}+s\sigma_{s}^{2}}{\beta_{s}}$, we can solve the above the SDE to obtain the reference process as

which implies that $Y_{s}\stackrel{{\scriptstyle d}}{{=}}\alpha_{s}x_{0}+\beta_{s}z+\sigma_{s}W_{s}\sim{\sf N}(\alpha_{s}x_{0},\beta_{s}^{2}+s\sigma_{s}^{2})$; this matches our previous calculation (91).

It remains to show that the process $X^{\text{\sc F}}\equiv X^{g^{\rm F}}$ defined by solution to the SDE (57) with $g_{s}=g^{\text{\sc F}}_{s}$ given in (70) is the Föllmer process associated with the process $Y$ defined by the solution to (93). To this end, recall that the KL divergence between two path measures is invariant under time-reversal, so that

where $\check{X}$ is the process defined by the solution to the SDE (81), and where $\check{X}^{\text{\sc R}}$ is its time-reversal. We can now use following decomposition, known as “disintegration” (Léonard, 2014), of the Kullback-Leibler divergence

where $\check{X}^{{\text{\sc R}},x}$ and $Y^{{\text{\sc R}},x}$ denote, respectively, the processes $\check{X}^{\text{\sc R}}$ and $Y^{\text{\sc R}}$ conditioned to start from $x$ (i.e. on $\check{X}^{{\text{\sc R}},x}_{s=0}=Y^{{\text{\sc R}},x}_{s=0}=x$). In addition, we used the fact that $Y^{\text{\sc R}}_{s=0}\sim{\sf N}(0,\text{\it Id})$ whereas $\check{X}^{\text{\sc R}}_{s=0}\sim\rho_{c}(\cdot|x_{0})$ by the constraint imposed in the minimization problem (82). The second term at the right hand side of (98) is fixed due to this constraint; the first term is always non-negative, but we can make it zero if we take $\check{X}^{\text{\sc R}}=X^{\text{\sc R}}$ with $X^{\text{\sc R}}$ defined as the solution to (84) since this process is the same as the one defined by the solution to (86) if we condition both on $X^{\text{\sc R}}_{s=0}=Y^{\text{\sc R}}_{s=0}=x$. Therefore $\check{X}^{\text{\sc R}}=X^{\text{\sc R}}$ minimizes (98), which means that its time-reversal $X^{g^{\rm F}}$ minimizes (82). ∎

When $\rho_{c}(x|x_{0})$ is a Gaussian mixture model (GMM), the drift $b_{s}$ is available analytically:

Inserting this expression in (102) we obtain

where in the first and second identities, we used the fact that $s\sigma^{2}_{s}\text{\it Id}=\text{\it Id}-\beta^{2}_{s}C_{j}\overline{C}_{j}^{-1}(s)$ and $x-\alpha_{s}x_{0}=x-\overline{m}_{j}(s)+\beta_{s}m_{j}$.

Now, using $b_{s}(x,x_{0})=\dot{\alpha}_{s}x_{0}+\dot{\beta}_{s}\beta_{s}^{-1}(x-\alpha_{s}x_{0})+s\sigma_{s}^{2}(\dot{\beta}_{s}\beta_{s}^{-1}-\dot{\sigma}_{s})\nabla\log\rho_{s}(x|x_{0})$, we get the final formula. ∎

Appendix B Details of Numerical Experiments

We create the synthetic example by starting with a Gaussian component ${\sf N}(m_{0},C_{0})$ with $m_{0}=$ and

We rotate this distribution counterclockwise by $2\pi/5$ four times to obtain the remaining four Gaussian modes. We assign each of the five modes equal weights to obtain our 2D Gaussian mixture model with density $\rho_{\text{GMM}}(x)$

The 2D particle jump-diffusion dynamics is constructed as follows. Between the jumps, the particle moves according to the Langevin dynamics

At jump times specified by a Poisson process with rate $\lambda=2$, the particle is rotated counterclockwise by an angle $2\pi/5$.

We simulate this dynamics using the following scheme with $\delta t=0.01$:

where $\xi\sim{\sf N}(0,I_{2\times 2})$ and $R_{2\pi/5}$ is the counterclockwise rotation operator in 2D with angle $2\pi/5$. We integrate this dynamics long enough to reach equilibrium and get enough data.

We keep the data at a regular time interval of $\Delta t=0.5$ and use paired $(x_{t},x_{t+0.5})$ as training data for learning the conditional distribution at lag $\tau=0.5$. In total we store $10^{5}$ training data pairs. We use a fully connected neural network with $5$ hidden layes with hidden dimension $500$ to approximate the velocity field $b_{s}(x,x_{0})$ in the SDEs. The input to the net is of dimension $5$ and the output is of dimension $2$. We train the network using a batch size of $10^{4}$, default AdamW optimizer with base learning rate $l=10^{-3}$ and cosine scheduler that decreases in each epoch the learning rate eventually to after $300$ epochs. We test the SDEs for new simulated trajectory data.

B.2 2D Stochastic Navier-Stokes Example

We set $\nu=10^{-3},\alpha=0.1,\varepsilon=1$. We consider the following random forcing

where $W_{i}(t)$, $1\leq i\leq 8$ are independent Wiener processes. For such random forcing, the NSE (17) has a unique invariant measure, as proved in (Hairer & Mattingly, 2006). Note that (106) is a Gaussian random field with covariance function $C(t,t^{\prime},x,y,x^{\prime},y^{\prime})=\min(t,t^{\prime})[\cos(6(x-x^{\prime}))+\cos(7(x-x^{\prime}))+\cos(5(x-x^{\prime}+y-y^{\prime})+\cos(8(x-x^{\prime}+y-y^{\prime}))]$ which is translation invariant in space. The damping term $-\alpha\omega$ is used to accelerate the mixing of the dynamics as it damps the vorticity at large length-scales/small Fourier modes to avoid all energy accumulating in large vortices.

We use a pseudo-spectral solver with Euler-Maruyama time-stepping scheme to solve the stochastic PDE in time and we use the jax-cfd package (Dresdner et al., 2022) for the mesh generation and domain discretization. We perform simulations with a grid sizes $256\times 256$ which is fine enough for resolving the numerical solutions of the 2D stochastic Navier-Stokes equations with our specifications.

We provide some additional numerical results on the 2D stochastic Navier-Stokes example below.

We consider superresolution, i.e., predicting the field with a resolution of $128\times 128$, from the downsized version with a resolution of $32\times 32$. In Fig. 7, we show a $32\times 32$ resolution field, as well as the true $128\times 128$ resolution field and the mean of samples drawn from our SDEs. Our method achieves an outstanding recovery. Additionally, we calculate the standard deviation of the samples. The spatial distribution of this standard deviation serves as a tool for uncertainty quantification; notably, there exists a pronounced correlation between the pattern of standard deviation and the vorticity field.

Forecasting efficiency

We conducted a comparative analysis of prediction speed between our models and the SPDE solver utilized for the 2D stochastic Navier-Stokes equation. For a scenario involving a short time lag of $\tau=0.5$, wherein our model accurately captures the resulting conditional distribution, sampling from our SDE forecaster with $200$ Euler-Maruyama steps takes $0.05$ seconds, while executing the SPDE solver on the same Nvidia RTX8000 GPU requires $8$ seconds. This observation highlights that our method accelerates forecasting by over 100 times without sacrificing physical information inherent in the conditional distribution. Furthermore, our approach could stand to gain even more acceleration in scenarios with lower viscosity, which typically necessitates a substantially reduced time step size for the SPDE solver.

Comparison with a deterministic approach

Comparisons between different SDE generative models

All the above experiments on NSE is done with the interpolant $\alpha_{s}=1-s,\beta_{s}=s^{2},\sigma_{s}=1-s$, as it performs the best. Below in Figures 8 and 9, we also post the loss and gradient norm curves, as well as the enstrophy spectra of the generated samples, for other choices of interpolants and Föllmer processes (trained with the same network, data, and number of epochs). We observe that $\beta_{s}=s^{2}$, namely $\dot{\beta}_{0}=0$, is important to ensure a stable gradient norm curve. This experiment demonstrates the superiority of the interpolant $\alpha_{s}=1-s,\beta_{s}=s^{2},\sigma_{s}=1-s$. Moreover, by changing the diffusion coefficients from $\sigma_{s}dW_{s}$ to the optimal $g_{s}^{\rm F}dW_{s}$ does not influence the spectrums significantly in this example.

We fix one initial vorticity field and compare the KL divergence between the true conditional distribution and the generated distribution. Here the KL divergence is calculated for the 1D conditional distributions of total enstrophy and energy, for a fixed $\omega_{t}$ and $\tau=1$. The goal here is to test whether changing the diffusion coefficient from $\sigma_{s}$ to $g_{s}^{\rm F}$ could improve the KL divergence accuracy of the generated distribution at $s=1$. Note that our theory in Theorem 3.3 shows that changing from $\sigma_{s}$ to $g_{s}^{\rm F}$ could improve the path KL divergence, which is an upper bound on the KL divergence of the generated marginal distribution at $s=1$.

We consider the generative SDEs corresponding to the stochastic interpolant with $\alpha_{s}=1-s,\beta_{s}=s^{2},\sigma_{s}=1-s$. We can vary the diffusion coefficient of the generative SDEs; we choose $\sigma_{s}=1-s$, or $g_{s}^{\rm F}=\sqrt{(3-s)(1-s)}$ that corresponds to a Föllmer process. We also compare the results with the Gaussian base ODE generative model (equivalent to flow matching (Lipman et al., 2022; Liu et al., 2022; Albergo & Vanden-Eijnden, 2022), corresponding to $\alpha_{s}=1-s,\beta_{s}=s$, and $x_{0}$ in the interpolant obeys an $\mathsf{N}(0,\text{\it Id})$ distribution). The quantitative KL results are reported in Table 3 and the densities are shown in Figure 10. We observe that SDE approaches lead to better KL accuracy. Moreover, changing the diffusion coefficients could potentially improve the KL accuracy, justifying the flexibility and usefulness of the interpolant approach.

B.3 Video generation