Riepilogo Tecnico: E2Former-V2: Attenzione Equivariante "On-the-Fly" con Memoria di Attivazione Lineare

1. Definizione del Problema

Le Reti Neurali a Grafo Equivarianti (EGNN) sono diventate lo standard per la modellazione di sistemi atomistici 3D grazie alla loro capacità di preservare le simmetrie rotazionali e traslazionali. Tuttavia, le architetture mainstream (ad es. MACE, EquiformerV2, eSCN) soffrono di critici colli di bottiglia di efficienza radicati nel loro design centrato sugli archi (edge-centric). Questi modelli costruiscono esplicitamente caratteristiche geometriche o eseguono prodotti tensoriali densi su ogni arco, portando a:

• Complessità di Memoria: Memoria di attivazione $O(|E|)$, dove $|E| \approx kN$ (con $k$ che rappresenta il numero medio di vicini). Per sistemi estesi, ciò si traduce in un enorme footprint di memoria che supera la capacità della GPU.
• Inefficienza dell'Hardware: La materializzazione esplicita dei tensori degli archi forza il movimento dei dati tra l'High Bandwidth Memory (HBM) e la SRAM on-chip, creando un "muro della memoria" che limita il throughput.
• Limiti di Scalabilità: All'aumentare della dimensione del sistema ($N$), la latenza scala negativamente, rendendo le simulazioni di dinamica molecolare (MD) su larga scala computazionalmente proibitive su hardware comune.

Sebbene i progressi recenti nei Transformer (ad es. FlashAttention) abbiano risolto colli di bottiglia simili per i dati sequenziali tramite l'esecuzione a streaming e basata su tile, questi principi sono rimasti inesplorati nelle architetture SO(3)-equivarianti.

2. Metodologia

E2Former-V2 affronta queste problematiche integrando la sparsificazione algebrica con un modello di esecuzione centrato sui nodi (node-centric) consapevole dell'hardware. L'architettura consiste in due innovazioni principali:

2.1. Sparificazione Equivariante Allineata all'Asse (EAAS)

Le contrazioni tensoriali SO(3) standard comportano somme dense su coefficienti di Clebsch–Gordan, risultando in una complessità $O(L^6)$ e pattern di accesso alla memoria irregolari. E2Former-V2 introduce EAAS per trasformare queste operazioni dense in operazioni sparse basate su permutazioni:

