Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Posterior Sampling with Denoising Oracles via Tilted Transport

Joan Bruna, Jiequn Han

TL;DR

This work tackles posterior sampling in linear inverse problems with priors learned via score-based diffusion. It introduces tilted transport, a time-varying quadratic-tilt framework that transforms the posterior into a boosted distribution \nu_T that is provably easier to sample under Bakry-Émery/Log-Sobolev criteria. A key contribution is a concrete condition involving the prior susceptibility \chi_t(\pi) and the tilt's condition number that guarantees strong log-concavity of the boosted posterior, enabling efficient Langevin sampling; the approach also provides insights for Ising models and Gaussian mixtures. The method is plug-and-play with existing samplers, supports iterated tilts, and is demonstrated through case studies (Gaussian mixtures, Ising, φ^4) and numerical experiments, highlighting its potential to improve uncertainty quantification in high-dimensional inverse problems.

Abstract

Score-based diffusion models have significantly advanced high-dimensional data generation across various domains, by learning a denoising oracle (or score) from datasets. From a Bayesian perspective, they offer a realistic modeling of data priors and facilitate solving inverse problems through posterior sampling. Although many heuristic methods have been developed recently for this purpose, they lack the quantitative guarantees needed in many scientific applications. In this work, we introduce the \textit{tilted transport} technique, which leverages the quadratic structure of the log-likelihood in linear inverse problems in combination with the prior denoising oracle to transform the original posterior sampling problem into a new `boosted' posterior that is provably easier to sample from. We quantify the conditions under which this boosted posterior is strongly log-concave, highlighting the dependencies on the condition number of the measurement matrix and the signal-to-noise ratio. The resulting posterior sampling scheme is shown to reach the computational threshold predicted for sampling Ising models [Kunisky'23] with a direct analysis, and is further validated on high-dimensional Gaussian mixture models and scalar field $\varphi^4$ models.

Posterior Sampling with Denoising Oracles via Tilted Transport

TL;DR

This work tackles posterior sampling in linear inverse problems with priors learned via score-based diffusion. It introduces tilted transport, a time-varying quadratic-tilt framework that transforms the posterior into a boosted distribution \nu_T that is provably easier to sample under Bakry-Émery/Log-Sobolev criteria. A key contribution is a concrete condition involving the prior susceptibility \chi_t(\pi) and the tilt's condition number that guarantees strong log-concavity of the boosted posterior, enabling efficient Langevin sampling; the approach also provides insights for Ising models and Gaussian mixtures. The method is plug-and-play with existing samplers, supports iterated tilts, and is demonstrated through case studies (Gaussian mixtures, Ising, φ^4) and numerical experiments, highlighting its potential to improve uncertainty quantification in high-dimensional inverse problems.

Abstract

Score-based diffusion models have significantly advanced high-dimensional data generation across various domains, by learning a denoising oracle (or score) from datasets. From a Bayesian perspective, they offer a realistic modeling of data priors and facilitate solving inverse problems through posterior sampling. Although many heuristic methods have been developed recently for this purpose, they lack the quantitative guarantees needed in many scientific applications. In this work, we introduce the \textit{tilted transport} technique, which leverages the quadratic structure of the log-likelihood in linear inverse problems in combination with the prior denoising oracle to transform the original posterior sampling problem into a new `boosted' posterior that is provably easier to sample from. We quantify the conditions under which this boosted posterior is strongly log-concave, highlighting the dependencies on the condition number of the measurement matrix and the signal-to-noise ratio. The resulting posterior sampling scheme is shown to reach the computational threshold predicted for sampling Ising models [Kunisky'23] with a direct analysis, and is further validated on high-dimensional Gaussian mixture models and scalar field models.
Paper Structure (50 sections, 16 theorems, 119 equations, 6 figures, 2 tables, 1 algorithm)

This paper contains 50 sections, 16 theorems, 119 equations, 6 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Notations:
  4. Acknowledgements:
  5. Preliminaries
  6. Problem Setup
  7. Denoising Oracles and Score-Based Diffusion
  8. Log-Sobolev Inequality and Fast Relaxation of Langevin Dynamics
  9. Evidence of Computational Hardness in the Generic Case
  10. Posterior Sampling via Tilted Transport
  11. A Motivating Example
  12. Hamilton-Jacobi Equation and Quadratic Tilts
  13. Tilt Transport Equation
  14. Solution to \ref{['eq:boost_ode']}
  15. Quantitative Conditions for Provable Sampling
  16. ...and 35 more sections

Key Result

Proposition 4

Assume a posterior sampler exists for the smooth prior with TV error $\epsilon$ and $\delta = o(d^{-1/2})$. Then there exists a sampler for the associated Ising model with TV error $1.1 \epsilon$.

Figures (6)

  • Figure 1: Schematic plot of tilted transport boosting posterior sampling with a 2D Gaussian mixture example. The density plot shows the first variable's density, and the scatter plot displays the samples.
  • Figure 2: Phase diagram for the boosted posterior ${\nu}_{T}$ being strongly log-concave in \ref{['coro:gaussmixt']}.
  • Figure 3: Scheme plot of iterated titled transport in an example with $d=3, k^\star=3$. The tilted transport from $T_{k+1}$ to $T_k$ transforms samples of $\nu_{k+1}$ (corresponding to eigenvalues $\bar{\lambda}_j(T_{k+1})$ to $\tilde{\nu}_k$ (corresponding to eigenvalues $\tilde{\lambda}_j(T_k)$), and marginalization/thermalization (denoted by the dashed lines) further transforms samples of $\tilde{\nu}_k$ toward ${\nu}_k$.
  • Figure 4: Comparison of Langevin and boosted Langevin for Gaussian mixture prior. We generate the prior, measurement and sample the posterior under 20 different instances in each setting. The sliced Wasserstein distances are plotted with the median in the middle, and the 25th and 75th percentiles indicated by the error bars.
  • Figure 5: Comparison of autocorrelation in Langevin dynamics between the original posterior and boosted posterior distributions. The autocorrelation is computed for each site individually and then averaged.
  • ...and 1 more figures

Theorems & Definitions (32)

  • Proposition 4: Hardness extends to smooth priors
  • Theorem 5: Tilted Transport under OU Semigroup
  • Corollary 6: Posterior Sampling via Tilted Transport
  • Remark 7: Tilted Transport under Heat Semigroup
  • Remark 8: Covariance Decomposition of bauerschmidt2019very and Polchinsky Flow
  • Corollary 9: Posterior Sampling via Tilted Transport, Heat Semigroup Setting
  • Proposition 10: Strong Log-Concavity of $\nu_{T}$
  • Proposition 11: Sufficient Condition for Finite Susceptibility
  • Proposition 12: Susceptibility Blow-up
  • Proposition 13: Properties of $\chi_t(\pi)$
  • ...and 22 more