CosmoFlow: Deep Learning for Cosmology

• CosmoFlow is a computational framework that combines deep generative models, scalable high-performance techniques, and Python tools to analyze cosmological fields and parameters.
• It employs flow-matching objectives and scale-aware latent masking to achieve high-fidelity reconstruction of CDM maps and efficient extraction of physical scales.
• High-performance training paradigms and ODE-based correlator methods enable precise modeling of primordial non-Gaussianity and robust cosmological parameter inference.

CosmoFlow refers to multiple advanced computational frameworks and algorithms for cosmology, spanning deep learning models for extracting cosmological parameters from large-scale structure simulations, scalable high-performance implementations for massive datasets, and Python packages for first-principles calculation of primordial cosmological correlators. The term encompasses distinct developments in both the physical modeling and data-driven inference of cosmological fields, including cold dark matter (CDM) maps and primordial non-Gaussianity.

1. Flow-Matching Generative Models for Cosmological Field Representation

CosmoFlow provides a generative, unsupervised framework that utilizes flow-matching objectives to learn compressed, interpretable latent representations of 2D cold dark matter fields without supervision. The model leverages a conditional flow-matching loss, seeking a time-dependent velocity field $u_\theta(x, t)$ that transforms samples from a Gaussian prior $p_0$ to the data distribution $p_1$ along straight-line interpolants $X_t = (1–t) X_0 + t X_1$. The objective is to minimize the mean squared error between the modeled velocity and the true displacement $(X_1 - X_0)$ across interpolated states:

$$L(\theta) = \mathbb{E}_{t, X_0, X_1} [\| u_\theta(X_t, t) - (X_1 - X_0) \|^2]$$

The network comprises an encoder $E_\phi$ that compresses $256\times256$ input CDM fields to $8\times16\times16$ latent maps and an 8-dimensional global summary, and a U-Net–based decoder $D_\theta$ that reconstructs sample fields or simulates field evolution by integrating the learned ODE $p_0$0 from $p_0$1 to $p_0$2. The architecture employs sinusoidal time embeddings, adaptive group normalization (conditioning on masks and summaries), and self-attention in key layers (Kannan et al., 16 Jul 2025).

2. Scale Awareness and Progressive Masking in Latent Space

A salient innovation in CosmoFlow’s architecture is the incorporation of scale-aware masking priors on the latent space. This progressive channel-wise masking enforces disentanglement of physical scales: at early integration times ($p_0$3), all latent channels are active, capturing large-scale (low-$p_0$4) structure, while at later times, channels are successively masked such that the final channel is forced to encode only the smallest-scale (high-$p_0$5) information. This results in each of the 8 latent channels specializing to different bands in the three-dimensional field power spectrum. The model thus compresses $p_0$6 pixels to $p_0$7 latents (a $p_0$8 reduction), with further reduction to an $p_0$9-dimensional summary supporting accurate low-dimensional inference (Kannan et al., 16 Jul 2025).

3. High-Performance Training Paradigms and Scaling

CosmoFlow has also designated a family of highly-optimized, parallelizable implementations for large-scale 3D convolutional neural networks on cosmological simulation data (Mathuriya et al., 2018, Oyama et al., 2020). The original application targets regression of cosmological parameters $p_1$0 from full 3D density fields, mapping $p_1$1 voxel subvolumes to parameter triplets using stacks of 3D convolution, pooling, and dense layers without batch normalization or augmentation (beyond duplication).

Key aspects of training at HPC scale include:

• Data-parallel synchronous stochastic gradient descent (SSGD) with global batch sizes up to $p_1$2 on tens of petaflops hardware.
• MPI-based all-reduce for gradient averaging, custom MKL-DNN kernels for optimized convolutions, and parallel I/O pipelines sustaining per-node bandwidth $p_1$3 MB/s.
• Enhanced scaling efficiency (up to $p_1$4 at $p_1$5 nodes), possible via the Cray PE Machine Learning Plugin and in-memory distributed sample caching.
• Empirically, parameter inference achieves mean relative errors of $p_1$6 ($p_1$7), $p_1$8 ($p_1$9), and $X_t = (1–t) X_0 + t X_1$0 ($X_t = (1–t) X_0 + t X_1$1), matching or exceeding Planck 2016 uncertainties (Mathuriya et al., 2018).

Extended hybrid-parallel strategies combine data and spatial parallelism to train on inputs up to $X_t = (1–t) X_0 + t X_1$2 voxels, reducing per-GPU memory demands from $X_t = (1–t) X_0 + t X_1$3 GiB to $X_t = (1–t) X_0 + t X_1$4 GiB and improving test mean squared error by an order of magnitude (Oyama et al., 2020).

4. Software for First-Principles Cosmological Correlator Calculation

A distinct usage of CosmoFlow names an open-source Python package automating the calculation of tree-level cosmological correlation functions via the cosmological flow (time-flow) method (Werth et al., 2024, Philcox et al., 24 Nov 2025). Operating at the level of canonical fields and conjugate momenta, the approach transforms the computation of equal-time $X_t = (1–t) X_0 + t X_1$5-point correlators from nested in-in integrals to closed ODE systems governed by Hamiltonian couplings.

