PaperlayerView exact revision on arXiv

Source-linked AI summary

Improved sampling via learned diffusions

Lorenz Richter, Julius Berner

arXiv:2307.01198v2cs.LGmath.OCmath.PRstat.ML

TL;DR

Sampling from unnormalized target densities is challenging when their normalizing constants are intractable and target samples are unavailable, while learned-diffusion approaches may face instability or mode collapse. The paper unifies these approaches through a path-space formulation that allows divergences beyond reverse KL and proposes the log-variance loss. It reports improved numerical stability and performance, and identifies log-variance divergence as particularly useful for sampling from unnormalized densities.

  • Problem

    Sampling from an unnormalized target density is challenging when its normalizing constant is intractable and target samples are unavailable; learned-diffusion methods may also face instability or mode collapse.

    Source: §1 · p. 1

  • Method

    A path-space formulation of time-reversed diffusions unifies learned-diffusion sampling methods and permits objectives based on divergences beyond reverse KL, including the log-variance loss.

    Source: §1 · p. 2

  • Results

    The log-variance loss significantly improves numerical stability and performance, and improves DIS and PIS metrics across the considered benchmark problems.

    Source: §1 · p. 2, §4.2 · p. 9

  • Takeaways & Limitations

    Log-variance divergence is particularly useful for sampling from unnormalized densities, with lower variance, computational efficiency, and exploration-exploitation trade-offs.

    Source: §5 · p. 9

    The general bridge framework remains less efficient and numerically unstable; high-dimensional divergence computation may be prohibitive, and Hutchinson trace estimation adds variance.

    Source: §4.2 · pp. 8–9

Abstract

from arXiv · show

Recently, a series of papers proposed deep learning-based approaches to sample from target distributions using controlled diffusion processes, being trained only on the unnormalized target densities without access to samples. Building on previous work, we identify these approaches as special cases of a generalized Schr\"odinger bridge problem, seeking a stochastic evolution between a given prior distribution and the specified target. We further generalize this framework by introducing a variational formulation based on divergences between path space measures of time-reversed diffusion processes. This abstract perspective leads to practical losses that can be optimized by gradient-based algorithms and includes previous objectives as special cases. At the same time, it allows us to consider divergences other than the reverse Kullback-Leibler divergence that is known to suffer from mode collapse. In particular, we propose the so-called log-variance loss, which exhibits favorable numerical properties and leads to significantly improved performance across all considered approaches.

1 INTRODUCTION

The paper addresses sampling from unnormalized target densities, whose normalizing constants are typically intractable, and the limits of learned-diffusion samplers trained without target samples. It presents a path-space framework that unifies prior methods and motivates objectives beyond reverse KL, including the log-variance loss.

1 INTRODUCTION

  • Learned-diffusion samplers target unnormalized densities without samples, but may fall short of state-of-the-art sampling performance and face instability or mode collapse.The normalization constant is typically intractable.Source: §1 · p. 1
  • A path-space perspective on time-reversed controlled diffusions unifies Schrödinger bridges, DIS, DDS, and PIS, and permits objectives beyond reverse KL.The paper notes that reverse-KL objectives are prone to mode collapse.Source: §1 · p. 2
  • The log-variance divergence avoids differentiation through the SDE solver and balances exploration and exploitation, with improved numerical stability and performance.The introduction points to improved convergence on a double-well problem.Source: §1 · p. 2, Figure 1 · p. 2

1.1 RELATED WORK

  • The framework builds on variational bridge formulations and extends them to other divergences; its path-space derivation connects existing sampling methods and identifies non-uniqueness in general bridge losses.The paper says special cases recover unique objectives for DIS, DDS, and PIS.Source: §1.1 · p. 2

1.2 OUTLINE OF THE ARTICLE

  • The article develops the path-space and time-reversed-SDE framework, derives connections to Schrödinger bridges, DIS, PIS, and DDS, then presents numerical evidence for the log-variance divergence.The outline places the numerical examples in Section 4.Source: §1.2 · p. 2

2 DIFFUSION-BASED SAMPLING

The paper frames diffusion-based sampling as matching time-reversed controlled processes through divergences between path-space measures. It compares KL and log-variance objectives, highlighting the latter’s flexibility and lower gradient-estimator variance at the optimum.

