⚛️ phenomenology

OmniCosmos: Transferring Particle Physics Knowledge Across the Cosmos

This paper demonstrates that the OmniLearned foundation model, originally trained on collider physics data, can successfully generalize across scientific fields to improve predictions of cosmological parameters and halo/galaxy velocities in CosmoBench datasets, marking the first instance of a collider physics model transferring knowledge to cosmology.

Original authors: Vinicius Mikuni, Ibrahim Elsharkawy, Benjamin Nachman

Published 2026-09-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Vinicius Mikuni, Ibrahim Elsharkawy, Benjamin Nachman

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The universe is a vast, dark place, filled with invisible matter that holds galaxies together and a mysterious energy that pushes them apart. To understand how this cosmic drama unfolds, scientists rely on massive computer simulations that recreate the history of the cosmos, from the first moments after the Big Bang to the present day. These simulations generate enormous amounts of data, essentially creating digital universes populated by billions of points representing dark matter and galaxies. The challenge for researchers is to look at these complex, shifting clouds of points and figure out the specific rules that govern them—such as how much matter exists in the universe or how fast it is expanding. Traditionally, scientists have had to compress this rich information into simple summaries to make sense of it, but this approach throws away a lot of the detail that could lead to deeper discoveries.

A new study introduces a clever way to solve this problem by borrowing a tool from a completely different field of physics. Researchers have taken a powerful artificial intelligence model, originally designed to analyze the chaotic sprays of particles created when atoms collide at high speeds, and taught it to understand the structure of the universe. This model, which had already learned to recognize patterns in particle collisions, was adapted to read the "point clouds" of cosmological simulations. The team found that this borrowed intelligence could predict cosmic properties with remarkable accuracy, often needing far fewer examples to learn the rules than models built from scratch. This suggests that the mathematical patterns describing how particles interact in a collider share a deep, hidden similarity with the way gravity shapes the cosmos, allowing knowledge to jump across the divide between the very small and the very large.

The researchers, working with data from a benchmark called CosmoBench, tested their adapted model on two different sets of simulated universes. One set, known as CAMELS-SAM, focused on detailed galaxy formation within a specific volume, while the other, Quijote, covered a much larger area with slightly less detail. The goal was to see if the model could look at the positions of dark matter halos and galaxies and correctly guess the values of key cosmic parameters, such as the density of matter in the universe and the strength of its fluctuations. They also asked the model to predict the velocities of these halos, essentially guessing how fast and in what direction these invisible structures are moving.

The results were striking. When the team used the full amount of available training data, their adapted model outperformed previous methods, including other advanced machine learning approaches and traditional statistical tools. However, the true power of the model revealed itself when the amount of data was limited. In situations where only a hundred simulated universes were available for training, the adapted model matched or exceeded the performance of other methods that had been trained on thousands of examples. This efficiency is crucial because running these cosmic simulations is incredibly expensive in terms of computing power; being able to learn from fewer examples means scientists can achieve high precision without needing to generate as many digital universes.

To ensure that this success was due to the specific knowledge gained from particle physics and not just the architecture of the neural network itself, the researchers ran several control tests. They compared their model to one that was built with the exact same structure but trained from scratch on random data, and another that was pre-trained on a simple, synthetic dataset of points that had no connection to real physics. The model trained on particle collisions consistently outperformed both. The model trained from scratch required much more data to reach the same level of accuracy, and the one trained on the simple synthetic points performed poorly, often failing to learn the patterns at all. This confirmed that the model had not just learned a generic way to process lists of points, but had internalized specific, transferable insights about how objects in a cloud relate to one another.

The study also explored whether the model needed to be completely retrained to work on the new data. They found that they could keep the vast majority of the model's internal "brain" frozen, changing only a tiny fraction of its settings to fit the new task. Even with this minimal adjustment, the model retained its ability to make accurate predictions. This indicates that the fundamental geometric relationships the model learned while studying particle collisions—how particles cluster, how they move relative to each other, and how they form structures—are surprisingly similar to the relationships found in the distribution of galaxies. The researchers noted that while the physical laws governing particle collisions are different from those governing the expansion of the universe, the underlying mathematical structures of the data are compatible enough to allow this transfer of knowledge.

In practical terms, this approach offers a significant shortcut for cosmologists. Adapting the model to a new dataset took only a few hours of computing time on standard graphics processors, a fraction of the time required to train a new model from scratch or run the simulations themselves. The team demonstrated that for the Quijote simulations, the model could match the performance of the best existing benchmarks using less than ten percent of the available data. For the more complex CAMELS-SAM simulations, it improved upon all previous results using half the data. These findings suggest that foundation models, which are large, pre-trained systems capable of learning broad patterns, can be a powerful tool for scientific discovery across different fields. By leveraging the massive datasets already generated in particle physics, researchers can accelerate their understanding of the cosmos, potentially reducing the computational cost of future studies and opening new avenues for exploring the fundamental nature of our universe.

Technical Summary: OmniCosmos: Transferring Particle Physics Knowledge Across the Cosmos

