🤖 AI

E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

E2Former-V2 ist eine skalierbare äquivariante Graph-Neural-Network-Architektur, die durch die Einführung von Equivariant Axis-Aligned Sparsification und einem On-the-Fly Equivariant Attention-Mechanismus die Rechenengpässe gängiger Modelle überwindet und dabei eine 20-fache TFLOPS-Verbesserung bei gleichbleibender Vorhersgeleistung für atomistische 3D-Systeme erzielt.

Ursprüngliche Autoren: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Veröffentlicht 2026-06-08

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen vorherzusagen, wie sich eine riesige Menschenmenge (Atome) in einem Raum bewegen und interagieren wird. In der Welt der Chemie und Physik nennt man das „atomistische Modellierung“. Lange Zeit haben Wissenschaftler ein spezielles KI-Modell verwendet, ein äquivariantes Graph-Neurales Netzwerk (EGNN). Diese Netzwerke sind großartig, weil sie verstehen, dass sich die Physik im Raum nicht ändert, wenn man den gesamten Raum dreht.

Dennoch haben die aktuellen Versionen dieser Netzwerke ein großes Problem: Sie sind ungeschickt und langsam, wenn die Menge groß wird.

Hier ist eine einfache Aufschlüsselung der Lösung des Papers, E2Former-V2, unter Verwendung alltäglicher Analogien.

Das Problem: Der „kantenzentrierte“ Verkehrsstau

Stellen Sie sich eine Party vor, bei der jeder mit seinen Nachbarn sprechen möchte.

Die alte Methode (Traditionelle EGNNs): Das System erstellt eine riesige, physische Liste von jedem einzelnen Gespräch, das auf der Party stattfindet. Wenn es 1.000 Leute gibt und jeder mit 50 Nachbarn spricht, muss das System 50.000 separate Gesprächsnotizen auf einem riesigen Whiteboard (Speicher) aufschreiben, bevor es überhaupt rechnen kann.
Der Flaschenhals: Wenn die Party größer wird (mehr Atome), wird dieses Whiteboard so gewaltig, dass der Computer keinen Platz mehr hat. Er verbringt die ganze Zeit damit, diese Notizen hin und her zu schieben, anstatt tatsächlich nachzudenken. Deshalb stürzen alte Modelle ab oder werden extrem langsam, wenn sie große Moleküle wie Proteine simulieren sollen.

Die Lösung: E2Former-V2

Die Autoren haben ein neues System entwickelt, das aufhört, jede einzelne Gesprächsnotiz aufzuschreiben. Stattdessen nutzen sie zwei clevere Tricks, um den Prozess blitzschnell zu machen.

Trick 1: Der „Magische Kompass“ (Äquivariante achsenorientierte Sparsifizierung)

Im alten System ist die Berechnung der Interaktion zwischen Atomen wie der Versuch, ein Puzzle zu lösen, bei dem jedes Teil in jede beliebige Richtung gedreht werden kann. Das ist unordentlich und erfordert viel Mathematik.

Die Analogie: Stellen Sie sich vor, Sie haben ein komplexes 3D-Puzzle. Anstatt zu versuchen, es zu lösen, während es sich dreht, benutzen Sie einen „Magischen Kompass“, um das gesamte Puzzle so zu drehen, dass ein bestimmtes Teil immer direkt nach oben zeigt.
Das Ergebnis: Sobald das Puzzle ausgerichtet ist, werden die meisten Teile leer oder gleich null. Sie müssen keine Mathematik mehr an den leeren Teilen durchführen. Das Paper nennt dies EAAS. Es verwandelt eine unordentliche, dichte Berechnung in ein einfaches „Neu-Indizierungs“-Spiel (nur das Verschieben von Teilen), was viel schneller ist.

Trick 2: Der „Streaming-Koch“ (On-the-Fly Äquivariante Attention)

Dies ist der größte Durchbruch.