2.1 NOTATION AND SETTING

  • Path-space measures represent controlled processes as trajectory distributions, with diffusion-based sampling transporting a prior density toward the target.The formulation seeks controls for which one process is the time reversal of the other.Source: §2.1 · p. 3, §2.2 · p. 3

2.2 SAMPLING AS TIME-REVERSAL PROBLEM

  • The sampling objective minimizes a divergence between the path measures of two time-reversed controlled processes, allowing divergences beyond KL.The framework derives a path-measure likelihood identity to compute suitable divergences.Source: §2.2 · pp. 3–4

2.3 COMPARISON OF THE KL AND LOG-VARIANCE DIVERGENCE

  • The log-variance loss can balance exploration and exploitation by changing its reference process, including along trajectories where the current process has low probability.The paper proposes this flexibility as a way to counteract mode collapse.Source: §2.3 · p. 4
  • Unlike the KL-based loss, the log-variance loss does not require differentiating the process through an SDE solver or differentiating the unnormalized target density.The process can be simulated with a black-box SDE solver without tracking its computation for automatic differentiation.Source: §2.3 · pp. 4–5
  • At the optimum, log-variance gradient-estimator variances vanish, whereas the KL-based estimator’s variances do not.For w = u, the log-variance loss can also be interpreted as a control-variate version of the KL loss.Source: §2.3 · p. 5

3 CONNECTIONS AND EQUIVALENCES OF DIFFUSION-BASED SAMPLING

The framework connects diffusion-based sampling objectives to Schrödinger bridges and recovers DIS, PIS, and DDS under particular choices. These connections also show that the general objectives need not have unique solutions, while additional conditions can yield unique ones.

3.1 SCHR ¨ODINGER BRIDGE PROBLEM (SB)

  • The general path-space objectives can have infinitely many solutions, while adding the Schrödinger bridge entropy constraint selects a unique optimum.An optimal SB control uniquely determines its paired control through Nelson’s identity.Source: §3 · p. 5, §3.1 · p. 6

3.2 DIFFUSION-BASED GENERATIVE MODELING (DIS)

  • Setting the DIS control choice to zero can reduce the entropy constraint to zero, but feasible sampling requires an approximately known, easily sampled terminal distribution.The paper notes a Gaussian prior as a practical approximation for that distribution.Source: §3.2 · p. 6

3.3 TIME-REVERSAL OF REFERENCE PROCESSES (PIS & DDS)

  • Particular reference-process choices recover the PIS and DDS objectives, and the framework also gives corresponding log-variance losses.The likelihood construction requires a tractable reference density and an approximate reverse process close to the optimal one.Source: §3.3 · p. 7

4 NUMERICAL EXPERIMENTS

The experiments compare KL-based and log-variance losses across bridge, PIS, and DIS approaches on benchmark sampling problems. Log-variance improves PIS and DIS across the considered benchmarks and metrics, while the general bridge approach retains efficiency and stability limitations.

4 NUMERICAL EXPERIMENTS

  • The study compares KL-based and log-variance losses for the general bridge, PIS, and DIS, using the same settings and number of target evaluations for both losses.DDS is not evaluated separately because the paper treats it as a special case of DIS.Source: §4 · pp. 7–8

4.2 RESULTS

  • For the GMM and double-well examples, log-variance avoids the mode collapse seen with reverse KL, which recovers only the prior mode.The double-well illustration uses d = 5, m = 5, and δ = 4; similar behavior is reported for other marginals and a d = 1000 setting.Source: Figure 2 · p. 8, Figure 3 · p. 8

4.1 BENCHMARK PROBLEMS

  • The benchmarks include a nine-mode, well-separated GMM, a 10-dimensional Funnel distribution, and a double-well target with high energy barriers.The double-well density factorizes across dimensions, allowing reference solutions by numerical integration and ground-truth samples by rejection sampling.Source: §4.1 · p. 8

4.2 RESULTS

  • The general bridge approach still has reduced efficiency and numerical instabilities, with divergence computation potentially prohibitive in high dimensions.Automatic differentiation for the divergence and Hutchinson’s trace estimator introduce additional variance; the more challenging problems therefore focus on DIS and PIS.Source: §4.2 · pp. 8–9
  • Log-variance significantly improves both DIS and PIS across the considered benchmark problems and metrics.The comparison replaces the KL-based loss without tuning the log-variance hyperparameter; in the few cases where KL performs better, the difference seems insignificant.Source: §4.2 · p. 9

