Differentiable Stochastic Halo Occupation Distribution with Galaxy Intrinsic Alignments

1. Problem Statement

Cosmological surveys (e.g., Rubin Observatory, Euclid, Roman) require precise modeling of galaxy clustering and Intrinsic Alignments (IA) to extract unbiased constraints from weak gravitational lensing data.

• The Bottleneck: Traditional Halo Occupation Distribution (HOD) models rely on stochastic sampling (Bernoulli/Poisson distributions) to populate dark matter halos with galaxies. These discrete, non-differentiable operations prevent the use of efficient gradient-based inference methods (like Hamiltonian Monte Carlo or gradient descent).
• Current Limitations:
  • MCMC: Standard Markov Chain Monte Carlo methods are computationally expensive for high-dimensional parameter spaces (e.g., Stage IV surveys may require years of compute time).
  • Emulators: Neural network-based emulators (e.g., IAEmu) offer speed but are limited to predicting specific summary statistics (like 2-point correlation functions) and cannot generate full galaxy catalogs or handle arbitrary statistics.
  • Analytic Models: Perturbation theory-based IA models struggle to capture nonlinear effects on small scales.

2. Methodology: diffHOD-IA

The authors present diffHOD-IA, a fully differentiable implementation of the HOD framework that incorporates galaxy intrinsic alignments. It extends the halotools-IA framework (Van Alfen et al., 2024) using automatic differentiation (via JAX).

Core Technical Components:

• Differentiable Sampling (The "Trick"):
  • Central Galaxies: The binary presence of a central galaxy (Bernoulli distribution) is relaxed using the Gumbel-Softmax trick. This replaces the discrete argmax with a continuous, differentiable softmax function controlled by a temperature parameter ($\tau$).
  • Satellite Galaxies: The Poisson distribution for satellite counts is approximated by a Binomial distribution (limit of $N_{max} \to \infty$), which is then relaxed using Gumbel-Softmax.
  • Spatial Placement: Satellites are assigned to subhalos using a softmax-based prioritization based on subhalo mass, replacing deterministic sorting with a differentiable weighting scheme. If subhalos are exhausted, positions are sampled from an NFW profile using inverse CDF sampling via Newton iteration.
• Differentiable Intrinsic Alignments:
  • Galaxy orientations relative to their host halos are modeled using the Dimroth–Watson distribution.
  • The authors implement differentiable inverse CDF sampling for this distribution. They derive analytic inverse forms where possible (using error functions) and use numerical Newton iterations for cases involving the imaginary error function (erfi), ensuring gradients flow from the alignment parameters ($\mu_{cen}, \mu_{sat}$) to the galaxy orientation vectors.
• Differentiable Summary Statistics:
  • The framework computes 2-point correlation functions (Position-Position $\xi(r)$, Position-Orientation $\omega(r)$, Orientation-Orientation $\eta(r)$) in a differentiable manner.
  • While pair counting is discrete, the authors fix the neighbor structure (pair indices) based on galaxy positions and allow gradients to flow only through the orientation vectors, enabling the calculation of $\partial \omega / \partial \mu$.

3. Key Contributions

1. End-to-End Differentiability: The first implementation to connect HOD and IA parameters directly to galaxy catalogs and summary statistics via automatic differentiation.
2. Dimroth–Watson Sampling: A novel differentiable sampling procedure for the Dimroth–Watson distribution, crucial for modeling galaxy misalignments.
3. Gradient Verification: Rigorous validation of computed gradients against finite difference estimates and analytic derivatives, confirming accuracy.
4. Hybrid Inference Pipeline: Demonstrates the integration of diffHOD-IA into Hamiltonian Monte Carlo (HMC) pipelines, a method previously inaccessible to stochastic HOD models.

4. Results & Validation

The model was validated against the reference non-differentiable halotools-IA using the Bolshoi-Planck simulation and parameters derived from the TNG300 simulation.

• Accuracy:
  • Galaxy Counts: Agreement within sample variance (mean difference < 0.2%).
  • 2-Point Statistics: Excellent agreement in $\xi(r)$, $\omega(r)$, and $\eta(r)$ across all scales, with fractional biases of < 2% (specifically 0.28% for $\xi$, 0.16% for $\omega$, and 1.77% for $\eta$).
  • Gradients: Autodifferentiated gradients matched finite-difference estimates with high precision for both HOD and IA parameters.
• Optimization Experiments:
  • Moment Matching: Successfully recovered IA parameters ($\mu_{cen}, \mu_{sat}$) from mock catalogs using gradient descent on first and second moments of alignment angles.
  • Correlation Function Matching: Recovered parameters by minimizing the loss between simulated and target $\omega(r)$, achieving convergence within statistical uncertainties.
• Inference Speed (HMC vs. MCMC):
  • diffHOD-IA (HMC): Converged in ~5 minutes on a single NVIDIA A100 GPU.
  • halotools-IA (MCMC): Required ~1 day across 150 CPU cores.
  • IAEmu (HMC): Converged in ~1 minute (faster than diffHOD-IA) but lacks catalog-level differentiability.

5. Significance & Future Impact

• Computational Efficiency: diffHOD-IA offers a massive speedup (orders of magnitude) for Bayesian inference compared to traditional MCMC, making high-dimensional parameter space exploration feasible for Stage IV surveys.
• Flexibility: Unlike emulators (which are "black boxes" trained on specific statistics), diffHOD-IA operates at the catalog level. This allows it to be extended to arbitrary summary statistics (e.g., higher-order moments, field-level inference) without retraining.
• Pipeline Integration: Because it is written in JAX, it can be integrated into larger differentiable cosmological pipelines (e.g., linking to differentiable N-body solvers like jaxPM or halo finders like jFoF), enabling end-to-end gradient flow from cosmological parameters to observed galaxy fields.
• Scientific Utility: The ability to perform joint inference on HOD and IA parameters simultaneously is a critical step toward reducing degeneracies in weak lensing analyses, ultimately leading to more precise cosmological constraints.

In summary, diffHOD-IA bridges the gap between realistic, simulation-based galaxy modeling and modern, gradient-based machine learning inference techniques, providing a scalable solution for the next generation of cosmological surveys.