Stochastic gradient descent based variational inference for infinite-dimensional inverse problems

This paper proposes and theoretically validates two stochastic gradient descent-based variational inference methods for infinite-dimensional inverse problems, utilizing constant-rate iterations with randomization to efficiently sample from posterior distributions and demonstrating their effectiveness through preconditioning and numerical applications to linear and non-linear problems.

The Gist

Imagine you are a detective trying to solve a mystery, but you can't see the culprit directly. You only have a few blurry photos (the data) and a hunch about what the culprit might look like (the prior knowledge). Your goal is to reconstruct the culprit's face as accurately as possible. In the world of math and science, this is called an inverse problem.

The paper you're asking about is about a new, super-fast way for computers to solve these mysteries, especially when the "culprit" is something incredibly complex, like the temperature distribution inside a star or the flow of water through a sponge (which are infinite-dimensional problems).

Here is the breakdown of their solution, using simple analogies:

1. The Old Way: The Slow, Careful Hiker (MCMC)

Traditionally, to solve these mysteries, scientists used a method called Markov Chain Monte Carlo (MCMC).

• The Analogy: Imagine you are trying to find the highest peak in a massive, foggy mountain range. The old method is like a hiker who takes one tiny, careful step at a time. They check every direction, make sure they aren't falling off a cliff, and slowly wander around until they have mapped the whole mountain.
• The Problem: This is incredibly accurate, but it takes forever. If the mountain is huge (infinite dimensions), the hiker might spend their whole life just walking in circles. It's too slow for modern, large-scale problems.

2. The New Idea: The Drunken Skier (Stochastic Gradient Descent)

The authors propose a new method based on Stochastic Gradient Descent (SGD).

• The Analogy: Instead of a careful hiker, imagine a skier who is slightly drunk. They know the general direction of the peak (the gradient), but they also have a bit of a wobble (random noise). They slide down the hill quickly, zig-zagging a bit.
• The Magic: In machine learning, this "drunken" sliding is usually used just to find the best spot (the peak). But the authors realized something brilliant: If you keep the skier moving at a constant speed with just the right amount of wobble, they don't just find the peak; they end up exploring the entire mountain range in a way that perfectly matches the probability of where the peak actually is.
• The Result: Instead of taking 100 years to map the mountain, the skier does it in a few hours.

3. The Two Versions: The Basic Skier vs. The Preconditioned Skier

The paper introduces two versions of this "drunken skier" technique:

• cSGD-iVI (The Basic Skier): This is the standard version. The skier has a fixed amount of wobble. It's fast and works well for simple mountains.
• pcSGD-iVI (The Preconditioned Skier): This is the upgraded version. Imagine the skier is now wearing special skis (a "preconditioner") that are tuned to the shape of the mountain.
  • If the mountain has a steep cliff on one side and a gentle slope on the other, normal skis make you spin out. These special skis adjust your balance automatically.
  • The Result: The preconditioned skier doesn't just go fast; they go smart. They explore the mountain much more efficiently and give a much more accurate map of the "culprit's" location.

4. Why "Infinite-Dimensional" Matters

Most computer problems are like a grid of 100 pixels. You can count them. But real-world physics (like fluid flow or heat) happens in a continuous space—you can zoom in forever, and there's always more detail. This is "infinite-dimensional."

• The Challenge: You can't count the pixels if there are infinite of them.
• The Solution: The authors proved mathematically that their "drunken skier" method works even when the mountain is infinitely detailed. They showed that the skier's path converges to the true answer, provided you tune the "wobble" (the noise) and the "speed" (the learning rate) correctly.

5. The "Secret Sauce": Tuning the Wobble

The most important part of the paper is how they figure out exactly how much the skier should wobble.

• They treat the problem like a game of "Hot or Cold." They want the skier's path to match the true probability of where the culprit is.
• They use a mathematical tool called KL Divergence (think of it as a "distance meter" between two maps) to find the perfect speed and wobble amount.
• The Outcome: They derived a formula that tells the computer exactly how to set these knobs so the "drunken" path becomes a perfect statistical map of the truth.

