PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.26283v1cs.DCcs.LG

TL;DR

The paper addresses recovery of shared, unmeasurable quantities from heterogeneous datasets whose joint analysis is computationally demanding. It extends SAGIPS with dataset-specific operators and discriminators whose constraints update a shared generator, and validates the approach on Rutherford-scattering data. The reported results indicate parameter recovery, improved resolution as datasets are added, robustness to detector systematics, and multi-GPU scalability.

  • Problem

    Recovering shared unknown quantities from heterogeneous measurements requires joint analysis of complementary datasets with differing experimental conditions and feature-space coverage.

  • Method

    The framework extends SAGIPS by assigning each dataset its own forward operator and discriminator while synchronizing generator gradients across GPUs.

  • Results

    The framework recovers unknown physical parameters, improves resolution as datasets are added, remains robust to unknown detector systematics, and scales to multi-GPU systems.

  • Takeaways & Limitations

    The approach is presented as broadly applicable to real-world multi-dataset inverse analyses where experimental conditions vary across measurements.

Abstract

from arXiv · show

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.

I. INTRODUCTION

The paper frames multi-dataset inverse analysis as recovering shared unknowns from heterogeneous measurements while addressing complementary information and computational scale. It extends SAGIPS with distributed generative modeling and validates the framework on a multi-detector Rutherford-scattering benchmark.

  • Inverse problems require recovering high-dimensional or shared unknown quantities from measurements that do not directly measure them.
  • The framework extends SAGIPS to non-identically distributed datasets by using dataset-specific forward operators and discriminators with shared generator-gradient updates.
  • The study evaluates the approach with a controlled multi-detector Rutherford-scattering experiment designed to test recovery of true physical parameters as datasets are added.
  • The work also examines robustness across gradient-transport mechanisms and scalability on single-node and multi-node HPC systems with up to 120 GPUs.
  • The reported results show successful recovery, improved resolution with additional datasets, and robustness to unknown detector systematics.
  • Datasets probe different feature-space regions and carry complementary information, while their size and forward-operator cost can exceed a single GPU.

A. Distributed Multi-Dataset Analysis

The distributed workflow assigns each dataset its own operator and discriminator while maintaining a shared generator across GPU replicas. Local discriminator training supplies dataset-specific gradients, whereas generator gradients are aggregated to synchronize the replicas.

  • Each dataset is assigned to a dedicated GPU together with its own forward operator and discriminator.
  • A local generator replica transforms a noise vector into predictions that are passed through the dataset-specific forward operator to produce synthetic data.
  • Each local discriminator compares synthetic and real data and is updated using discriminator gradients that remain local to its GPU.
  • Generator gradients from all datasets are shared and summed into synchronized gradients, ensuring consistent generator replicas.

B. Validity of Non-IID Gradient Accumulation

The framework permits gradient accumulation across non-IID datasets when all datasets share the same unknown parameters, while distributing dataset-specific processing across GPUs. Its validity depends on shared physical unknowns and sufficiently informative datasets.

  • Validity conditions: Non-IID gradient aggregation remains valid when all datasets share a common set of unknown parameters.The datasets need not be identically distributed, but their shared generator must represent the same underlying unknowns.
  • Distributed design: Dataset-specific forward operators and discriminators provide separate gradient contributions to the shared generator.The distributed setup assigns each dataset to its own processing path while synchronizing generator updates.
  • Distributed design: Each dataset can be distributed across m GPUs, with standard data-parallel training applied independently within each dataset group.This requires N ×m GPUs and reduces the computational load seen by each GPU.
  • Validity conditions: The framework can fail when datasets are governed by different unknowns or when a dataset is statistically underrepresented.Contradictory accumulated gradients or uninformative local discriminator gradients may prevent meaningful convergence.
  • Validity conditions: Failure detection may require diagnostics such as discriminator-loss plateaus or cross-dataset consistency checks.These diagnostics are suggested for identifying problematic aggregation conditions in practice.
  • Benchmark: The Rutherford benchmark uses four fixed-angle detectors that report particle rates ν(θi), with detector settings and measurement uncertainty entering the forward process.The controlled setup treats a0 and a1 as unknowns despite their known literature values, enabling validation against ground truth.

A. Loop Closure Tests and Datasets