Der alte Weg (Das Buffet): Das alte System bereitet das gesamte Buffet (alle 50.000 Gesprächsnotizen) vor und stellt es auf einen Tisch, bevor jemand essen darf. Wenn der Tisch zu klein ist, verschüttet das Essen und die Party stoppt.
Der neue Weg (Der Streaming-Koch): Das neue System, E2Former-V2, fungt wie ein Koch, der jeden Teller einzeln direkt vor dem Kunden zubereitet und ihn sofort serviert.
- Es schreibt die Notizen nicht auf.
- Es speichert nicht die ganze Liste.
- Es berechnet die Interaktion zwischen zwei Personen, fügt das Ergebnis dem Gesamtergebnis hinzu und vergisst die Details sofort wieder, um Platz für das nächste Paar zu machen.
Das Ergebnis: Dies wird als „On-the-Fly“-Verarbeitung bezeichnet. Da es niemals das riesige Whiteboard erstellt, verbraucht es fast keinen Speicher. Es passt alles in den superschnellen internen Cache (SRAM) des Computers, was es 20 Mal schneller macht als die alten Methoden.

Was sie tatsächlich erreicht haben

Die Autoren haben dieses neue System an zwei Hauptpunken getestet:

Kleine Moleküle (SPICE): Sie testeten es an drogenähnlichen Molekülen. Das neue System war genauso genau wie die besten existierenden Modelle, aber viel schneller.
Riesige Moleküle (OMol25): Sie testeten es an massiven chemischen Strukturen.
- Die Grenze: Alte Modelle stürzten ab (Speicherplatzmangel), wenn sie versuchten, Systeme mit mehr als 10.000 Atomen zu simulieren.
- Der Sieg: E2Former-V2 konnte erfolgreich Systeme mit 100.000 Atomen simulieren, ohne abzustürzen.

Erwähnte reale Tests

Die Autoren haben nicht nur Zahlen laufen lassen; sie haben echte Physik simuliert:

Flüssiges Wasser: Sie simulierten eine Box mit Wassermolekülen. Das neue Modell sagte, wie sich die Wassermoleküle anordnen (die radiale Verteilungsfunktion), fast perfekt voraus und stimmte besser mit realen Experimenten überein als bisherige Modelle.
Proteine: Sie simulierten ein großes Protein (etwa 30.000 Atome), das in Wasser schwimmt. Das Modell hielt das Protein stabil und half ihm, sich in eine Form zu falten, die mit echten experimentellen Daten übereinstimmte, was beweist, dass es komplexe biologische Strukturen bewältigen kann.

Zusammenfassung

E2Former-V2 ist wie der Wechsel von einem System, das jedes einzelne Gespräch auf einem riesigen, langsamen Whiteboard aufschreibt, zu einem System, bei dem ein intelligenter Koch Gespräche einzeln kocht und sofort serviert. Dies ermöglicht es Wissenschaftlern, massive molekulare Systeme (wie riesige Proteine) auf Standardcomputern zu simulieren, ohne dass der Speicher überläuft, während die physikalische Genauigkeit erhalten bleibt.

Technisches Resümee: E2Former-V2: On-the-Fly Äquivariante Attention mit Linearem Aktivierungsgedächtnis

1. Problemstellung

Äquivariante Graph-Neuronale Netze (EGNNs) sind zum Standard für die Modellierung von ataomistischen 3D-Systemen geworden, da sie Rotations- und Translationssymmetrien bewahren können. Jedoch leiden gängige Architekturen (z. B. MACE, EquiformerV2, eSCN) unter kritischen Effizienzengpässen, die in ihrem kantenzentrierten Design wurzeln. Diese Modelle konstruieren explizit geometrische Merkmale oder führen dichte Tensorprodukte auf jeder Kante durch, was zu Folgendem führt:

Speicherkomplexität: $O(|E|)$ Aktivierungsgedächtnis, wobei $|E| \approx kN$ (mit $k$ als durchschnittliche Nachbaranzahl). Für große Systeme resultiert dies in massiven Speicherabdrücken, die die GPU-Kapazität überschreiten.
Hardware-Ineffizienz: Die explizite Materialisierung von Kanten-Tensoren erzwingt Datentransfers zwischen dem High Bandwidth Memory (HBM) und dem On-Chip-SRAM, was eine Speicherwand schafft, die den Durchsatz begrenzt.
Skalierbarkeitslimits: Mit zunehmender Systemgröße ( $N$ ) skaliert die Latenz schlecht, was groß angelegte Molekulardynamik-Simulationen (MD) auf handelsüblicher Hardware rechentechnisch prohibitiv macht.

