Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Provable Mixed-Noise Learning with Flow-Matching

Paul Hagemann, Robert Gruhlke, Bernhard Stankewitz, Claudia Schillings, Gabriele Steidl

TL;DR

This work tackles Bayesian inverse problems with mixed Gaussian noise, where the noise parameters $(a^2,b^2)$ are unknown. It introduces an EM framework in which the E-step is implemented via conditional flow matching using neural ODEs to approximate the posterior $p_{X|Y_θ}$ across observations, while the M-step updates the noise parameters. The authors prove population-level convergence to the true parameter $ heta^*$ under local concavity and stability assumptions, and validate the approach on high-dimensional imaging tasks such as MNIST denoising and a reaction–diffusion PDE reconstruction. By combining flow-based posterior sampling with a principled EM analysis, the paper offers a scalable, theory-grounded method for learning mixed noise in Bayesian inverse problems and demonstrates practical effectiveness through qualitative and quantitative imaging results.

Abstract

We study Bayesian inverse problems with mixed noise, modeled as a combination of additive and multiplicative Gaussian components. While traditional inference methods often assume fixed or known noise characteristics, real-world applications, particularly in physics and chemistry, frequently involve noise with unknown and heterogeneous structure. Motivated by recent advances in flow-based generative modeling, we propose a novel inference framework based on conditional flow matching embedded within an Expectation-Maximization (EM) algorithm to jointly estimate posterior samplers and noise parameters. To enable high-dimensional inference and improve scalability, we use simulation-free ODE-based flow matching as the generative model in the E-step of the EM algorithm. We prove that, under suitable assumptions, the EM updates converge to the true noise parameters in the population limit of infinite observations. Our numerical results illustrate the effectiveness of combining EM inference with flow matching for mixed-noise Bayesian inverse problems.

Provable Mixed-Noise Learning with Flow-Matching

TL;DR

This work tackles Bayesian inverse problems with mixed Gaussian noise, where the noise parameters are unknown. It introduces an EM framework in which the E-step is implemented via conditional flow matching using neural ODEs to approximate the posterior across observations, while the M-step updates the noise parameters. The authors prove population-level convergence to the true parameter under local concavity and stability assumptions, and validate the approach on high-dimensional imaging tasks such as MNIST denoising and a reaction–diffusion PDE reconstruction. By combining flow-based posterior sampling with a principled EM analysis, the paper offers a scalable, theory-grounded method for learning mixed noise in Bayesian inverse problems and demonstrates practical effectiveness through qualitative and quantitative imaging results.

Abstract

We study Bayesian inverse problems with mixed noise, modeled as a combination of additive and multiplicative Gaussian components. While traditional inference methods often assume fixed or known noise characteristics, real-world applications, particularly in physics and chemistry, frequently involve noise with unknown and heterogeneous structure. Motivated by recent advances in flow-based generative modeling, we propose a novel inference framework based on conditional flow matching embedded within an Expectation-Maximization (EM) algorithm to jointly estimate posterior samplers and noise parameters. To enable high-dimensional inference and improve scalability, we use simulation-free ODE-based flow matching as the generative model in the E-step of the EM algorithm. We prove that, under suitable assumptions, the EM updates converge to the true noise parameters in the population limit of infinite observations. Our numerical results illustrate the effectiveness of combining EM inference with flow matching for mixed-noise Bayesian inverse problems.

Paper Structure

This paper contains 23 sections, 8 theorems, 92 equations, 5 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. EM Algorithm
  4. Flow Matching
  5. Convergence Analysis of the Population EM Operator
  6. Local Strong Concavity.
  7. First-Order Stability (FOS).
  8. Gradient Descent Variant.
  9. Sufficient Assumptions
  10. FOS Framework
  11. Gradient Descent Based Framework
  12. Step Size Choice Revisited
  13. Verification of Assumptions for Mixed Noise Model
  14. Assumption (A1)
  15. Assumption (A2)
  16. ...and 8 more sections

Key Result

Theorem 1

Let $\varepsilon > 0$, and suppose that $\mathcal{Q}(\cdot, \theta^*)$ is locally $\lambda$-strongly concave, and that the $\mathrm{FOS}(\gamma)$ condition holds on $B_\varepsilon(\theta^*)$ with $0 \leq \gamma < \lambda$. Then the EM operator $M$ is a contraction on $B_\varepsilon(\theta^*)$, i.e., In particular, for any $\theta^{(0)} \in B_\varepsilon(\theta^*)$, the sequence satisfies linear co

Figures (5)

  • Figure 1: Comparison of $c=\frac{2\mu\lambda}{\mu+\lambda}$ and $\lambda$ for fixed values of $\lambda =0.2,1,2$ for a range of $\mu > \lambda$.
  • Figure 2: Conditions and resulting flow matching samples.
  • Figure 3: Convergence of a and b for the MNIST denoising imaging example.
  • Figure 4: Ground truth, measurement, posterior mean and standard deviations of the model provided for the $u$- and $v$-component.
  • Figure 5: Convergence of a and b for the PDE example.

Theorems & Definitions (17)

  • Theorem 1: Convergence of Population EM em_conv
  • Theorem 2: Convergence of Gradient Descent EM em_conv
  • Proposition 3: Strong Concavity
  • proof
  • Proposition 4: FOS Condition
  • proof
  • Remark 5
  • Proposition 6
  • proof
  • Proposition 7: Gradient Smoothness Condition
  • ...and 7 more