Key features:

• Direct numerical integration of master flow equations for two- and three-point functions, with theory dependence encapsulated in Hamiltonian-derived $X_t = (1–t) X_0 + t X_1$6-tensors.
• Modularity: Users specify arbitrary field content, time-varying backgrounds, and interaction tensors for custom inflationary or late-time models through Python classes, abstracting numerics from physics.
• Cross-platform with only numpy, scipy, and matplotlib as core dependencies; parallelization supported via joblib or mpi4py.
• Capable of calculating non-factorizable primordial bispectra for models where analytic Schwinger–Keldysh (in-in) integrations are infeasible, enabling direct contact with CMB data pipelines (Philcox et al., 24 Nov 2025).

5. Downstream Applications and Performance Benchmarks

CosmoFlow’s generative model leverages learned representations for three principal downstream tasks (Kannan et al., 16 Jul 2025):

• Field-Level Reconstruction: Superior to variational autoencoders (VAEs), preserving $X_t = (1–t) X_0 + t X_1$7 of the power spectrum even at high $X_t = (1–t) X_0 + t X_1$8 and achieving $X_t = (1–t) X_0 + t X_1$9 lower MSE and $(X_1 - X_0)$0 higher SSIM scores relative to VAE baselines.
• Synthetic Sample Generation: Samples for unobserved cosmological parameter combinations maintain power spectra fidelity to within $(X_1 - X_0)$1 over $(X_1 - X_0)$2Mpc.
• Parameter Inference: Through an attached single-layer regressor on the $(X_1 - X_0)$3-dimensional summary, achieves mean relative errors of $(X_1 - X_0)$4 ($(X_1 - X_0)$5) and $(X_1 - X_0)$6 ($(X_1 - X_0)$7) with $(X_1 - X_0)$8 latents ($(X_1 - X_0)$9, $$L(\theta) = \mathbb{E}_{t, X_0, X_1} [\| u_\theta(X_t, t) - (X_1 - X_0) \|^2]$$0 with $$L(\theta) = \mathbb{E}_{t, X_0, X_1} [\| u_\theta(X_t, t) - (X_1 - X_0) \|^2]$$1 latents, though with less scale disentanglement).

The Python CosmoFlow for correlator calculation enables robust, high-fidelity bispectrum evaluation (e.g., for collider-type inflation models), with Cython-accelerated core routines providing $$L(\theta) = \mathbb{E}_{t, X_0, X_1} [\| u_\theta(X_t, t) - (X_1 - X_0) \|^2]$$2 speed-up, hierarchical neural-network–based separable approximations reducing CMB estimator complexity from thousands to just three MLP basis functions (with $$L(\theta) = \mathbb{E}_{t, X_0, X_1} [\| u_\theta(X_t, t) - (X_1 - X_0) \|^2]$$3 cosine correlation), and direct interfacing to PolySpec for CMB bispectrum constraints (Werth et al., 2024, Philcox et al., 24 Nov 2025).

6. Integration with Cosmology Pipelines and Comparative Context

CosmoFlow models and software enable several novel scientific pipelines:

• Compression and semantic understanding of high-resolution field data for fast simulation and analysis (enabling survey generation and inference with compressed latents).
• Parameter inference directly from simulated or observed maps without the need for summary statistics, leveraging the expressive power of deep 3D CNNs and unsupervised representations.
• Ab initio calculation of primordial correlators for arbitrary inflationary or multi-field models, relevant for constraining early-universe physics through CMB and large-scale structure observations.
• In CMB non-Gaussianity analysis, CosmoFlow–derived bispectra supplement or replace analytic templates that may be unavailable for strongly-mixed or featureful scenarios, as in collider models.

Compared to baselines:

• VAEs capture large-scale structure but lose small-scale fidelity.
• ResNet regressors provide strong parameter inference if trained directly on raw data, but lack generative or compression capabilities.
• Traditional modal bispectrum decompositions require thousands of basis functions for comparable accuracy, while CosmoFlow-based neural factorizations suffice with only three functions in non-factorizable cases (Philcox et al., 24 Nov 2025).

7. Limitations, Extensibility, and Future Directions

CosmoFlow’s flow-matching generative models are currently implemented for 2D fields; expansion to 3D, more realistic survey geometries, or additional physical fields (e.g., weak lensing) is tractable. At extreme HPC scale, convergence sensitivity to batch size and learning rate decay, as well as I/O bottlenecks, remain key challenges. In the ODE-based correlator package, computational time for highly squeezed configurations or large field-content models scales steeply but is mitigated by parallelization and modular extensions (e.g., to four-point functions or higher spins) (Werth et al., 2024).

A plausible implication is that CosmoFlow’s frameworks are positioned to unify data-driven and first-principles approaches in cosmological field and correlator analysis, providing both efficient inference pipelines and precise theoretical predictions for next-generation observational cosmology.