A Taxonomy of Loss Functions for Stochastic Optimal Control

TL;DR

This paper classifies various loss functions used in stochastic optimal control, revealing their shared optimization landscapes and differences in gradient variance, with experiments illustrating their relative strengths and weaknesses.

Contribution

It provides a taxonomy of SOC loss functions, showing their equivalence in expectation and analyzing their variance, thus clarifying their relationships and performance.

Findings

SOC loss functions share the same expected gradient

Loss functions differ mainly in gradient variance

Experimental results highlight strengths and weaknesses of different loss functions

Abstract

Stochastic optimal control (SOC) aims to direct the behavior of noisy systems and has widespread applications in science, engineering, and artificial intelligence. In particular, reward fine-tuning of diffusion and flow matching models and sampling from unnormalized methods can be recast as SOC problems. A recent work has introduced Adjoint Matching (Domingo-Enrich et al., 2024), a loss function for SOC problems that vastly outperforms existing loss functions in the reward fine-tuning setup. The goal of this work is to clarify the connections between all the existing (and some new) SOC loss functions. Namely, we show that SOC loss functions can be grouped into classes that share the same gradient in expectation, which means that their optimization landscape is the same; they only differ in their gradient variance. We perform simple SOC experiments to understand the strengths and weaknesses of different loss functions.