Vision Transformers and Graph Neural Networks for Charged Particle Tracking in the ATLAS Muon Spectrometer

Questo articolo presenta due approcci basati sull'apprendimento automatico per migliorare il tracciamento delle particelle cariche nello spettrometro a muoni di ATLAS: l'integrazione di reti neurali grafiche per il rifiuto dei segnali di fondo, che riduce i tempi di ricostruzione del 15%, e un prototipo di tracciamento end-to-end basato su Vision Transformer in grado di raggiungere un'efficienza del 98% in soli 2,3 ms.

L'essenziale

Immagina di essere in una stanza enorme e affollatissima, come uno stadio durante un concerto rock. C'è un'esplosione di suoni, luci e persone che si muovono in tutte le direzioni. Il tuo compito è trovare due amici specifici (i muoni) che stanno camminando insieme attraverso questa folla, distinguendoli da tutti gli altri rumori e movimenti casuali.

Questo è esattamente il problema che affronta il Large Hadron Collider (LHC) al CERN, e in particolare l'esperimento ATLAS. Con l'arrivo della nuova era "High-Luminosity" (dopo il 2030), la folla diventerà ancora più densa: invece di 60 collisioni al secondo, ne avremo fino a 200. Trovare i nostri "amici" (le particelle) in mezzo a così tanto caos diventa un incubo per i computer attuali.

Questo documento presenta due nuove idee, basate sull'Intelligenza Artificiale (AI), per risolvere questo problema in modo più veloce ed efficiente.

Ecco come funzionano, spiegate con parole semplici:

1. Il Primo Approccio: Il "Filtro Intelligente" (GNN)

L'Analogia: Il Controllore di Folla

Immagina che i dati grezzi che arrivano dai sensori siano un mucchio di biglietti da visita gettati a terra. La maggior parte sono falsi (rumore di fondo, errori elettronici, raggi cosmici), ma alcuni appartengono davvero ai tuoi amici.
Il sistema attuale prova a raccogliere tutti i biglietti e a metterli in ordine, il che richiede molto tempo e fatica.

Il primo metodo proposto usa una Rete Neurale a Grafo (GNN).

• Come funziona: Invece di guardare ogni singolo biglietto da solo, l'AI raggruppa i biglietti vicini in "mucchietti" (chiamati Bucket). Poi, usa la sua intelligenza per dire: "Ehi, questo mucchietto sembra solo rumore, buttalo via subito!".
• Il risultato: L'AI agisce come un controllore di folla super-veloce che rimuove il 97% dei biglietti falsi prima che il sistema principale inizi a lavorare.
• Il vantaggio: Il computer principale deve lavorare su molto meno materiale. Il tempo per analizzare un evento scende da 255 millisecondi a 217 millisecondi. È un miglioramento del 15%, che sembra poco, ma in fisica delle particelle significa poter salvare più eventi interessanti senza bloccare il sistema.

2. Il Secondo Approccio: L'"Occhio Magico" (Vision Transformers)

L'Analogia: Il Detective che guarda una foto intera

Il secondo metodo è ancora più audace. Immagina di non voler raccogliere i biglietti uno per uno, ma di voler guardare l'intera stanza come se fosse una fotografia e dire: "Vedo due linee che attraversano la stanza, quelle sono i miei amici!".

Questo metodo usa le Vision Transformers (ViT), la stessa tecnologia che fa funzionare le app di riconoscimento immagini sui tuoi smartphone o che permette alle auto a guida autonoma di vedere la strada.