• Cambio di Base da SO(3) a SO(2): Il metodo allinea il sistema di coordinate globale con il vettore locale $\vec{r}$ (l'asse che connette due atomi). In questo frame allineato, le armoniche sferiche solide diventano massimamente sparse (Lemma 4.1), dove solo la componente $m=0$ è non nulla.
• Operatore di Re-indexing: Commutando le rotazioni attraverso il prodotto tensoriale, la contrazione densa collassa in una regola di re-indexing deterministica (Eq. 13). Questo operatore mappa le feature di input sulle feature di output tramite mappe lineari sparse per blocchi, eliminando la necessità di una sommatoria densa sugli indici magnetici.
• Risultato: Ciò riduce la complessità della fase di convoluzione e ottiene un'accelerazione di circa $\sim 1.5\times$ rispetto alle implementazioni standard.

2.2. Attenzione Equivariante "On-the-Fly"

Per eliminare il collo di bottiglia della memoria $O(|E|)$, gli autori propongono un meccanismo di attenzione completamente fuso e centrato sui nodi, implementato tramite un kernel Triton personalizzato:

• Fattorizzazione Centrata sui Nodi: Invece di materializzare le feature degli archi, il meccanismo di attenzione decouple la computazione. I vettori Query e Key sono costruiti utilizzando proiezioni scalari SO(3)-invarianti (Eq. 6), mentre l'informazione geometrica equivariante viene propagata solo attraverso il percorso dei valori (value path).
• Riduzione in Streaming: Il kernel elabora i vicini in modalità streaming. Itera sui vicini $j \in \mathcal{N}(i)$, calcolando i punteggi di attenzione e accumulando gli aggiornamenti dei valori on-the-fly all'interno della SRAM.
• Eliminazione della Materializzazione degli Archi: I tensori intermedi di forma $N \times K \times H \times D$ non vengono mai scritti in HBM. Il kernel mantiene statistiche correnti (max, normalizzazione e accumulatori di valore) on-chip, riducendo la complessità della memoria di attivazione da $O(|E|)$ a $O(|V|)$.
• Softmax Online: Per garantire la stabilità numerica durante lo streaming, viene utilizzata una formulazione di softmax online per aggiornare gli accumulatori in modo incrementale.

3. Contributi Chiave

1. Sparsificazione Equivariante Allineata all'Asse (EAAS): Una nuova tecnica algebrica che unifica il recoupling Wigner-6j con l'allineamento dell'asse SO(2), convertendo le contrazioni tensoriali dense SO(3) in operazioni di re-indexing sparse.
2. Attenzione Equivariante On-the-Fly: Un kernel Triton completamente fuso che evita la materializzazione degli archi, massimizzando il riutilizzo della SRAM e raggiungendo fino a 20$\times$ più TFLOPS rispetto alle implementazioni standard.
3. Prestazioni Scalabili: Capacità dimostrata di gestire sistemi fino a 100.000 atomi su una singola GPU, un regime in cui i precedenti Transformer equivarianti falliscono a causa di errori Out-Of-Memory (OOM).
4. Fedeltà Fisica: Validazione nella dinamica molecolare che mostra la riproduzione accurata delle strutture dell'acqua liquida e il raffinamento stabile delle strutture proteiche predette da AlphaFold ($\sim$30k atomi).

4. Risultati Sperimentali

Gli autori hanno valutato E2Former-V2 sui dataset SPICE (chimica medicinale) e OMol25 (scala modello di fondazione), nonché nelle simulazioni di MD.

• Accuratezza Predittiva:

  • Su SPICE, E2Former-V2 ha ottenuto i più bassi Errori Assoluti Medi (MAE) di Energia e Forza su la maggior parte dei subset, superando MACE-Large del 48% in Energy MAE su Dimers.
  • Su OMol25, la variante conservativa ha raggiunto un Energy MAE aggregato di 1.27 meV/atomo, eguagliando eSCN-small e superando significativamente UMA-S (3.62 meV/atomo).
  • Su OC20 S2EF-2M, ha raggiunto un'accuratezza comparabile a forti baseline come EquiformerV2 e eSCN.

• Efficienza e Scalabilità:

  • Throughput: Il kernel fuso ha ottenuto un'accelerazione media di 20$\times$ in TFLOPS rispetto alle implementazioni naive in PyTorch.
  • Memoria: Il picco di utilizzo della memoria è stato ridotto di 16.9$\times$ fino a 19.5$\times$ rispetto alle baseline.
  • Scalabilità: Mentre i modelli precedenti (ad es. EquiformerV2, MACE-Large) incontravano fallimenti OOM oltre i 10k–50k atomi, E2Former-V2 è scalato in modo affidabile fino a 100k atomi.
  • Velocità di Inferenza: A 100k atomi, E2Former-V2 (Conservative) ha raggiunto 0.12 step/s, quasi 3$\times$ più veloce di UMA-S. Nel setting Direct, è stato un ordine di grandezza più veloce di Allegro ed EquiformerV2 su scale minori.

• Dinamica Molecolare:

  • Le simulazioni di acqua bulk periodica (216 molecole) hanno mostrato che E2Former-V2 corrisponde strettamente alle funzioni di distribuzione radiale (RDF) sperimentali, superando MACE-OFF.
  • Una simulazione di 2 ns di una proteina solvatata ($\sim$30k atomi) ha mantenuto alti punteggi TM rispetto alla struttura nativa, dimostrando una dinamica stabile.

5. Significato e Rivendicazioni

Il paper sostiene che E2Former-V2 dimostra la fattibilità pratica di grandi transformer equivarianti su GPU commerciali. Spostando il paradigma di esecuzione da centrato sugli archi a centrato sui nodi e sfruttando lo streaming consapevole dell'hardware, il metodo risolve il fondamentale collo di bottiglia della memoria che ha limitato la scalabilità dei modelli equivarianti.

Gli autori sottolineano che non si tratta semplicemente di un miglioramento algoritmico, ma di una realizzazione a livello di sistema che colma il divario tra la fattorizzazione teorica (riduzione della complessità da $O(|E|)$ a $O(|V|)$) ed l'efficienza dell'hardware. Il lavoro suggerisce che i modelli di fondazione molecolare su larga scala e le simulazioni MD a lungo termine, precedentemente proibitive dal punto di vista computazionale, sono ora accessibili utilizzando l'infrastruttura GPU standard. Il codice è rilasciato per facilitare ulteriore adozione nella comunità.