Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Score matching for bridges without learning time-reversals

Elizabeth L. Baker, Moritz Schauer, Stefan Sommer

TL;DR

The paper addresses sampling from bridges of diffusion processes conditioned on an endpoint, where the score term $\nabla_x \log p(t, x; T, y)$ is generally intractable. It proposes a score-matching approach based on adjoint diffusions to learn the forward bridge score directly, avoiding the need to learn time-reversals. The authors derive a KL-based loss that ensures the learned drift produces a bridged process and show that this yields training-time savings and competitive performance relative to time-reversal-based methods across linear and nonlinear SDEs; they also provide code. The approach broadens the applicability of diffusion-bridge methods to non-linear dynamics and end-point distributions, with potential impact for shape analysis and related domains.

Abstract

We propose a new algorithm for learning bridged diffusion processes using score-matching methods. Our method relies on reversing the dynamics of the forward process and using this to learn a score function, which, via Doob's $h$-transform, yields a bridged diffusion process; that is, a process conditioned on an endpoint. In contrast to prior methods, we learn the score term $\nabla_x \log p(t, x; T, y)$ directly, for given $t, y$, completely avoiding first learning a time-reversal. We compare the performance of our algorithm with existing methods and see that it outperforms using the (learned) time-reversals to learn the score term. The code can be found at https://github.com/libbylbaker/forward_bridge.

Score matching for bridges without learning time-reversals

TL;DR

The paper addresses sampling from bridges of diffusion processes conditioned on an endpoint, where the score term is generally intractable. It proposes a score-matching approach based on adjoint diffusions to learn the forward bridge score directly, avoiding the need to learn time-reversals. The authors derive a KL-based loss that ensures the learned drift produces a bridged process and show that this yields training-time savings and competitive performance relative to time-reversal-based methods across linear and nonlinear SDEs; they also provide code. The approach broadens the applicability of diffusion-bridge methods to non-linear dynamics and end-point distributions, with potential impact for shape analysis and related domains.

Abstract

We propose a new algorithm for learning bridged diffusion processes using score-matching methods. Our method relies on reversing the dynamics of the forward process and using this to learn a score function, which, via Doob's -transform, yields a bridged diffusion process; that is, a process conditioned on an endpoint. In contrast to prior methods, we learn the score term directly, for given , completely avoiding first learning a time-reversal. We compare the performance of our algorithm with existing methods and see that it outperforms using the (learned) time-reversals to learn the score term. The code can be found at https://github.com/libbylbaker/forward_bridge.
Paper Structure (32 sections, 4 theorems, 46 equations, 8 figures, 1 table, 1 algorithm)

This paper contains 32 sections, 4 theorems, 46 equations, 8 figures, 1 table, 1 algorithm.

Table of Contents

  1. INTRODUCTION
  2. Contributions
  3. BACKGROUND
  4. Doob's h-transform
  5. Adjoint processes with reversed dynamics
  6. Score matching bridge sampling
  7. RELATED WORK
  8. Score matching methods
  9. MCMC methods
  10. Adjoint methods
  11. METHOD
  12. A loss function
  13. Optimising the loss function
  14. Training on multiple end points
  15. Distributions of endpoints
  16. ...and 17 more sections

Key Result

Theorem 4.1

Let $X$ be a diffusion process as defined in eq: unconditioned process and suppose further that $f, \sigma$ are $C^2$ and that $X$ admits $C^2$ transition densities. Let $\Bar{X}$ be the conditioned diffusion satisfying eq: conditioned process. Let $\Bar{\mathbb{P}}$ be the path probability of $\Bar where the infimum is over functions $s\colon [0, T]\times \mathbb{R}^d \to \mathbb{R}^d$.

Figures (8)

  • Figure 1: We learn the SDE in \ref{['eq: kunita sde']} conditioned to hit the shape in black at time $t=1$. We see that, starting the trajectory at the shape in red at time $t=0$, we get a path between the two emojis, where each time point in the path represents a shape.
  • Figure 2: For varying $x, y$ and different times $t \in [0.25, 0.5, 0.75]$, we plot the true score $\nabla_x \log p(t, x; 1, y)$ and the learned score $s_\theta(t, x; y)$, as described in \ref{['sec: ou experiment']}. The learned score takes a fixed end time $T=1$, but takes both $x$ and $y$ as input values. The absolute error between the true and learned score is reported in \ref{['fig: ou error']}.
  • Figure 3: We plot the square error between the true and learned score averaged over times from $t=0$ to $t=1$. For the proposed method and DB (fw) we learn $\nabla_{x} \log (t, x; T, y)$, and for DB (bw) we learn $\nabla_x \log (0, x_0; t, x)$. Top: Trained for $x_0=1$, $y=1$, with $1$-dimension and varying end times. Bottom: Trained for $x_0=1.0$, $y=1.0$, with $T=1.0$. See \ref{['sec: ou experiment']} for details.
  • Figure 4: We train a 2D Wiener process process to hit a circle of radius $3$ at time $T=1.0$. On the top we plot $20$ trajectories started from the centre, and on the bottom we plot $20$ trajectories started from evenly spaced points on a circle of radius $5$.
  • Figure 5: We plot $100$ conditioned trajectories with the methods Proposed, Guided MCMC and Forward bridge (DB) and the time-reversed bridge with Time-reversed bridge. In grey, we plot some unconditioned trajectories that end in a given window about the point $(1.5, 0.2)$. For details see \ref{['sec: cell']}.
  • ...and 3 more figures

Theorems & Definitions (8)

  • Theorem 4.1
  • proof
  • Theorem 4.2
  • proof
  • Theorem
  • proof
  • Theorem
  • proof