Distributional Shrinkage II: Higher-Order Scores Encode Brenier Map

TL;DR

This paper introduces a hierarchy of denoisers based on higher-order score functions that approximate the Brenier map for Gaussian additive models, connecting optimal transport, Fisher information, and combinatorics.

Contribution

It develops a novel hierarchy of denoisers using higher-order score functions that recover the Brenier map without prior knowledge of the signal distribution.

Findings

Higher-order denoisers achieve faster convergence rates in Wasserstein distance.

The hierarchy is characterized by Bell polynomial recursions.

Rates of score estimation are established via kernel density and score matching methods.

Abstract

Consider the additive Gaussian model $Y = X + \sigma Z$, where $X \sim P$ is an unknown signal, $Z \sim N(0,1)$ is independent of $X$, and $\sigma > 0$ is known. Let $Q$ denote the law of $Y$. We construct a hierarchy of denoisers $T_0, T_1, \ldots, T_\infty \colon \mathbb{R} \to \mathbb{R}$ that depend only on higher-order score functions $q^{(m)}/q$, $m \geq 1$, of $Q$ and require no knowledge of the law $P$. The $K$-th order denoiser $T_K$ involves scores up to order $2K{-}1$ and satisfies $W_r(T_K \sharp Q, P) = O(\sigma^{2(K+1)})$ for every $r \geq 1$; in the limit, $T_\infty$ recovers the monotone optimal transport map (Brenier map) pushing $Q$ onto $P$. We provide a complete characterization of the combinatorial structure governing this hierarchy through partial Bell polynomial recursions, making precise how higher-order score functions encode the Brenier map. We further establish rates of convergence for estimating these scores from $n$ i.i.d.\ draws from $Q$ under two complementary strategies: (i) plug-in kernel density estimation, and (ii) higher-order score matching. The construction reveals a precise interplay among higher-order Fisher-type information, optimal transport, and the combinatorics of integer partitions.