Flinch: A Differentiable Framework for Field-Level Inference of Cosmological parameters from curved sky data

Abstract: We present Flinch, a fully differentiable and high-performance framework for field-level inference on angular maps, developed to improve the flexibility and scalability of current methodologies. Flinch is integrated with differentiable cosmology tools, allowing gradients to propagate from individual map pixels directly to the underlying cosmological parameters. This architecture allows cosmological inference to be carried out directly from the map itself, bypassing the need to specify a likelihood for intermediate summary statistics. Using simulated, masked CMB temperature maps, we validate our pipeline by reconstructing both maps and angular power spectra, and we perform cosmological parameter inference with competitive precision. In comparison with the standard pseudo-$C_\ell$ approach, Flinch delivers substantially tighter constraints, with error bars reduced by up to 40%. Among the gradient-based samplers routinely employed in field-level analyses, we further show that MicroCanonical Langevin Monte Carlo provides orders-of-magnitude improvements in sampling efficiency over currently employed Hamiltonian Monte Carlo samplers, greatly reducing computational expense.

• The paper demonstrates that a differentiable, field-level inference framework using AD and gradient-based samplers achieves up to 40% tighter cosmological constraints than traditional summary-statistics methods.
• It introduces a hierarchical Bayesian model with spherical harmonic transforms to effectively handle noisy, masked angular data while enabling end-to-end gradient propagation.
• Empirical results validate significantly improved sampling efficiency and precision in map and power spectrum reconstruction, highlighting the method’s robustness for next-generation cosmological surveys.

Flinch.jl: Differentiable Field-Level Inference for Cosmological Parameter Estimation on the Sphere

Introduction and Motivation

The Flinch.jl framework addresses the challenge of extracting maximal cosmological information from high-dimensional, masked, and noisy angular map data, such as those produced by CMB and large-scale structure surveys. Traditional approaches compress data into summary statistics (e.g., two-point functions), which, while effective for Gaussian fields, discard higher-order information and can be suboptimal in the presence of survey masks, inhomogeneous noise, and non-Gaussianity. Field-level inference (FLI) circumvents these limitations by operating directly on the pixelized data, enabling statistically optimal extraction of information under a correct generative model. However, FLI introduces significant computational and algorithmic complexity due to the high dimensionality and intricate posterior geometry.

Flinch.jl is designed to overcome these challenges by leveraging the Julia programming language and automatic differentiation (AD) to enable flexible, scalable, and fully differentiable field-level inference. The framework supports end-to-end propagation of gradients from map pixels to cosmological parameters, facilitating the use of advanced gradient-based samplers and seamless integration with differentiable cosmological theory modules.

Hierarchical Bayesian Model and Spherical Harmonic Representation

Flinch.jl adopts a hierarchical Bayesian model, where the observed map $\mathbf{d}$ is modeled as a noisy, masked realization of a latent field with spherical harmonic coefficients $\mathbf{a}$, whose covariance is determined by the angular power spectrum $C_\ell$ (or, in the cosmological hierarchy, by cosmological parameters $\boldsymbol{\theta}$ via $C_\ell(\boldsymbol{\theta})$). The generative process is:

1. Draw $C_\ell$ (or $\boldsymbol{\theta}$) from a prior.
2. Draw $\mathbf{a}$ from a zero-mean Gaussian with covariance $C_\ell$.
3. Generate the observed map via a linear response (mask, beam, pixel window) and additive noise.

This structure is formalized in a directed acyclic graph (DAG), with the flexibility to switch between power-spectrum-level and cosmological-parameter-level inference.

The spherical harmonic transform (SHT) is central to the model, enabling efficient representation and manipulation of fields on the sphere. The forward model incorporates all relevant observational effects (masking, pixelization, beam), and the likelihood is evaluated only on unmasked pixels.

Automatic Differentiation and Differentiable SHT

A key innovation in Flinch.jl is the pervasive use of AD for gradient computation. Unlike frameworks with hand-coded gradients, Flinch.jl leverages Julia's Zygote.jl to automatically propagate derivatives through the entire pipeline, including the SHT and cosmological theory modules. This enables rapid prototyping, model extension, and direct inference on cosmological parameters.

The implementation of differentiable SHT is achieved via adjoint-aware wrappers, where the forward pass computes the map from harmonic coefficients, and the pullback reconstructs gradients in harmonic space, respecting the symmetries and reality conditions of the coefficients.

Figure 1: Example of an adjoint-aware wrapper for a spherical-harmonic transform, illustrating how AD support is added with minimal code once the mathematical rules are derived.

This approach ensures that the computational overhead of gradient evaluation is modest (approximately twice the cost of the forward pass), and that gradients can flow seamlessly from the data to the highest levels of the hierarchy.

