CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators

• URL: http://arxiv.org/abs/2305.06347v2
• Date: Thu, 22 Jun 2023 14:29:11 GMT
• Title: CosmoPower-JAX: high-dimensional Bayesian inference with differentiable cosmological emulators
• Abstract summary: We present CosmoPower-JAX, a JAX-based implementation of the CosmoPower framework, which accelerates cosmological inference. We show how, using the automatic differentiation, batch evaluation and just-in-time compilation features of JAX, parameter estimation can be accelerated by orders of magnitude. We show how CosmoPower-JAX provides converged posterior contours in 3 days, as opposed to the estimated 6 years required by standard methods.
• Score: 0.0
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We present CosmoPower-JAX, a JAX-based implementation of the CosmoPower framework, which accelerates cosmological inference by building neural emulators of cosmological power spectra. We show how, using the automatic differentiation, batch evaluation and just-in-time compilation features of JAX, and running the inference pipeline on graphics processing units (GPUs), parameter estimation can be accelerated by orders of magnitude with advanced gradient-based sampling techniques. These can be used to efficiently explore high-dimensional parameter spaces, such as those needed for the analysis of next-generation cosmological surveys. We showcase the accuracy and computational efficiency of CosmoPower-JAX on two simulated Stage IV configurations. We first consider a single survey performing a cosmic shear analysis totalling 37 model parameters. We validate the contours derived with CosmoPower-JAX and a Hamiltonian Monte Carlo sampler against those derived with a nested sampler and without emulators, obtaining a speed-up factor of $\mathcal{O}(10^3)$. We then consider a combination of three Stage IV surveys, each performing a joint cosmic shear and galaxy clustering (3x2pt) analysis, for a total of 157 model parameters. Even with such a high-dimensional parameter space, CosmoPower-JAX provides converged posterior contours in 3 days, as opposed to the estimated 6 years required by standard methods. CosmoPower-JAX is fully written in Python, and we make it publicly available to help the cosmological community meet the accuracy requirements set by next-generation surveys.

Related papers

• Cosmos 3: Omnimodal World Models for Physical AI [321.3056180868173]
  We introduce Cosmos 3, a family of omnimodal world models designed to jointly process and generate language, image, video, audio, and action sequences.<n>Our evaluation demonstrates that Cosmos 3 unifies a new state-of-the-art across a diverse suite of understanding and generation tasks.
  arXiv  Detail & Related papers  (2026-06-01T19:12:30Z)
• Infinite-World: Scaling Interactive World Models to 1000-Frame Horizons via Pose-Free Hierarchical Memory [101.2076718776139]
  We propose a robust interactive world model capable of maintaining coherent visual memory over 1000+ frames in complex real-world environments.<n>We introduce a Pose-free Memory (HPMC) that distills historical latents into a fixed-budget geometric representation.<n>We also propose an Uncertainty-aware Action Labeling module that discretizes continuous motion into a tri-state logic.
  arXiv  Detail & Related papers  (2026-02-02T17:52:56Z)
• Cosmo-FOLD: Fast generation and upscaling of field-level cosmological maps with overlap latent diffusion [1.9116784879310027]
  We introduce a novel technique to rapidly generate accurate and arbitrarily large cosmological and astrophysical 3-dimensional fields.<n>We are able to generate TNG300-2 dark matter density and gas temperature fields from a model trained only on 1% of the volume.<n>These results are obtained within a small fraction of the original simulation cost and produced on a single GPU.
  arXiv  Detail & Related papers  (2026-01-20T19:00:04Z)
• Dark Energy Survey Year 3 results: Simulation-based $w$CDM inference from weak lensing and galaxy clustering maps with deep learning. I. Analysis design [0.2854985824260583]
  This work presents the first simulation-based inference ( SBI) pipeline that combines weak lensing and galaxy clustering maps.<n>We develop a scalable forward model based on the CosmoGridV1 suite of N-body simulations to generate over one million self-consistent mock realizations of DES Y3.<n>We train deep graph convolutional neural networks on the full survey footprint in spherical geometry to learn low-dimensional features that approximately maximize mutual information with target parameters.
  arXiv  Detail & Related papers  (2025-11-06T18:59:59Z)
