# Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 50 references that are in the map.
- First posted on arXiv 29 Jul 2026.

## Builds on (most cited first, 11 of 11)

- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Cellular automata as convolutional neural networks](https://mlsanity.com/paper/1809.02942/) (2018)
- [Learning to Predict the Cosmological Structure Formation](https://mlsanity.com/paper/1811.06533/) (2018)
- [Learning Graph Cellular Automata](https://mlsanity.com/paper/2110.14237/) (2021)
- [Field Level Neural Network Emulator for Cosmological N-body Simulations](https://mlsanity.com/paper/2206.04594/) (2022)
- [Translation and Rotation Equivariant Normalizing Flow (TRENF) for Optimal Cosmological Analysis](https://mlsanity.com/paper/2202.05282/) (2022)
- [Multiscale Flow for Robust and Optimal Cosmological Analysis](https://mlsanity.com/paper/2306.04689/) (2023)
- [Neural cellular automata: applications to biology and beyond classical AI](https://mlsanity.com/paper/2509.11131/) (2025)
- [A volumetric deep Convolutional Neural Network for simulation of mock dark matter halo catalogues](https://mlsanity.com/paper/1805.04537/) (2018)
- [Predicting the Thermal Sunyaev-Zel'dovich Field using Modular and Equivariant Set-Based Neural Networks](https://mlsanity.com/paper/2203.00026/) (2022)
- [COmoving Computer Acceleration (COCA): N-body simulations in an emulated frame of reference](https://mlsanity.com/paper/2409.02154/) (2024)

## Abstract (arXiv)

Field-level inference of cosmological initial conditions from galaxy surveys requires a forward model that is simultaneously accurate in the non-linear regime, computationally efficient, and fully differentiable. Traditional N-body simulations are accurate but computationally prohibitive for iterative inference, while approximate solvers like Lagrangian Perturbation Theory (LPT) fail to capture the knotty halo-forming dynamics of the cosmic web at late times. We introduce the textitLagrangian Neural Cellular Automaton (LNCA), a hybrid deep learning framework that can be applied to emulate structure formation as a local, iterative dynamical process on a comoving lattice. Unlike convolutional emulators which map fixed grids in a single pass, the LNCA operates in the Lagrangian frame, iteratively advecting the computational graph itself to follow the flow of mass. By training the network to learn only the textitresidual displacement corrections to the Zeldovich approximation, we achieve high-fidelity emulation of the non-linear physics while guaranteeing accuracy at large scales. We further constrain our model to produce complete trajectories, not just final states, by adopting an equivariant cellular automaton architecture, which recurrently iterates on its internal states to yield a dynamic history. The resulting update rule acts on a single fixed neighborhood, is translationally and rotationally equivariant, and is integrated over discrete substeps to yield trajectories, making it a differentiable forward model for reconstructing the initial conditions of the universe from lightcone data. We train our model to emulate the real-space matter distributions of a diverse set of N-body simulations from the Quijote suite, for a range of redshifts and cosmologies, given only the initial conditions.