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

TL;DR

This paper introduces an Eulerian discretization approach combined with low-rank tensor decompositions to efficiently solve high-dimensional Bayesian inverse problems using a regularized JKO scheme, overcoming the curse of dimensionality.

Contribution

It proposes a novel Eulerian method with low-rank tensor trains for Bayesian inversion, enabling gradient-free, scalable computations and effective sampling in high dimensions.

Findings

Achieves comparable or better performance than MCMC methods.

Reduces computational cost via caching and low-rank tensor formats.

Enables approximate sampling and model modification for related problems.

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.