The loop-closure tests use analytically defined toy datasets with known pipelines and a two-parameter solution space. Combining datasets supplies the constraints needed to identify the unknowns, while controlled uncertainties mimic imperfect measurements.

  • Loop closure tests: Toy datasets Ri are generated with known pipelines Pi so generator predictions can be compared directly with the input unknowns.The generator predictions A0 and A1 replace a0 and a1, providing clean performance metrics.
  • Loop closure tests: The natural logarithm transforms products and fractions in the Rutherford expression into summations and differences for the pipeline.This preserves mathematical equivalence while simplifying the toy inverse problem.
  • Datasets: Six datasets are created, each containing 10^6 values of ln(ν) and covering a different region of feature space.Their distributions have similar O(10) magnitudes but differ in centers and spreads.
  • Datasets: Table I specifies the settings used to generate the six datasets and their sequential analysis order.The dataset indices in the table determine how SAGIPS incorporates the datasets.
  • Identifiability: At least two datasets are required to uniquely determine a0 and a1 because each single dataset yields a one-dimensional solution manifold.A single dataset therefore leaves infinitely many parameter pairs consistent with the observations.
  • Benchmark design: The toy setup intentionally provides an analytically tractable benchmark for visualizing convergence and validating the generative framework.The two-dimensional solution space makes the recovery behavior directly interpretable.

B. SAGIPS Setup

The SAGIPS experiments incrementally add datasets and use independently initialized GAN ensembles to quantify uncertainty. Single-GPU baselines compare shared-discriminator and dataset-specific-discriminator configurations.

  • Experimental protocol: Experiments add datasets incrementally, progressing from R0 alone to all six datasets.Each experiment starts from a freshly initialized ensemble to avoid contamination from earlier runs.
  • Experimental protocol: An ensemble of independently initialized GANs provides uncertainty quantification across the inverse-solver experiments.The ensemble is trained separately for each dataset-count configuration.
  • Run groups: Run Group A uses one GPU per dataset, including single-GPU experiments and multi-GPU configurations with multiplicity m = 1.This grouping isolates experiments without within-dataset sharding.
  • Run groups: Run Group B shards each of N datasets across m > 1 GPUs.This grouping studies the effects of assigning multiple GPUs to each dataset.
  • Run groups: Separating the run groups isolates the effects of dataset count from the effects of GPU multiplicity.The comparison is designed to distinguish data-scaling behavior from within-dataset sharding behavior.
  • SAGIPS configurations: The single-GPU baseline either concatenates datasets for one discriminator or assigns each dataset its own discriminator.In the multi-discriminator setting, discriminator gradients are combined before updating the shared generator.

D. Multi-GPU Runs

Multi-GPU runs evaluate dataset-specific parallelism, several gradient-transport strategies, and within-dataset sharding. Increasing multiplicity lowers per-GPU memory and compute demands but also reduces discriminator throughput per step.

  • Multi-GPU architecture: Each dataset is assigned to its own GPU, enabling parallel evaluation of its forward operator and discriminator.This configuration implements the distributed architecture used by the framework.
  • Gradient transport: The study compares four gradient-transport methods implemented with mpi4py for inter-GPU communication.The methods are conventional ARAR, grouped ARAR, Strong ARAR, and Double Binary Tree.
  • Gradient transport: Conventional ARAR maximizes cross-GPU communication but depends heavily on inter-node communication speed.Grouped ARAR restricts outer-group communication to one GPU per node, while Strong ARAR combines broader synchronization with lower communication cost.
  • Multiplicity: When datasets are sharded, the generator’s predicted sample count scales inversely with multiplicity m.The default generated-sample count is 1 k when data is not sharded.
  • Multiplicity: Increasing m reduces discriminator throughput per step while decreasing memory and compute demands on each GPU.The experiment tests whether this trade-off accelerates training without degrading inference quality.
  • Computing systems: Run Group A uses a single Jefferson Lab node with 8 NVIDIA A800 GPUs, whereas Run Group B uses Polaris nodes with 4 NVIDIA A100 GPUs each.The two run groups therefore use different computing environments for single-GPU and sharded experiments.

IV. EVALUATION AND RESULTS

