Scalable Dark Siren Cosmology with gwcosmo: GPU Acceleration, Validation and Systematics

Abstract: As the number of confident gravitational-wave detections grows, population-level hierarchical analyses face increasing computational costs. Dark-siren cosmological inference integrates over the localisation volume of each gravitational-wave source. To remain feasible without discarding information from the quieter but more numerous sources in the catalogue, significant efficiency improvements are vital for analysis pipelines. In this work, we present an upgraded version of the cosmological inference pipeline gwcosmo, which leverages vectorisation on graphics processing units to process the entire gravitational-wave catalogue in parallel with each iteration. This new implementation achieves a speed-up of 1000 times over the previous version, facilitating analyses of O5-like numbers of GW events on wall-clock timescales of hours. Our results demonstrate the scalability of the gwcosmo pipeline, specifically its ability to handle the increasing computational load of expanding event catalogues, positioning it as a vital tool for future advances in dark-siren cosmology.

• The paper presents a GPU-accelerated gwcosmo pipeline that reduces likelihood evaluation times by nearly three orders of magnitude.
• It leverages batched processing of 3D tensors and custom CUDA kernels to enhance throughput and maintain constant scaling with large event catalogs.
• Validation against CPU implementations demonstrates robust posterior consistency and effective downsampling with minimal impact on inference accuracy.

Scalable Dark Siren Cosmology with gwcosmo: GPU Acceleration, Validation and Systematics

Introduction and Context

The hierarchical Bayesian inference of cosmological parameters using gravitational-wave (GW) standard sirens represents a pivotal methodology for probing the expansion rate of the universe without recourse to traditional electromagnetic distance ladders. The gwcosmo framework has played a central role in enabling population and cosmological inference from compact binary coalescence (CBC) event catalogs, especially under the dark siren paradigm where electromagnetic counterparts are absent. However, the rapid growth in GW event detections presents a critical computational challenge: scaling traditional likelihood evaluation approaches becomes infeasible due to their serial and memory-intensive nature. The work "Scalable Dark Siren Cosmology with gwcosmo: GPU Acceleration, Validation and Systematics" (2605.23538) systematically addresses these scalability and performance bottlenecks, introducing a fundamentally re-engineered, GPU-accelerated version of gwcosmo.

Hierarchical Likelihood Reformulation and Parallelization

The core methodological development hinges on the hierarchical likelihood formulation for GW population data, which includes event- and pixel-level marginalizations, GW selection functions, and reweighting by astrophysical and cosmological priors. The standard CPU-based implementation evaluates the likelihood through nested event-pixel-sample loops, resulting in computational cost that scales as $$\mathcal{O}(N_\text{events} N_\text{pix} N_\text{samples})$$. This serial evaluation approach, visualized in Figure 1, limits practical analyses to modest catalog sizes due to prohibitive wall times and energy consumption.

Figure 1: The legacy likelihood evaluation in gwcosmo as a sequential, CPU-bound scheme over events, sky pixels, and posterior samples, with each computational element waiting on completion of its predecessors.

To break this bottleneck, the upgraded gwcosmo architecture leverages massively parallel graphics processing units (GPUs), refactoring the entire evaluation into a vectorized, batched process. Events, pixels, and samples are packed into three-dimensional dense tensors, padding as necessary to maintain memory and compute efficiency, and processed using custom CUDA kernels and PyTorch backends. This reimagined architecture eliminates all explicit Python loops in likelihood calculations, and parallelizes both KDE construction and marginalization over all events and pixels, as illustrated in Figure 2.

Figure 2: The GPU-accelerated approach where event, pixel, and sample computations are vectorized and dispatched to many parallel threads, yielding orders of magnitude speedup.

This transformation shifts the computational scaling toward a (near) constant cost with input size, up to hardware-GPU memory limits. Non-parallelized remainders and memory-bound components introduce modest scaling at the upper bounds of catalog size, but the practical gains are decisive.

Memory Management and Downsampling Approximations

A major constraint for any parallel implementation is GPU memory: the dense 3D tensor approach can be infeasible for large catalogs (thousands of events, many pixels, and posterior samples) due to VRAM limitations. The pipeline introduces an optional random downsampling of posterior samples per pixel, reducing memory loads while preserving statistical fidelity. Further, the likelihood can be evaluated in batches across events to fit hardware limits at the expense of some throughput. The systematic implications of such approximations are explicitly validated in the results section.

Quantitative Performance and Scalability

The paper presents rigorous benchmarking of likelihood evaluation timings for CPU and GPU backends. Figure 3 demonstrates a speedup of nearly three orders of magnitude (close to a factor of 1000) for the full likelihood evaluation—including both numerator and denominator—when using the GPU-accelerated pipeline on catalogs up to 2000 GW events. The legacy CPU scaling is strictly linear with event number, while GPU evaluations remain flat until hardware saturation.

