Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Score Operator Newton transport

Nisha Chandramoorthy, Florian Schaefer, Youssef Marzouk

TL;DR

The approach is an infinite-dimensional Newton method, involving a linear PDE, for finding a zero of a ``score-residual'' operator, and it is proved sufficient conditions for convergence to a valid transport map.

Abstract

We propose a new approach for sampling and Bayesian computation that uses the score of the target distribution to construct a transport from a given reference distribution to the target. Our approach is an infinite-dimensional Newton method, involving a linear PDE, for finding a zero of a ``score-residual'' operator. We prove sufficient conditions for convergence to a valid transport map. Our Newton iterates can be computed by exploiting fast solvers for elliptic PDEs, resulting in new algorithms for Bayesian inference and other sampling tasks. We identify elementary settings where score-operator Newton transport achieves fast convergence while avoiding mode collapse.

Score Operator Newton transport

TL;DR

The approach is an infinite-dimensional Newton method, involving a linear PDE, for finding a zero of a ``score-residual'' operator, and it is proved sufficient conditions for convergence to a valid transport map.

Abstract

We propose a new approach for sampling and Bayesian computation that uses the score of the target distribution to construct a transport from a given reference distribution to the target. Our approach is an infinite-dimensional Newton method, involving a linear PDE, for finding a zero of a ``score-residual'' operator. We prove sufficient conditions for convergence to a valid transport map. Our Newton iterates can be computed by exploiting fast solvers for elliptic PDEs, resulting in new algorithms for Bayesian inference and other sampling tasks. We identify elementary settings where score-operator Newton transport achieves fast convergence while avoiding mode collapse.
Paper Structure (14 sections, 3 theorems, 33 equations, 8 figures, 1 algorithm)

This paper contains 14 sections, 3 theorems, 33 equations, 8 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Infinite-dimensional score matching
  3. Learning a zero of the score-residual operator
  4. Score operator Newton (SCONE) method
  5. SCONE transport algorithm
  6. Related work
  7. Convergence proof
  8. Numerical results
  9. Discussion
  10. Other iterations based on the contraction mapping principle
  11. SCONE example
  12. Additional experiments
  13. 1D unbounded target: bimodal Gaussian with equally weighted modes
  14. 1D unbounded target: bimodal Gaussian with unequally weighted modes

Key Result

Theorem 1

Let $\Omega$ be a bounded and open subset of $\mathbb{R}^d$ with a smooth boundary. Let $L(x,D) u = f$ be a second-order strongly elliptic system, with $L(x,D) = \sum_{|\alpha| \leq 2} a_\alpha(x) D^\alpha$, and zero Dirichlet boundary conditions. If the coefficients $a_\alpha$ and the right hand s where $K$ only depends on $\|a_\alpha\|_{s,\gamma}$ and $d$.

Figures (8)

  • Figure 1: A graphical overview of the construction of transport maps. The score of the source and target densities are given as inputs. The method outputs samples from the target distribution. Each iteration involves solving an elliptic PDE that gradually transports the score of the source to that of the target. The PDE solutions across iterations are combined via a simple composition operation to obtain the transport map from the source to the target.
  • Figure 2: Left: the solution $v_n$ after 1 (top) and 5 (bottom) iterations computed using second order finite difference with O(500) grid points. Center: the transformed empirical density. Right: $p_1$ (top) and $p_5$ (bottom).
  • Figure 3: Convergence of SCONE. Left: the convergence of $\|v_n\|$. Center: transformed empirical density after 15 iterations. Right: $T_n$ after 15 iterations.
  • Figure 4: Left: SVGD with RBF kernels (median heuristic for bandwidth) and 512 particles. Center: parameterized monotone transport map parno2022mpart, with polynomial degree 10, 512$^2$/11 samples, optimized using gradient descent + line search. Right: SCONE transport with ODE updates solved with 512 grid points.
  • Figure 5: Numerical validation of SCONE transport on 1D densities: the first row depicts results after 1 iteration of our Newton method and the second row after 5 iterations. The scalar field $v,$ the histogram approximations of the target density and the target score are plotted for the two different targets described in section \ref{['sec:appx-num']}. The results of the computed transformed densities and scores from the SCONE iteration are compared against the target densities and scores in columns 2 and 3.
  • ...and 3 more figures

Theorems & Definitions (9)

  • Remark 1: Availability of scores
  • Remark 2: Validity of transport maps
  • Definition 1: Hölder space of order $k$ and exponent $\gamma$
  • Theorem 1
  • Theorem 2: Score-matching
  • proof
  • Theorem 3: SCONE construction of transport
  • proof
  • Remark 3