Skip to content

Dynamic Diffusion Schrödinger Bridge in Astrophysical Observational Inversions

Conference: NeurIPS 2025 arXiv: 2506.08065 Code: Available Area: Physics Keywords: Schrödinger Bridge, Astrophysical Inverse Problems, Giant Molecular Clouds, Probabilistic Generative Models, Out-of-Distribution Generalization

TL;DR

This paper proposes Astro-DSB, a Diffusion Schrödinger Bridge-based framework for modeling astrophysical inverse problems. It directly learns a probabilistic mapping from observables to true physical distributions, requires only 25% of the training cost of conditional DDPM, demonstrates significant generalization advantages in out-of-distribution (OOD) evaluation, and is successfully applied to real observational data from Taurus B213.

Background & Motivation

Giant Molecular Clouds (GMCs) are key components of the interstellar medium, and understanding their internal physical distributions (density, magnetic field) is critical for studying star formation. The core task is to infer the true physical state \(x_1\) (e.g., volume density, magnetic field strength) from limited observables \(y\) (e.g., column density maps)—a classical observational inverse problem.

Limitations of existing methods: - Traditional astrophysical statistical methods (power-law fitting, DCF method): rely on overly simplistic physical assumptions with poor generalization - Discriminative ML methods (U-Net): performance degrades severely in OOD scenarios - Conditional DDPM: requires an unnatural Gaussian prior assumption, slow training convergence (~400 epochs), and is misaligned with the underlying physical process

Core motivation: The Gaussian prior \(p_0 = \mathcal{N}(0, I)\) in diffusion models is mismatched with astrophysical systems. The Schrödinger Bridge framework can directly model the mapping from observables to physical distributions without requiring artificially imposed prior structures.

Method

Overall Architecture

Astro-DSB is built upon the paired Diffusion Schrödinger Bridge (I2SB) framework, with three key improvements:

  • M0: Pairwise Matching — exploits the deterministic pairing between observables \(y\) and physical states \(x_1\)
  • M1: Noise Alignment — aligns ML noise with observational noise
  • M2: Observable Enhancement — injects observables as additional conditioning at each timestep

Training uses a patch-based strategy; inference employs aggregated scalable sampling to handle large-scale observational data.

Key Designs

M0: Pairwise Matching

Adopts the analytic posterior form from I2SB: $\(p(x_t | y, x_1) = \mathcal{N}\!\left(x_t;\, \frac{\bar\sigma_t^2}{\bar\sigma_t^2 + \sigma_t^2} y + \frac{\sigma_t^2}{\bar\sigma_t^2 + \sigma_t^2} x_1,\, \frac{\bar\sigma_t^2 \sigma_t^2}{\bar\sigma_t^2 + \sigma_t^2} I\right)\)$

This eliminates the complexity of wave-function coupling in standard DSB and renders training tractable. Unlike vision-domain I2SB, \(p_0\) and \(p_1\) here correspond to observables and physical states, respectively—not arbitrary image domains.

M1: Noise Perturbation Alignment

Explicitly models observational noise \(\varepsilon \sim \lambda\mathcal{N}(0,I)\) (\(\lambda=0.1\)): - Incorporates the noise in the observation model \(y = F(x_1) + \varepsilon\) into the DSB framework - The optimization target shifts to learning \(p(x_1 \mid y - \varepsilon)\) rather than \(p(x_1 \mid y)\) - Simple yet critical—a fundamental distinction from conventional vision tasks that assume clean inputs

M2: Observable Enhancement

Injects the clean observable \(y\) as additional conditioning into the score function at each timestep \(t\): - Effectively learns a conditional Doob's h-transform - Continuously "reminds" the model of the observational constraints - Empirically found to be crucial for both convergence speed and predictive accuracy

Patch-based Training + Aggregated Sampling

  • Training: Astrophysical data have large resolutions; images are cropped into \(128 \times 128\) patches
  • Inference: For large-scale observations (e.g., Taurus B213), the data are split into overlapping patches, inferred independently, and then aggregated
  • Unlike padding-based schemes (Li et al., 2025), no boundary effects need to be handled

Loss & Training

Reparameterized training objective: $\(\mathcal{F}_{\theta,t} = \left\|\varepsilon_{\theta,t}(x_t, y, t;\,\theta) - \frac{y - \varepsilon - x_1}{\sigma_t}\right\|^2\)$

  • U-Net backbone, ~81M parameters (unified with baselines)
  • 1000-step time discretization
  • Batch size 16, AdamW, lr=\(5\times10^{-5}\), \(4\times\) NVIDIA T4 GPUs
  • Convergence in only ~100 epochs (vs. ~400 epochs for conditional DDPM), reducing training cost by 75%
  • Inference: 6–30 seconds per test sample (depending on whether step skipping is enabled)

Key Experimental Results

Main Results

Molecular Density Prediction (ID + OOD)