Problem Statement
Current cosmological surveys aim to revolutionize the understanding of dark matter and dark energy by moving beyond low-dimensional summaries to field-level analysis. However, this approach faces two primary bottlenecks: the computational cost of simulating entire universes (limiting the availability of large synthetic datasets) and the need for machine learning architectures capable of handling variable-length sets of objects (point clouds) while remaining permutation-invariant. While transformer-based models have shown success in particle physics, they typically require massive datasets to leverage their capacity, a resource often unavailable in cosmology. Furthermore, existing foundation models in particle physics often rely on specialized tokenization that requires redefinition for cosmological data, complicating knowledge transfer.

Methodology
The authors propose OmniCosmos, a strategy to transfer knowledge from a particle physics foundation model, OmniLearned, to cosmological inference tasks. OmniLearned is a Point-Edge-Transformer (PET) pre-trained on one billion labeled and unlabeled hadronic jets from real and simulated particle collisions.

The adaptation process involves three key modifications to bridge the domain gap between collider jets and cosmological halos:

Input Representation Adaptation: While OmniLearned inputs are particle four-momenta, cosmological inputs are 3D spatial coordinates of dark matter halos. The authors replace the initial feature embedding with a strategy that includes:
- A global feature embedding mapping inputs to a latent space.
- A local embedding based on $k$ -nearest neighbors ( $k=10$ for CAMELS-SAM, $k=20$ for Quijote).
- New pairwise geometric features specific to cosmology: absolute Euclidean distance, cosine distance between vectors, and cosine distance between a halo and the vector difference to its neighbor.
Task Adaptation: The pre-training tasks (classification and generation) are replaced with regression tasks for cosmological parameters ( $\Omega_m$ and $\sigma_8$ ) and halo velocity prediction. The output layers are modified to match the dimensionality of the target parameters or velocity vectors.
Fine-Tuning Strategy: To prevent "catastrophic forgetting" of the geometric priors learned during pre-training while accommodating new layers, the authors employ a differential learning rate strategy. Newly initialized layers (input embeddings and output heads) are assigned a higher learning rate, while the pre-trained transformer weights are updated with a lower rate. In some ablation studies, transformer blocks are frozen entirely.

The model is evaluated on the CosmoBench datasets, specifically the CAMELS-SAM (semi-analytic galaxy formation, lower volume, higher resolution) and Quijote (larger volume, lower resolution) simulation suites.

Key Contributions

Cross-Domain Transfer: This work demonstrates, for the first time, that a foundation model trained on particle collider data can successfully generalize to cosmological data, despite the physical laws governing the two domains being distinct.
Efficient Adaptation: The authors show that only a negligible fraction of the model weights (0.09%) need to be randomly initialized and retrained, while the vast majority of the pre-trained knowledge is retained.
Architecture Validation: The study validates that a point-cloud architecture designed for unordered sets is effective for cosmological inference without requiring explicit ordering or new tokenization schemes.

Results
The performance of OmniCosmos was evaluated using the coefficient of determination ( $R^2$ ) for predicting $\Omega_m$ , $\sigma_8$ , and halo velocities.

Data Efficiency: OmniCosmos significantly outperforms models trained from scratch, particularly in low-data regimes. On the CAMELS-SAM dataset, OmniCosmos matches or exceeds previous benchmarks using only 100 training examples, whereas scratch-trained models require the full dataset to reach similar performance.
Parameter Regression: For the Quijote dataset, OmniCosmos matches the performance of the two-point correlation function (2PCF) for $\Omega_m$ and greatly improves upon the previous deep learning model (GNN) for both $\Omega_m$ and $\sigma_8$ , often using less than 10% of the available training data.
Velocity Prediction: OmniCosmos improves upon previous benchmarks for halo velocity prediction, surpassing linear least squares (LLS) models using less than 10% of the available dataset.
Ablation Studies:
- Frozen vs. Full Fine-tuning: Models with frozen transformer blocks performed similarly to fully fine-tuned models, suggesting the pre-trained representations are highly transferable.
- Synthetic Pre-training: A model pre-trained on a synthetic Gaussian point cloud dataset performed significantly worse than OmniCosmos and, in some cases, worse than a model trained from scratch. This indicates that the specific structure learned from particle physics data is crucial, and generic pre-training is insufficient.

Significance and Claims
The paper claims that OmniCosmos represents a breakthrough in cross-disciplinary scientific machine learning. By leveraging the "effective representation" of data learned from high-energy physics, the model achieves superior precision in cosmological parameter estimation with fewer simulations. The authors argue that this approach offers a practical solution to the computational bottleneck of cosmological simulations, potentially reducing the number of required simulations by an order of magnitude to achieve precision cosmology. The work establishes that the geometric and structural patterns learned from particle collisions are sufficiently universal to be applicable to the distribution of dark matter halos in the universe.

Technical Summary: OmniCosmos: Transferring Particle Physics Knowledge Across the Cosmos

More like this