Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Control, Transport and Sampling: Towards Better Loss Design

Qijia Jiang, David Nabergoj

TL;DR

Novel objective functions are proposed that can be used to transport $\nu$ to $\mu$ and consequently sample from the target $\mu$ via optimally controlled dynamics, via optimally controlled dynamics.

Abstract

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.

Control, Transport and Sampling: Towards Better Loss Design

TL;DR

Novel objective functions are proposed that can be used to transport to and consequently sample from the target via optimally controlled dynamics, via optimally controlled dynamics.

Abstract

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.
Paper Structure (28 sections, 6 theorems, 148 equations, 6 figures, 5 tables, 1 algorithm)

This paper contains 28 sections, 6 theorems, 148 equations, 6 figures, 5 tables, 1 algorithm.

Table of Contents

  1. INTRODUCTION
  2. GENERAL FRAMEWORK
  3. Setup
  4. Goal and Related Approaches
  5. Forward KL
  6. Reverse KL
  7. METHODOLOGY
  8. Training Loss Proposal
  9. Discretization & Implementation
  10. NUMERICS & COMPARISONS
  11. Algorithm Specification
  12. Comparison between discretized losses
  13. Result
  14. CONCLUSION
  15. ADDITIONAL MOTIVATION & PRACTICAL CONSIDERATION
  16. ...and 13 more sections

Key Result

Proposition 1

For the problem of eqn:bridge_problem, the following losses are valid: (a) $\arg\min_{\nabla \phi,\nabla \psi}\, D_{KL}(\overrightarrow{\mathbb{P}}^{\nu,f+\sigma^2 \nabla \phi}\Vert\overleftarrow{\mathbb{P}}^{\mu,f-\sigma^2\nabla \psi})+\lambda \cdot \mathbb{E}_{X\sim \overrightarrow{\mathbb{P}}^{\n (c) $\arg\min_{\phi,\psi}\, \text{Var}_{X\sim \overrightarrow{\mathbb{P}}^{\nu,f+\sigma^2 \nabla \p

Figures (6)

  • Figure 1: Weighted Marginal for Double Well
  • Figure 2: Optimal Transport Between Fixed Marginals
  • Figure 3: Training loss plot for GMM with TD and Separate Control
  • Figure 4: GMM marginal before (left) and after (right) importance weighting
  • Figure 5: Standard normal marginal (left) and Funnel marginal (right) with importance weighting
  • ...and 1 more figures

Theorems & Definitions (29)

  • Remark 1: Nelson's identity
  • Remark 2
  • Proposition 1: Control Training Objective
  • Remark 3: Stochastic gradient w.r.t controls at optimality
  • Remark 4: Comparisons
  • Proposition 2: Importance Sampling
  • Lemma 1: Discretized Loss and Estimator
  • Remark 5: Optimal $\log Z$ estimator
  • Remark 6: Langevin
  • Remark 7: Bound on optimal objective
  • ...and 19 more