| Method | ID: \(\mu\)(|↓|) | ID: \(\sigma\)(|↓|) | OOD: \(\mu\)(|↓|) | OOD: \(\sigma\)(|↓|) | |--------|--------------|-----------------|----------------|-----------------| | 3PLF (Traditional) | -0.77 | 4.44 | -3.23 | 4.21 | | U-Net | -0.25 | 1.29 | 5.01 | 7.08 | | cDDPMs | -0.05 | 0.61 | 2.88 | 5.33 | | Astro-DSB | -0.02 | 0.71 | 0.51 | 2.32 |

Astro-DSB achieves the best \(\mu\) in the ID setting and comprehensively outperforms all baselines in the OOD setting (\(\sigma\) reduced from 5.33 to 2.32).

Magnetic Field Strength Prediction (ID + OOD)

| Method | ID: \(\mu\)(|↓|) | ID: \(\sigma\)(|↓|) | OOD: \(\mu\)(|↓|) | OOD: \(\sigma\)(|↓|) | |--------|--------------|-----------------|----------------|-----------------| | DCF (Traditional) | 11.80 | 21.78 | 5.19 | 9.64 | | U-Net | -0.06 | 0.38 | -0.32 | 0.74 | | cDDPMs | 0.07 | 0.38 | 0.12 | 0.72 | | Astro-DSB | 0.15 | 0.55 | 0.19 | 0.67 |

Astro-DSB also demonstrates OOD generalization advantages for magnetic field prediction, with only 25% of the training cost of cDDPMs.

Ablation Study

Ablation of Key Components (Density Prediction)

Configuration ID: \(\mu\) ID: \(\sigma\) OOD: \(\mu\) OOD: \(\sigma\)
w/o M1 & M2 -0.19 1.75 4.24 14.71
w/o M1 \(-3\times10^{-3}\) 0.65 -1.84 5.06
w/o M2 \(7\times10^{-3}\) 1.63 -14.66 54.95
Full -0.02 0.71 0.51 2.32

M2 (Observable Enhancement) is critical for OOD generalization—its removal causes \(\sigma\) to surge from 2.32 to 54.95. M1 (Noise Alignment) contributes primarily to in-distribution accuracy.

Key Findings

  • Probabilistic generative models > discriminative models: the advantage is especially pronounced in OOD scenarios
  • Distribution-level learning yields better generalization: DSB learns a probabilistic mapping \(p_0 \to p_1\) rather than a pixel-level mapping, resulting in greater robustness
  • Successful application to real Taurus B213 observational data, with predictions consistent with HC3N density tracer measurements
  • The 75% training cost reduction stems from \(p_0\) (observables) being substantially closer to \(p_1\) (physical states) than \(\mathcal{N}(0,I)\) (Gaussian noise), enabling faster convergence

Highlights & Insights

  1. Physics–ML alignment: Removing the Gaussian prior assumption and directly modeling the observable-to-physical-state probabilistic mapping yields stronger interpretability
  2. Substantially improved training efficiency: The 75% training cost reduction arises from a better prior choice, with clear theoretical intuition
  3. Scientifically meaningful OOD generalization: In astrophysics, OOD is defined via different initial conditions and physical processes, making it more rigorous than OOD in vision tasks
  4. Validation on real observations: Predictions on Taurus B213 are consistent with independent observational constraints, demonstrating practical applicability
  5. Interdisciplinary contribution: Showcases the potential of diffusion models for scientific discovery beyond conventional visual synthesis

Limitations & Future Work

  • The complete forward physical equations governing GMC evolution are not explicitly modeled; the approach relies on data-driven mappings
  • Future work could consider integrating multiple intermediate physical states or encoding physical laws directly into the generative process
  • Inference speed (6–30 seconds per sample) still requires optimization for large-scale survey data processing
  • The current work validates only two physical quantities (density and magnetic field); extension to temperature, velocity, and other quantities is warranted
  • I2SB (Liu et al., 2023) provides the foundational theoretical framework for paired DSB; Astro-DSB adds noise alignment and observable enhancement
  • Distinction from Li et al. (2025) on physical field reconstruction: single-stage learning (no pretraining) with systematic OOD evaluation
  • MHD simulation data are drawn from the ENZO and ORION2 codes, with physically meaningful ID/OOD splits
  • Implication for the vision community: probabilistic models exhibit greater robustness than discriminative models under OOD conditions

Rating

  • Novelty: ⭐⭐⭐⭐ (first application of DSB to astrophysical inverse problems; M1/M2 designs are innovative)
  • Technical Depth: ⭐⭐⭐⭐ (Schrödinger Bridge theory combined with physical constraints is substantive; implementation is relatively concise)
  • Experimental Thoroughness: ⭐⭐⭐⭐⭐ (dual ID/OOD evaluation, real observations, complete ablations, fair comparisons)
  • Value: ⭐⭐⭐⭐ (practical value in astrophysics; 75% training cost reduction is appealing)
  • Writing Quality: ⭐⭐⭐⭐ (astrophysical background is well introduced and accessible to ML readers)