PaperlayerView exact revision on arXiv

Source-linked AI summary

Inverse Problems Conditioned on Observation Ensembles: Applications and Methods

Zhengyan Huan, Camila Pazos, Martin Klassen, Vincent Croft, Pierre-Hugues Beauchemin, Shuchin Aeron

arXiv:2601.22029v2cs.LGstat.ML

TL;DR

The paper addresses inference when multiple truth–observation datasets share an unknown forward model and new observations come from an unknown prior. Its ensemble inverse generative models condition posterior sampling on both each observation and its observation set; EI-FM shows superior recovery among compared methods and performs similarly to cFM-γ across the reported γ range.

  • Problem

    EIP requires recovering an unknown prior or its posterior from observations when the shared forward model is unknown and the observation ensemble must supply information about the prior.

    Source: §1 · pp. 2–3

  • Method

    Ensemble inverse generative models use a permutation-invariant representation of the observation set alongside each observation to form a non-iterative, distribution-dependent posterior sampler.

    Source: §1.2 · p. 4, §2.1 · p. 5

  • Results

    EI-FM displays superior recovery among compared methods and performs similarly to cFM-γ for γ ∈[−1, 1].

    Source: §3.2 · p. 8

  • Takeaways & Limitations

    The observation ensemble can provide information about an unknown prior that supports posterior inference and an inductive bias for generalization to unseen priors.

    Source: §2.1 · p. 5, §4 · p. 12

    Reliable inference depends on an informative observation set; weakly separated induced observation distributions require a larger set size N.

    Source: §1 · p. 3

Abstract

from arXiv · show

We introduce a new multivariate statistical problem that we refer to as the Ensemble-conditioned Inverse Problem (EIP). The aim of EIP is to invert for an ensemble that is distributed according to the pushforward of a prior under a forward process. In high energy physics (HEP), this is related to a widely known problem called unfolding, which aims to reconstruct the true physics distribution from observations that are distorted by detector effects. The EIP also arises in full waveform inversion (FWI) and inverse imaging with unknown priors. We propose non-iterative inference-time methods that construct posterior samplers based on a new class of conditional generative models, which we call ensemble inverse generative models. For the posterior modeling, these models additionally use the ensemble information contained in the observation set on top of single observations. Unlike existing methods, our proposed methods avoid explicit and iterative use of the forward model at inference time via training across several sets of truth-observation pairs that are consistent with the same forward model, but originate from a wide range of priors. We empirically demonstrate that this training procedure can implicitly encode the likelihood model, enabling direct posterior inference for unseen priors to some degree. We benchmark the proposed method on several synthetic and real datasets in inverse imaging, HEP, and FWI. Our code is available at https://github.com/ZhengyanHuan/EIP.

1 Introduction

The paper formalizes the Ensemble-conditioned Inverse Problem (EIP), where truth–observation pairs from multiple priors share an unknown forward model, and proposes non-iterative posterior sampling conditioned on individual observations and their ensemble. It positions this approach among existing inverse-problem methods and reports results across synthetic and real applications.

1 Introduction

  • EIP uses multiple truth–observation datasets generated under different priors but a shared, unknown forward model to infer a new prior or its posterior without explicitly using that model.The observation set must contain enough distributional information to distinguish candidate priors; weakly separated induced observation distributions require larger sets.Source: §1 · pp. 2–3
  • EIP-I targets recovery of the prior distribution, while EIP-II targets posterior samples for a given observation; EIP-I is further restricted to generate samples through posterior sampling.Source: §1 · pp. 2–3

1.1 Related Works

  • Prior work includes iterative unfolding and posterior-sampling methods, while transformer theory and empirical-Bayes methods leave unseen-measure generalization unaddressed or focus on one-dimensional Poisson settings.GDDPM also uses observation moments to support generalization across physics processes, and this work targets EIP-II with generative models.Source: §1.1 · pp. 3–4

