Differentiable Fuzzy Cosmic-Web for Field Level Inference

A comprehensive analysis of the cosmological large-scale structure derived from galaxy surveys involves field-level inference, which requires a forward modelling framework that simultaneously accounts for structure formation and tracer bias. While structure formation models are well-understood, the development of an effective field-level bias model remains challenging within Bayesian reconstruction methods, which we address in this work. To bridge this gap, we have developed a differentiable model that integrates augmented Lagrangian perturbation theory, nonlinear, nonlocal, and stochastic biasing. At the core of our approach is the HICOBIAN model, which provides a description of a field with a positive number of tracers while incorporating a long- and short-range nonlocal framework and deviations from Poissonity in the likelihood. A key insight of our model is that transitions between cosmic-web regions are inherently smooth, which we implement using sigmoid-based gradient operations. This enables a fuzzy, and, hence, differentiable hierarchical cosmic-web description, making the model well-suited for machine learning frameworks. We test the practical implementation of this model through GPU-accelerated computations implemented in JaX, the BRIDGE code, enabling scalable evaluation of complex biasing. Our approach accurately reproduces the primordial density field within associated error bars derived from Bayesian posterior sampling within a self-specified setting as validated by two- and three-point statistics in Fourier space. Furthermore, we demonstrate that the methodology approaches the maximum encoded information consistent with Poisson noise. We also demonstrate that the bias parameters of a higher-order nonlocal bias model can be accurately reconstructed within the Bayesian framework for bias models with eight parameters.

💡 Research Summary

The paper presents a fully differentiable, GPU‑accelerated framework called BRIDGE for Bayesian field‑level inference of the primordial density field from galaxy‑type tracers. The core innovation lies in coupling an augmented Lagrangian perturbation theory (ALPT) gravity solver with a novel hierarchical, non‑local bias model named HICOBIAN. ALPT combines second‑order Lagrangian perturbation theory (2LPT) for long‑range tidal displacements with a spherical‑collapse (SC) solution for short‑range dynamics, blended through a smooth kernel of fixed radius (≈4 h⁻¹ Mpc). This hybrid approach retains analytic dependence on cosmological parameters, suppresses shell‑crossing, and reproduces N‑body halo statistics down to k ≈ 0.4 h Mpc⁻¹ while being far more computationally efficient than full particle‑mesh or N‑body simulations.

HICOBIAN addresses a long‑standing difficulty in perturbative bias modelling: truncation at a fixed order often yields unphysical negative densities and oscillatory behaviour. By first classifying each grid cell into one of the four cosmic‑web environments (node, filament, sheet, void) and then assigning a local, non‑linear bias function specific to that environment, the model isolates dominant contributions and preserves physical positivity. Transitions between environments are rendered smooth using sigmoid functions, producing a “fuzzy” cosmic‑web representation that is fully differentiable. The bias model also incorporates stochasticity and deviations from Poissonian shot noise by employing log‑normal priors on strictly positive bias parameters and allowing for super‑Poissonian dispersion.

All components are implemented in JAX, enabling automatic differentiation and seamless execution on GPUs. The forward model maps a white‑noise field ν (the latent initial conditions) to a linear density field δ₀ via the linear power spectrum, evolves δ₀ to a non‑linear matter overdensity δ with ALPT, and finally transforms δ into an expected tracer field (\bar n) using HICOBIAN. Observed tracer counts n* are assumed to follow an independent discrete distribution (Poisson or super‑Poisson) with mean (\bar n). The posterior (P(ν,b,p|n^*)) is sampled using gradient‑based Hamiltonian Monte Carlo, specifically the No‑U‑Turn Sampler (NUTS) from NumPyro, after an initial Wiener‑filtering step to accelerate convergence.

The authors validate the pipeline on mock data generated with the same forward model (“self‑specified” setting). At grid resolutions of 5 and 10 h⁻¹ Mpc, the reconstructed primordial density field matches the input power spectrum and bispectrum within the 1‑σ Bayesian credible intervals. Moreover, an eight‑parameter bias model (including linear, quadratic, cubic, and both long‑ and short‑range non‑local terms) is recovered with high fidelity; posterior means align with true values and uncertainties reflect the imposed log‑normal priors. Information‑theoretic analysis shows that the method extracts nearly the maximum possible information given Poisson shot noise, confirming that the fuzzy, differentiable web representation does not sacrifice statistical efficiency.

In summary, the BRIDGE‑HICOBIAN framework delivers (1) accurate reconstruction of the initial density field, (2) robust recovery of complex, high‑order bias parameters, and (3) a scalable, fully differentiable pipeline suitable for modern large‑scale surveys such as DESI, Euclid, and the Roman Space Telescope. Its combination of analytic perturbation theory, fuzzy cosmic‑web biasing, and modern autodiff tools positions it as a powerful bridge between theoretical cosmology and next‑generation observational data.