Summary of Results

The authors tested this on two real-world scenarios:

1. A simple smooth equation: Like finding a hidden shape in a smooth fog.
2. Darcy Flow: Like figuring out how water moves through a complex underground rock layer (very messy and non-linear).

The Verdict:

• The Basic Skier (cSGD) was much faster than the old "Careful Hiker" (MCMC) but sometimes missed the edges of the mountain.
• The Preconditioned Skier (pcSGD) was the winner. It was fast, accurate, and its map of the mountain was almost identical to the slow, perfect map made by the old method, but it did it in a fraction of the time.

In a nutshell: This paper teaches computers how to solve incredibly complex, infinite puzzles by turning a slow, careful search into a fast, slightly chaotic slide that magically lands on the right answer every time.

Technical Summary

Here is a detailed technical summary of the paper "Stochastic gradient descent based variational inference for infinite-dimensional inverse problems."

1. Problem Statement

The paper addresses infinite-dimensional inverse problems governed by Partial Differential Equations (PDEs), where the goal is to recover an unknown function $u$ (e.g., a source term or permeability) from noisy measurements $d$.

• Context: These problems are inherently ill-posed and defined in infinite-dimensional Hilbert spaces ($H_u$).
• Challenge: Traditional Bayesian methods often rely on Markov Chain Monte Carlo (MCMC) sampling (e.g., pCN). While rigorous, MCMC is computationally prohibitive for large-scale problems due to slow mixing and high iteration counts.
• Gap: Existing Variational Inference (VI) methods for PDE-constrained problems often follow a "Discretize-then-Bayesianize" approach, which introduces discretization errors and lacks uniform convergence guarantees. There is a lack of efficient, infinite-dimensional VI methods that operate under a "Bayesianize-then-discretize" framework using Stochastic Gradient Descent (SGD).

2. Methodology

The authors propose two novel methods: cSGD-iVI (constant-step SGD Variational Inference) and pcSGD-iVI (preconditioned cSGD-iVI). These methods treat the SGD iteration not just as an optimization tool, but as a sampling mechanism to approximate the posterior distribution.

A. Core Framework

1. Problem Setup: The inverse problem is formulated as $d = H(u) + \epsilon$, with a Gaussian prior $\mu_0 = \mathcal{N}(0, C_0)$. The posterior $\mu$ is defined via Bayes' theorem.
2. Randomization Strategy: Instead of mini-batching (which is unsuitable for PDE-constrained cost functions), the authors introduce stochastic gradient noise directly into the gradient update.
   • The stochastic gradient is defined as $\tilde{G}(u) = G(u) - \frac{1}{\sqrt{S}}\Delta G(u)$, where $\Delta G(u) \sim \mathcal{N}(0, C_{GN})$.
   • $S$ acts as a scaling parameter analogous to batch size.
3. Discrete-Time Process: The constant-step SGD iteration is viewed as a discrete-time stochastic process:
   $$u_{k+1} = u_k - \eta \tilde{G}(u_k)$$
   The stationary distribution of this process, denoted $\nu$, is used as the approximate posterior.

B. Theoretical Derivations

• Covariance Characterization: By analyzing the iteration in the eigenbasis of the prior covariance $C_0$, the authors derive the stationary mean and covariance of the process. The covariance operator $C_\nu$ depends explicitly on the learning rate $\eta$ and the noise scale $S$.
• Optimal Learning Rate: The optimal learning rate $\eta^\dagger$ is determined by minimizing the Kullback-Leibler (KL) divergence between the estimated posterior $\nu$ and the true posterior $\mu$. This yields a closed-form expression for $\eta^\dagger$ that balances the data fidelity and prior information.
• Regularization & Error Bounds: The paper establishes that the cSGD formulation possesses regularization properties. It provides a theoretical bound on the discretization error between the estimated posterior mean and the true background truth, showing the error is controlled by the learning rate and the truncation level $M$ of the eigenmodes.

