Sampling from the Random Linear Model via Stochastic Localization Up to the AMP Threshold

TL;DR

This paper demonstrates that sampling from the linear inverse problem's posterior using stochastic localization matches the mean estimation threshold, with convergence proven under specific noise conditions.

Contribution

It establishes the equivalence of sampling and mean estimation thresholds in the linear inverse problem and provides convergence proofs in KL and Wasserstein distances.

Findings

Sampling threshold coincides with mean estimation threshold.

Convergence in smoothed KL divergence below the AMP threshold.

Wasserstein convergence under a dimension-free operator norm bound.

Abstract

Recently, Approximate Message Passing (AMP) has been integrated with stochastic localization (diffusion model) by providing a computationally efficient estimator of the posterior mean. Existing (rigorous) analysis typically proves the success of sampling for sufficiently small noise, but determining the exact threshold involves several challenges. In this paper, we focus on sampling from the posterior in the linear inverse problem, with an i.i.d. random design matrix, and show that the threshold for sampling coincides with that of posterior mean estimation. We give a proof for the convergence in smoothed KL divergence whenever the noise variance $\Delta$ is below $\Delta_{\rm AMP}$, which is the computation threshold for mean estimation introduced in (Barbier et al., 2020). We also show convergence in the Wasserstein distance under the same threshold assuming a dimension-free bound on the operator norm of the posterior covariance matrix, a condition strongly suggested by recent breakthroughs on operator norm bounds in similar replica symmetric systems. A key observation in our analysis is that phase transition does not occur along the sampling and interpolation paths assuming $\Delta<\Delta_{\rm AMP}$.