{"id":"86cb253f-9aa9-4a34-8880-212fa0483182","arxiv_id":"2510.26851","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"jFoF runs friends-of-friends halo finding entirely on GPUs in JAX and adds frozen-assignment and REINFORCE-based gradient modes for differentiable cosmology.","lead":"jFoF is a friends-of-friends halo finder that runs entirely on GPUs using JAX, and it adds three ways to push gradients through the discrete halo-finding step. It gives cosmologists a building block for end-to-end differentiable simulations, with an order-of-magnitude speedup over a 128-core CPU baseline in single-GPU tests that shrinks or reverses at the largest tested scale.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The 'fully equivalent to standard FoF' catalog claim (Algorithms 1–2, §2.2) is unproven: neighbor-list caps kmax/max_per_cell can silently drop true FoF links in overdense regions; no cap values or truncation analysis are reported, and r_c>0.995 doesn't test individual membership.","rationale":"The reader's weakest_assumption precisely identifies the cap-truncation issue in Algorithms 1–2, and I agree that this is the most load-bearing concern. The central claim has two parts: (i) jFoF attains order-of-magnitude speedups over CPU FoF, and (ii) it maintains consistent halo catalogs and enables gradient propagation. The speedup part is honestly benchmarked in part and has a disclosed slower 1024^3 case; the gradient part is explicitly presented as proof-of-concept with limitations. The catalog-equivalence part, however, is stated categorically ('fully equivalent', 'exact equivalence') and is not independently verified at the level needed. The cap-truncation issue is correct in the direction of failure—exactly where FoF matters most, in overdense halos—and the reported validation is too coarse to rule it out. This does not warrant rejection: the code is public, the caps are likely tunable, and a cap audit could settle the question. But it does warrant maintaining the CONDITIONAL verdict until the cap behavior is characterized. I do not see a need to move beyond the reader's assessment, hence UNCHANGED.","tokens_in":18091,"tokens_out":4966,"duration_ms":45811,"concrete_test":"Run jFoF-LC (Algorithm 2) on the 512^3 fastPM box with max_per_cell and kmax swept over {8, 16, 32, 64, 128, 256} and compare the resulting per-particle halo labels to a no-cap exact reference (e.g., serial cFoF or a direct scipy cKDTree range query with no neighbor cap). Compute the fraction of particles whose connected-component label changes at each cap setting, and report the empirical distribution of within-l_link neighbor counts for the densest 1% of particles. If any label differences are nonzero or the neighbor-count tail exceeds kmax, the caps are binding and the 'fully equivalent' claim fails.","verdict_should_be":"UNCHANGED","load_bearing_attack":"Section 2.2 states that the linked-cell grid maintains 'exact equivalence with the canonical FoF connectivity criterion,' and the abstract promises 'consistent halo catalogs.' But Algorithm 2 caps candidate lists at max_per_cell per neighbor cell and kmax filtered neighbors, while Algorithm 1 queries only the kmax nearest neighbours. In a dense halo core, the number of particles within the linking length l_link can exceed any fixed cap, so true FoF edges are silently discarded. Because FoF connected components depend on the full edge set, dropping edges can split or merge halos; the claimed equivalence holds only if the caps never bind. The paper reports no cap values, no distribution of neighbor counts, and no sensitivity analysis. The validation (r_c > 0.995 on a coarse mesh and 'small differences at the very low mass end') is a field-level correlation coefficient, not a per-particle label comparison, so it cannot detect this truncation. This is the load-bearing soft spot in the central claim: the speedups may be real, but the 'consistent halo catalogs' underpinning the scientific claim are not established unless caps provably do not truncate true friends.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The manuscript presents jFoF, a JAX-based GPU-native Friends-of-Friends halo finder. Two neighbor-search strategies are implemented — a k-d tree and a linked-cell grid — followed by iterative label propagation and group compaction. A slab-parallel extension is also described. The authors benchmark jFoF against the C-based MPI halo finder cFoF on fastPM boxes up to 512^3 particles, reporting up to an order-of-magnitude speedup and catalog agreement measured by a field-level cross-correlation r_c > 0.995. A 1024^3 parallel run is disclosed as slower on GPU than CPU. For differentiability, jFoF provides frozen-assignment gradients, a decorated variant using particle variance as a mass proxy, and a REINFORCE-based stochastic-edge topological mode. Demonstrations include σ8 field-level inference, halo power-spectrum derivative estimation, and a toy halo-assembly reconstruction.","tokens_in":18310,"tokens_out":7775,"duration_ms":76442,"significance":"If the central claims are substantiated, this is a valuable contribution to GPU-resident cosmological analysis and differentiable inference. The code is public, the REINFORCE identity (Eq. 6) is textbook-correct, and the separation of an exact forward pass from approximate backward gradients is a sensible design. The main barriers are the unproven catalog-equivalence claim under bounded neighbor lists and the hardware-unfair benchmark framing. The paper's ability to back-propagate through discrete halo operations is plausible but needs a more precise estimator statement and stronger validation.","major_comments":[{"comment":"The text states the linked-cell formulation maintains 'exact equivalence with the canonical FoF connectivity criterion,' but Algorithm 2 caps candidate lists at 27×max_per_cell and then keeps only kmax filtered neighbors per particle; Algorithm 1 likewise queries only the kmax nearest neighbors. In overdense halo cores, a true friend beyond these caps is silently discarded, and FoF components depend on the full transitive closure of the edge set, so dropped edges can split or merge halos. The paper reports no cap values used in the benchmarks, no distribution of neighbor counts, and no sensitivity analysis. The validation via Eq. (2) (r_c > 0.995 on a coarse mesh) is a field-level statistic and cannot detect per-particle mislabeling, especially the acknowledged 'small differences at the very low mass end.' This undermines the abstract's 'consistent halo catalogs' claim. Please report the","section":"§2.2, Algorithm 2 (also Algorithm 1)"},{"comment":"The headline 'order-of-magnitude speedup' compares jFoF on a single GPU to cFoF on 128 CPU cores, and the y-axis of Figure 2 is labeled 'core time,' not wall-clock time. This is not an equal-hardware or standard production-cost comparison. The paper's own largest test, the 1024^3 run in §2.3, shows jFoF slower than cFoF (4038.0 GPU-seconds vs 3218.8 CPU-seconds). Please qualify the abstract accordingly, report wall-clock timings and hardware specifications, and, if possible, compare against a GPU FoF baseline or a realistic CPU core count. The current presentation overstates the generality of the speed advantage.","section":"§2.3, Figure 2"},{"comment":"The REINFORCE estimator in Eq. (7) is unbiased for ∇θ E[L(y)] only when L does not depend on θ except through the sampled y. In the field-level demonstrations, however, the loss is evaluated on painted halo fields whose values depend on halo positions, and those positions are explicit functions of particle positions θ. If the deterministic pathwise dependence is not also included in the gradient, the implemented estimator is biased/incomplete. The paper should state the full estimator — including any direct ∇θ terms for the deterministic parts of the loss — and verify the topological-mode gradients against a finite-difference check on a field-level target. The current finite-difference comparisons in Figures 5 and 7 are for frozen assignment, not for the topological mode.","section":"§3.2, Eq. (7), §3.2.2"}],"minor_comments":[{"comment":"The k-d tree version is marked non-periodic, yet it is compared against cFoF on periodic fastPM boxes. Please clarify how periodicity is handled (or why it is not needed for the benchmark), and whether the r_c comparison includes only the linked-cell/jFoF-LC results for periodic cases.","section":"§2.1, Algorithm 1"},{"comment":"The 6D distance metric for subhalo finding is written with ambiguous notation. Please define the vector norms and the normalizations σ_x and σ_v, which appear without specification.","section":"Eq. (3)"},{"comment":"The text notes a 'systematic offset' between the autodiff and finite-difference derivatives. Please quantify this offset and discuss its origin more explicitly; as presented, the reader cannot judge whether the offset is physically expected or a limitation of the frozen-assignment approximation.","section":"Figures 5 and 7"},{"comment":"The figure axis label 'core time (sec)' is ambiguous. State whether this is wall-clock time, core-seconds, or another metric, and define it in the caption.","section":"General"}],"recommendation":"major_revision","confidential_remarks":"The paper addresses an important problem and the public code is a strength. However, the exact-equivalence claim is currently unsupported by the bounded-neighbor algorithms, and the performance claims need reframing after the disclosed slower 1024^3 case. These are fixable with additional analysis and careful reporting, so I recommend major revision rather than rejection. The REINFORCE estimator issue also needs a precise statement of the full gradient estimator."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"jFoF is a useful, honest paper. The genuinely new thing is the packaged combination: a JAX-native k-d tree and linked-cell FoF plus three gradient modes (frozen assignment, decorated frozen, REINFORCE) in one public tool. The REINFORCE estimator is textbook-correct, the code runs, and the catalog comparison to cFoF is done against an external, independently implemented baseline. r_c > 0.995 is respectable. The frozen-assignment demo converging on sigma8 in a few steps is a nice proof of concept.\n\nThe soft spot that bothers me is the 'fully equivalent' catalog claim. Algorithm 2 caps candidates at max_per_cell per cell and kmax filtered neighbors; Algorithm 1 queries only the kmax nearest neighbors. In a dense halo core, the true number of friends within l_link can exceed any fixed cap, so real FoF edges are silently dropped and the connected components can change — splitting or merging halos. The paper reports no cap values, no distribution of neighbor counts, and no truncation analysis. The r_c > 0.995 test is a coarse field-level correlation, not a per-particle label comparison, so it can't catch this. This isn't a fatal flaw — caps can be set generously or justified — but it is untested, and it sits under the paper's central scientific claim.\n\nThe speedup claim needs boundary conditions too. The 'order of magnitude' compares one GPU to 128 MPI CPU cores. No GPU FoF baseline is tested, and the paper's own 1024^3 parallel run is slower than cFoF (4038 GPU-seconds vs 3218.8 CPU-seconds). That is honestly reported, but it means the abstract's 'up to an order-of-magnitude speedup' overstates what is actually established.\n\nThe gradient modes are clearly marked as proofs of concept, which I appreciate. The frozen-assignment derivatives show a systematic offset from finite differences, the mass decoration is calibrated on the same boxes it is validated against, and the REINFORCE convergence is fragile — the authors say so themselves in Section 4. None of this is hidden; it just means these modes are not yet drop-in tools.\n\nMy take: this deserves peer review. The code is public, the core idea is solid, and the limitations are addressable. I would ask referees to require (1) cap values and a truncation analysis or a proof that caps never bind, (2) a GPU-vs-GPU or matched-resource benchmark, (3) a per-particle catalog comparison rather than only r_c, and (4) at least a toy-problem variance characterization for REINFORCE. With those, I would be comfortable citing jFoF as the differentiable FoF baseline.","headline":"jFoF is a real contribution — GPU-native FoF with gradient modes — but the 'equivalent to canonical FoF' claim rests on untested neighbor caps, and the speedup needs a same-baseline comparison.","tokens_in":18983,"tokens_out":2481,"would_cite":true,"duration_ms":22678,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"A GPU-native Friends-of-Friends halo finder claims up to a tenfold speedup over CPU codes and opens gradient paths through discrete clustering.","keywords":["Friends-of-Friends","GPU halo finding","differentiable simulation","automatic differentiation","REINFORCE gradient estimator","halo catalogs","field-level inference","cosmological structure"],"falsifier":"Run jFoF on a dense synthetic or simulated cluster where at least one particle has more than kmax true neighbors within the linking length, then compare the resulting connected components to an exact, uncapped union-find FoF; if the components differ, the 'consistent/canonical' catalog claim fails. Reporting the cap values used in the benchmarks and counting how often particles hit the caps would also settle the severity of truncation.","tokens_in":17768,"feed_emoji":"🌌","tokens_out":4411,"duration_ms":51470,"temperature":0.7,"pith_summary":"The paper presents a halo finder built to run entirely on a GPU, eliminating the usual CPU–GPU data transfers in cosmological simulation analysis. It claims two advances: the GPU code produces halo catalogs that match a standard CPU Friends-of-Friends implementation closely while running up to an order of magnitude faster, and it makes the discrete group-finding operation differentiable. Differentiability is achieved in three modes: frozen assignment, where halo membership is fixed during back-propagation; decorated frozen assignment, where an analytic property proxy like particle variance stands in for halo mass; and topological optimization, where halo links are treated as stochastic edges and gradients are estimated with a REINFORCE rule. If these claims hold, cosmological inference pipelines could optimize initial conditions or model parameters all the way from the simulation to the final halo catalog with gradient descent instead of derivative-free sampling.","feed_headline":"GPU halo finder runs 10x faster and learns gradients","feed_subtitle":"Keeps halo finding on the GPU and lets gradients flow through discrete cluster grouping for end-to-end inference.","key_machinery":"The mechanical core is the neighbor-search plus label-propagation pipeline: particles are sorted into either a k-d tree or a periodic linked-cell grid, each particle retrieves a bounded candidate list (capped by parameters kmax and max_per_cell), filters candidates by the linking length, and then iteratively adopts the minimum label among linked neighbors until convergence. For differentiability, the paper introduces a stochastic edge model in which each candidate link is a Bernoulli variable with probability sigma(alpha * (b^2 - d^2)/b^2), so that the expected halo connectivity is a smooth function of particle positions; a REINFORCE estimator then turns sampled graph connectivity into unbia","core_discovery":"The central claim is that a fully GPU-resident implementation of Friends-of-Finds, using either a flat-array k-d tree or a linked-cell grid for candidate neighbor search and iterative min-label propagation for group assembly, produces halo catalogs effectively equivalent to a canonical CPU FoF output (cross-correlation coefficient exceeding 0.995, with small differences only reported at the lowest masses) while achieving up to an order-of-magnitude speedup on a single GPU relative to a 128-core MPI CPU baseline. On top of performance, the paper claims gradients can be propagated through discrete halo finding: positions via frozen assignment, mass proxies via a decorated assignment, and halo","pith_inferences":["The order-of-magnitude speedup is demonstrated on a single GPU versus a 128-core CPU baseline; the paper itself reports that a 1024^3 slab-parallel run was slower than the CPU code, so the speed claim is regime-dependent and likely weakens at high particle counts or in multi-GPU configurations without further optimization.","The catalog-equivalence claim is only as strong as the fixed caps on neighbor candidates; in dense halo cores where friend counts exceed kmax or max_per_cell, connectivity would be silently altered, so 'equivalent' should be read as statistically consistent (r_c > 0.995) rather than particle-identical.","Because REINFORCE gradients are unbiased only in expectation across many samples, practical use in Hamiltonian Monte Carlo or high-dimensional inference will likely need variance reduction (baselines, multiple samples, annealing) and careful tuning of alpha and the linking length, which the paper acknowledges as a convergence fragility.","The decorated frozen assignment depends on a calibration between particle variance and halo mass on a specific simulation; applying it to a different resolution, redshift, or cosmology would require revalidating that surrogate relationship before the gradients are trustworthy."],"forward_implications":["Halo finding can be co-located with a GPU simulation, removing expensive host–device movement and enabling in-situ halo catalogs in GPU-accelerated pipelines.","Halo position gradients from frozen assignment are stable enough to drive gradient-descent inference of a cosmological parameter (e.g., sigma_8) at fixed initial phases, converging in a few evaluations.","Decorated frozen assignment yields low-variance gradients of the halo mass function with respect to sigma_8, matching finite-difference trends with less noise.","The REINFORCE-based topological mode allows optimization of halo number, membership, and connectivity in expectation, demonstrated by reconstructing a target halo field from random initial conditions in both a 2D toy and a 3D cosmological setup.","The forward pass remains exact standard FoF; only the backward pass is approximate, so the forward model itself is unchanged when used inside gradient-based samplers or optimizers.","The k-d tree FoF extends naturally to 6D phase-space subhalo finding by adding a velocity term to the distance metric, with qualitative agreement on substructure morphology."],"fun_headline_variants":["GPU halo finder: 10x faster, gradients flow through FoF","jFoF: 10x GPU speedup for halo finding, now differentiable","Differentiable FoF on GPU: 10x faster, JAX-based","10x faster GPU FoF with gradient flow for halo finding"],"cache_read_input_tokens":2304,"weakest_assumption_plain":"The claim that jFoF produces 'fully equivalent' FoF catalogs assumes that the fixed caps on how many neighbor candidates each particle can return never exclude a true friend; in dense halo cores, where friend counts can exceed any fixed cap, the algorithm silently changes connectivity and stops being canonical FoF.","fun_headline_variants_meta":{"raw":{"variants":["GPU halo finder: 10x faster, gradients flow through FoF","jFoF: 10x GPU speedup for halo finding, now differentiable","Differentiable FoF on GPU: 10x faster, JAX-based","10x faster GPU FoF with gradient flow for halo finding"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000956,"raw_usage":{"total_tokens":3900,"prompt_tokens":720,"completion_tokens":3180,"prompt_tokens_details":{"cached_tokens":256},"prompt_cache_hit_tokens":256,"prompt_cache_miss_tokens":464,"completion_tokens_details":{"reasoning_tokens":3108}},"tokens_in":464,"tokens_out":3180,"duration_ms":20625,"temperature":1.0,"reasoning_tokens":3108,"cache_read_input_tokens":256,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-04T07:10:10.045662+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Run jFoF on a dense synthetic or simulated cluster where at least one particle has more than kmax true neighbors within the linking length, then compare the resulting connected components to an exact, uncapped union-find FoF; if the components differ, the 'consistent/canonical' catalog claim fails. Reporting the cap values used in the benchmarks and counting how often particles hit the caps would also settle the severity of truncation.","supporting_citations":[],"review_version":1}