Rethinking Losses for Diffusion Bridge Samplers

2025-06-12

AI Summary

This paper addresses the theoretical deficiency of the widely adopted Log Variance (LV) loss in diffusion bridge samplers—namely, its lack of rigorous derivation and misalignment between optimization objective and sampling performance. We propose rKL-LD, a novel loss function grounded in the reverse KL divergence and leveraging the log-derivative trick, which combines theoretical rigor with practical superiority. We first demonstrate that LV cannot be justified via the data processing inequality within the diffusion bridge framework, whereas rKL-LD emerges naturally and admits a clear variational interpretation. Experiments across multiple challenging benchmarks show that rKL-LD significantly improves sample quality (FID reduced by 12–28%), enhances training stability, and exhibits greater robustness to hyperparameter choices. Our work establishes a new theoretical standard and practical paradigm for diffusion bridge modeling.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• Comparing Log Variance and rKL losses for diffusion bridges
• Analyzing gradient equivalence in diffusion bridge samplers
• Evaluating rKL-LD loss performance in diffusion bridge training

Innovation

• Uses Log Variance loss for diffusion bridges
• Employs rKL loss with log-derivative trick
• Reduces hyperparameter optimization needs significantly