Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

An Eulerian approach to regularized JKO scheme with low-rank tensor decompositions for Bayesian inversion

Vitalii Aksenov, Martin Eigel

TL;DR

The possibility of using the Eulerian discretization for the problem of modelling high-dimensional distributions and sampling is studied and the resulting method requires only pointwise computations of the unnormalized posterior and is, in particular, gradient-free.

Abstract

The possibility of using the Eulerian discretization for the problem of modelling high-dimensional distributions and sampling, is studied. The problem is posed as a minimization problem over the space of probability measures with respect to the Wasserstein distance and solved with entropy-regularized JKO scheme. Each proximal step can be formulated as a fixed-point equation and solved with accelerated methods, such as Anderson's. The usage of low-rank Tensor Train format allows to overcome the \emph{curse of dimensionality}, i.e. the exponential growth of degrees of freedom with dimension, inherent to Eulerian approaches. The resulting method requires only pointwise computations of the unnormalized posterior and is, in particular, gradient-free. Fixed Eulerian grid allows to employ a caching strategy, significally reducing the expensive evaluations of the posterior. When the Eulerian model of the target distribution is fitted, the passage back to the Lagrangian perspective can also be made, allowing to approximately sample from it. We test our method both for synthetic target distributions and particular Bayesian inverse problems and report comparable or better performance than the baseline Metropolis-Hastings MCMC with same amount of resources. Finally, the fitted model can be modified to facilitate the solution of certain associated problems, which we demonstrate by fitting an importance distribution for a particular quantity of interest.

An Eulerian approach to regularized JKO scheme with low-rank tensor decompositions for Bayesian inversion

TL;DR

The possibility of using the Eulerian discretization for the problem of modelling high-dimensional distributions and sampling is studied and the resulting method requires only pointwise computations of the unnormalized posterior and is, in particular, gradient-free.

Abstract

The possibility of using the Eulerian discretization for the problem of modelling high-dimensional distributions and sampling, is studied. The problem is posed as a minimization problem over the space of probability measures with respect to the Wasserstein distance and solved with entropy-regularized JKO scheme. Each proximal step can be formulated as a fixed-point equation and solved with accelerated methods, such as Anderson's. The usage of low-rank Tensor Train format allows to overcome the \emph{curse of dimensionality}, i.e. the exponential growth of degrees of freedom with dimension, inherent to Eulerian approaches. The resulting method requires only pointwise computations of the unnormalized posterior and is, in particular, gradient-free. Fixed Eulerian grid allows to employ a caching strategy, significally reducing the expensive evaluations of the posterior. When the Eulerian model of the target distribution is fitted, the passage back to the Lagrangian perspective can also be made, allowing to approximately sample from it. We test our method both for synthetic target distributions and particular Bayesian inverse problems and report comparable or better performance than the baseline Metropolis-Hastings MCMC with same amount of resources. Finally, the fitted model can be modified to facilitate the solution of certain associated problems, which we demonstrate by fitting an importance distribution for a particular quantity of interest.

Paper Structure

This paper contains 30 sections, 5 theorems, 85 equations, 7 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Sampling approaches
  3. Gradient flow structure of sampling approaches
  4. Low-rank Tensor decompositions
  5. Contribution and paper organization
  6. Metric space of measures with respect to the Wasserstein distance
  7. Wasserstein distance and its dynamic reformulation
  8. Minimizing movements scheme
  9. Entropic regularization of the Minimizing Movements scheme
  10. Derivation of the method
  11. Fixed-point formulation of the proximal step
  12. Convergence of the Minimizing Movements with regularized proximal
  13. Numerical solution
  14. Discretization and Tensor Train compression
  15. Solution of the heat equation in TT format
  16. ...and 15 more sections

Key Result

Theorem 2.1

Continuity equation ambrosio2005gradient There exists a Borel vector field $v_t \in L^2(\rho_t, X):\ \|v_t\|_{L^2(\rho_t, X)} \leq |\rho_t'|$ and the continuity equation holds in the sense of distributions, i.e.

Figures (7)

  • Figure 1: Schematic of proposed fixed-point algorithm.
  • Figure 2: Contour lines of the test distribution and the samples generated with the deterministic method (left) and the combination of deterministic and stochastic method (right).
  • Figure 3: Reconstruction for a Gaussian target distribution.
  • Figure 4: Contour lines of two-dimensional marginals of the test distributions shown with the reference sample, the TT-generated sample, and an MCMC sample using a shorter chain with an equivalent number of posterior calls as the TT method.
  • Figure 5: Posterior distribution and reconstruction of the initial solution for the hyperbolic problem.
  • ...and 2 more figures

Theorems & Definitions (10)

  • Definition 2.1
  • Theorem 2.1
  • Definition 2.2
  • Corollary 2.1.1
  • proof
  • Theorem 2.2
  • Corollary 2.2.1
  • Definition 4.1
  • Theorem 4.1: qin2022error
  • Remark