C. Preconditioning (pcSGD-iVI)

To improve sampling efficiency and convergence, a preconditioned version is introduced:

• The update rule becomes $u_{k+1} = u_k - \eta T \tilde{G}(u_k)$, where $T$ is a bounded linear symmetric preconditioning operator.
• Theoretical results are extended to show how $T$ modifies the optimal learning rate and the stationary covariance, allowing for better alignment with the true posterior geometry.

D. Noise Operator Selection

The paper discusses how to choose the operator $Q$ (defining the noise covariance $C_{GN} = C_0 Q$). It proposes using a random projection operator $P$ to approximate the gradient noise, deriving the necessary conditions on the projection dimension $p$ to ensure the stability of the Lyapunov equation governing the stationary distribution.

3. Key Contributions

1. Infinite-Dimensional VI via SGD: The first development of a constant-step SGD-based variational inference framework specifically for infinite-dimensional inverse problems, bridging the gap between optimization and Bayesian sampling in function spaces.
2. Theoretical Validation: Rigorous proof that the stationary distribution of the cSGD iteration approximates the true posterior. The authors derive the exact form of the estimated covariance operator and prove the existence of an optimal learning rate that minimizes KL divergence.
3. Error Analysis: Establishment of discretization error bounds between the approximated posterior mean and the true solution, dependent on the learning rate and the spectral truncation level.
4. Preconditioned Variant: Introduction of pcSGD-iVI, which significantly accelerates convergence and improves the accuracy of the posterior covariance estimation compared to the unpreconditioned version.
5. Noise Strategy: A novel approach to defining stochastic gradient noise in PDE-constrained settings without relying on mini-batches, utilizing random projections to determine the noise structure.

4. Numerical Results

The methods were tested on two problems: a linear elliptic equation and a nonlinear steady-state Darcy flow equation.

• Linear Elliptic Problem:

  • Comparison: Compared against the standard pCN (preconditioned Crank-Nicolson) MCMC method.
  • Accuracy: pcSGD-iVI produced posterior mean functions and covariance operators nearly identical to pCN (relative error $\approx 0.5\%$), while cSGD-iVI showed larger deviations ($\approx 4.7\%$), particularly near boundaries.
  • Uncertainty Quantification: The 95% credibility regions of pcSGD-iVI successfully contained the ground truth, whereas cSGD-iVI failed to fully capture the uncertainty in certain regions.
  • Efficiency: Both SGD methods required significantly fewer PDE solves (4,000–6,600) compared to pCN ($5 \times 10^5$) to achieve comparable accuracy.

• Nonlinear Darcy Flow Problem:

  • Comparison: Compared against SVGD (Stein Variational Gradient Descent).
  • Performance: pcSGD-iVI provided the most accurate uncertainty quantification, with credibility regions fully containing the truth. SVGD had the lowest relative error in the mean but failed to capture the full uncertainty structure (credibility regions missed the truth). cSGD-iVI performed poorly in both mean accuracy and uncertainty.
  • Cost: pcSGD-iVI was computationally more expensive per step than cSGD-iVI but significantly cheaper than SVGD (which required 2,000 steps and ~11,000s runtime).

5. Significance

• Scalability: The proposed methods offer a scalable alternative to MCMC for high-dimensional and infinite-dimensional inverse problems, reducing computational costs by orders of magnitude.
• Theoretical Rigor: By operating in the infinite-dimensional setting first and then discretizing, the methods avoid the "curse of dimensionality" and discretization artifacts common in finite-dimensional VI approaches.
• Practical Utility: The introduction of preconditioning (pcSGD-iVI) makes the method viable for complex, real-world PDE-constrained problems where both accurate point estimates and reliable uncertainty quantification are critical.
• New Paradigm: The work successfully reinterprets SGD as a variational inference tool for Bayesian inverse problems, providing a new theoretical foundation for using optimization-based dynamics for probabilistic sampling in function spaces.