Während jüngste Fortschritte bei Transformern (z. B. FlashAttention) ähnliche Engpässe für Sequenzdaten durch Streaming und kachelbasierte Ausführung gelöst haben, blieben diese Prinzipien in SO(3)-äquivarianten Architekturen bisher unerschlossen.

2. Methodik

E2Former-V2 adressiert diese Probleme durch die Integration von algebraischer Sparsifizierung mit einem hardware-bewussten, knotenzentrierten Ausführungsmodell. Die Architektur besteht aus zwei Kerninnovationen:

2.1. Äquivariante Achsen-ausgerichtete Sparsifizierung (EAAS)

Standardmäßige SO(3)-Tensorkontraktionen beinhalten dichte Summationen über Clebsch–Gordan-Koeffizienten, was zu einer Komplexität von $O(L^6)$ und irregulären Speicherzugriffsmustern führt. E2Former-V2 führt EAAS ein, um diese dichten Operationen in dünnbesetzte, permutationsbasierte Operationen zu transformieren:

SO(3) $\to$ SO(2) Basiswechsel: Die Methode richtet das globale Koordinatensystem an dem lokalen Vektor $\vec{r}$ (der Achse, die zwei Atome verbindet) aus. In diesem ausgerichteten Frame werden sphärische Kugelfunktionen maximal dünnbesetzt (Lemma 4.1), wobei nur die $m=0$ Komponente ungleich Null ist.
Re-Indexing-Operator: Durch das Kommutieren von Rotationen durch das Tensorprodukt kollabiert die dichte Kontraktion in eine deterministische Re-Indexing-Regel (Gl. 13). Dieser Operator bildet Eingangsmerkmale auf Ausgangsmerkmale via dünnbesetzter blockweiser linearer Abbildungen ab, wodurch die Notwendigkeit einer dichten Summation über magnetische Indizes entfällt.
Ergebnis: Dies reduziert die Komplexität der Konvolutionsphase und erzielt eine $\sim 1,5\times$ Beschleunigung gegenüber Standardimplementierungen.

2.2. On-the-Fly Äquivariante Attention

Um den $O(|E|)$ Speicherengpass zu eliminieren, schlagen die Autoren einen vollständig fusionierten, knotenzentrierten Attention-Mechanismus vor, der mittels eines benutzerdefinierten Triton-Kernels implementiert ist:

Knotenzentrierte Faktorisierung: Anstatt Kantenmerkmale zu materialisieren, entkoppelt der Attention-Mechanismus die Berechnung. Query- und Key-Vektoren werden unter Verwendung von SO(3)-invarianten skalaren Projektionen (Gl. 6) konstruiert, während äquivariante geometrische Informationen nur über den Value-Pfad propagiert werden.
Streaming-Reduktion: Der Kernel verarbeitet Nachbarn in einem Streaming-Verfahren. Er iteriert über Nachbarn $j \in \mathcal{N}(i)$ , berechnet Attention-Scores und akkumuliert Value-Updates on-the-fly innerhalb des SRAM.
Eliminierung der Kanten-Materialisierung: Zwischenliegende Tensoren der Form $N \times K \times H \times D$ werden niemals in das HBM geschrieben. Der Kernel hält laufende Statistiken (Maximalwert, Normalisierung und Value-Akkumulatoren) on-chip aufrecht, wodurch die Aktivierungs-Speicherkomplexität von $O(|E|)$ auf $O(|V|)$ reduziert wird.
Online-Softmax: Um die numerische Stabilität während des Streamings zu gewährleisten, wird eine Online-Softmax-Formulierung verwendet, um Akkumulatoren inkrementell zu aktualisieren.

3. Kernbeiträge

