Control, Transport and Sampling: Towards Better Loss Design
This work addresses the challenge of improving loss design for sampling and transport tasks in machine learning, offering incremental advancements through a unified framework.
The paper tackles the problem of designing better training losses for diffusion-based sampling and optimal transport by leveraging connections to the Schrödinger bridge problem, proposing novel objective functions that enable controlled dynamics to transport and sample from target distributions with numerical advantages in implementation.
Leveraging connections between diffusion-based sampling, optimal transport, and stochastic optimal control through their shared links to the Schrödinger bridge problem, we propose novel objective functions that can be used to transport $ν$ to $μ$, consequently sample from the target $μ$, 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ödinger bridge comes with the additional practical capability of baking in inductive bias when it comes to Neural Network training.