jFoF: GPU Cluster Finding with Gradient Propagation

Technical Summary of "jFoF: GPU Cluster Finding with Gradient Propagation"

Problem Statement
Cosmological simulations rely heavily on the Friends-of-Friends (FoF) algorithm to identify dark matter halos, the fundamental building blocks of cosmic structure. While effective, traditional FoF implementations face two primary limitations in modern astrophysical workflows:

1. Performance Bottlenecks: Standard implementations often rely on CPU-based spatial data structures (e.g., $k$-d trees) or distributed-memory MPI frameworks. These approaches incur significant overhead from host-device data transfers and struggle to exploit the massive parallelism of modern GPU architectures, particularly when integrated into end-to-end GPU-native simulation pipelines (e.g., PMWD, GRAM-X).
2. Non-Differentiability: Halo finding involves discrete operations—thresholding, neighbor selection, and graph connectivity—that are inherently non-differentiable. This prevents the propagation of gradients through the halo identification step, hindering the use of gradient-based optimization and inference methods (such as Hamiltonian Monte Carlo) to calibrate cosmological models or reconstruct initial conditions directly from observed halo catalogs.

Methodology
The authors present jFoF, a fully GPU-native Friends-of-Friends halo finder implemented in JAX. The code is designed to operate entirely on the device, eliminating costly CPU-GPU data transfers and enabling integration into differentiable astrophysical pipelines.

• Neighbor Search Strategies: jFoF implements two complementary neighbor-search algorithms optimized for GPU execution:
  1. $k$-d Tree: A hierarchical spatial partitioning method adapted for GPUs using flat arrays and fixed shapes to avoid recursion and pointer indirection. It utilizes the jaxkd library, offering both pure JAX and CUDA-compiled versions.
  2. Linked-Cell Grid: A method that discretizes the simulation volume into a uniform grid of cubic cells. This approach restricts neighbor searches to a local stencil of adjacent cells, reducing complexity from $O(N^2)$ to $O(N)$ and naturally handling periodic boundary conditions via minimum-image wrapping.
• Group Construction: Both methods employ an iterative label-propagation scheme. Particles adopt the minimum label among themselves and their linked neighbors, converging to connected components that define the halos.
• Gradient Propagation: To address the non-differentiability of discrete clustering, jFoF introduces three modes for gradient estimation:
  1. Frozen Assignment: Halo membership is treated as fixed during backpropagation. Gradients flow only through continuous variables (e.g., particle positions), allowing for the differentiation of halo positions and properties defined as analytical functions of particle coordinates (e.g., "decorated" proxies like particle variance for mass).
  2. Topological Optimization (REINFORCE): A stochastic approach where the linking criterion is relaxed into a probabilistic Bernoulli process. Using the REINFORCE (score-function) estimator, gradients are propagated through the discrete topology of the halo graph. This allows the optimization of halo connectivity, membership, and number in expectation, bridging continuous fields with discrete catalogs.

Key Results

• Performance Benchmarks: The authors benchmarked jFoF against cFoF, a parallel CPU-based implementation using MPI. On simulations ranging from $64^3$ to $512^3$ particles, the GPU-native jFoF implementations (both $k$-d tree and linked-cell) achieved speedups of up to an order of magnitude in core-hours compared to the CPU baseline, particularly at large particle counts.
• Accuracy: The halo catalogs generated by jFoF show near-complete agreement with the standard cFoF implementation. Cross-correlation coefficients for halo positions exceed $0.995$ across all scales. Small discrepancies were noted only at the very low-mass end ($\leq 40$ particles), attributed to floating-point precision differences.
• Gradient Validation:
  • In frozen assignment tests, the authors demonstrated that gradients of halo positions and proxy mass properties (particle variance) could be computed efficiently. These gradients successfully guided a gradient-descent optimizer to recover a cosmological parameter ($\sigma_8$) from a fixed initial phase simulation, converging to the true value within a few function evaluations.
  • In topological optimization tests, the REINFORCE-based approach successfully reconstructed a target halo field distribution from random initial conditions. The optimizer learned to fragment halos outside the target region and merge particles within it, recovering large-scale structure morphology with high correlation coefficients.

Significance and Claims
The paper claims that jFoF establishes a foundation for differentiable structure identification in cosmology. By enabling end-to-end, gradient-based optimization of structure formation models within GPU-accelerated pipelines, jFoF allows for:

• Direct Inference: The ability to perform inference on cosmological parameters or subgrid physics models by differentiating directly through the halo finding process, rather than relying on approximate bias models or emulators.
• Unified Workflows: The integration of simulation, analysis, and inference into a single differentiable computational graph, facilitating joint constraints on physical parameters.
• Scalability: The demonstration that FoF halo finding can be reimplemented for modern accelerator architectures without sacrificing accuracy, paving the way for analysis of simulations exceeding $10^9$ particles.

The authors maintain a modest tone regarding the topological optimization approach, noting that while it provides an unbiased gradient estimate, it exhibits high variance and requires careful tuning of relaxation parameters and sampling strategies. They position jFoF not as a final solution for all dynamical problems but as a proof-of-concept tool that bridges the gap between discrete observational catalogs and continuous physical fields.