Differentiable Fuzzy Cosmic-Web for Field Level Inference

1. Problem Statement

The reconstruction of the primordial density field from present-day galaxy surveys is a fundamental challenge in cosmology. It requires field-level inference, which involves sampling the posterior distribution of initial conditions given observed tracer distributions (galaxies, quasars, etc.).

• Limitations of Inverse Methods: Traditional inverse reconstruction methods fail in the non-linear regime because gravitational evolution is non-invertible once shell-crossing occurs (multiple matter streams overlap), leading to irreversible loss of phase-space information.
• Limitations of Forward Modeling: While forward modeling (evolving initial conditions to the present day) is physically principled, it faces two major hurdles:
  1. Computational Cost: High-fidelity N-body simulations are too expensive for the millions of evaluations required for Bayesian posterior sampling.
  2. Bias Modeling: Developing an effective field-level bias model is difficult. Standard perturbation theory bias models often produce unphysical negative densities or oscillatory behavior when truncated. Furthermore, existing cosmic-web classification methods (e.g., based on eigenvalue signs) rely on hard, non-differentiable thresholds, making them incompatible with gradient-based sampling algorithms like Hamiltonian Monte Carlo (HMC).

2. Methodology

The authors introduce BRIDGE (Bayesian Reconstruction and Inference of Data-driven Generative Environments), a differentiable, GPU-accelerated framework implemented in JAX. The methodology integrates three core components:

A. Forward Model Architecture

The forward model $\mathcal{M}$ maps white noise (initial conditions), bias parameters, and cosmological parameters to the expected tracer field:

1. Initial Conditions: A linear transformation maps white noise to the initial overdensity field using the linear matter power spectrum.
2. Gravity Solver: Uses Augmented Lagrangian Perturbation Theory (ALPT). This hybrid approach combines Second-Order LPT (2LPT) for long-range tidal displacements and a Spherical Collapse (SC) solution for short-range dynamics. ALPT is computationally efficient, suppresses shell-crossing, and reproduces N-body statistics down to $k \approx 0.4 \, h \, \text{Mpc}^{-1}$.
3. Bias Model: Maps the evolved dark matter field to the mean tracer field using a non-linear, non-local, and stochastic bias prescription.

B. The HICOBIAN Bias Model

The core innovation is the Hierarchical Cosmic-Web Biasing Nonlocal (HICOBIAN) model:

• Hierarchical Classification: It combines a long-range classification ($\Phi$-web, based on the tidal field) and a short-range classification ($\delta$-web, based on the density Hessian). This partitions space into 16 distinct regions (e.g., voids, sheets, filaments, knots, and their sub-structures).
• Local Bias per Region: Within each region, a specific local non-linear bias model is applied. The model uses a power-law component modulated by low-pass and high-pass filters to ensure positive tracer densities and account for halo exclusion.
• Stochasticity: The model incorporates a Negative Binomial distribution (via a Gamma-Poisson mixture) to account for over-dispersion in tracer counts, moving beyond simple Poisson noise.

C. Differentiable "Fuzzy" Cosmic Web

To enable gradient-based inference (HMC), the authors replaced hard cosmic-web thresholds with sigmoid-based transitions.

• Instead of a step function (Heaviside) that creates discontinuous gradients, they use a sigmoid function to calculate "fuzzy" membership weights ($p^{(K)}, p^{(F)}, p^{(S)}, p^{(V)}$) for each voxel.
• Bias parameters for a voxel are computed as a weighted sum of the parameters for each cosmic-web type. This ensures the entire forward model is end-to-end differentiable, allowing for efficient gradient computation required by HMC.

D. Inference Engine

• Sampling: The framework uses Hamiltonian Monte Carlo (HMC), specifically the No-U-Turn Sampler (NUTS), to sample the high-dimensional posterior of initial conditions and bias parameters.
• Initialization: Chains are initialized using a Wiener-filtered reconstruction of the observed data to accelerate convergence.
• Redshift-Space Distortions (RSD): RSDs are modeled at the dark-matter field level in a fully differentiable manner, mapping real-space coordinates to redshift-space coordinates based on peculiar velocities.

3. Key Contributions

1. Differentiable Cosmic Web: The first implementation of a cosmic-web classification scheme that is fully differentiable, enabling the use of gradient-based Bayesian inference methods for complex, non-local bias models.
2. HICOBIAN Model: A unified bias framework that simultaneously handles non-linearity, non-locality (long and short-range), and stochasticity across hierarchical cosmic-web environments.
3. BRIDGE Code: A scalable, GPU-accelerated codebase written in JAX that performs field-level inference orders of magnitude faster than CPU-based alternatives, capable of handling complex bias models with up to 64 parameters.
4. Unified Framework: A single pipeline that integrates structure formation (ALPT), complex biasing, and RSDs, mirroring the forward models used to generate mock catalogs.

4. Results

The authors validated BRIDGE through three numerical tests (TEST1, TEST2, TEST3) with varying resolutions (5–10 $h^{-1}$ Mpc) and bias complexities (8 to 64 parameters).

• Reconstruction Accuracy: The framework successfully reconstructed the primordial density field with high fidelity. Two-point (power spectrum) and three-point (bispectrum) statistics in Fourier space matched the ground truth within error bars.
• Information Limit: The reconstruction reached the theoretical limit of encoded information consistent with Poisson noise (shot noise), as verified by the propagator (cross-correlation between reconstructed and true fields).
• Parameter Recovery:
  • In simpler models (TEST1 & TEST2), the code accurately recovered both the initial density field and the bias parameters (power-law $\alpha$ and negative binomial $\beta$) with sub-percent to few-percent precision.
  • In the complex HICOBIAN model (TEST3, 64 parameters), the density field was accurately recovered, but the specific bias parameters were not fully recovered when allowed to vary freely. This suggests a degeneracy in complex bias models that may require additional constraints or intra-halo physics to resolve.
• Performance: Single-gradient calculations took less than one second on an NVIDIA A100 GPU. A full inference run for a $128^3$ grid required approximately 100 GPU hours.

5. Significance

This work represents a significant step toward unified, interpretable, and fully differentiable cosmological data analysis.

• Scalability: By leveraging JAX and GPUs, BRIDGE makes field-level inference feasible for large-scale surveys (e.g., DESI, Euclid, Roman) where brute-force sampling was previously impossible.
• Physical Consistency: The use of a forward model identical to mock catalog generation ensures that the inference is consistent with the physics used to simulate the data, reducing systematic errors.
• Future Applications: The framework is designed to be extended to light-cone evolution, selection effects, and more complex tracers (Lyman-$\alpha$ forest, intensity mapping). It paves the way for joint inference of cosmological parameters and complex astrophysical bias models, potentially breaking degeneracies in our understanding of dark matter and galaxy formation.