Control, Transport and Sampling: Towards Better Loss Design

TL;DR

This paper introduces new loss functions for diffusion-based sampling by connecting optimal transport, stochastic control, and Schr"odinger bridges, enhancing sampling efficiency and training bias incorporation.

Contribution

It proposes novel objective functions based on Schr"odinger bridges that improve transport and sampling in diffusion models, emphasizing the pathwise perspective and optimality conditions.

Findings

New loss functions improve sampling accuracy.

Pathwise perspective offers numerical advantages.

Incorporating Schr"odinger bridges introduces beneficial inductive biases.

Abstract

Leveraging connections between diffusion-based sampling, optimal transport, and stochastic optimal control through their shared links to the Schr\"odinger bridge problem, we propose novel objective functions that can be used to transport $\nu$ to $\mu$, consequently sample from the target $\mu$, via optimally controlled dynamics. We highlight the importance of the pathwise perspective and the role various optimality conditions on the path measure can play for the design of valid training losses, the careful choice of which offer numerical advantages in implementation. Basing the formalism on Schr\"odinger bridge comes with the additional practical capability of baking in inductive bias when it comes to Neural Network training.