Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Han Cui, Zhiyuan Yu, Jingbo Liu

TL;DR

This work addresses sampling from the posterior in a high-dimensional linear Gaussian model with i.i.d. Gaussian design by combining diffusion-based stochastic localization with AMP to approximate the posterior mean. The authors introduce a discretized SDE whose drift is the AMP-based posterior mean and show that sampling converges in a smoothed KL sense for Delta below the AMP threshold, with Wasserstein convergence under a dimension-free Lipschitz bound on the mean-field map. They show the sampling threshold matches the mean-estimation threshold and that no phase transitions occur along the diffusion path in this regime, aided by state-evolution analysis and recent operator-norm insights. A corollary indicates that for large measurement rate alpha, sampling succeeds for any fixed Delta, and numerical experiments corroborate practical efficiency and accuracy advantages over DPS-Gaussian baselines.

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 $Δ$ is below $Δ_{\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 $Δ<Δ_{\rm AMP}$.

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

TL;DR

This work addresses sampling from the posterior in a high-dimensional linear Gaussian model with i.i.d. Gaussian design by combining diffusion-based stochastic localization with AMP to approximate the posterior mean. The authors introduce a discretized SDE whose drift is the AMP-based posterior mean and show that sampling converges in a smoothed KL sense for Delta below the AMP threshold, with Wasserstein convergence under a dimension-free Lipschitz bound on the mean-field map. They show the sampling threshold matches the mean-estimation threshold and that no phase transitions occur along the diffusion path in this regime, aided by state-evolution analysis and recent operator-norm insights. A corollary indicates that for large measurement rate alpha, sampling succeeds for any fixed Delta, and numerical experiments corroborate practical efficiency and accuracy advantages over DPS-Gaussian baselines.

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 is below , 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 .
Paper Structure (13 sections, 20 theorems, 125 equations, 2 figures, 1 table, 1 algorithm)

This paper contains 13 sections, 20 theorems, 125 equations, 2 figures, 1 table, 1 algorithm.

Table of Contents

  1. INTRODUCTION
  2. Main contributions
  3. Related works
  4. METHOD
  5. CONVERGENCE GUARANTEES
  6. NUMERICAL EXPERIMENTS
  7. DISCUSSION
  8. PROOFS
  9. Proof of Theorem \ref{['thm1']}
  10. Proof of Theorem \ref{['theorem 7']}
  11. BOUNDING THE WASSERSTEIN DISTANCE BETWEEN $z_T/T$ AND $\tilde{z}_T/T$
  12. Connection between overlap concentration and Lipschitz constant of $m_t$
  13. ADDITIONAL EXPERIMENTS

Key Result

Theorem 1

Consider the linear observation model in model1, with fixed measurement rate $\alpha$ and noise level $\Delta$, and suppose that $\boldsymbol{\theta}$ has a product prior compactly supported on $[-L_\theta,L_\theta]^{N}$. Suppose that $\Delta<\Delta_{\rm AMP}$, where $\Delta_{\rm AMP}$ is the thresh

Figures (2)

  • Figure 1: log Algorithm MSE and log DPS-Gaussian MSE for different values of $N$. For this experiment, we let the prior be standard normal distribution and set $\Delta=0.01, \alpha=2, K=50, T=300, \delta=0.1$.
  • Figure 2: Trajectories of the first and second coordinates of the samples generated by Algorithm \ref{['Alg1']}. In this experiment, the prior distribution is $P_0(\theta)=\frac{1}{2}\delta_{-1}+\frac{1}{2}\delta_{1}$, and we set $M=1000, N=1250, K=20, T=200, \delta=0.1$. The three rows correspond to $\Delta$=0.01, 1, 10 respectively.

Theorems & Definitions (39)

  • Theorem 1
  • Remark 1
  • Remark 2
  • Remark 3
  • Corollary 1
  • Proposition 1
  • proof
  • Lemma 1
  • Lemma 2
  • Theorem 2
  • ...and 29 more