Differentiable Fuzzy Cosmic-Web for Field Level Inference

TL;DR

This paper introduces BRIDGE, a GPU-accelerated, differentiable framework for field-level inference of the cosmological large-scale structure that combines ALPT-based gravity with a hierarchical, nonlocal HICOBIAN bias. By treating bias as a morphology-dependent, stochastic, and smoothly varying quantity across the cosmic web, the authors enable gradient-based Bayesian inference of the primordial density field from tracer data, including redshift-space distortions. The approach is validated with numerical tests, showing accurate recovery of initial and final density fields and bias parameters under self-consistent forward-model conditions, and approaching the shot-noise information limit in summary statistics such as the power spectrum and bispectrum. While highly effective in recovering fields and many bias parameters, the work identifies challenges in recovering parameters for highly complex bias models, motivating future work on intra-halo physics, light-cone evolution, and direct application to real survey data like DESI.

Abstract

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.

Figures (15)

• Figure 1: Schematic overview of the BRIDGE pipeline. The framework combines a scalable, differentiable structure formation model with a flexible field-level bias prescription, all implemented in a GPU-accelerated JAX environment. This enables efficient Bayesian inference of the primordial density field from tracer observations, with support for complex, nonlocal bias models and multi-resolution analysis. The modular design allows for seamless integration of physical models while maintaining end-to-end differentiability and high computational performance.
• Figure 2: Example of fuzzy cosmic-web classification with $N=256$ and $\Delta L=1.7\,h^{-1}\mathrm{Mpc}$. Top left: evolved dark matter density contrast field. Bottom left: "hard" $\Phi$-web classification. Right panels: fuzzy membership weights $p_i^{(\mathrm{V})}$, $p_i^{(\mathrm{S})}$, $p_i^{(\mathrm{F})}$, and $p_i^{(\mathrm{K})}$ for voids, sheets, filaments, and knots, respectively, as defined in Eq. \ref{['eq:soft_cweb_p']}.
• Figure 3: TEST1: Mean autocorrelation $\xi(n)$ as a function of lag $n$ for a chain of length $M\approx1000$, evaluated over $10^4$ randomly selected voxels. The red curve shows the ensemble average $\langle\xi(n)\rangle$ across parameters, with the shaded band indicating the $\pm1\sigma$ dispersion. The horizontal dashed line marks the threshold $\xi=0.1$ used to assess mixing efficiency.
• Figure 4: Computing times for single gradient evaluations of the forward model employing ALPT evolution and the HICOBIAN bias model for different mesh sizes, run on a single NVIDIA A100-SXM4 GPU equipped with 40 GB of on-board HBM2 memory. For reference, the full burn-in and sampling procedure for the $128^3$ case in this study required approximately 100 GPU-hours.
• Figure 5: Gelman-Rubin convergence diagnostic computed for three independently initialized chains of 250 samples after convergence. This statistic is calculated for all cells in Fourier space, and we show their distribution with scale $k$. Values close to one indicate that the drawn samples have converged to the same distribution.