🤖 machine learning

Zeroth-Order Non-Log-Concave Sampling with Variance Reduction and Applications to Inverse Problems

This paper introduces a variance-reduced zeroth-order Langevin sampling method that achieves the first non-asymptotic convergence guarantees for high-dimensional non-log-concave distributions without gradient access, and applies this framework to develop ZO-APMC, a novel algorithm for solving black-box inverse problems with pre-trained score-based priors.

Original authors: M. Berk Sahin, Behzad Sharif, Abolfazl Hashemi

Published 2026-06-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: M. Berk Sahin, Behzad Sharif, Abolfazl Hashemi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the best spot to set up a campsite in a vast, foggy, and mountainous wilderness. Your goal is to find the "perfect" spot (the target distribution), but the terrain is tricky: it has deep valleys, flat plateaus, and hidden peaks (non-log-concave).

Usually, to navigate this, you would need a map with a compass that tells you exactly which way is "downhill" (the gradient). But in this paper, the authors are dealing with a situation where you have no compass and no map. You can only look at your current location and ask, "If I take a tiny step in this random direction, is the ground higher or lower?" This is called Zeroth-Order information. It's like feeling the ground with your feet instead of reading a topographic map.

Here is the breakdown of their solution, the ZO-APMC method, using simple analogies:

1. The Problem: The "Noisy Flashlight"

When you can't see the gradient (the slope), you have to guess it by taking random steps.

The Old Way: To get a decent guess of the slope, you had to take thousands of random steps in every single move to average out the noise. This is like trying to find a path by taking 1,000 random steps, measuring the height, and then repeating that 1,000-step process for every single step you take. It's incredibly slow and expensive.
The Paper's Insight: The authors realized that if you take a few steps, then take a few more steps, the change in the ground between those two moments is usually small and predictable. You don't need to re-measure the whole world every time.

2. The Solution: The "Variance-Reduced Flashlight"

The authors invented a new way to guess the slope that is like a smart, memory-equipped flashlight.

How it works: Instead of taking thousands of steps every time, they take a big batch of steps occasionally (to get a fresh, accurate baseline) and then, for the next few moves, they just measure the difference between where they were and where they are now.
The Analogy: Imagine you are walking in the dark. Instead of stopping every 10 feet to shine a giant, expensive spotlight on the whole forest to see the path (the old way), you shine a small flashlight once to get your bearings, and then just watch how the ground shifts under your feet as you walk. You only turn on the big spotlight occasionally to correct your course.
The Result: This reduces the "noise" (variance) in your guess without needing thousands of measurements. It makes the process fast and stable, even in high-dimensional spaces (like navigating a 1,000-dimensional wilderness).

3. The Application: Solving "Black Box" Puzzles

The paper applies this to Inverse Problems.

The Scenario: Imagine you have a photo of a blurry black hole or a damaged MRI scan. You want to reconstruct the original clear image.
The "Black Box": Usually, to fix the image, you need to know exactly how the camera or MRI machine works (the math behind the blur). But sometimes, that machine is a "black box"—it's a closed system, a legacy computer program, or a complex physics simulation where you can't see the math inside. You can only put an image in and get a blurry result out.
The Paper's Approach: They combine their "smart flashlight" (the sampling method) with a pre-trained AI that knows what "good images" look like (a Score-Based Generative Model, or SGM).
- The AI acts as a guide saying, "This looks like a real brain scan."
- The "smart flashlight" checks the black box: "If I change this pixel, does the blurry output get closer to what I measured?"
- Together, they navigate the fog to find the most likely original image without ever needing to open the black box or see its internal math.

4. The Proof: Why It Works

The authors didn't just guess; they proved mathematically that this method works.

Convergence: They showed that even with this noisy, "blind" navigation, the method is guaranteed to eventually find the correct solution (the target distribution).
Efficiency: They proved that you don't need the massive number of steps the old methods required. You can get accurate results with a constant, small number of checks per step, regardless of how complex the problem is.

Summary of Results

