Active Diffusion-Based Inference for Ill-Posed Inverse Problems under Incomplete Priors

AI summary

Overview

Research area: Machine learning for scientific inverse problems — specifically diffusion models (DMs), active learning, and Bayesian uncertainty quantification applied to ill-posed inverse problems.

Technical level: Intermediate. The paper assumes familiarity with diffusion models, conditional generation, posterior covariance, and Kullback–Leibler divergence, but the algorithmic loop and results are described concretely enough for a reader with a general ML background.

Scope (one sentence): The paper proposes and validates an active, diffusion-based inverse solver that iteratively expands its parameter domain when posterior uncertainty or ensemble disagreement signals that the true parameters lie outside the initial training support.

What This Paper Is About

Many scientific problems require recovering unknown parameters from measured observables, but the inverse mapping is usually ambiguous, unstable, and noise-sensitive — and generative solvers typically assume the range of plausible parameter values is known in advance. When the true parameters fall outside that assumed range, diffusion models silently extrapolate and can report misleadingly confident answers. This paper's goal is an inference framework that detects such out-of-domain situations from the model's own uncertainty and actively grows its training domain until the true parameter region is captured.

Key Contributions

1. An active diffusion-based inverse solver for incomplete priors. A conditional diffusion model learns the mapping between parameter space and observable space, and an active learning loop expands the training domain whenever the posterior uncertainty exceeds a threshold τ, with no requirement that the initial bounds contain the true parameters.

2. Posterior variance as a diagnostic, with a Bayesian justification. The paper argues formally that if θ* ∉ Θ_k then Tr(Σ_θ) becomes large, so the posterior covariance acts as an estimator of model misspecification error, and shows that expanding the domain toward the MAP estimate minimizes the KL divergence between the true and learned posteriors.

3. An ensemble-based correction for overconfident extrapolation. Because a single diffusion model can extrapolate confidently with narrow variance, the authors train K independently initialized models and use disagreement among their posterior predictions (e.g., variance of posterior means, pairwise Wasserstein distances) as a more robust trigger for domain expansion.

4. Validation on two problems plus a model-agnostic comparison. The framework is demonstrated on a toy ill-posed problem with infinitely many solutions and on a QCD proxy problem parameterizing quark correlation functions, and the active loop is shown to improve MDN, INN, NF, and DM solvers alike.

Main Findings

• Toy problem, in-domain: With training domain Θ₀ = {(x,y) : x² + y² ≤ 4} (radii r ∈ [0, 2]) and target z* = 4 (r = 2), four independently trained diffusion models recovered the correct posterior, with samples spread along the target circle of radius 2 and low pairwise Wasserstein distances (mean W̄ = 0.12).

• Toy problem, out-of-domain: For z* = 25 (r = 5), the four models produced inconsistent posteriors scattered across parameter space, with much larger pairwise Wasserstein distances (mean W̄ = 3.00), signalling extrapolation beyond training support.

• Toy problem, after active learning: Applying the active loop for z* = 25 brought the predictions back onto the target circle of radius 5 and reduced the mean pairwise Wasserstein distance to W̄ = 0.19, comparable to the in-domain case.

• QCD, in-domain: When ground-truth parameters were drawn from the interior of the prior, the trained diffusion model recovered all six parameters with posterior standard deviations ranging from 0.03 to 0.06.

• QCD, out-of-domain without active learning: For true values N_u = 2.0, a_u = 0.6, b_u = −0.5, N_d = 1.5, a_d = −0.4, b_d = 1.8 — all outside the initial bounds — the four-model ensemble showed inflated posterior standard deviations and ensemble disagreement σ_ens ranging from 0.04 to 0.18.

• QCD, out-of-domain with active learning: After iterative domain expansion, posterior means were N̂_u = 1.92, â_u = 0.65, b̂_u = −0.53, N̂_d = 1.41, â_d = −0.42, b̂_d = 1.80, with every ground-truth value inside three standard deviations of its posterior mean and uncertainties contracting to σ̂ ≈ 0.03–0.05, similar to in-domain levels.

• Model-agnosticism (toy circle problem, mean radial error): All methods trained on z ∈ [0, 4] and tested at z* = 4 and z* = 25. MDN: 0.067 in-domain, 2.497 OOD fixed, 0.082 OOD +AL. INN: 0.003, 1.679, 0.025. NF: 0.045, 107.243, 0.079. DM: 0.015, 1.279, 0.032. All fixed-domain solvers degrade sharply out-of-domain, while the active learning framework restores near in-domain accuracy for every model class.

