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 · Mathematics and Computer Science Division, Argonne National Laboratory, Lemont, IL, USA · Department of Computer Science, Old Dominion University, Norfolk, VA, USA · Theory Division, Thomas Jefferson National Accelerator Facility, Newport News, VA 23606, USA
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 R0 ,… RN−1 . Each of the subsets is connected via an operator Pi to a set of M unknowns: Ri=Pi[a0,…,aM−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Φ . 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 θ2 and registered by detector 2. The observable reported to the experimentalist is the particle rate ν(θ2) .
Dataset Ri
θi [Deg]
ln(Cexp(i))
T [a.u.]
δν(i)
R0
0.5
50.0
5.0
0.1
R1
5.0
57.0
5.0
0.08
R2
10.0
54.0
5.0
0.15
R3
15.0
55.0
5.0
0.12
R4
25.0
65.0
5.0
0.2
R5
50.0
67.0
5.0
0.07
Table 1: Settings to create 6 datasets Ri , following Eq. 10 . The dataset indices Ri presented in the left column reflect the order in which the datasets are analyzed by SAGIPS.
Figure 4: Left: Datasets Ri , created via Eq. 10 in combination with the settings in Tab. 1 . Right: Solution space for each dataset Ri . For visualization purposes, the δv(i) -term has been omitted for each plot. The dashed horizontal and dashed-dot vertical lines represent the true values a0 and a1 respectively.
Figure 5: Top Row: Relative residual from Eq. 14 for unknown a0 (left) and a1 (right) as a function of the number of datasets fed into SAGIPS. Bottom Row: Uncertainties associated with r^0 (left) and 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 a0 (left) and a1 (right) as a function of the number of datasets fed into SAGIPS. Bottom Row: Uncertainties associated with r^0 (left) and 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 a0 (left) and a1 (right) as a function of the number of datasets fed into SAGIPS. Bottom Row: Uncertainties associated with r^0 (left) and 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
r^0±Δr^0
0.034±0.107
−0.033±0.145
0.697±1.536
0.063±0.303
r^1±Δr^1
−0.001±0.007
−0.003±0.009
−0.033±0.121
−0.0013±0.024
Elapsed Time [s]
6,334
9,209
2,077
2,243
GPU Utilization [%]
8.42
1.97
28.73
20.16
Elapsed Time (m=20)Elapsed Time (m=4)
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
Figure 10: Top Row: Densities ρ0(x,y) ( left ) and ρ1(x,y) ( right ), determined w.r.t Eq. 29 and 30 . Each density was created with a resolution of 100×100 pixels. Bottom Row: Datasets R0,R1,R2 generated according to Eq. 28 .
Figure 11: Density projection ρ0(y∣x∈[a,b]) ( top row ) and ρ1(y∣x∈[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 μ of the GAN before and after training respectively. The shaded blue areas represent the uncertainty σ of the trained GAN ensemble. The results shown here have been obtained by training the GAN on dataset R0 only.
Figure 12: Density projection ρ0(y∣x∈[a,b]) (top row) and ρ1(y∣x∈[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 μ of the GAN before and after training respectively. The shaded blue areas represent the uncertainty σ of the trained GAN ensemble. The results shown here have been obtained by training the GAN on datasets R0,R1 and R2 simultaneously.
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.
Chenglong Bao, Yiming Dang, Chenguang Duan +2
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 models have emerged as powerful generative priors for solving inverse imaging problems. However, their practical deployment is hindered by the substantial computational cost of slow, multi-step sampling. Although Consistency Models (CMs) address this limitation by enabling high-quality generation in only one or a few steps, their direct application to inverse problems has remained largely unexplored. This paper introduces a modified consistency sampling framework specifically designed for inverse problems. The proposed approach regulates the sampler's stochasticity through a measurement-consistency mechanism that leverages the degradation operator, thereby enforcing fidelity to the observed data while preserving the computational efficiency of consistency-based generation. Comprehensive experiments on the Fashion-MNIST and LSUN Bedroom datasets demonstrate consistent improvements across both perceptual and pixel-level metrics, including the Fréchet Inception Distance (FID), Kernel Inception Distance (KID), peak signal-to-noise ratio (PSNR), and structural similarity index measure (SSIM), compared with baseline consistency and diffusion-based sampling methods. The proposed method achieves competitive or superior reconstruction quality with only a small number of sampling steps.
Amirreza Tanevardi, Pooria Abbas Rad Moghadam, Seyed Mohammad Eshtehardian +2
Department of Electrical Engineering Sharif University of Technology Tehran, Iran · Electronics Research Institute Sharif University of Technology Tehran, Iran
Building datasets remains one of the most manual and brittle parts of AI development. In this technical report, we focus on the most extreme but also most prevalent setting real world practitioners face: a zero data regime. Here, practitioners don't have any data for the capability they want to learn. We introduce Invent-A-Dataset which is a prompt based system to go from dataset description to realistic and large scale post-training datasets. We evaluate Invent-A-Dataset against five frontier model APIs including Anthropic, Google, Open AI, DeepSeek, Zai. Across eight task types and dataset sizes up to 20K samples, Invent-A-Dataset significantly outperforms with both the highest quality (17% relative gains) while simultaneously producing the most diverse samples (19% relative gains). Its diversity advantage widens with scale of training dataset size (from parity at 200 samples to 37% relative gains at 20K samples). This translates into considerable downstream training gains, resulting in far more performant post-trained models. Invent-A-Dataset fine-tune consistently ranks higher compared to other generator fine-tunes across different post-trained model architectures.