Emulating Cosmic Structure Formation with a Lagrangian Neural Cellular Automaton

Technische Samenvatting: Emulatie van Kosmische Structuurvorming met een Lagrangiaanse Neurale Cellulaire Automaat

Probleemstelling
De volgende generatie sterrenstelsel-roodverschuivingssurveys (bijv. DESI, Euclid, LSST) vereist field-level inferentiemethoden om de volledige fysieke informatie van kosmologische kaarten te extraheren, waarbij verder wordt gekeken dan standaard samenvattende statistieken zoals het vermogensspectrum. Dit inverse probleem vereist een voorwaarts model dat lineaire begincondities verbindt met niet-lineaire structuren in een latere tijd. Een dergelijk model moet drie concurrerende beperkingen voldoen:

1. Fysieke Nauwkeurigheid: Het moet de niet-lineaire dynamica correct modelleren, inclusief shell-crossing en halo-vorming, waar traditionele Lagrangiaanse Perturbatietheorie (LPT) faalt.
2. Computationele Efficiëntie: Het moet snel genoeg zijn om miljoenen bemonsteringsstappen te ondersteunen in Bayesiaanse inferentie-pipelines, in tegenstelling tot standaard $N$-body simulaties die te computationeel kostbaar zijn voor iteratieve lussen.
3. Differentieerbaarheid: Het moet volledig differentieerbaar zijn om gradiëntgebaseerde optimalisatie mogelijk te maken voor het reconstrueren van begincondities en kosmologische parameters.

