Rethinking Losses for Diffusion Bridge Samplers

TL;DR

This paper critically examines loss functions for diffusion bridge samplers, revealing that the reverse KL loss with the log-derivative trick (rKL-LD) outperforms the commonly used Log Variance loss, especially for learned diffusion coefficients.

Contribution

The authors demonstrate that rKL-LD is a more conceptually sound and practically effective loss function for diffusion bridges, outperforming LV loss in stability and performance.

Findings

rKL-LD outperforms LV loss on benchmarks

rKL-LD requires less hyperparameter tuning

Training with rKL-LD is more stable

Abstract

Diffusion bridges are a promising class of deep-learning methods for sampling from unnormalized distributions. Recent works show that the Log Variance (LV) loss consistently outperforms the reverse Kullback-Leibler (rKL) loss when using the reparametrization trick to compute rKL-gradients. While the on-policy LV loss yields identical gradients to the rKL loss when combined with the log-derivative trick for diffusion samplers with non-learnable forward processes, this equivalence does not hold for diffusion bridges or when diffusion coefficients are learned. Based on this insight we argue that for diffusion bridges the LV loss does not represent an optimization objective that can be motivated like the rKL loss via the data processing inequality. Our analysis shows that employing the rKL loss with the log-derivative trick (rKL-LD) does not only avoid these conceptual problems but also consistently outperforms the LV loss. Experimental results with different types of diffusion bridges on challenging benchmarks show that samplers trained with the rKL-LD loss achieve better performance. From a practical perspective we find that rKL-LD requires significantly less hyperparameter optimization and yields more stable training behavior.