• Come funziona: L'AI prende tutti i segnali (i "punti" lasciati dalle particelle) e li tratta come pixel di un'immagine. Usa un meccanismo chiamato "Attenzione" (come se l'AI potesse concentrarsi solo sulle parti importanti della foto ignorando il resto) per collegare i punti e disegnare le linee (le traiettorie delle particelle).
• Il trucco: Per non impazzire con la quantità di dati, l'AI prima usa un filtro veloce per pulire la "foto" dal rumore, lasciando solo i punti importanti. Poi, disegna le traiettorie.
• Il risultato: È incredibilmente veloce. Su una scheda video normale (quella che usi per giocare), riesce a ricostruire l'intero percorso delle particelle in 2,3 millisecondi.
• L'efficacia: Riesce a trovare il 98% delle particelle giuste e a collegare correttamente i punti quasi il 93% delle volte. È come se il detective avesse un occhio infallibile per trovare i suoi amici nella folla, anche se a volte non riesce a dire esattamente quanto velocemente stavano camminando (la precisione dei parametri è ancora da migliorare).

Perché è importante?

Attualmente, i computer dell'esperimento ATLAS sono come un gruppo di persone che cercano di risolvere un puzzle gigante usando solo un pezzo alla volta, molto lentamente.
Queste nuove tecniche AI sono come dare a quelle persone:

1. Un super-filtro che scarta i pezzi sbagliati del puzzle prima ancora di iniziare (Metodo 1).
2. Un super-occhio che guarda l'intero puzzle e vede subito la soluzione (Metodo 2).

In sintesi

Con l'aumento della folla di particelle nel futuro, i metodi attuali non saranno più sufficienti. Questi due approcci mostrano che l'Intelligenza Artificiale può:

• Risparmiare tempo: Elaborando i dati molto più velocemente.
• Risparmiare soldi: Funzionando bene anche su hardware economico (schede video da gaming) invece che su supercomputer costosissimi.
• Migliorare la fisica: Permettendo agli scienziati di vedere meglio i fenomeni rari e misteriosi che potrebbero cambiare la nostra comprensione dell'universo.

È un passo fondamentale per trasformare il caos di un'esplosione di particelle in una storia chiara e leggibile.

Sommario tecnico

Contesto e Problema

Il lavoro affronta le sfide critiche poste dall'avvento dell'era dell'High-Luminosity LHC (HL-LHC), prevista dopo il 2030. In questa fase, il numero di collisioni protone-protone per incrocio di fascio (pileup, $\langle\mu\rangle$) aumenterà da circa 60 (Run 3) fino a 200. Questo incremento drastico della densità di interazione comporta:

1. Alta occupazione dei rivelatori: Un aumento significativo dei "hit" (segnali) nel Muon Spectrometer dell'esperimento ATLAS.
2. Rapporto segnale-rumore estremamente basso: In alcune regioni (es. New Small Wheel), il rapporto segnale-rumore scende fino allo 0,6%.
3. Vincoli di latenza: Il sistema di trigger (in particolare l'Event Filter) deve processare eventi a 1 MHz con vincoli di tempo stringenti per selezionare gli eventi fisici interessanti.
4. Complessità geometrica: Il rivelatore integra cinque tecnologie diverse con campi magnetici non omogenei, rendendo difficile l'uso di algoritmi tradizionali basati su trasformate di Hough e fit $\chi^2$, che attualmente richiedono circa 255 ms per evento su CPU.

L'obiettivo è sviluppare strategie di elaborazione dati in tempo reale più efficienti e robuste per il recupero delle tracce di muoni.

Metodologia

Gli autori propongono due approcci distinti basati sul Machine Learning (ML) per affrontare il problema del recupero delle tracce:

1. Reti Neurali su Grafi (GNN) per il Rifiuto dei Fondi

• Concetto: Integrare le GNN nella catena di ricostruzione esistente (baseline) per filtrare i "hit" di fondo prima della ricerca delle tracce.
• Architettura:
  • I dati non sono trattati come singoli hit, ma come cluster di livello superiore chiamati "Muon Buckets" (raggruppamenti di hit entro 30 cm).
  • Si costruisce un grafo dinamico dove i nodi sono i bucket e gli archi connettono bucket in settori adiacenti o identici, rispettando criteri di prossimità spaziale ($|\Delta z| < 15$ m, $\sqrt{\Delta x^2 + \Delta y^2} < 6.8$ m).
  • Viene utilizzato un modello basato su EdgeConv (architettura che propaga informazioni tra nodi vicini).
• Implementazione: Il modello è stato integrato nel framework C++ di ATLAS (Athena) convertendo le operazioni PyTorch Geometric in formato ONNX per l'inferenza.

2. Vision Transformers (ViT) per il Recupero End-to-End

• Concetto: Un approccio "end-to-end" che utilizza architetture di visione artificiale per risolvere il problema combinatorio di associazione hit-traccia in un'unica fase ML.
• Architettura: Basata su una modifica di Mask2Former (originariamente per segmentazione panottica di immagini).
  • Tokenizzazione: Ogni singolo hit del rivelatore è trattato come un "token" (simile a un patch di pixel).
  • Fase di Filtraggio: Utilizza un meccanismo di Flash Attention a finestra (windowed Flash Attention) ordinato secondo l'angolo azimutale $\phi$. Questo riduce la complessità computazionale da quadratica $O(N^2)$ a lineare $O(W \times N)$, permettendo di gestire l'alto rumore.
  • Fase di Tracking: Utilizza un decoder Transformer con cross-attention condizionata da maschere e self-attention tra query (candidati traccia).
  • Output: Assegnazione hit-traccia e regressione dei parametri cinematici ($p_T, \eta, \phi, q$).
• Dati: Addestrato su simulazioni ATLAS Run 4 con $\langle\mu\rangle=200$, utilizzando eventi $Z \to \mu\mu$, $J/\psi \to \mu\mu$ e $t\bar{t}$.

Risultati Chiave

Performance GNN (Filtro Bucket)

• Velocità: Riduzione del tempo di ricostruzione per evento del 15% (da 255 ms a 217 ms) su hardware CPU/GPU ibrido.
• Efficienza: Raggiunge un tasso di rifiuto dei bucket di fondo del 97% a $\langle\mu\rangle=60$.
• Qualità Fisica: Non introduce bias negativi; l'efficienza di ricostruzione e la precisione dei muoni sono preservate in tutto lo spazio cinematico.

Performance ViT (Tracking End-to-End)

• Filtraggio Hit:
  • AUC (Area Under Curve) di 0.9997.
  • A un punto di lavoro con efficienza del segnale del 99%, la purezza degli hit passa dallo 0,6% iniziale al 66,5%, rifiutando il 99,7% del fondo.
  • Riduzione dell'occupazione media da ~6.900 hit/evento a soli 55 hit/evento.
• Tracking:
  • Efficienza di rilevamento segnale: 98,0% (per muoni con $p_T \ge 5$ GeV e $\ge 3$ hit).
  • Efficienza di associazione hit: 92,9%.
  • Purezza di associazione hit: 88,90%.
  • Efficienza di doppio matching (efficienza e purezza > 50%): 94,59%.
  • Accuratezza nella classificazione del segno di carica: 96,35%.
• Velocità di Inferenza:
  • Su GPU consumer (NVIDIA RTX 3090): 2,3 ms per evento (in batch di 200).
  • Su GPU H100: 0,9 ms per evento.
  • Nota: La fase di regressione dei parametri è molto veloce (0,07 ms), ma attualmente limitata dall'overhead di lancio dei kernel GPU.

Contributi e Significatività

1. Validazione di Architetture ML Avanzate: Dimostra che GNN e Vision Transformers sono applicabili con successo al Muon Spectrometer, un ambiente noto per la sua bassa densità di segnali e alta complessità geometrica, dove i metodi ML erano poco esplorati.
2. Efficienza Computazionale: L'approccio ViT offre una riduzione di due ordini di grandezza nella latenza per evento rispetto alla catena CPU attuale (2,3 ms vs 255 ms), rendendo fattibile l'uso di GPU consumer per il trigger in tempo reale.
3. Gestione del Rumore: L'uso di meccanismi di attenzione (Flash Attention) permette di gestire efficacemente il rapporto segnale-rumore estremamente basso (0,6%) tipico dell'HL-LHC, riducendo drasticamente il rumore combinatorio prima della ricostruzione.
4. Scalabilità: Sebbene la latenza singola sembri bassa, il vero vantaggio risiede nella capacità di elaborare batch di eventi in parallelo su GPU, sfruttando la parallelizzazione massiccia disponibile nei moderni acceleratori.
5. Futuro: Il lavoro apre la strada all'integrazione diretta di queste tecniche nel sistema di trigger dell'Event Filter, potenzialmente sostituendo o affiancando gli algoritmi tradizionali per gestire la complessità dei dati dell'HL-LHC.

In sintesi, lo studio conferma che le tecniche di intelligenza artificiale di ultima generazione possono risolvere i colli di bottiglia computazionali e di rumore previsti per la fisica delle particelle di prossima generazione, offrendo soluzioni veloci, robuste e scalabili.