Gradient-Based Samplers and Performance

Flinch.jl supports several advanced gradient-based samplers:

• Hamiltonian Monte Carlo (HMC): Standard gradient-based MCMC, efficient in high dimensions but sensitive to tuning and mass matrix adaptation.
• No-U-Turn Sampler (NUTS): An adaptive variant of HMC that eliminates the need to pre-specify trajectory length, improving robustness.
• MicroCanonical Langevin Monte Carlo (MCLMC): A microcanonical, unadjusted sampler that operates on constant-energy surfaces, offering superior scaling in high dimensions at the cost of a controllable bias.
• Pathfinder: A quasi-Newton variational inference method used for rapid initialization of MCMC chains.

Empirical results demonstrate that MCLMC achieves orders-of-magnitude improvements in sampling efficiency over HMC and NUTS, particularly as the dimensionality increases.

Figure 2: Sampler efficiency (ESS per gradient evaluation) as a function of dimensionality, showing that MCLMC exhibits markedly better scaling than HMC and NUTS.

Map and Power Spectrum Reconstruction

Flinch.jl is validated on simulated, masked CMB temperature maps. The framework accurately reconstructs both the underlying map and its angular power spectrum, even in the presence of complex masks and noise. The field-level approach enables partial recovery of large-scale modes beneath the mask due to global mode coupling, while small-scale modes remain unconstrained.

Figure 3: Pixel residuals normalized by the recovered per-pixel uncertainty, demonstrating that NUTS and MCLMC achieve residuals consistent with a standard normal, indicating high convergence.

Figure 4: Power spectrum comparison between the fiducial theory and reconstructed spectra, with residuals normalized by their standard deviations, confirming consistency across scales.

Comparison with the standard pseudo-$C_\ell$ pipeline (e.g., NaMaster) reveals that Flinch.jl achieves substantially tighter constraints, with error bars reduced by up to 40%. This improvement is purely methodological, as both approaches use identical data and masks.

Figure 5: Binned power spectra recovered by HMC, NUTS, MCLMC, and the standard pseudo-$C_\ell$ pipeline, highlighting the suboptimality of the latter.

Cosmological Parameter Inference

By interfacing with differentiable cosmological emulators (e.g., Capse.jl), Flinch.jl enables direct inference on cosmological parameters from the map, bypassing summary statistics. Posterior constraints on parameters such as $\log(10^{10}A_s)$, $n_s$, $H_0$, $\omega_b$, and $\omega_c$ are systematically tighter than those obtained from pseudo-$C_\ell$ analyses, with error bars reduced by 20–40%.

Figure 6: Marginalized posterior contours for key cosmological parameters, showing that Flinch.jl yields noticeably tighter constraints than the NaMaster-based pseudo-$C_\ell$ approach.

This demonstrates that field-level likelihoods retain additional statistical information, leading to sharper and more precise cosmological constraints. The framework is readily extensible to more complex models, including non-Gaussian fields and higher-order statistics.

Convergence Diagnostics and Sampler Validation

Extensive convergence diagnostics are performed, including the Gelman-Rubin statistic, integrated autocorrelation time, and bias analysis. NUTS and MCLMC achieve excellent convergence, with $\hat{R} < 1.01$ for most parameters. MCLMC's bias is shown to be negligible compared to the Monte Carlo error of NUTS.

Figure 7: Histogram of Gelman-Rubin statistics for the reconstructed $C_\ell$'s, indicating excellent convergence for NUTS and MCLMC.

Figure 8: Relative bias of the cumulative mean of each sampler with respect to the NUTS reference, confirming the reliability of MCLMC.

Implications and Future Directions

Flinch.jl establishes a new standard for flexible, scalable, and statistically optimal field-level inference on the sphere. The combination of AD-enabled modeling and high-performance samplers (notably MCLMC) makes previously intractable analyses computationally viable, even at high resolution. The framework is well-positioned for application to upcoming high-precision datasets from CMB and LSS surveys.

Future developments include extension to spin fields, integration with differentiable 3x2pt statistics codes (e.g., Blast.jl), and application to real survey data for constraints on parameters such as $f_{\mathrm{NL}}$. The methodology also opens avenues for advanced systematics detection via cross-validation and for incorporating non-Gaussian likelihoods.

Conclusion

Flinch.jl demonstrates that fully differentiable, field-level inference frameworks can deliver substantial gains in cosmological parameter estimation by retaining the full information content of angular map data. The use of AD and advanced samplers enables both methodological flexibility and computational efficiency, with strong empirical evidence for improved statistical precision over traditional summary-statistics-based pipelines. This approach is poised to play a central role in the analysis of next-generation cosmological surveys.