Figure 3: Likelihood evaluation timings across catalog sizes for CPU and GPU versions; GPU cost remains roughly constant, while CPU cost grows linearly, particularly for selection effects.

Furthermore, Figure 4 highlights that by capping the maximum number of posterior samples per pixel, both computation times (upper panel) and memory utilization (lower panel) can be reduced by an additional factor of two for large catalogs, without significant impact on inference quality.

Figure 4: Top—wall time reduction from sample downsampling; bottom—downsampling significantly decreases peak GPU memory use.

Direct comparisons of final energy consumption and wall time show an order-of-magnitude reduction in total energy cost and elapsed time for the GPU implementation (15 hours and 10.5 kWh) relative to a 32-core CPU deployment (32 days and 230.4 kWh).

Validation and Posterior Consistency

To establish scientific reliability, the new implementation was validated through two critical experiments:

1. Simulated Population Recovery: Analysis of 2000 mock GW events confirms exact recovery of injected cosmological and population hyperparameters across fiducial and downsampled GPU settings. For instance, spectral siren analysis constrains Hubble constant $H_0=64.01^{+5.20}_{-5.30}$ (16% uncertainty) from this simulated catalog, with only a 21-hour GPU run time. Figure 5 presents the corner plot for the key parameters, showing consistency between full and downsampled posterior sample runs.

   Figure 5: Posterior constraints on cosmological and selected population hyperparameters from 2000 simulated GW events, with both 'all samples' and downsampled configurations recovering the true values.

2. Cross-Implementation Agreement: Comparison between legacy CPU and new GPU versions on the GWTC-4.0 catalog (141 real events, FullPop-4.0 model) indicates that the GPU posteriors are in excellent agreement with the CPU reference as measured by the marginal KL divergence, remaining well within bootstrap-derived noise thresholds. The downsampling option (max 2500 samples per pixel) leads to a minimal additional deviation but also matches CPU results. Figure 6 details these correspondence results.

   Figure 6: Direct comparison of marginalized posteriors for the Hubble parameter and population parameters—legacy CPU (blue) matches both GPU full (orange) and GPU-downsampled (green) runs.

Notably, the GPU analysis with all samples converges in four days (96 hours), and the downsampled configuration in just 15 hours, compared to 32 days on the CPU.

Systematics and Precision Tests

The availability of rapid evaluation enables comprehensive systematic validation. Figure 7 compiles the impact of varying key implementation hyperparameters:

• Cap on maximum samples per pixel: No significant degradation until aggressive reduction ($\ll 2500$ samples/pixel).
• KDE bandwidth choice: No observable effect on $H_0$ posterior.
• KDE grid resolution: Deviations emerge for grid size $\leq 128$ points; 256+ points are stable.
• CPU/GPU floating-point precision: Single precision is as precise as double but about twice as fast.
• Minimum $N_\text{eff,PE}$ stability criterion: Thresholds above 10 introduce significant variance and posterior shifts.

  Figure 7: Impact of pipeline configuration choices on Hubble parameter posterior (left) and residuals against fiducial (right); only aggressive downsampling or grid coarsening degrade results.

Across these tests, KL divergences with respect to the fiducial analysis remain below the empirically derived threshold for agreement unless extreme parameter values are adopted.

Theoretical and Practical Implications

The results establish that hierarchical Bayesian inference for GW cosmology can be rendered computationally tractable for expanding event catalogs with thousands of entries. The practical implications include:

• Enabling population-level and cosmological inference for current and anticipated future catalogs (e.g., O5, O6, ET/LISA era) that would have been otherwise intractable under CPU-based pipelines.
• Supporting rapid systematic studies, model selection, and robustness testing due to drastic reduction in wall time and energy expenditure.
• Providing a baseline for further developments such as ragged tensor representations for improved memory utilization, deeper integration of galaxy catalog incompleteness, and more complex multi-dimensional selection functions.

For the broader methodological landscape, the vectorized, parallelized architecture described outperforms or matches the best alternative pipelines (e.g., icarogw [icarogw2023], CHIMERA [borghi2024]) in both runtime and flexibility. Furthermore, the modular PyTorch acceleration provides a generalizable template for other hierarchical inference frameworks in astroparticle and cosmological data analysis.

Conclusion

This work presents a comprehensive redesign of the hierarchical Bayesian dark siren cosmology pipeline, delivering acceleration factors of over 1000, stable and validated posteriors, and robust systematics control for GW event catalogs of realistic O5 scope and beyond. As GW detection catalogs continue to expand, these developments render high-precision, model-flexible, and energy-efficient inference in gravitational-wave cosmology scientifically and computationally sustainable.