jFoF: GPU Cluster Finding with Gradient Propagation

Sintesi Tecnica di "jFoF: GPU Cluster Finding with Gradient Propagation"

Problematica
Le simulazioni cosmologiche si affidano pesantemente all'algoritmo Friends-of-Friends (FoF) per identificare gli aloni di materia oscura, i mattoni fondamentali della struttura cosmica. Sebbene efficace, le implementazioni tradizionali di FoF affrontano due limitazioni primarie negli odierni flussi di lavoro astrofisici:

1. Colli di bottiglia delle prestazioni: Le implementazioni standard spesso si affidano a strutture dati spaziali basate su CPU (ad esempio, alberi $k$-d) o framework a memoria distribuita MPI. Questi approcci comportano un overhead significativo dovuto ai trasferimenti dati host-device e faticano a sfruttare il massic l'parallelismo delle moderne architetture GPU, in particolare quando integrati in pipeline di simulazione end-to-end native per GPU (ad esempio, PMWD, GRAM-X).
2. Non differenziabilità: La ricerca degli aloni comporta operazioni discrete — sogliatura, selezione dei vicini e connettività dei grafi — che sono intrinsecamente non differenziabili. Ciò impedisce la propagazione dei gradienti attraverso la fase di identificazione dell'alone, ostacolando l'uso di metodi di ottimizzazione e inferenza basati sul gradiente (come il Monte Carlo Hamiltoniano) per calibrare i modelli cosmologici o ricostruire le condizioni iniziali direttamente dai cataloghi di aloni osservati.

Metodologia
Gli autori presentano jFoF, un cercatore di aloni Friends-of-Friends completamente nativo per GPU, implementato in JAX. Il codice è progettato per operare interamente sul dispositivo, eliminando i costosi trasferimenti dati CPU-GPU e consentendo l'integrazione in pipeline astrofisiche differenziabili.

• Strategie di ricerca dei vicini: jFoF implementa due algoritmi complementari di ricerca dei vicini ottimizzati per l'esecuzione su GPU:
  1. $k$-d Tree: Un metodo di partizionamento spaziale gerarchico adattato per le GPU utilizzando array piatti e forme fisse per evitare la ricorsione e l'indirezione dei puntatori. Utilizza la libreria jaxkd, offrendo versioni sia in puro JAX che compilate in CUDA.
  2. Griglia a celle collegate (Linked-Cell Grid): Un metodo che discretizza il volume di simulazione in una griglia uniforme di celle cubiche. Questo approccio restringe la ricerca dei vicini a uno stencil locale di celle adiacenti, riducendo la complessità da $O(N^2)$ a $O(N)$ e gestendo naturalmente le condizioni al contorno periodiche tramite l'avvolgimento dell'immagine minima (minimum-image wrapping).
• Costruzione dei gruppi: Entrambi i metodi impiegano uno schema iterativo di propagazione delle etichette. Le particelle adottano l'etichetta minima tra se stesse e i propri vicini collegati, convergendo verso componenti connesse che definiscono gli aloni.
• Propagazione del gradiente: Per affrontare la non differenziabilità del clustering discreto, jFoF introduce tre modalità per la stima del gradiente:
  1. Assegnazione congelata (Frozen Assignment): L'appartenenza all'alone è trattata come fissa durante la backpropagation. I gradienti fluiscono solo attraverso variabili continue (ad esempio, le posizioni delle particelle), consentendo la differenziazione delle posizioni e delle proprietà degli aloni definite come funzioni analitiche delle coordinate delle particelle (ad esempio, proxy "decorati" come la varianza delle particelle per la massa).
  2. Ottimizzazione topologica (REINFORCE): Un approccio stocastico in cui il criterio di collegamento viene rilassato in un processo Bernoulli probabilistico. Utilizzando lo stimatore REINFORCE (score-function), i gradienti vengono propagati attraverso la topologia discreta del grafo dell'alone. Ciò consente l'ottimizzazione della connettività, dell'appartenenza e del numero di aloni in senso atteso, colmando il divario tra campi continui e cataloghi discreti.

Risultati Chiave

• Benchmark delle prestazioni: Gli autori hanno confrontato jFoF con cFoF, un'implementazione parallela su CPU che utilizza MPI. Su simulazioni che vanno da $64^3$ a $512^3$ particelle, le implementazioni di jFoF native per GPU (sia $k$-d tree che cella collegata) hanno ottenuto accelerazioni fino a un ordine di grandezza in termini di ore-core rispetto alla baseline CPU, particolarmente con un elevato numero di particelle.
• Accuratezza: I cataloghi di aloni generati da jFoF mostrano un accordo quasi completo con l'implementazione standard cFoF. I coefficienti di correlazione incrociata per le posizioni degli aloni superano lo $0.995$ su tutte le scale. Piccole discrepanze sono state notate solo alla scala di massa molto bassa ($\leq 40$ particelle), attribuite alle differenze di precisione in virgola mobile.
• Validazione del gradiente:
  • Nei test di assegnazione congelata, gli autori hanno dimostrato che i gradienti delle posizioni degli aloni e delle proprietà proxy della massa (varianza delle particelle) potevano essere calcolati efficientemente. Questi gradienti hanno guidato con successo un ottimizzatore a discesa del gradiente per recuperare un parametro cosmologico ($\sigma_8$) da una simulazione a fase fissa, convergendo al valore reale entro poche valutazioni della funzione.
  • Nei test di ottimizzazione topologica, l'approccio basato su REINFORCE ha ricostruito con successo una distribuzione target di campi di aloni da condizioni iniziali casuali. L'ottimizzatore ha imparato a frammentare gli aloni al di fuori della regione target e a fondere le particelle al suo interno, recuperando la morfologia della struttura su larga scala con alti coefficienti di correlazione.

Significatività e Rivendicazioni
Il paper sostiene che jFoF stabilisca le fondamenta per l'identificazione della struttura differenziabile in cosmologia. Abilitando l'ottimizzazione end-to-end basata sul gradiente della formazione delle strutture all'interno di pipeline accelerate da GPU, jFoF permette:

• Inferenza Diretta: La capacità di eseguire l'inferenza su parametri cosmologici o modelli di fisica sub-grid differenziando direttamente attraverso il processo di ricerca degli aloni, anziché affidarsi a modelli di bias approssimativi o emulatori.
• Workflow Unificati: L'integrazione di simulazione, analisi e inferenza in un unico grafo computazionale differenziabile, facilitando i vincoli congiunti sui parametri fisici.
• Scalabilità: La dimostrazione che la ricerca degli aloni FoF può essere reimplementata per le moderne architetture acceleratrici senza sacrificare l'accuratezza, aprendo la strada all'analisi di simulazioni che superano i $10^9$ particelle.

Gli autori mantengono un tono modesto riguardo all'approccio di ottimizzazione topologica, notando che, sebbene fornisca una stima del gradiente non distorta, presenta un'alta varianza e richiede una calibrazione attenta dei parametri di rilassamento e delle strategie di campionamento. Posizionano jFoF non come una soluzione finale per tutti i problemi dinamici, ma come uno strumento di prova di concetto che colma il divario tra cataloghi osservativi discreti e campi fisici continui.