# CosmoFlow: Using Deep Learning to Learn the Universe at Scale

arXiv: https://arxiv.org/abs/1808.04728 · HTML: https://mlsanity.com/paper/1808.04728/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 14 Aug 2018.

## Built on by (most cited first, 14 of 14)

- [Neural Implicit Flow: a mesh-agnostic dimensionality reduction paradigm of spatio-temporal data](https://mlsanity.com/paper/2204.03216/) (2022)
- [MLPerf HPC: A Holistic Benchmark Suite for Scientific Machine Learning on HPC Systems](https://mlsanity.com/paper/2110.11466/) (2021)
- [Clairvoyant Prefetching for Distributed Machine Learning I/O](https://mlsanity.com/paper/2101.08734/) (2021)
- [HammingMesh: A Network Topology for Large-Scale Deep Learning](https://mlsanity.com/paper/2209.01346/) (2022)
- [Using Machine Learning at Scale in HPC Simulations with SmartSim: An Application to Ocean Climate Modeling](https://mlsanity.com/paper/2104.09355/) (2021)
- [A Comprehensive Performance Study of Large Language Models on Novel AI Accelerators](https://mlsanity.com/paper/2310.04607/) (2023)
- [RINAS: Training with Dataset Shuffling Can Be General and Fast](https://mlsanity.com/paper/2312.02368/) (2023)
- [I/O in Machine Learning Applications on HPC Systems: A 360-degree Survey](https://mlsanity.com/paper/2404.10386/) (2024)
- [An Oracle for Guiding Large-Scale Model/Hybrid Parallel Training of Convolutional Neural Networks](https://mlsanity.com/paper/2104.09075/) (2021)
- [The Disharmony between BN and ReLU Causes Gradient Explosion, but is Offset by the Correlation between Activations](https://mlsanity.com/paper/2304.11692/) (2023)
- [SOLAR: A Highly Optimized Data Loading Framework for Distributed Training of CNN-based Scientific Surrogates](https://mlsanity.com/paper/2211.00224/) (2022)
- [SAIH: A Scalable Evaluation Methodology for Understanding AI Performance Trend on HPC Systems](https://mlsanity.com/paper/2212.03410/) (2022)
- [GPU-centric Communication Schemes for HPC and ML Applications](https://mlsanity.com/paper/2503.24230/) (2025)
- [LUMOS: Democratizing SciML Workflows with L0-Regularized Learning for Unified Feature and Parameter Adaptation](https://mlsanity.com/paper/2602.22537/) (2026)

## Abstract (arXiv)

Deep learning is a promising tool to determine the physical model that describes our universe. To handle the considerable computational cost of this problem, we present CosmoFlow: a highly scalable deep learning application built on top of the TensorFlow framework. CosmoFlow uses efficient implementations of 3D convolution and pooling primitives, together with improvements in threading for many element-wise operations, to improve training performance on Intel(C) Xeon Phi(TM) processors. We also utilize the Cray PE Machine Learning Plugin for efficient scaling to multiple nodes. We demonstrate fully synchronous data-parallel training on 8192 nodes of Cori with 77% parallel efficiency, achieving 3.5 Pflop/s sustained performance. To our knowledge, this is the first large-scale science application of the TensorFlow framework at supercomputer scale with fully-synchronous training. These enhancements enable us to process large 3D dark matter distribution and predict the cosmological parameters OmegaM, σ8 and ns with unprecedented accuracy.