5 CONCLUSION

The paper unifies diffusion-based methods through path-space measures of time-reversed processes and permits objectives based on arbitrary divergences. Its log-variance divergence improves sampling quality across multiple numerical examples, with further performance gains from problem- and approach-specific tuning left as a possibility.

5 CONCLUSION

  • A path-space perspective on time-reversed diffusion processes connects SB, DIS, PIS, and DDS.The paper presents this as a unifying perspective on diffusion-based generative modeling.Source: §5 · p. 9
  • The framework allows arbitrary divergences between path measures as objectives, while the KL divergence recovers known methods..Source: §5 · p. 9
  • The log-variance divergence has lower variance, computational efficiency, and exploration-exploitation trade-offs, and improves sampling quality across multiple numerical examples.The paper suggests problem- and approach-specific finetuning might further enhance its performance.Source: §5 · p. 9

A.1 ASSUMPTIONS

The proofs assume sufficient regularity of the coefficients of all appearing SDEs. This supports Novikov’s condition and unique strong solutions with smooth, strictly positive densities at intermediate times.

A.1 ASSUMPTIONS

  • The proof assumptions require sufficiently regular coefficient functions for all appearing SDEs.The stated conditions ensure Novikov’s condition and unique strong solutions with smooth, strictly positive densities for t ∈ (0, T).Source: §A.1 · p. 14
  • Novikov’s condition is assumed to hold for the appearing SDEs.Source: §A.1 · p. 14
  • The SDEs are assumed to have unique strong solutions with smooth, strictly positive densities at intermediate times.Source: §A.1 · p. 14

A.2 PROOFS

The proofs use path-space likelihoods and stochastic calculus to derive objectives and gradient properties. They also connect the log-variance gradient to a control-variate interpretation and discuss computational alternatives.

A.2 PROOFS

  • Backward stochastic integrals define a divergence-free objective that may be beneficial in high dimensions, where automatic differentiation of the divergence is expensive.The backward integral uses the right endpoint when discretizing on a time grid.Source: §A.2 · p. 15
  • Even at optimal controls, the corresponding Monte Carlo estimator’s variance can usually remain nonzero.Source: §A.2 · p. 16
  • Choosing w = u makes the log-variance gradient with respect to u interpretable as a control-variate version of the KL gradient.The authors describe this as promising reduced variance for the corresponding Monte Carlo estimators.Source: §A.2 · p. 16
  • An importance-sampling view in path space yields an expression equivalent to the one in Lemma 3.1 without requiring time-reversals.Source: §A.2 · p. 16

A.3 SAMPLING VIA LEARNED DIFFUSIONS

Controlled-diffusion samplers learn a control that drives a process toward the target while additional trajectory costs enforce constraints. PIS, DIS, and DDS correspond to different choices of reference processes and divergences.

A.3 SAMPLING VIA LEARNED DIFFUSIONS

  • PIS, DIS, and DDS arise from choices of reference process and divergence within the controlled-diffusion sampling formulation.The terminal cost penalizes ending in low-target-density regions, while other costs constrain process trajectories.Source: §A.3 · p. 17
  • PIS minimizes reverse KL to the uncontrolled process, whereas DIS and DDS minimize reverse KL to the time-reversed noising process.PIS corresponds to a classical Schrödinger bridge; DIS and DDS enforce denoising of samples from the reversed noising process.Source: §A.3 · p. 17
  • The KL-based methods can also be derived from the underlying PDEs or BSDE systems.The section uses a general KL-based loss because most previous methods are special cases of it.Source: §A.3 · p. 17

A.4 THE SCHR ¨ODINGER BRIDGE PROBLEM

The classical Schrödinger bridge selects, among processes connecting prescribed endpoint distributions, the one closest in KL divergence to a reference process. Its optimality conditions can be expressed through coupled PDEs or a corresponding BSDE system, which also recovers the sampling losses.

