Conditioning Diffusions Using Malliavin Calculus

Jakiw Pidstrigach, Elizabeth Louise Baker, Carles Domingo-Enrich, George Deligiannidis, Nikolas Nüsken

TL;DR: We propose a Malliavin-calculus-based framework for steering diffusions under singular rewards (e.g.\ diffusion bridges), enabling stable training and outperforming existing gradient-based methods that fail when the reward is non-differentiable.

Abstract:

In generative modelling and stochastic optimal control, a central computational task is to modify a reference diffusion process to maximise a given terminal-time reward.
Most existing methods require this reward to be differentiable, using gradients to steer the diffusion towards favourable outcomes.
However, in many practical settings, like diffusion bridges, the reward is singular, taking an infinite value if the target is hit and zero otherwise.
We introduce a novel framework, based on Malliavin calculus and centred around a generalisation of the Tweedie score formula to nonlinear stochastic differential equations, that enables the development of methods robust to such singularities.
This allows our approach to handle a broad range of applications, like diffusion bridges, or adding conditional controls to an already trained diffusion model.
We demonstrate that our approach offers stable and reliable training, outperforming existing techniques.
As a byproduct, we also introduce a novel score matching objective. Our loss functions are formulated such that they could readily be extended to manifold-valued and infinite dimensional diffusions.

Lay Summary:

In many stochastic systems, one is interested in simulating outcomes that satisfy a prescribed condition. For instance, one might wish to generate chemical reaction trajectories that terminate in a specific state, or simulate weather patterns consistent with a particular scenario. In generative artificial intelligence, this includes tasks such as generating images with a given depth map or designing proteins that exhibit certain desired properties.
We introduce a new method rooted in Malliavin calculus to address these challenges. This mathematical framework yields a highly general formulation, enabling the control of diffusion processes conditioned on a wide variety of events.

Primary Area: Probabilistic Methods->Monte Carlo and Sampling Methods

Keywords: Bridges, Optimal control, stochastic differential equations, diffusion processes, Malliavin calculus, conditioning

Submission Number: 12474