jFoF: GPU Cluster Finding with Gradient Propagation

Technische Samenvatting van "jFoF: GPU Cluster Finding met Gradiëntpropagatie"

Probleemstelling
Kosmologische simulaties vertrouwen zwaar op het Friends-of-Friends (FoF) algoritme om donkere materie-halo's te identificeren, de fundamentele bouwstenen van de kosmische structuur. Hoewel effectief, kampen traditionele FoF-implementaties in moderne astrofysische workflows met twee primaire beperkingen:

1. Prestatieknelpunten: Standaardimplementaties vertrouwen vaak op CPU-gebaseerde ruimtelijke datastructuren (bijv. $k$-d bomen) of gedistribueerde geheugen MPI-frameworks. Deze benaderingen veroorzaken aanzienlijke overhead door host-device dataoverdrachten en hebben moeite om de massale parallelliteit van moderne GPU-architecturen uit te buiten, vooral wanneer ze zijn geïntegreerd in end-to-end GPU-native simulatiepipelines (bijv. PMWD, GRAM-X).
2. Niet-differentieerbaarheid: Halo-identificatie omvat discrete operaties—drempelwaarden, nabuurschapselectie en graafconnectiviteit—die inherent niet-differentieerbaar zijn. Dit verhindert de propagatie van gradiënten door de stap van halo-identificatie, wat het gebruik van gradiëntgebaseerde optimalisatie- en inferentiemethoden (zoals Hamiltonian Monte Carlo) belemmert om kosmologische modellen te kalibreren of initiële condities direct te reconstrueren vanuit geobserveerde halo-catalogi.

Methodologie
De auteurs presenteren jFoF, een volledig GPU-native Friends-of-Friends halo-finder geïmplementeerd in JAX. De code is ontworpen om volledig op het device te opereren, waardoor kostbare CPU-GPU dataoverdrachten worden geëlimineerd en integratie in differentieerbare astrofysische pipelines mogelijk wordt gemaakt.

• Strategieën voor Nabuurschapzoekopdrachten: jFoF implementeert twee complementaire nabuurschapzoekalgoritmen die geoptimaliseerd zijn voor GPU-executie:
  1. $k$-d Boom: Een hiërarchische ruimtelijke partitioneringsmethode aangepast voor GPU's met behulp van platte arrays en vaste vormen om recursie en pointer-indirection te vermijden. Het maakt gebruik van de jaxkd bibliotheek, die zowel pure JAX- als CUDA-gecompileerde versies biedt.
  2. Linked-Cell Grid: Een methode die het simulatievolume discretiseert in een uniform rooster van kubieke cellen. Deze benadering beperkt nabuurschapzoekopdrachten tot een lokale stencil van aangrenzende cellen, wat de complexiteit reduceert van $O(N^2)$ naar $O(N)$ en periodieke randvoorwaarden natuurlijk afhandelt via minimum-image wrapping.
• Groepsvorming: Beide methoden maken gebruik van een iteratief label-propagatieschema. Deeltjes nemen het minimale label aan van zichzelf en hun gekoppelde buren, wat convergeert naar verbonden componenten die de halo's definiëren.
• Gradiëntpropagatie: Om de niet-differentieerbaarheid van discrete clustering aan te pakken, introduceert jFoF drie modi voor gradiëntschatting:
  1. Frozen Assignment: Halo-lidmaatschap wordt tijdens backpropagation als vast beschouwd. Gradiënten stromen alleen door continue variabelen (bijv. deeltjesposities), waardoor de differentiatie van halo-posities en eigenschappen die gedefinieerd zijn als analytische functies van deeltjescoördinaten (bijv. "gedecoreerde" proxies zoals deeltjesvariantie voor massa) mogelijk is.
  2. Topologische Optimalisatie (REINFORCE): Een stochastische benadering waarbij het koppelingscriterium wordt versoepeld tot een probabilistisch Bernoulli-proces. Met behulp van de REINFORCE (score-function) estimator worden gradiënten gepropageerd door de discrete topologie van de halo-graaf. Dit maakt de optimalisatie van halo-connectiviteit, lidmaatschap en aantal in verwachting mogelijk, waardoor continue velden worden overbrugd met discrete catalogi.

Belangrijkste Resultaten

• Prestatiebenchmarks: De auteurs hebben jFoF getest tegenover cFoF, een parallelle CPU-gebaseerde implementatie met gebruik van MPI. Op simulaties variërend van $64^3$ tot $512^3$ deeltjes bereikten de GPU-native jFoF-implementaties (zowel $k$-d boom als linked-cell) snelheden die tot een orde van grootte hoger liggen in core-uren vergeleken met de CPU-baseline, met name bij grote aantallen deeltjes.
• Nauwkeurigheid: De door jFoF gegenereerde halo-catalogi vertonen bijna volledige overeenstemming met de standaard cF-FoF implementatie. De cross-correlatiecoëfficiënten voor halo-posities overstijgen $0,995$ over alle schalen. Kleine discrepanties werden alleen opmerkt bij de zeer lage massa-zijde ($\leq 40$ deeltjes), toegeschreven aan verschillen in floating-point precisie.
• Gradiëntvalidatie:
  • In frozen assignment tests hebben de auteurs aangetoond dat gradiënten van halo-posities en proxy massa-eigenschappen (deeltjesvariantie) efficiënt berekend kunnen worden. Deze gradiënten leidden een gradiënt-afdaling optimizer succesvol naar het herstellen van een kosmologische parameter ($\sigma_8$) uit een vaste initiële fase-simulatie, waarbij de ware waarde binnen enkele functie-evaluaties werd bereikt.
  • In topologische optimalisatie tests slaagde de REINFORCE-gebaseerde benadering erin om een doel-halo veldverdeling te reconstrueren uit willekeurige initiële condities. De optimizer leerde halo's te fragmenteren buiten de doelregio en deeltjes te fuseren binnen deze regio, waarbij de grootschalige structuur morfologie met hoge correlatiecoëfficiënten werd hersteld.

Betekenis en Claims
Het artikel claimt dat jFoF een fundament legt voor differentieerbare structuuridentificatie in de kosmologie. Door end-to-end, gradiëntgebaseerde optimalisatie van structuurvormingsmodellen binnen GPU-versnelde pipelines mogelijk te maken, staat jFoF toe dat:

• Directe Inferentie: Het vermogen om inferentie uit te voeren op kosmologische parameters of subgrid fysica-modellen door direct te differentiëren door het halo-zoekproces, in plaats van te vertrouwen op benaderde bias-modellen of emulatoren.
• Verenigde Workflows: De integratie van simulatie, analyse en inferentie in één enkel differentieerbaar computationeel grafisch model, wat gezamenlijke beperkingen op fysieke parameters faciliteert.
• Schaalbaarheid: De demonstratie dat FoF halo-zoeken kan worden geherimplementeerd voor moderne accelerator-architecturen zonder nauwkeurigheid op te offeren, wat de weg vrijmaakt voor de analyse van simulaties die meer dan $10^9$ deeltjes bevatten.

De auteurs hanteren een bescheiden toon met betrekking tot de topologische optimalisatie benadering, waarbij zij opmerken dat hoewel het een onbevooroordeelde gradiëntschatting biedt, het een hoge variantie heeft en zorgvuldige afstemming van relaxatieparameters en samplingstrategieën vereist. Zij positioneren jFoF niet als een definitieve oplossing voor alle dynamische problemen, maar als een proof-of-concept tool die de kloof overbrugt tussen discrete observationele catalogi en continue fysieke velden.