Multi-Dataset Inverse Problem Solving with Distributed Generative AI

Daniel Lersch; Steven Goldenberg; Johann Rudi; Markus Diefenthaler; Kevin Brager; Xingfu Wu; Yaohang Li; Nobuo Sato

cs.DCAug 26, 2026

Multi-Dataset Inverse Problem Solving with Distributed Generative AI

Authors: Daniel Lersch, Steven Goldenberg, Johann Rudi, Markus Diefenthaler, Kevin Brager, Xingfu Wu, Yaohang Li, Nobuo Sato

Organizations: Data Science Department, Thomas Jefferson National Accelerator Facility, Newport News, VA, USA · Department of Mathematics, Virginia Tech, Blacksburg, VA, USA · Electron Ion Collider Group, Thomas Jefferson National Accelerator Facility, Newport News, VA, USA · Department of Physics, William & Mary, Williamsburg, VA, USA ·

Abstract

Extracting a shared set of unknown, not directly measurable quantities from multiple, heterogeneous datasets is a common challenge across scientific domains. A prominent example is the combination of datasets obtained from different measurements with different settings (e.g. varying detector resolutions). Analyzing such datasets jointly, rather than independently or after naive merging, is essential for obtaining precise and unbiased estimates of the unknowns, but requires careful treatment of dataset heterogeneity and is computationally demanding. We present a generalized framework for simultaneously analyzing multiple heterogeneous datasets in the context of generative AI-based inverse problem solvers. Building on our recent Scalable Asynchronous Generative Inverse Problem Solver (SAGIPS) framework, we extend the well-established distributed data-parallel training paradigm to non-identically distributed datasets, where each dataset is controlled by the same set of unknown inference parameters but covers a different region of the available feature space. Each dataset is processed through its own forward operator and discriminator, providing complementary constraints that collectively guide a shared generator toward global parameter consistency. We validate the approach using a controlled setup inspired by a multi-detector scattering experiment. We provide numerical evidence that our framework is robust to different data fidelities, which arise from unknown detector systematics in the Rutherford experiment, and we show the scaling behavior on multi-GPU leadership computing systems. The results show that our approach is well suited for real-world multi-dataset analyses in which experimental conditions vary across measurements.

Figures & tables

Figure 1: A large volume dataset consists of $N$ subsets $R_{0}$ ,… $R_{N-1}$ . Each of the subsets is connected via an operator $P_{i}$ to a set of M unknowns: $R_{i}=P_{i}[a_{0},\ldots,a_{M-1}]$ .

Figure 2: Schematic presentation of the proposed approach to analyze $N$ datasets in parallel through the SAGIPS inverse problem solver. Displayed here are, for simplicity, the first and last out of $N$ GPUs. The red (blue) colors indicate the operator, discriminator and real dataset that are bound to the first (last) GPU. The generator is identical on each GPU and updated through the shared gradients $g_{\Phi}$ . The discriminator on each GPU is trained locally and thus shares on information with adjacent GPUs.

Figure 3: Schematic representation of the Rutherford experiment: A particle gun fires a particle on a target foil. The scattered particle will be detected by one of the four detectors that are mounted at fixed angles. Each detector counts the number of particles (per time) that are leaving the reaction chamber at that specific angle. The scenario presented here displays a particle being scattered at angle $\theta_{2}$ and registered by detector 2. The observable reported to the experimentalist is the particle rate $\nu(\theta_{2})$ .

Dataset $R_{i}$	$\theta_{i}$ [Deg]	$\ln(C_{exp}(i))$	$T$ [a.u.]	$\delta_{\nu}(i)$
$R_{0}$	$0.5$	$50.0$	$5.0$	$0.1$
$R_{1}$	$5.0$	$57.0$	$5.0$	$0.08$
$R_{2}$	$10.0$	$54.0$	$5.0$	$0.15$
$R_{3}$	$15.0$	$55.0$	$5.0$	$0.12$
$R_{4}$	$25.0$	$65.0$	$5.0$	$0.2$
$R_{5}$	$50.0$	$67.0$	$5.0$	$0.07$