1.2 Contributions

  • Ensemble inverse generative models provide a non-iterative EIP-II framework that conditions posterior sampling on both an individual observation and its observation set.Source: §1.2 · p. 4
  • A permutation-invariant encoder extracts ensemble information so the posterior sampler depends on the observation distribution and is reported to generalize to unseen priors.Source: §1.2 · p. 4
  • Across synthetic settings and real HEP, FWI, and image-inversion applications, the proposed methods outperform baselines without explicit prior or forward-model knowledge.Source: §1.2 · p. 4

2 Method

The method performs non-iterative posterior sampling by conditioning generative models on both a single observation and its observation ensemble. EI-DDPM and EI-FM learn a permutation-invariant ensemble representation alongside the posterior sampler, providing an inductive bias for inference under unseen priors.

2 Method

  • For EIP-II, non-iterative posterior samplers condition on both the single observation y and the observation set Y.The ensemble information in Y is intended to provide an inductive bias for unseen priors.Source: §2 · p. 4

2.1 Ensemble Inverse Generative Models for EIP-II

  • The permutation-invariant network ϕw maps observation set Y to a k-dimensional representation of ensemble information.It treats Y as a set, ignoring sample order; deep set and set transformer are listed as implementation choices.Source: §2.1 · p. 5
  • EI-DDPM and EI-FM use εθ to model p(x|y, Y), with ϕw(Y) supplying ensemble information and y identifying the sample to invert.The models are conditional-DDPM and conditional-flow-matching frameworks, respectively.Source: §2.1 · p. 5
  • Inference set size should match training size N; larger sets can be covered by repeated subsets, while smaller sets require duplication that may perform poorly when N′ ≪ N.The authors attribute this concern to duplicated samples reflecting incorrect ensemble information.Source: §2.1 · p. 5
  • Conditioning on ϕw(Y) makes the sampler distribution-dependent, unlike a vanilla model that averages sample-wise posteriors across training priors.The jointly learned representation can capture nonlinear and cross-variable dependencies that prescribed low-order moments may miss, without domain-specific statistic selection.Source: §2.1 · pp. 5–6

3 Experiments

Experiments compare ensemble-conditioned methods with generative and task-specific baselines across synthetic recovery, image inversion, particle unfolding, and seismic inversion. The proposed methods show strong recovery on unseen priors and structural families in these evaluations.

3.1 Baselines for Comparison

  • Baselines include single-observation conditional generative models, moment-conditioned GDDPM, reweighting-based OmniFold, Schrödinger-bridge SBUnfold, and Sourcerer.Sourcerer requires a differentiable forward model or surrogate, while OmniFold reweights an initial distribution toward the prior.Source: §3.1 · p. 7

3.2 Synthetic EIP

  • Across γ ∈ [−1, 1], EI-FM has the strongest synthetic prior recovery and performs similarly to cFM-γ, including in the unseen central range.OmniFold-best has low SWD where its initial distribution matches a training prior, but weaker generalization for γ ∈ [−0.25, 0.25]; at unseen γ = 0.9, EI-FM also recovers a thin prior more closely than cFM.Source: Figure 3 · p. 8, §3.2 · p. 8, Figure 2 · p. 8

3.3 Image Inversion of MNIST Digits Mixture

  • On MNIST mixtures, EI-FM and EI-DDPM outperform baselines in MSE and SSIM across t and preserve the truth’s S-shaped structure at t = 0.5.The regions near t = 0 and t = 1 are excluded from training; accumulated transport-path errors toward t = 1 increase recovery difficulty.Source: Figure 4 · p. 9, Figure 5 · p. 9, §3.3 · p. 10

3.4 Particle Physics Data Unfolding

  • On four unseen physics processes, EI-DDPM and EI-FM show superior recovery performance for pT, E, and px using 1-D Wasserstein distances to truth-level data, while achieving comparable or superior performance to GDDPM.GDDPM outperforms GDDPM-v, suggesting redundant moment information in GDDPM-v impairs recovery.Source: Table 2 · p. 11, §3.4 · p. 11

3.5 Seismic Full Waveform Inversion

  • EI-DDPM and EI-FM achieve the best results in most FWI family–metric comparisons and generalize strongly to the two families excluded from training.The comparisons use MSE, MAE, and SSIM; prior-only methods are not applicable because this task requires posterior recovery.Source: Table 3 · p. 12, §3.5 · p. 12

