🤖 machine learning

Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems

This paper presents the first unified analysis of variance reduction techniques for sampling from high-dimensional, non-log-concave distributions, establishing improved non-asymptotic convergence rates and demonstrating their effectiveness in enhancing sample quality for inverse problems with score-based generative priors.

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

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: M. Berk Sahin, Ahmet Ege Tanriverdi, 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 perfect spot to set up a campsite in a vast, foggy, and incredibly complex mountain range. This mountain range represents a probability distribution. Your goal is to find the "valleys" (the most likely places) where the campsite should be.

In the world of machine learning, this is called sampling. You want to generate random points that accurately represent the shape of this mountain range.

The Problem: The Foggy Hike with a Bad Map

Usually, you have a map (a mathematical formula) that tells you the slope of the terrain. If you follow the slope downhill, you eventually find a valley. This is like optimization.

However, in many modern AI problems (like creating realistic images or solving medical scans), the map is too big to read all at once. You can only peek at a tiny, blurry piece of the map at a time. This is called a stochastic gradient.

The problem is that these "peeked" maps are noisy. Sometimes they tell you to go left, sometimes right, even if the true slope is straight down. This noise is called variance.

The Old Way (SGLD): To get a clear picture, the old method said, "Take a huge group of hikers (a large batch size), look at the map together, and average their opinions." This works, but it's slow and requires a massive team (huge memory and computing power) for every single step. If you try to do it with a small team, you get lost in the fog (the samples are bad).

The Solution: The Smart Hiking Team (Variance Reduction)

This paper introduces two new hiking strategies, named ML-VRLD and SL-VRLD. Think of them as "Smart Hiking Teams" that use Variance Reduction.

Instead of just asking a new group of hikers for directions every single step, these teams use memory:

Remembering the Past: They remember where they looked before.
Smoothing the Noise: If the new "peek" at the map says "Go Left!" but the memory says "We were just there and it was flat," the team knows the new peek was probably a fluke. They combine the new info with the old info to get a much clearer, more accurate direction.

The paper claims these teams can achieve the same (or better) accuracy as the massive "Old Way" team, but they only need one hiker (or a tiny, fixed-size team) at a time. They get the job done with O(1) effort per step, meaning the cost doesn't explode as the problem gets bigger.

The Terrain: Bumpy Mountains (Non-Log-Concave)

Most previous research assumed the mountains were nice, smooth bowls (Log-Concave). But real-world data is messy. It has multiple peaks, deep valleys, and weird shapes (Non-Log-Concave).

The Paper's Claim: These new "Smart Teams" are the first to prove mathematically that they can navigate these bumpy, messy mountains efficiently without needing a massive team at every step. They prove the hikers will eventually find the right valleys, even in this difficult terrain.

The Application: Fixing Blurry Photos (Inverse Problems)

The paper also applies this to Inverse Problems. Imagine you have a blurry photo of a crime scene (the measurements) and you want to reconstruct the original clear image (the solution).

The Challenge: There are infinite ways to make a blurry photo. You need a "prior" (a guess of what a normal photo looks like) to guide you. Modern AI uses "Score-Based Generative Models" (SGMs) as this guess.
The Bottleneck: Usually, to fix the photo, the AI has to look at every single pixel of the blurry image at once to figure out the next step. This is impossible for huge 3D medical scans (like MRI or CT) because it requires too much computer memory.
The Paper's Fix: They adapted their "Smart Hiking Teams" to work with these photo-fixing AI models. Now, the AI can fix the photo by looking at just a small slice of the image at a time (a mini-batch), while still using its memory to stay on track.

The Results: Sharper Photos, Less Memory

The authors tested this on:

Synthetic Mountains: They proved the new teams found the valleys faster and more accurately than the old method.
MRI Scans: They reconstructed brain images from incomplete data. The new method produced sharper images with fewer artifacts (blurry spots) compared to the old method, using the same amount of computing power per step.
CT Scans: They reconstructed images from very few X-ray angles. Again, the new method preserved fine details better than the competition.

Summary

In short, this paper says: "We found a way to navigate complex, messy mathematical landscapes using a tiny, memory-efficient team instead of a massive, expensive one. We proved this works mathematically, and we showed it creates better medical images (MRI and CT) by fixing blurry scans more accurately without needing supercomputers."

Technical Summary: Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems

Problem Statement
Sampling from high-dimensional, non-log-concave distributions with unnormalized densities is a fundamental challenge in machine learning, particularly when exact gradients of the potential function are unavailable. In such scenarios, gradients must be approximated via stochastic estimates, which often exhibit high variance under a fixed computational budget per iteration. While variance reduction techniques (e.g., SGD with momentum, STORM, PAGE) have improved convergence in non-convex optimization, their application to sampling from non-log-concave distributions remains largely unexplored. Furthermore, existing posterior sampling methods for inverse problems using score-based generative models (SGMs) typically operate in a full-batch regime, limiting their applicability to large-scale datasets common in 3D and dynamic imaging.

Methodology
The authors propose a unified framework that adapts variance reduction techniques from non-convex optimization to the sampling domain. They introduce two unified estimators:

Mean-Lipschitz Variance-Reduced Langevin Diffusion (ML-VRLD): This estimator unifies SGD with momentum and its stochastic counterpart. It assumes Lipschitz continuity of the exact gradients (Assumption 1) and provides a unified treatment where the estimator is a weighted average of a mini-batch gradient and a momentum-based update.
Sample-Lipschitz Variance-Reduced Langevin Diffusion (SL-VRLD): This estimator unifies PAGE and STORM. It assumes Lipschitz continuity of the stochastic gradients in the mean-squared sense (Assumption 4). It incorporates a correction term based on the difference between current and past stochastic gradients, which improves estimation accuracy under the sample-Lipschitz condition.

Both methods are designed to achieve $\epsilon$ -relative Fisher Information (FI) convergence with $O(1)$ stochastic gradient computations per iteration, contrasting with the $O(Ld\sigma^2/\epsilon)$ requirement of standard Stochastic Gradient Langevin Dynamics (SGLD).

The framework is extended to inverse problems (e.g., MRI and CT reconstruction) by integrating these estimators into an annealed Langevin Monte Carlo (LMC) sampling process with SGM priors. The posterior sampling formulation utilizes a sequence of intermediate target distributions with annealing and smoothing schedules to guide the iterates from a smoothed prior to the true posterior. Theoretical guarantees are established under bounded SGM estimation error assumptions.

Key Contributions

Unified Analysis: The paper provides the first unified analysis of variance-reduced estimators (specifically ML-VRLD and SL-VRLD) for sampling from non-log-concave distributions.
Improved Complexity:
- Both ML-VRLD and SL-VRLD achieve $\epsilon$ -relative FI convergence with $O(1)$ gradient computations per iteration, significantly improving upon the per-iteration complexity of SGLD.
- SL-VRLD further reduces the total computational complexity from $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD) to $O(L^2d^2\sigma^2/\epsilon^3)$ .
Convergence Guarantees: The authors establish non-asymptotic convergence rates in squared total variation (TV) distance under a Poincaré inequality assumption and prove weak convergence to the target distribution for both algorithms.
Inverse Problem Extension: The framework is extended to posterior sampling with SGM priors, establishing the same convergence guarantees under bounded SGM estimation error. This addresses the limitation of existing full-batch posterior sampling methods.

Results

Synthetic Experiments: On Gaussian mixture models (GMM) and compressed sensing inverse problems, the proposed methods demonstrate faster convergence to lower relative FI values compared to SGLD. SGLD with a fixed batch size converges to a suboptimal FI level, whereas ML-VRLD and SL-VRLD achieve the target accuracy with the same fixed batch size. SL-VRLD consistently reaches the target distribution faster than ML-VRLD.
Statistical Validation: In a posterior sampling task with a known ground truth, SGLD produced samples with significantly inflated variance. ML-VRLD and SL-VRLD significantly reduced this variance, with SL-VRLD achieving the closest match to the ground truth posterior variance.
Imaging Applications:
- MRI Reconstruction: On the fastMRI dataset with acceleration factors of 8 and 12, ML-VRLD and SL-VRLD produced sharper anatomical details and fewer artifacts compared to baselines (including DPS, PnP-DM, and SGLD). SL-VRLD achieved the highest PSNR and SSIM scores.
- Sparse-Angle CT Reconstruction: On the 2DeteCT dataset with 180 and 360 projection views, the proposed methods outperformed baselines in preserving fine structural details, demonstrating the benefit of variance reduction in high-noise, undersampled regimes.

Significance
The paper claims to provide the first non-asymptotic convergence guarantees for non-log-concave sampling using variance-reduced estimators with $O(1)$ gradient computations per iteration. By extending these techniques to inverse problems with SGM priors, the work enables principled posterior sampling for large-scale imaging tasks where full-batch methods are computationally prohibitive. The empirical results validate that variance reduction consistently improves sample quality and statistical accuracy under fixed memory and computational constraints, offering a scalable solution for high-dimensional Bayesian inference.