Äquivariante Achsen-ausgerichtete Sparsifizierung (EAAS): Eine neuartige algebraische Technik, die Wigner-6j-Rekombination mit SO(2)-Achsausrichtung vereint und dichte SO(3)-Tensorkontraktionen in dünnbesetzte Re-Indexing-Operationen umwandelt.
On-the-Fly Äquivariante Attention: Ein vollständig fusionierter Triton-Kernel, der die Kanten-Materialisierung vermeidet, die SRAM-Wiederverwendung maximiert und bis zu 20 $\times$ höhere TFLOPS im Vergleich zu Standardimplementierungen erreicht.
Skalierbare Performance: Nachgewiesene Fähigkeit, Systeme mit bis zu 100.000 Atomen auf einer einzelnen GPU zu handhaben – ein Regime, in dem vorherige äquivariante Transformer aufgrund von Out-of-Memory (OOM) Fehlern scheitern.
Physikalische Treue: Validierung in der Molekulardynamik, die eine exakte Reproduktion von Strukturen von flüssigem Wasser sowie eine stabile Verfeinerung von AlphaFold-prädizierten Proteinstrukturen ( $\sim$ 30k Atome) zeigt.

4. Experimentelle Ergebnisse

Die Autoren evaluierten E2Former-V2 auf den Datensätzen SPICE (Medizinalchemie) und OMol25 (Foundation-Modell-Skala) sowie in MD-Simulationen.

Prädiktive Genauigkeit:
- Auf SPICE erreichte E2Former-V2 die niedrigsten Energie- und Kraft-MAE (Mean Absolute Error) über die meisten Teilmengen und übertraf MACE-Large um 48 % beim Energie-MAE bei Dimeren.
- Auf OMol25 erreichte die konservative Variante einen aggregierten Energie-MAE von 1,27 meV/Atom, was eSEN-small entspricht und UMA-S (3,62 meV/Atom) signifikant übertrifft.
- Auf OC20 S2EF-2M erzielte es eine Genauigkeit, die mit starken Baselines wie EquiformerV2 und eSCN vergleichbar ist.
Effizienz und Skalierbarkeit:
- Durchsatz: Der fusionierte Kernel erreichte eine durchschnittliche 20 $\times$ Beschleunigung in den TFLOPS gegenüber naiven PyTorch-Implementierungen.
- Speicher: Der Spitzen-Speicherverbrauch wurde im Vergleich zu den Baselines um den Faktor 16,9 $\times$ bis 19,5 $\times$ reduziert.
- Skalierbarkeit: Während vorherige Modelle (z. B. EquiformerV2, MACE-Large) bei mehr als 10k–50k Atomen OOM-Fehler erlitten, skalierte E2Former-V2 zuverlässig auf 100k Atome.
- Inferenzgeschwindigkeit: Bei 100k Atomen erreichte E2Former-V2 (Conservative) 0,12 Schritte/s, was fast 3 $\times$ schneller ist als UMA-S. Im Direct-Setting war es bei kleineren Skalen um eine Größenordnung schneller als Allegro und EquiformerV2.
Molekulardynamik:
- Simulationen von periodischem Bulk-Wasser (216 Moleküle) zeigten, dass E2Former-V2 die experimentellen Radialverteilungsfunktionen (RDF) eng nachbildet und dabei MACE-OFF übertrifft.
- Eine 2-ns-Simulation eines solvatisierten Proteins ( $\sim$ 30k Atome) behielt hohe TM-Scores relativ zur nativen Struktur bei, was eine stabile Dynamik demonstriert.

5. Bedeutung und Ansprüche

Das Paper behauptet, dass E2Former-V2 die praktische Machbarkeit großer äquivarianter Transformer auf handelsüblichen GPUs demonstriert. Durch den Wechsel des Ausführungsparadigmas von kantenzentriert zu knotenzentriert und die Nutzung von hardware-bewusstem Streaming löst die Methode den fundamentalen Speicherengpass, der die Skalierbarkeit äquivarianter Modelle begrenzt hat.

Die Autoren betonen, dass dies nicht bloß eine algorithmische Verbesserung ist, sondern eine systemische Realisierung, die die Lücke zwischen theoretischer Faktorisierung (Reduktion der Komplexität von $O(|E|)$ auf $O(|V|)$ ) und Hardware-Effizienz schließt. Die Arbeit legt nahe, dass groß angelegte molekulare Foundation-Modelle und langzeitige MD-Simulationen, die zuvor rechentechnisch prohibitiv waren, nun mit Standard-GPU-Infrastruktur zugänglich sind. Der Code wird veröffentlicht, um die weitere Nutzung in der Community zu fördern.