The evaluation measures reconstruction accuracy, uncertainty, and computational behavior as datasets are added to SAGIPS. Multi-GPU configurations improve training time while maintaining distributed memory requirements.

  • IV. Evaluation and Results: The evaluation uses relative residuals and ensemble uncertainty, with rank-replica averaging applied before computing ensemble means.Training time is measured from timestamped generator checkpoints written by all ensemble members and GPUs.
  • A. Results from Run Group A: Residuals are consistent with zero within uncertainties, indicating successful reconstruction of the unknowns a0 and a1 across the analyzed approaches.
  • A. Results from Run Group A: Adding datasets drives residuals toward zero and shrinks their uncertainties, whereas one dataset poorly constrains the two unknowns.A single dataset defines a one-dimensional manifold in the solution space for two unknowns.
  • A. Results from Run Group A: The single-GPU single-discriminator baseline performs weakest, particularly for the uncertainty associated with the second residual.The passage attributes this to difficulty handling varying feature ranges across datasets.
  • A. Results from Run Group A: Both multi-GPU approaches achieve a factor of ≳2 improvement in training time through distributed computation.The single-GPU multi-discriminator approach takes longest, while the single-GPU single-discriminator setup is approximately 100 s faster.
  • A. Results from Run Group A: Multi-GPU configurations require near-constant memory per GPU as the number of analyzed datasets increases.Single-GPU approaches require the datasets to reside simultaneously in one A800 GPU’s memory.

3. Assessing possible systematic Bias from the Input Data

The study tests whether dataset ordering and distributed execution affect reconstruction as heterogeneous measurements are added. Low-resolution subsets can reduce precision, while full-data multi-GPU analyses preserve performance and expose communication–convergence trade-offs.

  • 3. Assessing possible systematic Bias from the Input Data: All ordering schemes converge to the same performance once all six datasets are analyzed because they then access identical information.
  • 3. Assessing possible systematic Bias from the Input Data: The framework evaluates six datasets through successive ordering schemes defined by dataset-specific resolution δν(i).Up uses ascending resolution, Down descending resolution, and Mix a random permutation of the default order.
  • 3. Assessing possible systematic Bias from the Input Data: The Down ordering worsens residual quality at n = 3 when only the three lowest-resolution datasets are analyzed.Residuals remain statistically consistent with zero, but uncertainties become substantially larger.
  • 3. Assessing possible systematic Bias from the Input Data: Performance recovers markedly after adding a fourth dataset, while Up and Mix remain consistent with the baseline across tested points.The results indicate that higher-resolution data can disproportionately improve precision.
  • B. Results from Run Group B: With conventional ARAR transport, performance is poor for one dataset but improves noticeably as more datasets are added under multiplicity m = 20.Additional data sharding has no negative impact on inverse-solver performance.
  • B. Results from Run Group B: Grouped ARAR shows inferior accuracy, whereas conventional ARAR and double binary-tree achieve the best reconstruction with training times of multiple hours.Strong-ARAR trains faster while offering competitive convergence quality.

1. Scaling with Data Multiplicity

The analysis examines how dataset multiplicity and gradient-transport choices affect training time, generator consistency, and parameter recovery in distributed multi-dataset inference.

  • Scaling with Data Multiplicity: Increasing multiplicity can reduce training time, with speedup evaluated by comparing conventional ARAR runs at m = 4 and m = 20.The speedup ratio is defined from the two training times and reported for conventional and strong ARAR.
  • Model Drift Analysis: Weight and output drift are monitored as complementary diagnostics of generator consistency during SAGIPS training.The metrics are based on pairwise replica differences and are summarized across ensemble members.
  • Model Drift Analysis: All drift curves start at zero because replicas share identical weights after initial synchronization.The initial pairwise distance is therefore exactly zero by construction.
  • Model Drift Analysis: Conventional ARAR preserves perfect consistency, strong ARAR has bounded weight drift near 0.01 with negligible output drift, and ARAR has the largest drift.The reported drift patterns align with differences in convergence quality and training time.
  • Model Drift Analysis: The proposed drift metrics are intended as general diagnostics, while improved synchronization and posttraining weight averaging remain future mitigation strategies.These mitigation techniques are not developed in the present work.
  • Scaling with Data Multiplicity: The framework uses synchronized generator gradients across GPUs while assigning dataset-specific processing to separate devices.This design supports simultaneous analysis of heterogeneous datasets.
  • Scaling with Data Multiplicity: Parameter recovery and inference quality improve as more datasets are incorporated, although performance also depends on which datasets are available.Combining only the lowest-resolution datasets can transiently degrade precision.

