🔭 astrophysics

Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

The paper introduces the Lagrangian Neural Cellular Automaton (LNCA), a highly efficient and fully differentiable hybrid deep learning framework that emulates non-linear cosmic structure formation by iteratively learning residual displacement corrections to the Zeldovich approximation, thereby enabling accurate, percent-level reconstruction of cosmological initial conditions from galaxy surveys with significantly fewer parameters than existing models.

Original authors: Cooper Jacobus, Beatriz Tucci, Oliver Philcox

Published 2026-07-31

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Cooper Jacobus, Beatriz Tucci, Oliver Philcox

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

Imagine the universe as a giant, invisible ocean of dark matter, stretching across billions of light-years. In the very beginning, this ocean was calm and smooth, like a glassy pond. But over billions of years, gravity acted like a giant hand, pulling the water into swirling whirlpools, stretching it into long, thin threads, and crushing it into dense, clumpy islands where galaxies are born. This process is called "cosmic structure formation." Scientists want to understand exactly how this happened because the pattern of these clumps and threads holds the secret to the universe's history, its ingredients, and its ultimate fate.

To figure this out, researchers usually run massive computer simulations. They start with a smooth, computer-generated version of the early universe and let the laws of physics run the show, watching how gravity pulls everything together. However, these simulations are like trying to simulate every single water molecule in the ocean; they are incredibly accurate but take so much computing power that they are too slow to use for certain types of detective work. On the other hand, there are faster, simpler math shortcuts, but they are like using a sketch instead of a photograph—they miss the messy, tangled details where the real action happens. The big challenge is finding a way to get the accuracy of the heavy simulations with the speed of the simple shortcuts, all while being able to work backward from what we see today to guess what the universe looked like at the start.

This paper introduces a clever new tool called the Lagrangian Neural Cellular Automaton (LNCA) to solve this problem. Think of the universe not as a fixed grid of pixels (like a digital photo) but as a swarm of invisible, dancing particles. Traditional computer models try to update a fixed grid, which gets messy and blurry when things move fast. The LNCA, however, treats the universe like a game of "follow the leader" where the grid itself moves with the particles. It's like having a set of sticky notes that you can slide around on a table; instead of trying to paint the picture on a fixed wall, you move the notes to where the action is.

The researchers trained this AI on a massive library of super-accurate, but slow, universe simulations. They didn't ask the AI to learn the whole story from scratch. Instead, they taught it a specific trick: "Start with a simple, fast guess (called the Zeldovich approximation), and then just learn the tiny, messy corrections needed to make it perfect." It's like teaching a student the basic rules of a sport, and then only having them practice the tricky, high-speed moves that the basic rules miss.

The results are impressive. In their tests, the LNCA managed to recreate the complex, tangled web of the universe with a precision of about 1% in the non-linear regime (specifically for scales up to k ≲ 0.5 hMpc−1), which is where the universe gets really messy and clumpy. It did this while using roughly 10,000 times fewer learning parameters than other similar AI models. Perhaps most importantly, because the AI learns the "rules of motion" step-by-step rather than just guessing the final picture, it can generate a smooth, continuous history of the universe. This means scientists can pause the simulation at any moment to see what the universe looked like at any point in time, or even build a "lightcone" (a view of the universe as it would appear to an observer looking out from Earth), something that previous fast models struggled to do.

The authors show that this method works well not just for the specific universe they trained on, but also for a wide variety of different cosmic settings, suggesting it has learned the general laws of how gravity builds structure. While it doesn't quite match the absolute peak accuracy of the slowest, most expensive simulations at the very smallest scales, it comes remarkably close while being vastly more efficient and flexible. By combining the speed of a simple math shortcut with the brainpower of a neural network, the LNCA offers a new, powerful way to rewind the clock on the universe and understand how we got here.

Technical Summary: Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

Problem Statement
The next generation of galaxy redshift surveys (e.g., DESI, Euclid, LSST) demands field-level inference methods to extract the full physical information content of cosmological maps, moving beyond standard summary statistics like the power spectrum. This inverse problem requires a forward model that connects linear initial conditions to non-linear late-time structures. Such a model must satisfy three competing constraints:

Physical Accuracy: It must correctly model non-linear dynamics, including shell-crossing and halo formation, where traditional Lagrangian Perturbation Theory (LPT) fails.
Computational Efficiency: It must be fast enough to support millions of sampling steps in Bayesian inference pipelines, unlike standard $N$ -body simulations which are too computationally prohibitive for iterative loops.
Differentiability: It must be fully differentiable to enable gradient-based optimization for reconstructing initial conditions and cosmological parameters.

Existing solutions face trade-offs: $N$ -body simulations are accurate but slow; LPT and Effective Field Theory (EFT) models are fast but break down in the non-linear regime; and standard Eulerian Convolutional Neural Networks (CNNs) struggle with mass advection over large distances, lack Galilean invariance, and often fail to provide temporally continuous trajectories required for lightcone construction.

