Τα Ισοβαρή Γραφοδικτυακά Νευρωνικά Δίκτυα (Equivariant Graph Neural Networks – EGNNs) έχουν καθιερωθεί ως μια ευρέως χρησιμοποιούμενη μέθοδος για τη μοντελοποίηση τρισδιάστατων ατομικών συστημάτων. Ωστόσο, οι επικρατούσες αρχιτεκτονικές αντιμετωπίζουν σημαντικά προβλήματα κλιμακωσιμότητας, τα οποία οφείλονται στην ρητή κατασκευή γεωμετρικών χαρακτηριστικών ή σε πυκνά τανυστικά γινόμενα σε κάθε ακμή.
Προς αντιμετώπιση αυτής της πρόκλησης, παρουσιάζεται η αρχιτεκτονική E2Former-V2, μια κλιμακωτή λύση που ενσωματώνει αλγεβρική αραιότητα με εκτέλεση προσαρμοσμένη στο υλικό.
Πρόταση Ισοβαρούς Αξονικά Ευθυγραμμισμένης Αραίωσης (EAAS)
Αρχικά, προτείνεται η Ισοβαρής Αξονικά Ευθυγραμμισμένη Αραίωση (Equivariant Axis-textbf{A}ligned textbf{S}parsification – EAAS). Η EAAS βασίζεται στη συνέλιξη Wigner-$6j$ και στην αλλαγή βάσης $mathrm{SO}(3) rightarrow mathrm{SO}(2)$, μετατρέποντας τις υπολογιστικά δαπανηρές πυκνές τανυστικές συστολές σε αποδοτικές, αραιές λειτουργίες ανα-ευρετηρίασης ισοτιμίας.
Μηχανισμός Προσοχής On-the-Fly Equivariant Attention
Βασιζόμενοι σε αυτή την αναπαράσταση, εισάγεται ο μηχανισμός On-the-Fly Equivariant Attention, ο οποίος είναι πλήρως κόμβο-κεντρικός και υλοποιείται μέσω ενός προσαρμοσμένου, συγχωνευμένου πυρήνα Triton. Με τη μεγιστοποίηση της χρήσης της μνήμης SRAM, ο εν λόγω πυρήνας επιτυγχάνει μια 20πλάσια βελτίωση σε TFLOPS σε σύγκριση με τις τυπικές υλοποιήσεις.
Αποτελέσματα Πειραμάτων και Συμπεράσματα
Εκτεταμένα πειράματα που διεξήχθησαν στα σύνολα δεδομένων SPICE και OMol25 καταδεικνύουν ότι η E2Former-V2 διατηρεί συγκρίσιμη προβλεπτική απόδοση, ενώ παράλληλα επιταχύνει σημαντικά την εξαγωγή συμπερασμάτων. Αυτή η εργασία αποδεικνύει ότι μεγάλοι ισοβαροί μετασχηματιστές μπορούν να εκπαιδευτούν αποτελεσματικά χρησιμοποιώντας ευρέως προσβάσιμες πλατφόρμες GPU. Ο κώδικας είναι διαθέσιμος στη διεύθυνση https://github.com/IQuestLab/UBio-MolFM/tree/e2formerv2.