Table 1: Settings to create 6 datasets

R_{i}

, following Eq. 10 . The dataset indices

R_{i}

presented in the left column reflect the order in which the datasets are analyzed by SAGIPS.

Panel of Figure 4 — Figure 4: Left: Datasets $R_{i}$ , created via Eq. 10 in combination with the settings in Tab. 1 . Right: Solution space for each dataset $R_{i}$ . For visualization purposes, the $\delta_{v}(i)$ -term has been omitted for each plot. The dashed horizontal and dashed-dot vertical lines represent the true values $a_{0}$ and $a_{1}$ respectively.

Figure 5: Top Row: Relative residual from Eq. 14 for unknown $a_{0}$ (left) and $a_{1}$ (right) as a function of the number of datasets fed into SAGIPS. Bottom Row: Uncertainties associated with $\hat{r}_{0}$ (left) and $\hat{r}_{1}$ (right) as a function of the number of analyzed datasets. The different curves in each panel present the various analysis approaches explored within this run group.

Figure 6: Total training ( top panel ), GPU memory fraction ( center panel ) and GPU utilization ( bottom panel ) as a function of the number of datasets used for the SAGIPS analysis.

Figure 7: Top Row: Relative residual from Eq. 14 for unknown $a_{0}$ (left) and $a_{1}$ (right) as a function of the number of datasets fed into SAGIPS. Bottom Row: Uncertainties associated with $\hat{r}_{0}$ (left) and $\hat{r}_{1}$ (right) as a function of the number of analyzed datasets. The different curves in each panel present the different ordering schemes (explained in main text) for successively combining and analyzing the data. Each analysis was carried out with the conventional ARAR gradient transport.

Figure 8: Top Row: Relative residual from Eq. 14 for unknown $a_{0}$ (left) and $a_{1}$ (right) as a function of the number of datasets fed into SAGIPS. Bottom Row: Uncertainties associated with $\hat{r}_{0}$ (left) and $\hat{r}_{1}$ (right) as a function of the number of analyzed datasets. All curves shown here have been obtained by using the conventional ARAR gradient transport method.

	conv-ARAR	Double Binary-Tree	ARAR	strong-ARAR
$\hat{r}_{0}\pm\Delta\hat{r}_{0}$	$0.034\pm 0.107$	$-0.033\pm 0.145$	$0.697\pm 1.536$	$0.063\pm 0.303$
$\hat{r}_{1}\pm\Delta\hat{r}_{1}$	$-0.001\pm 0.007$	$-0.003\pm 0.009$	$-0.033\pm 0.121$	$-0.0013\pm 0.024$
Elapsed Time $[s]$	$6,334$	$9,209$	$2,077$	$2,243$
GPU Utilization $[\%]$	$8.42$	$1.97$	$28.73$	$20.16$
$\frac{\text{Elapsed Time }(m=4)}{\text{Elapsed Time }(m=20)}$	$1.521$	-	-	$3.354$

Table 2: Performance evaluation of the multi-dataset analyses conducted on the Polaris HPC, using a

m=20

multiplicity. The underlying analyses ran with all six datasets and different gradient transport methods. The last row represents the ratio between training times using a

m=4

and

m=20

multiplicity respectively. The ratio was determined for two gradient transport methods only.

Figure 9: Weight drifts ( top row ) and output drifts ( bottom row ) computed according to Eq. 16 and 17 as a function of the SAGIPS training times. Shown here are the median, interquantile range (IQR) and maximum determine from the rank pair $(i,j)$ averaged over the ensemble size. The columns (from left to right) represent the pair wise distances for conventional ARAR, strong ARAR and ARAR.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 10 — Figure 10: Top Row: Densities $\rho_{0}(x,y)$ ( left ) and $\rho_{1}(x,y)$ ( right ), determined w.r.t Eq. 29 and 30 . Each density was created with a resolution of $100\times 100$ pixels. Bottom Row: Datasets $R_{0},R_{1},R_{2}$ generated according to Eq. 28 .