• CosmoBench: A Multiscale, Multiview, Multitask Cosmology Benchmark for Geometric Deep Learning [4.340305187316021]
  Cosmological simulations provide a wealth of data in the form of point clouds and directed trees.<n>A crucial goal is to extract insights from this data that shed light on the nature and composition of the Universe.<n>We introduce CosmoBench, a benchmark dataset curated from state-of-the-art cosmological simulations.
  arXiv  Detail & Related papers  (2025-07-04T16:46:25Z)
• Cosmos-Transfer1: Conditional World Generation with Adaptive Multimodal Control [97.98560001760126]
  We introduce Cosmos-Transfer, a conditional world generation model that can generate world simulations based on multiple spatial control inputs. We conduct evaluations to analyze the proposed model and demonstrate its applications for Physical AI, including robotics2Real and autonomous vehicle data enrichment.
  arXiv  Detail & Related papers  (2025-03-18T17:57:54Z)
• Predictable Scale: Part I -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining [56.58170370127227]
  We show that optimal learning rate follows a power-law relationship with both model parameters and data sizes, while optimal batch size scales primarily with data sizes. This work is the first work that unifies different model shapes and structures, such as Mixture-of-Experts models and dense transformers.
  arXiv  Detail & Related papers  (2025-03-06T18:58:29Z)
• Conditional Diffusion-Flow models for generating 3D cosmic density fields: applications to f(R) cosmologies [0.0]
  Next-generation galaxy surveys promise unprecedented precision in testing gravity at cosmological scales. We explore conditional generative modelling to create 3D dark matter density fields via score-based (diffusion) and flow-based methods.
  arXiv  Detail & Related papers  (2025-02-24T12:06:23Z)
• AIvaluateXR: An Evaluation Framework for on-Device AI in XR with Benchmarking Results [55.33807002543901]
  We present AIvaluateXR, a comprehensive evaluation framework for benchmarking large language models (LLMs) running on XR devices.<n>We deploy 17 selected LLMs across four XR platforms: Magic Leap 2, Meta Quest 3, Vivo X100s Pro, and Apple Vision Pro, and conduct an extensive evaluation.<n>We propose a unified evaluation method based on the 3D Optimality theory to select the optimal device-model pairs from quality and speed objectives.
  arXiv  Detail & Related papers  (2025-02-13T20:55:48Z)
• GPS-Gaussian+: Generalizable Pixel-wise 3D Gaussian Splatting for Real-Time Human-Scene Rendering from Sparse Views [67.34073368933814]
  We propose a generalizable Gaussian Splatting approach for high-resolution image rendering under a sparse-view camera setting. We train our Gaussian parameter regression module on human-only data or human-scene data, jointly with a depth estimation module to lift 2D parameter maps to 3D space. Experiments on several datasets demonstrate that our method outperforms state-of-the-art methods while achieving an exceeding rendering speed.
  arXiv  Detail & Related papers  (2024-11-18T08:18:44Z)
• $\mathtt{emuflow}$: Normalising Flows for Joint Cosmological Analysis [0.0]
  We show that normalising flows can be used to efficiently combine cosmological constraints from independent datasets. We show that the method is able to accurately describe the posterior distribution of real cosmological datasets. The resulting joint constraints can be obtained in a fraction of the time it would take to combine the same datasets at the level of their likelihoods.
  arXiv  Detail & Related papers  (2024-09-02T18:04:14Z)