Bestaande oplossingen kampen met afwegingen: $N$-body simulaties zijn accuraat maar traag; LPT en Effective Field Theory (EFT) modellen zijn snel maar breken af in het niet-lineaire regime; en standaard Euleriaanse Convolutionele Neurale Netwerken (CNN's) hebben moeite met massa-advectie over grote afstanden, missen Galileïsche invariantie en falen vaak in het bieden van tijdelijk continue trajecten die vereist zijn voor lichtkegelconstructie.

Methodologie: De Lagrangiaanse Neurale Cellulaire Automaton (LNCA)
De auteurs introduceren de Lagrangiaanse Neurale Cellulaire Automaton (LNCA), een hybride deep learning-framework dat structuurvorming emuleert als een lokaal, iteratief dynamisch proces op een bewegend rooster (comoving lattice).

• Lagrangiaans Kader & Vaste Topologie: In tegenstelling tot Euleriaanse CNN's die opereren op vaste ruimtelijke roosters, werkt de LNCA in het Lagrangiaanse kader. Het definieert een regelmatig Cartesiaans rooster van Lagrangiaanse coördinaten ($q$) waarbij de node-connectiviteit vaststaat tot de initiële dichtstbijzijnde buren. De nodes (die vloeistofelementen vertegenwoordigen) worden verplaatst volgens update-regels, maar de topologie van de buurt-graaf blijft onveranderlijk. Dit maakt het gebruik van hooggeoptimaliseerde 3D-convoluties mogelijk terwijl de continue deeltjestracking van graafgebaseerde netwerken behouden blijft.
• Residuele Leren: Om grootschalige consistentie en efficiëntie te waarborgen, leert het netwerk niet de totale zwaartekracht, maar leert het enkel de residuele verplaatsingscorrecties ten opzichte van de Zeldovich-benadering (ZA). De totale verplaatsing wordt gedecomposeerd als $\Psi(t) = D(t)\Psi_{IC} + \Psi_{res}(t)$, waarbij het netwerk $\Psi_{res}$ voorspelt.
• Equivariante Architectuur: Het model dwingt strikte $E(3)$ (translatie en rotatie) equivariantie af door constructie. Het opereert niet op ruwe toestandsvectoren, maar op een basis van geometrische invarianten afgeleid van de afgeleiden van het verplaatsingsveld (Laplacians, Hessians, vectorproducten).
  • Scalar Pool ($I_i$): Een set $E(3)$-invariante scalairen geconstrueerd uit gradiënten en dot-producten.
  • Vector Menu ($E_i$): Een basis van equivariante vectoren (verplaatsing, snelheid, latente vectoren).
  • De update-regel brengt de scalar pool in kaart naar coëfficiënten die de vector menu lineair moduleren.
• Kolmogorov-Arnold Netwerken (KANs): De kern van de updatefunctie $f_\theta$ wordt geïmplementeerd als een geconstrueerde KAN, waarbij activatiefuncties geleerde polynomen zijn in plaats van standaard ReLU's. Dit garandeert analytische continuïteit, interpreteerbaarheid en een significante reductie in het aantal leerbare parameters.
• Temporele Continuïteit: Het model voorspelt de tijdsafgeleiden van verplaatsing en latente velden, die worden geïntegreerd over continue tijdstappen met behulp van een differentieerbaar schema. Dit genereert volledige deeltjest trajecten in plaats van single-shot mappingen, wat de reconstructie van tussenliggende roodverschuivingen en lichtkegels mogelijk maakt.
• Training: Het model wordt getraind op de Quijote simulatie-suite met een multi-objective loss functie die fouten bestraft in per-deeltje verplaatsing, momentum, dichtheid in de reële ruimte en halo-populatie-momenten. Het is geconditioneerd op kosmologische parameters ($\Omega_m, \Omega_b, h, n_s, \sigma_8$) die naar alle nodes worden uitgezonden.

Belangrijkste Resultaten
De getrainde LNCA werd gevalideerd tegen $N$-body simulaties over een bereik van roodverschuivingen ($z \in [0, 3]$) en kosmologieën (inclusclusief held-out Latin Hypercube punten):

• Vermogensspectrum & Cross-correlatie: Het model bereikt een precisie op percentniveau in het materie-vermogensspectrum ($P(k)$) en de cross-correlatie ($r(k)$) diep in het niet-lineaire regime, specifiek voor $k \lesssim 0.5 \, h \text{Mpc}^{-1}$. Bij $k \sim 1 \, h \text{Mpc}^{-1}$ blijven de afwijkingen onder de 5%.
• Verplaatsingsnauwkeurigheid: De mediane verplaatsingsfout is ongeveer $0.2 \, \text{Mpc}/h$, waarbij 95% van de deeltjes fouten $< 2 \, \text{Mpc}/h$ heeft. Het model reconstrueert succesvol complexe niet-lineaire kenmerken zoals shell-crossing caustics en bubbelachtige voids.
• Halo-vorming: Het model herstelt nauwkeurig de abundantie en ruimtelijke distributie van massieve halo's. Hoewel er een licht tekort is aan laag-massa halo's (door de diffuse aard van post-shell-crossing componenten), komt de hoog-massa zijde overeen met de $N$-body waarheid binnen de Poisson-spreiding.
• Generalisatie: Het model generaliseert robuust naar kosmologieën die ver van de trainingsdistributie liggen, waarbij het een hoge getrouwheid behoudt zonder hertraining.

Betekenis en Claims
Het artikel beweert dat de LNCA een significante vooruitgang vertegenwoordigt in kosmologische emulatie door tegelijkertijd nauwkeurigheid, efficiëntie en fysieke beperkingen aan te pakken:

1. Parameter Efficiëntie: De LNCA bereikt een nauwkeurigheid die vergelijkbaar is met state-of-the-art U-Net emulatoren (die $\sim 2 \times 10^7$ gewichten vereisen) met slechts $\sim 4 \times 10^3$ leerbare parameters. Deze vier-orde-grootte reductie wordt toegeschreven aan de recurrente toepassing van een enkele, compacte, fysiek beperkte update-regel.
2. Localiteit en Schaalbaarheid: Door te vertrouwen op lokale updates in plaats van globale Fourier-transformaties (vereist voor Poisson-solvers in $N$-body codes), vermijdt de LNCA globale synchronisatiebarrières. Dit zorgt ervoor dat communicatie schaalt met het oppervlak in plaats van het volume, wat superieure schaalbaarheid biedt voor grote volumes.
3. Differentieerbaarheid en Lichtkegels: In tegenstelling tot snapshot-gebaseerde emulatoren produceert de LNCA een continu, differentieerbaar traject. Dit ondersteunt van nature de constructie van lichtkegel-outputs (het simuleren van observaties langs de verleden lichtkegel) zonder extra computationele stappen, een cruciaal kenmerk voor survey-analyse.
4. Interpreteerbaarheid: Het gebruik van KANs en een equivariante basis maakt het mogelijk om de interne dynamica te inspecteren als expliciete algebraïsche expressies over fysieke invarianten, wat een niveau van interpreteerbaarheid biedt dat afwezig is in "black-box" deep learning modellen.

De auteurs concluderen dat hoewel de LNCA niet noodzakelijkerwijs het snelste voorwaartse model is in termen van pure wall-clock tijd vergeleken met gespecialiseerde solvers zoals FastPM of pmwd, de unieke combinatie van localiteit, recursie, differentieerbaarheid en fysieke beperkingen het een betrouwbaar en schaalbaar instrument maakt voor de volgende generatie field-level kosmologische inferentie.