Figure 11: Density projection $\rho_{0}(y|x\in[a,b])$ ( top row ) and $\rho_{1}(y|x\in[a,b])$ ( bottom row ), following Eq. 33 for different x-intervals, as a function of y. The dashed, black lines represent the true density projections whereas the dot-dashed red and blue curves represent the mean $\mu$ of the GAN before and after training respectively. The shaded blue areas represent the uncertainty $\sigma$ of the trained GAN ensemble. The results shown here have been obtained by training the GAN on dataset $R_{0}$ only.

Figure 12: Density projection $\rho_{0}(y|x\in[a,b])$ (top row) and $\rho_{1}(y|x\in[a,b])$ ( bottom row ), following Eq. 33 for different x-intervals, as a function of y. The dashed, black lines represent the true density projections whereas the dot-dashed red and blue curves represent the mean $\mu$ of the GAN before and after training respectively. The shaded blue areas represent the uncertainty $\sigma$ of the trained GAN ensemble. The results shown here have been obtained by training the GAN on datasets $R_{0},R_{1}$ and $R_{2}$ simultaneously.

Inverse Problems

Inverse problems infer hidden causes, parameters, or signals from indirect, noisy measurements using physical models, priors, and statistical methods.

See all 246 papers

Diffusion Based Unpaired Data Learning for Inverse Problems
Sep 1, 2026Chenglong Bao, Yiming Dang, Chenguang Duan +2Diffusion-Based Inverse Problems Inverse Problems
Yau Mathematical Sciences Center, Tsinghua University, Beijing, China · Department of Mathematical Sciences, Tsinghua University, Beijing, China · Institut für Geometrie und Praktische Mathematik, RWTH Aachen University, Aachen, Germany +2
MACS: Measurement-Aware Consistency Sampling for Inverse Problems
Oct 2, 2025Amirreza Tanevardi, Pooria Abbas Rad Moghadam, Seyed Mohammad Eshtehardian +2Diffusion Model Sampling Inverse Problems
Department of Electrical Engineering Sharif University of Technology Tehran, Iran · Electronics Research Institute Sharif University of Technology Tehran, Iran
Invent a Dataset: Measuring Dataset Generation Abilities with Zero Seed Data
Oct 1, 2026Shivalika Singh, Andrija Djurisic, Gbemileke Onilude +2Synthetic Data Generation Synthetic Data Generation for Language Models
Adaption

Sep 1, 2026J/K move · Enter open · S save

Diffusion Based Unpaired Data Learning for Inverse Problems

Chenglong Bao, Yiming Dang, Chenguang Duan, Yuling Jiao+1

Yau Mathematical Sciences Center, Tsinghua University, Beijing, China · Department of Mathematical Sciences, Tsinghua University, Beijing, China · Institut für Geometrie und Praktische Mathematik, RWTH Aachen University, Aachen, Germany+2

Diffusion-Based Inverse Problems Inverse Problems Diffusion Models Image Inverse Problems

Open paper AskPDF arXiv

Data is important in many deep learning-based inverse problem solvers. However, obtaining sufficient paired data in many scenarios remains highly challenging, while unpaired data is cheap. To maximize data utilization, this paper proposes LUD-DIF, a diffusion-based approach for solving inverse problems with unpaired data. Starting from the evidence lower bound (ELBO) of the joint distribution, we decouple it into two independent diffusion processes under the weak-coupling assumption. The method provides theoretical support from a variational inference perspective, derives the loss function, quantitatively analyzes the error bound introduced by the assumption, and offers a theorem-motivated heuristic for hyperparameter selection. Experimental results demonstrate that LUD-DIF achieves outstanding performance on multiple image inverse problems, validating its effectiveness and generalization capability in unpaired inverse problem settings.

Abstract

Inverse Problems

Explore similar work

Diffusion Based Unpaired Data Learning for Inverse Problems

MACS: Measurement-Aware Consistency Sampling for Inverse Problems

Invent a Dataset: Measuring Dataset Generation Abilities with Zero Seed Data