E2Former‑V2: Schnellere, skalierbare Equivariante Attention‑Netzwerke

Equivariante Graph Neural Networks (EGNNs) sind inzwischen ein Standardwerkzeug zur Modellierung von 3‑D‑atomistischen Systemen. Trotz ihrer hohen Genauigkeit stoßen die gängigen Architekturen an Skalierbarkeitsgrenzen, weil sie für jede Kante explizite geometrische Features oder dichte Tensorprodukte erzeugen.

Mit E2Former‑V2 wird dieses Problem angegangen. Das neue System kombiniert algebraische Sparsität mit hardware‑sensibler Ausführung. Im Kern steht die Methode EAAS (Equivariant Axis‑Aligned Sparsification), die auf Wigner‑6j‑Convolution aufbaut und durch einen SO(3) → SO(2)-Basiswechsel die aufwändigen dichten Tensorkontraktionen in effiziente, spärliche Paritäts‑Re‑Indexierungsoperationen umwandelt.

Auf dieser sparsamen Basis wird ein vollständig node‑zentrierter Mechanismus namens On‑the‑Fly Equivariant Attention eingeführt. Durch einen maßgeschneiderten Triton‑Kernel werden Kanten‑Tensoren vermieden und die SRAM‑Nutzung maximiert, was zu einer beeindruckenden Steigerung von 20‑fachen TFLOPS im Vergleich zu herkömmlichen Implementierungen führt.

Umfangreiche Tests auf den Datensätzen SPICE und OMol25 zeigen, dass E2Former‑V2 die gleiche Vorhersageleistung wie herkömmliche Modelle beibehält, die Inferenzzeit jedoch deutlich verkürzt. Damit beweist die Arbeit, dass große equivariant Transformer‑Modelle effizient auf allgemein verfügbaren GPU‑Plattformen trainiert werden können.

Der Quellcode ist frei verfügbar unter https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2.