Function-Space Decoupled Diffusion for Forward and Inverse Modeling in Carbon Capture and Storage

Paper: 2602.12274 Authors: Xin Ju, Jiachen Yao, Anima Anandkumar, Sally M. Benson, Gege Wen Categories: cs.LG, physics.geo-ph

Abstract

Carbon Capture and Storage (CCS) requires accurate characterization of subsurface flow dynamics, but faces fundamental challenges from sparse observational data and ill-posed inverse problems. This paper introduces Fun-DDPS (Function-space Decoupled Diffusion Posterior Sampling), a generative framework that decouples the learning of geological parameter distributions from physics-based dynamics modeling. The approach trains a diffusion model on geological parameters (geomodels) while using a Local Neural Operator (LNO) surrogate to provide differentiable physics guidance for cross-field conditioning. This architecture enables robust recovery of missing information in parameter space while maintaining physical consistency in the dynamics field through efficient gradient-based data assimilation.

Key Contributions

• Decoupled Architecture: Separates prior learning in parameter space from physics-based guidance in dynamics space, enabling single-channel diffusion models to handle complex multi-field problems
• Extreme Sparsity Handling: Achieves 7.7% relative error with only 25% observations versus 86.9% for standard surrogates—an 11x improvement demonstrating capability where deterministic methods completely fail
• Rigorous Validation Framework: First work to validate diffusion-based inverse solvers against asymptotically exact Rejection Sampling posteriors, achieving Jensen-Shannon divergence<0.06
• Physical Consistency: Produces artifact-free realizations while joint-state baselines suffer from high-frequency artifacts, with 4x better sample efficiency than rejection sampling

Methodology: Function-Space Decoupling

The core innovation lies in architectural decoupling. Traditional joint-state approaches like DPS condition diffusion models on both geological parameters and flow dynamics simultaneously, requiring multi-channel models that struggle with cross-field dependencies. Fun-DDPS instead:

1. Prior Learning: Trains a single-channel diffusion model $p_\theta(x)$ exclusively on geological parameter fields $x$ (permeability, porosity), learning the natural distribution of subsurface structures
2. Physics Surrogate: Employs a pre-trained Local Neural Operator $\mathcal{G}_\phi: x \rightarrow y$ mapping parameters to dynamics (pressure, saturation fields)
3. Guided Sampling: During inference, uses the surrogate to compute gradients $\nabla_x \log p(y_{obs}|x)$ for data likelihood, guiding the diffusion process toward observations

The sampling process follows:

$x_{t-1} = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}_t}}\epsilon_\theta(x_t, t)\right) - \zeta\nabla_x\|\mathcal{G}_\phi(x_t) - y_{obs}\|^2 + \sigma_t z$

where $\zeta$ controls guidance strength and $z \sim \mathcal{N}(0, I)$ maintains stochasticity.

Results: Forward and Inverse Performance

Forward Modeling Under Sparsity: When reconstructing full dynamics fields from 25% sparse observations, Fun-DDPS demonstrates remarkable robustness. The relative $L_2$ error of 7.7% contrasts sharply with 86.9% for standard neural operator surrogates. This isn’t incremental improvement—it’s the difference between usable predictions and complete failure. The decoupled architecture allows the diffusion prior to “fill in” missing spatial information using learned geological structure, while the surrogate ensures physical plausibility.

Inverse Problem Validation: The paper provides unprecedented rigor in validating generative inverse solvers. By comparing against Rejection Sampling—which asymptotically recovers the true posterior but requires ~1000 forward simulations per sample—the authors establish ground truth. Both Fun-DDPS and Fun-DPS achieve JS divergence<0.06, confirming they approximate the true posterior distribution. However, Fun-DDPS produces physically consistent samples without the high-frequency noise artifacts plaguing joint-state methods, while requiring only 250 function evaluations versus 1000 for rejection sampling.

Computational Efficiency: The Local Neural Operator surrogate enables gradient computation in milliseconds versus minutes for full physics simulators. This 1000x speedup makes iterative refinement practical, allowing Fun-DDPS to generate diverse posterior samples in reasonable time while maintaining physical fidelity through differentiable physics constraints.

Implications for Subsurface Modeling

This work addresses a critical bottleneck in CCS deployment: uncertainty quantification under sparse monitoring. Traditional deterministic inversion methods produce single point estimates that ignore posterior uncertainty, while Monte Carlo methods are computationally prohibitive. Fun-DDPS offers a middle path—generating diverse, physically plausible scenarios that capture epistemic uncertainty while remaining computationally tractable.

The decoupling principle extends beyond CCS. Any inverse problem involving parameter-to-observation mappings through expensive physics simulators could benefit: seismic inversion, reservoir engineering, climate model calibration, or medical imaging. The key insight is that learning priors in natural parameter space (where data is abundant) and applying physics constraints through surrogates (where gradients are needed) is more effective than joint modeling.

The validation against rejection sampling sets a new standard for generative inverse modeling. Too often, diffusion-based methods are evaluated only on reconstruction metrics without verifying they actually sample from the posterior. This work demonstrates that with proper architecture and guidance, diffusion models can provably approximate Bayesian posteriors while being orders of magnitude more efficient than traditional sampling methods.

Takeaways

1. Decoupling prior learning from physics guidance enables diffusion models to handle extreme data sparsity (75% missing observations) where deterministic methods fail completely
2. Function-space architectures with neural operator surrogates provide 11x improvement in reconstruction accuracy while maintaining 4x better sample efficiency than rejection sampling
3. First rigorous validation shows diffusion-based inverse solvers can provably approximate true Bayesian posteriors (JS divergence<0.06) when properly designed
4. Physical consistency emerges naturally from the decoupled architecture—Fun-DDPS produces artifact-free realizations while joint-state methods suffer from high-frequency noise
5. The framework generalizes beyond CCS to any inverse problem with expensive physics simulators and sparse observations, offering a practical path to uncertainty quantification in scientific computing