Appendix A: Distributed Data Parallel Training - DDP

Appendix A explains standard distributed data-parallel training and extends its gradient aggregation to heterogeneous datasets that share common unknown parameters.

  • Appendix A: Distributed Data Parallel Training - DDP: Distributed data parallelism minimizes a shared loss by updating replicated network parameters from aggregated gradients.The parameter update uses gradients and a learning rate, with optimizer-specific modifications possible.
  • Appendix A: Distributed Data Parallel Training - DDP: A batch can be divided into k equal shards, allowing k smaller gradient estimates to be computed in parallel and locally aggregated.The shard size satisfies B = ks.
  • Appendix A: Distributed Data Parallel Training - DDP: Correct DDP requires synchronized initial network parameters and optimizer states across all ranks.Each rank holds a copy of the network before training begins.
  • Appendix A: Distributed Data Parallel Training - DDP: All ranks must receive the accumulated gradient through all-reduce or ring-all-reduce to update local networks consistently.Without gradient accumulation, replicas would follow different optimization trajectories.
  • Appendix A: Distributed Data Parallel Training - DDP: Standard DDP assumes IID shards so averaged local gradients remain unbiased estimates of the same underlying gradient.Under- or over-represented features can invalidate the stochastic gradient estimate.
  • Using DDP for distributed Multi-Dataset Analysis: For multi-dataset analysis, the joint objective sums dataset-specific losses and aggregates their local gradients across GPUs.The gradient identities hold regardless of the joint distribution of the datasets.
  • Using DDP for distributed Multi-Dataset Analysis: Equal normalization by 1/N prevents any dataset from dominating the update, while the aggregated gradient forms a consensus over complementary constraints.Per-dataset weighting can be added by modifying the normalization term.

Appendix B: Teaser: Multi-Dataset Analysis in 2D

Appendix B constructs a two-density, three-dataset toy problem to demonstrate multi-dataset analysis beyond scalar parameter recovery.

  • Appendix B: Teaser: Multi-Dataset Analysis in 2D: The teaser problem uses two densities, ρ0(x, y) and ρ1(x, y), to define three datasets through mixtures with coefficients βji.The mixture coefficients satisfy β0i + β1i = 1 for each dataset.
  • Appendix B: Teaser: Multi-Dataset Analysis in 2D: The densities ρ0 and ρ1 are specified by distinct functions over x and y, each including the shared factor (1 + 0.75 · xy).Their exponents differ across x, 1 − x, y, and 1 − y.
  • Appendix B: Teaser: Multi-Dataset Analysis in 2D: Three datasets containing 10^7 (x, y)-pairs are generated from the densities and fitted with a GAN tasked with recovering both density functions.The datasets are visualized in Fig. 10 and generated using the mixture pipeline.

1. Setting up the GAN

The 2D teaser adapts the GAN to predict two density images and evaluates recovery through density projections before and after training with one or three datasets.

  • Setting up the GAN: The generator predicts two color channels representing the densities at the resolution used for the density images.The discriminator retains the earlier architecture with two numerical inputs instead of one.
  • Setting up the GAN: An additional integral loss enforces proper density normalization, with the tunable scale set to λI = 0.008.The integrals can correspond to expected cross sections or theoretical values in a nuclear-physics application.
  • Experiments and Results: The experiments assess GAN performance using projections of ρ0 and ρ1 across different x-intervals.The one-dataset and three-dataset experiments are summarized in Figs. 11 and 12.
  • Experiments and Results: With dataset R0 alone, the GAN cannot recover both densities because one mixture equation constrains two unknown densities.Predictions tend toward ρ0 because R0 is dominated by that density.
  • Experiments and Results: With R0, R1, and R2, the generator properly represents both densities, although disagreement remains most pronounced for ρ1 at high x.The discrepancy may be reduced through greater capacity, longer training, or stronger regularization.
  • Experiments and Results: The toy experiments support applying the multi-dataset framework to complex pipelines and to proton 3D imaging.The appendix presents the problem as closely related to a real QuantOm application.
Loading 2608.26283v1…