• Computational cost: In the QCD proxy experiment, each training step required approximately 2.6 seconds of forward simulation for 64 parameter samples with 10,000 events per sample, while the diffusion model training update took approximately 0.03 seconds per step.

Methodology in Plain English

The authors train a conditional diffusion model to generate parameters θ given an observable y, using synthetic pairs produced by running the known (but expensive) forward model F on parameters sampled from a guessed domain Θ₀. Given a target observable y*, they draw posterior samples by reverse diffusion and measure how spread out those samples are.

If the spread (trace of the posterior covariance) stays below a threshold τ, the answer is accepted. If it exceeds τ, the inferred samples are treated as proposals indicating where training data are missing: the parameter domain is enlarged to include them, new forward simulations are generated there, and the model is fine-tuned on the augmented dataset. This loop repeats until the posterior contracts.

Because a single diffusion model can occasionally extrapolate confidently with small variance, the authors also train an ensemble of independently initialized models and look at how much they disagree — measured, for example, by the variance of their posterior means or by pairwise Wasserstein distances between predicted distributions. Large disagreement means the inverse problem is underconstrained in that region and more training data are needed there.

The authors draw an analogy to the classical shooting method: as shooting adjusts unknown initial conditions based on a boundary mismatch, this method adjusts the support of the parameter distribution based on posterior uncertainty, converging when the distribution stabilizes in a region consistent with the observations.

Why This Matters

Impact on research: Generative inverse solvers are usually only as good as the prior domain they were trained on, and expressive models give no reliable warning when they leave that domain. This work turns the model's own uncertainty — and, more robustly, disagreement across an ensemble — into an actionable signal, and provides a Bayesian justification for using that signal to drive domain augmentation. It also shows that the loop is not specific to diffusion models, since MDN, INN, and normalizing flow solvers all improved under the same active framework.

Real-world applications (as indicated by the paper's framing):

• Medical imaging, where reconstructing underlying structure from measurements is ill-posed and noisy.
• Remote sensing, where observations must be inverted to recover surface or atmospheric parameters.
• Non-destructive inspection, where internal properties are inferred from external measurements.
• Nuclear physics, specifically extracting quark correlation functions from event-level observables in deep inelastic scattering experiments, where the valid parameter domain is not known a priori.

Industry relevance: The framework targets the common engineering situation where the plausible parameter range is unknown, effectively unbounded, or too high-dimensional to cover uniformly, while each forward simulation is expensive. It offers a way to spend a limited simulation budget only where the model's own uncertainty says information is missing, which is directly useful in any design, calibration, or characterization pipeline built on costly physics or simulation codes.

Future Directions

• Scaling the active diffusion framework to higher-dimensional parameter spaces and more sophisticated forward models.
• Investigating surrogate or approximate forward modeling to reduce the cost of generating new training data, since forward simulation dominates the compute budget (approximately 2.6 seconds per step versus approximately 0.03 seconds for the model update in the QCD proxy experiment).
• Incorporating physics-informed constraints into the solver to improve efficiency and stability.
• Applying the method to real experimental data where data uncertainties, systematic uncertainties, and model uncertainties coexist, as a step toward practical uncertainty-quantified generative inference for large-scale scientific discovery.

Target Audience

Researchers and practitioners working at the intersection of generative modeling and scientific inference: physicists and domain scientists facing expensive forward models and unknown parameter bounds, machine learning researchers interested in uncertainty quantification and active learning for diffusion models, and method developers building inverse solvers for imaging, remote sensing, or non-destructive evaluation who need reliable behavior when the prior is incomplete.

Authors’ abstract

Many scientific and engineering applications require estimating unknown parameters from experimentally observable data -- an inverse problem that is inherently challenging due to nonlinearity, noise, and ill-posedness. In this paper, we propose an active diffusion-based inverse problem solver. A DM is trained to learn the mapping between the parameter space and the observable space. By iteratively detecting and correcting model misspecification through posterior uncertainty, the method discovers and learns the correct region of parameter space, even when initial training bounds exclude the true parameters. This provides a principled, Bayesian justification for adaptive domain augmentation and ensures robust inference for inverse problems under incomplete prior knowledge. We demonstrate the effectiveness of our inverse solver for a toy inverse problem with infinite solutions, and for the parameterization of the quantum correlation functions to event observables in a Quantum Chromodynamics analysis of nucleon structure.