Methodology: The Lagrangian Neural Cellular Automaton (LNCA)
The authors introduce the Lagrangian Neural Cellular Automaton (LNCA), a hybrid deep learning framework that emulates structure formation as a local, iterative dynamical process on a comoving lattice.

Lagrangian Frame & Fixed Topology: Unlike Eulerian CNNs that operate on fixed spatial grids, the LNCA operates in the Lagrangian frame. It defines a regular Cartesian grid of Lagrangian coordinates ( $q$ ) where node connectivity is fixed to initial nearest neighbors. The nodes (representing fluid elements) are displaced according to update rules, but the neighborhood graph topology remains immutable. This allows the use of highly optimized 3D convolutions while preserving the continuous particle tracking of graph-based networks.
Residual Learning: To ensure large-scale consistency and efficiency, the network does not learn the total gravitational force. Instead, it learns only the residual displacement corrections to the Zeldovich Approximation (ZA). The total displacement is decomposed as $\Psi(t) = D(t)\Psi_{IC} + \Psi_{res}(t)$ , where the network predicts $\Psi_{res}$ .
Equivariant Architecture: The model enforces strict $E(3)$ $E (3)$ (translation and rotation) equivariance by construction. It does not operate on raw state vectors but on a basis of geometric invariants derived from the displacement field's derivatives (Laplacians, Hessians, vector products).
- Scalar Pool ( $I_i$ ): A set of $E(3)$ -invariant scalars constructed from gradients and dot products.
- Vector Menu ( $E_i$ ): A basis of equivariant vectors (displacement, velocity, latent vectors).
- The update rule maps the scalar pool to coefficients that linearly modulate the vector menu.
Kolmogorov-Arnold Networks (KANs): The core update function $f_\theta$ is implemented as a constrained KAN, where activation functions are learned polynomials rather than standard ReLUs. This ensures analytic continuity, interpretability, and a significant reduction in learnable parameters.
Temporal Continuity: The model predicts time derivatives of displacement and latent fields, which are integrated over continuous time steps using a differentiable scheme. This generates complete particle trajectories rather than single-shot mappings, enabling the reconstruction of intermediate redshifts and lightcones.
Training: The model is trained on the Quijote simulation suite using a multi-objective loss function that penalizes errors in per-particle displacement, momentum, real-space density, and halo population moments. It is conditioned on cosmological parameters ( $\Omega_m, \Omega_b, h, n_s, \sigma_8$ ) broadcast to all nodes.

Key Results
The trained LNCA was validated against $N$ -body simulations across a range of redshifts ( $z \in [0, 3]$ ) and cosmologies (including held-out Latin Hypercube points):

Power Spectrum & Cross-Correlation: The model achieves percent-level precision in the matter power spectrum ( $P(k)$ ) and cross-correlation ( $r(k)$ ) well into the non-linear regime, specifically for $k \lesssim 0.5 \, h \text{Mpc}^{-1}$ . At $k \sim 1 \, h \text{Mpc}^{-1}$ , deviations remain below 5%.
Displacement Accuracy: The median displacement error is approximately $0.2 \, \text{Mpc}/h$ , with 95% of particles having errors $< 2 \, \text{Mpc}/h$ . The model successfully reconstructs complex non-linear features like shell-crossing caustics and bubble-like voids.
Halo Formation: The model accurately recovers the abundance and spatial distribution of massive halos. While there is a slight deficit in low-mass halos (due to the diffuse nature of post-shell-crossing components), the high-mass end matches the $N$ -body truth within Poisson scatter.
Generalization: The model generalizes robustly to cosmologies far from the training distribution, maintaining high fidelity without retraining.

Significance and Claims
The paper claims that the LNCA represents a significant advance in cosmological emulation by simultaneously addressing accuracy, efficiency, and physical constraints:

Parameter Efficiency: The LNCA achieves accuracy comparable to state-of-the-art U-Net emulators (which require $\sim 2 \times 10^7$ weights) using only $\sim 4 \times 10^3$ learnable parameters. This four-order-of-magnitude reduction is attributed to the recurrent application of a single, compact, physically constrained update rule.
Locality and Scalability: By relying on local updates rather than global Fourier transforms (required for Poisson solvers in $N$ -body codes), the LNCA avoids global synchronization barriers. This allows communication to scale with surface area rather than volume, offering superior scalability for large volumes.
Differentiability and Lightcones: Unlike snapshot-based emulators, the LNCA produces a continuous, differentiable trajectory. This natively supports the construction of lightcone outputs (simulating observations along the past lightcone) without additional computational passes, a critical feature for survey analysis.
Interpretability: The use of KANs and an equivariant basis allows the internal dynamics to be inspected as explicit algebraic expressions over physical invariants, offering a level of interpretability absent in "black-box" deep learning models.

The authors conclude that while the LNCA may not be the fastest forward model in terms of raw wall-clock time compared to specialized solvers like FastPM or pmwd, its unique combination of locality, recurrence, differentiability, and physical constraints makes it a reliable and scalable tool for the next generation of field-level cosmological inference.

Technical Summary: Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

More like this