A.4 THE SCHR ¨ODINGER BRIDGE PROBLEM

  • The Schrödinger bridge selects the solution of the general bridge problem that minimizes path-space KL divergence to a chosen reference process.The traditional choice is the uncontrolled process, corresponding to reference control r = 0.Source: §A.4 · p. 17
  • The optimal control is characterized by coupled PDEs with endpoint conditions set by the prior and target distributions.The PDE perspective uses the Hopf-Cole transform to obtain optimality conditions.Source: §A.4 · pp. 17–18
  • Applying Itô’s lemma yields a BSDE system, whose combined identity recovers the KL-based loss and the log-variance loss.At the optimal controls, the variance of the relevant difference is zero, motivating the log-variance objective.Source: §A.4 · pp. 18–19

A.4.1 SCHR ¨ODINGER HALF-BRIDGES (PIS)

For a point-mass prior, the Schrödinger bridge equations decouple, yielding the PIS objective. The optimal control can be expressed using the target-linked PDE solution, and the log-variance loss has a corresponding PIS form.

A.4.1 SCHR ¨ODINGER HALF-BRIDGES (PIS)

  • PIS assumes a point-mass prior, which decouples the bridge PDE system and permits an alternative derivation of its objective.The stated assumption is pprior = δx0 for a fixed initial point x0.Source: §A.4.1 · p. 20
  • Under this assumption, the optimal control is obtained from the bridge solution, and the log-variance loss can also be written in PIS form.The derivation uses the decoupled process and the boundary condition at the initial point.Source: §A.4.1 · p. 20

A.5 TRACTABLE SDES

The appendix specifies tractable SDEs and implementation choices, then reports additional experiments on partial trajectories and broader comparisons. These experiments show improved mode coverage and robust performance for the log-variance loss, while partial-trajectory training can recover all modes on a 40-mode benchmark.

A.5 TRACTABLE SDES

  • The appendix presents variance-exploding and variance-preserving SDE examples and a generalized sampler-training loop that selects either KL-based or log-variance loss.The loop initializes the process, simulates discretized paths, and approximates costs.Source: §A.5 · pp. 20–21

A.6 COMPUTATIONAL DETAILS

  • Experiments use fixed Euler-Maruyama discretization and specify method-specific SDEs and control parametrizations for PIS, DIS, and the general bridge.The implementation also reports gradient-step counts by dimension and uses a decaying learning rate.Source: §A.6 · pp. 21–22

A.6.1 COMPUTATION OF LOG-NORMALIZING CONSTANTS

  • With approximate controls, the proposed expression gives a lower bound on log Z, which the experiments use as an estimator.The bound follows from Jensen’s inequality; for optimal controls, the integrand is almost surely constant.Source: §A.6.1 · p. 23

A.7 PARTIAL TRAJECTORY OPTIMIZATION

  • Partial-trajectory training uses shorter intervals and selectable starting states, which can support faster simulation and exploration; it is available for the log-variance loss, not the KL-based loss.The paper reports that this approach can achieve competitive results without target-gradient information.Source: §A.7 · pp. 23–24, Table 3 · p. 24
  • Partial-trajectory optimization identifies all 40 modes, outperforming DIS even with the log-variance loss and target-gradient information in its parametrization.The comparison is shown for the 40-mode Gaussian-mixture benchmark.Source: Figure 4 · p. 24, §A.7 · p. 24

A.8 FURTHER EXPERIMENTS AND COMPARISONS

  • The log-variance loss consistently performs better with batch size 512 and reduces mode collapse in both moderate and very high dimensions.The paper attributes the smaller-batch benefit to reduced variance in gradient estimators; Figures 6 and 7 show marginal mode coverage.Source: Table 4 · p. 25, Figure 5 · p. 25, §A.8 · p. 25
  • The general bridge improves on DIS or PIS with KL divergence in selected benchmarks, while log-variance methods are competitive with CRAFT across reported metrics.The comparisons report medians over five independent runs.Source: Table 5 · p. 27, Table 6 · p. 27
  • On the Wu et al. target, the log-variance loss improves performance over KL-based losses; its gains are robust across seeds and reduce average time per gradient step.The KL-based methods place 86% of PIS samples and 67% of DIS samples outside the domain; the log-variance comparison is described as competitive with stochastic normalizing flows.Source: Figure 8 · p. 27, Figure 9 · p. 28

Related papers

Loading 2307.01198v2…