4 Conclusions and Limitations

The paper presents ensemble inverse generative models for posterior inference using both individual observations and information from their ensemble. It reports strong inference and generalization across three domains, while identifying observation-set informativeness and encoder quality as limitations.

4 Conclusions and Limitations

  • EI-DDPM and EI-FM demonstrate strong posterior inference and generalization across inverse imaging, HEP unfolding, and FWI.The method conditions on individual observations and ensemble information extracted from the observation set.Source: §4 · p. 12
  • The approach uses a permutation-invariant neural network to extract ensemble information and turn sample-wise posterior sampling into distribution-dependent sampling.This information is used as an inductive bias for generalization to unseen priors.Source: §4 · p. 12
  • Performance depends on observation-set informativeness and the quality, capacity, and stability of the permutation-invariant encoder.The work does not address the encoder’s optimal design.Source: §4 · p. 12
  • Future work includes validation on directly measured real-world data, theoretical guarantees on recovery discrepancy, and improved ensemble-information architectures.Source: §4 · p. 12

A.1 Relationship between EIP-I and EIP-II

This appendix relates the EIP-I and EIP-II solutions and introduces DDPM and flow matching as conditional generative-model backgrounds. It also includes a notation summary.

A.1 Relationship between EIP-I and EIP-II

  • Integrating EIP-II’s solution corresponds to EIP-I’s solution, although integrating an incorrect posterior can still lead to the correct prior.Source: Figure 7 · p. 16

A.2 Summary of Notation

B An Introduction to DDPM and FM

  • DDPM learns to reverse gradual corruption of data, mapping Gaussian samples toward the target distribution through a learned reverse process.The forward process uses a predefined noise schedule; the reverse process is learned.Source: B An Introduction to DDPM and FM · p. 17
  • Flow matching learns velocity fields for continuous flows between an initial distribution and a conditional target distribution.Its objective can be optimized through an equivalent mean-squared-error objective, with a user-selected interpolant such as a linear one.Source: B An Introduction to DDPM and FM · p. 17

C.1 Model Configuration and Computational Cost

The section describes task-specific ensemble encoders and denoisers, then reports computational costs. Training is more expensive, while sampling incurs a moderate memory increase on the tested FWI subset.

C.1 Model Configuration and Computational Cost

  • Ensemble encoders use set-transformer structures for 2-D Gaussian and particle-physics inputs, while image and seismic inputs first pass through convolutional encoders.Source: §C.1 · pp. 17–18
  • The denoising networks are task-adapted: MLPs process Gaussian and particle-physics inputs, while U-nets process image and seismic inputs.Their inputs include intermediate data, a single observation, and the encoded observation ensemble.Source: §C.1 · p. 18
  • Inference sampling is nearly unchanged in time from vanilla DDPM/FM and substantially faster than iterative optimization or reweighting, with unfolding taking seconds versus several minutes.Training the ensemble models is approximately twice as expensive because it also learns the ensemble representation.Source: §C.1 · p. 18
  • 1004.41 MB for EI-FM versus 792.41 MB for FM is the peak memory use on a subset of FWI sampling, an additional cost of approximately 26.8%.The authors characterize this as moderate rather than prohibitive overhead.Source: §C.1 · p. 19, Table 5 · p. 19

C.2 Effect of N, N ′ in 2-D Gaussian EIP

In the 2-D Gaussian EIP, recovery improves as the observation-set size grows, while very small sets can fail to convey useful ensemble information. When inference has fewer observations than training, duplication can remain effective if the available set is sufficiently informative.

C.2 Effect of N, N ′ in 2-D Gaussian EIP

  • Training with fixed N can simplify the pipeline and stabilize optimization; random training sizes are recommended when inference set sizes vary near N.The passage assumes N′ ≫ N in most cases, so enough samples are available to form observation sets.Source: §C.2 · p. 19
  • For N ≤ 10, EI-FM performs worse than sample-wise cFM; its SWD improves as N grows and stabilizes at a sufficiently large size.The authors report that sets that are too small can mislead the model during training and inference.Source: §C.2 · p. 19
  • With a model trained at N = 4000, duplication of N′ observations yields SWD that decreases as N′ approaches 4000; at N′ = 1000 it is close to the N′ = 4000 case.Even at N′ = 10, EI-FM slightly outperforms sample-wise cFM.Source: §C.2 · p. 19

