Control, Transport and Sampling: Towards Better Loss Design (2405.13731v2)

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.