The paper tested this on three real-world "black box" challenges:

MRI Reconstruction: Rebuilding clear brain images from blurry, undersampled data. Their method produced images almost as good as methods that do have access to the internal math, but without needing that access.
Black Hole Imaging: Reconstructing images of black holes from telescope data. Their method produced the clearest images among all "black box" competitors, capturing fine details that others missed.
Fluid Dynamics (Navier-Stokes): Trying to figure out the starting conditions of a fluid flow based on a snapshot. Here, their method performed competitively, proving it works even on complex physics simulations.

In a nutshell: The authors created a "blind navigation" tool that is smart enough to remember its previous steps, allowing it to solve complex, high-dimensional puzzles (like medical imaging and astrophysics) without needing to see the internal math of the systems it's trying to reverse-engineer.

Technical Summary: Zeroth-Order Non-Log-Concave Sampling with Variance Reduction and Applications to Inverse Problems

1. Problem Statement

The paper addresses the fundamental challenge of sampling from high-dimensional, non-log-concave distributions when the potential function $f$ (where the target distribution is $\pi \propto \exp(-f)$ ) is only accessible via zeroth-order (ZO) evaluations. This "black-box" setting arises when gradient information is unavailable, computationally prohibitive, or undefined (e.g., in legacy code, proprietary systems, or PDE-based simulators).

Existing methods face two primary limitations in this context:

High Variance: Standard ZO gradient estimators based on finite differences along random Gaussian directions exhibit high variance. To control this variance, traditional approaches require batch sizes scaling linearly with the dimension $d$ ( $O(d)$ ), leading to prohibitive computational and memory costs in high-dimensional settings.
Lack of Theoretical Guarantees: While Langevin dynamics provides a framework for sampling with gradients, its extension to ZO settings lacks non-asymptotic convergence guarantees for non-log-concave distributions. Furthermore, existing black-box posterior sampling methods for inverse problems often rely on heuristic approximations without rigorous convergence guarantees under standard discrepancy measures like Fisher Information (FI) or Total Variation (TV) distance.

2. Methodology

Variance-Reduced Zeroth-Order Estimator

The core methodological contribution is a novel variance-reduced ZO gradient estimator, $g_k$ , designed to achieve accurate gradient approximation with only $O(1)$ function evaluations per iteration, independent of the ambient dimension $d$ . The estimator operates via a probabilistic update rule:

With probability $p$ , a large batch estimator is computed using $b$ function evaluations.
With probability $1-p$ , the estimate is updated recursively using a small batch ( $b' \ll b$ ) of fresh evaluations combined with the previous estimate $g_{k-1}$ .

This mechanism exploits the strong correlation between consecutive iterates ( $x_{k\gamma}$ and $x_{(k-1)\gamma}$ ) to approximate the gradient change, thereby reducing the per-iteration computational cost while maintaining variance control. The method assumes the potential function $f$ is Lipschitz continuous to bound the variance of the gradient variation.

ZO-APMC Algorithm for Inverse Problems

The authors extend this framework to solve ill-posed inverse problems ( $y = A(x) + \xi$ ) in a black-box setting using a pre-trained Score-Based Generative Model (SGM) as a prior. The proposed algorithm, ZO-APMC (Zeroth-Order Annealed Posterior Monte Carlo), combines:

Annealed Langevin Dynamics: A sequence of weighted posterior distributions where the prior is gradually smoothed and the likelihood weight is increased.
Variance-Reduced Estimation: Replacing the naive ZO gradient of the likelihood with the proposed variance-reduced estimator.
SGM Priors: Utilizing a pre-trained SGM to estimate the prior score $-\nabla h(x)$ , accounting for both the perturbation error (due to noise level $\sigma_k$ ) and the network estimation error.

The algorithm iterates through a continuous-time interpolation that incorporates the annealing schedule and the ZO estimator, allowing for sampling from complex, non-log-concave posteriors without requiring gradients of the forward operator $A(\cdot)$ .

3. Key Contributions

First Non-Asymptotic Guarantees for ZO Non-Log-Concave Sampling: The paper establishes the first non-asymptotic convergence guarantees for zeroth-order sampling from non-log-concave distributions.
- It achieves an $\epsilon$ -relative Fisher Information error after $O(d^7 L_m^4 / \epsilon^4)$ iterations.
- Under a Poincaré inequality assumption, it yields $\epsilon$ -accuracy in squared Total Variation distance.
- Crucially, these guarantees are achieved with $O(1)$ function evaluations per iteration, eliminating the $O(d)$ batch-size scaling required by standard ZO estimators.
Variance-Reduced Estimator: The proposal of a novel estimator that decouples the batch size from the dimension, substantially reducing memory and computational costs while enabling stable sampling.
Theoretically Grounded Black-Box Posterior Sampling: The development of ZO-APMC, the first posterior sampling method for black-box inverse problems with SGM priors that offers formal non-asymptotic and weak convergence guarantees. The analysis explicitly accounts for biases introduced by annealing schedules and SGM estimation errors.
Empirical Validation: Extensive experiments demonstrating that ZO-APMC outperforms existing black-box baselines (such as EnKG, DPG, SCG, and Forward/Central-GSG) in MRI reconstruction and black-hole imaging, while achieving competitive performance on the Navier–Stokes inverse problem.

4. Results

Synthetic Experiments: Numerical validation confirms the theoretical convergence rates. The variance-reduced mechanism ( $p < 1$ ) is shown to converge to lower Fisher Information values compared to naive ZO approaches ( $p=1$ ) under fixed per-iteration budgets. Statistical validation on a bimodal Gaussian mixture prior demonstrates that ZO-APMC accurately recovers posterior mode statistics (mean and variance), closely matching gradient-based counterparts (APMC).
MRI Reconstruction: On the FastMRI dataset, ZO-APMC achieves state-of-the-art performance among black-box methods, with PSNR (35.29 dB) and SSIM (0.966) metrics closely matching the gradient-based APMC (36.55 dB, 0.973) and significantly outperforming other black-box baselines. Visualizations show that ZO-APMC captures fine details without gradient information.
Black-Hole Imaging: On the InverseBench dataset, ZO-APMC produces reconstructions that most closely match ground truth, outperforming all black-box baselines across PSNR, closure-phase error, and closure-amplitude error metrics.
Navier–Stokes Inverse Problem: While ZO-APMC achieves performance comparable to the strong baseline DPG and outperforms other black-box methods, it trails slightly behind EnKG. The authors attribute this to the high nonlinearity of the PDE solver, where Gaussian perturbations (used in ZO) may move inputs off the data manifold, affecting solver stability.

5. Significance and Claims

The paper claims to establish a theoretically grounded framework for zeroth-order sampling in non-log-concave settings, addressing a fundamental gap in derivative-free probabilistic inference. By providing the first non-asymptotic convergence guarantees for both general sampling and black-box posterior sampling with SGM priors, the work enables principled probabilistic inference in scenarios where gradients are inaccessible, undefined, or prohibitively expensive.

The authors position ZO-APMC as a significant step toward scalable, black-box posterior sampling, offering an algorithm with formal convergence guarantees that can be applied to diverse fields including scientific simulators (PDEs), medical imaging (MRI), and astrophysics (black-hole imaging). The work emphasizes that while the theoretical iteration complexity has a high polynomial dependence on dimension, empirical results on high-dimensional tasks show convergence with substantially fewer iterations, matching the practical performance of gradient-based counterparts.

Limitations and Modesty: The paper acknowledges that the theoretical bounds depend on Lipschitz continuity assumptions and Poincaré inequalities, which may not hold globally for all distributions. It also notes that in highly nonlinear PDE settings (Navier–Stokes), the ZO perturbation strategy may be less robust than ensemble-based methods like EnKG. The authors do not claim to solve all black-box inverse problems but rather provide a rigorous foundation and a practical algorithm that bridges the gap between heuristic black-box methods and theoretically guaranteed gradient-based methods.