C.3 Extension of 2-D Gaussian EIP

The section extends the 2-D Gaussian EIP to priors with three independent parameters and evaluates recovery across parameter ranges. EI-FM performs comparably to cFM-γ, which has direct prior knowledge, and outperforms cFM.

C.3 Extension of 2-D Gaussian EIP

  • The extended prior has three independent parameters, µ1, µ2, and γ1, and its samples undergo the same forward process as equation 6.The recovery task is to infer the prior from its observation set when γ is unknown.Source: §C.3 · p. 19
  • Training uses γ1 values in [−0.75, −0.25]∪[0.25, 0.75] and µ1, µ2 values in [−1.5, −0.5]∪[0.5, 1.5], while inference tests wider ranges.At inference, γ1 ranges over [−1, 1] and µ1, µ2 over [−2, 2].Source: §C.3 · pp. 19–20
  • Across the evaluated parameter ranges, EI-FM matches cFM-γ and outperforms cFM on SWD.The comparison uses µ1, µ2, and γ1; cFM-γ is provided with the prior parameters directly.Source: Figure 9 · p. 20, §C.3 · p. 21

C.4 Detailed Results of Particle Physics Data Unfolding

The section reports 1-D Wasserstein distances for detector-level and recovered data across seven physics components. Detector-level distributions for η and ϕ are already close to the truth, so their best results can occur before unfolding.

C.4 Detailed Results of Particle Physics Data Unfolding

  • Detector-level distortion is small for η and ϕ, whose detector-level distributions are already close to the true prior.This explains why some of their best performances occur before unfolding.Source: §C.4 · p. 21
  • For some η and ϕ results, detector-level data perform best before unfolding.The section attributes this to the small detector-level distortion for these components.Source: §C.4 · p. 21
  • Table 6 reports 1-D Wasserstein distances between truth-level data and detector-level or recovered data for seven physics components.The components are pT, η, ϕ, E, px, py, and pz; results cover four unseen physics distributions.Source: Table 6 · p. 21, §C.4 · p. 21

C.5 TARP Coverage as an Additional Metric for Posterior Accuracy

The section uses TARP expected coverage to assess whether estimated posteriors capture the truth within credible regions. Across the Gaussian EIP and HEP unfolding settings, samplers using ensemble information show slightly smaller deviations from ideal coverage.

C.5 TARP Coverage as an Additional Metric for Posterior Accuracy

  • TARP tests posterior accuracy by checking how often estimated credible regions contain the truth for truth–observation pairs.It uses randomly drawn reference points and scans credibility levels.Source: §C.5 · p. 21
  • Under the stated regularity conditions and sufficiently rich reference-point positions, matching expected coverage at every credibility level characterizes the true posterior.The section also reports mean absolute deviation from ideal expected coverage over uniformly sampled α values.Source: §C.5 · p. 22
  • Posterior samplers using ensemble information show slightly smaller deviations from ideal coverage in the Gaussian EIP and HEP unfolding results.Figures 10 and 11 show ECP versus credibility level curves for these respective settings.Source: §C.5 · p. 22, Figure 10 · p. 22, Figure 11 · p. 22

C.6 Visualization of Full Waveform Inversion Results

The section visualizes recovered velocity maps from methods in each family. EI-DDPM preserves subsurface structural patterns more faithfully than baselines, which introduce stronger distortions or miss geometric features.

C.6 Visualization of Full Waveform Inversion Results

  • EI-DDPM preserves subsurface structural patterns more faithfully, while baselines introduce stronger distortions or miss key geometric features.Figure 12 visualizes the recovered velocity maps from the compared methods in each family.Source: §C.6 · p. 22, Figure 12 · p. 23
Loading 2601.22029v2…