• Outdoor Scene Extrapolation with Hierarchical Generative Cellular Automata [70.9375320609781]
  We aim to generate fine-grained 3D geometry from large-scale sparse LiDAR scans, abundantly captured by autonomous vehicles (AV) We propose hierarchical Generative Cellular Automata (hGCA), a spatially scalable 3D generative model, which grows geometry with local kernels following, in a coarse-to-fine manner, equipped with a light-weight planner to induce global consistency.
  arXiv  Detail & Related papers  (2024-06-12T14:56:56Z)
• The future of cosmological likelihood-based inference: accelerated high-dimensional parameter estimation and model comparison [5.79711608843279]
  We advocate for a new paradigm of cosmological likelihood-based inference, leveraging recent developments in machine learning and its underlying technology. We demonstrate the application of this paradigm on a simulated cosmic shear analysis for a Stage IV survey in 37- and 39-dimensional parameter spaces. We recover posterior contours and evidence estimates that are in excellent agreement with those computed by the traditional nested sampling approach.
  arXiv  Detail & Related papers  (2024-05-21T17:45:36Z)
• GPS-Gaussian: Generalizable Pixel-wise 3D Gaussian Splatting for Real-time Human Novel View Synthesis [70.24111297192057]
  We present a new approach, termed GPS-Gaussian, for synthesizing novel views of a character in a real-time manner. The proposed method enables 2K-resolution rendering under a sparse-view camera setting.
  arXiv  Detail & Related papers  (2023-12-04T18:59:55Z)
• A point cloud approach to generative modeling for galaxy surveys at the field level [0.5099081649205313]
  We introduce a diffusion-based generative model to describe the distribution of galaxies in our Universe. We demonstrate a first application to massive dark matter haloes in the Quijote simulation suite. This approach can be extended to enable a comprehensive analysis of cosmological data.
  arXiv  Detail & Related papers  (2023-11-28T19:00:00Z)
• ArchGym: An Open-Source Gymnasium for Machine Learning Assisted Architecture Design [52.57999109204569]
  ArchGym is an open-source framework that connects diverse search algorithms to architecture simulators. We evaluate ArchGym across multiple vanilla and domain-specific search algorithms in designing custom memory controller, deep neural network accelerators, and custom SOC for AR/VR workloads.
  arXiv  Detail & Related papers  (2023-06-15T06:41:23Z)
• Fast emulation of cosmological density fields based on dimensionality reduction and supervised machine-learning [0.0]
  We show that it is possible to perform fast dark matter density field emulations with competitive accuracy using simple machine-learning approaches. New density cubes for different cosmological parameters can be estimated without relying directly on new N-body simulations.
  arXiv  Detail & Related papers  (2023-04-12T18:29:26Z)
• Cosmology from Galaxy Redshift Surveys with PointNet [65.89809800010927]
  In cosmology, galaxy redshift surveys resemble such a permutation invariant collection of positions in space. We employ a textitPointNet-like neural network to regress the values of the cosmological parameters directly from point cloud data. Our implementation of PointNets can analyse inputs of $mathcalO(104) - mathcalO(105)$ galaxies at a time, which improves upon earlier work for this application by roughly two orders of magnitude.
  arXiv  Detail & Related papers  (2022-11-22T15:35:05Z)
• Satellite galaxy abundance dependency on cosmology in Magneticum simulations [101.18253437732933]
  We build an emulator of satellite abundance based on cosmological parameters. We find that $A$ and $beta$ depend on cosmological parameters, even if weakly. We also show that satellite abundance cosmology dependency differs between full-physics (FP) simulations, dark-matter only (DMO) and non-radiative simulations.
  arXiv  Detail & Related papers  (2021-10-11T18:00:02Z)
• Emulation of cosmological mass maps with conditional generative adversarial networks [0.0]
  We propose a novel conditional GAN model that is able to generate mass maps for any pair of matter density $Omega_m$ and matter clustering strength $sigma_8$. Our results show that our conditional GAN can interpolate efficiently within the space of simulated cosmologies. This contribution is a step towards building emulators of mass maps directly, capturing both the cosmological signal and its variability.
  arXiv  Detail & Related papers  (2020-04-17T09:34:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.