Sparse graph attention networks as efficient ionic liquid potentials

Abstract

Diese Masterarbeit verwendet Graphen-basierte neuronale Netzwerke zur Vorhersage der potentiellen Energieoberfläche von organischen Molekülen und ionischen Flüssigkeiten. Die Repräsentation von atomistischen Systemen durch Graphen ermöglicht die Verwendung sogenannter message passing neural networks, um Information über die Graphenstruktur an benachbarte Atome zu verbreiten. Dies ermöglicht es, die herkömmlich verwendeten und rechnerisch sehr teuren Fingerabdrücke zu vermeiden. Graph attention networks werden adaptiert, um einen dünnbesetzten molekularen Graphen zu generieren und durch den Attention-Mechanismus gewichtete Nachrichten zu versenden. Nach mehreren solchen Ebenen werden die Energiebeiträge durch eine hierarchischen Regressions-Kopf ausgelesen. Die Summe dieser atomistischen Beiträge ist die Vorhersage der potentiellen Energie des Moleküls. Die auf die Atome wirkenden Kräfte werden als Ableitung des Modells evaluiert. Eine effiziente Implementierung in JAX macht Nutzen von einer Grafikkarte sowie just-in-time Kompilierung und ermöglicht dadurch die Simulation der temporalen Entwicklung molekularer Systeme durch Integration der Kräfte. Die quadratische Komplexität des Attention-Mechanismus wird durch Berechnung auf dem dünnbesetzten Graphen vermieden, wodurch Skalierbarkeit zu größeren Systemen ermöglicht wird. Diese Arbeit adaptiert die Implementation von NeuralIL, einer Fingerabdruck-basierten Architektur zur Berechnung der Energie und Kräfte von ionischen Flüssigkeiten. Die Methode nutzt kugelförmige Bessel-Fingerabdrücke sowie ein neuronales Netzwerk und dient als Referenzmethode zum Vergleich der Genauigkeit und Geschwindigkeit des Graphen-basierten Modells. Das Modell erreicht chemische Genauigkeit bei der Vorhersage der Energie und Kräfte von ionischen Flüssigkeiten bei einem 4x Geschwindigkeitsvorteil beim Trainieren. Evaluierung auf der sehr großen Datenbank von organischen Molekülen ANI-1 zeigt die Fähigkeit des Modells, effizient die Energie von Molekülen in dissoziierten Zuständen und während der molekularen Vibrationen mit chemischer Genauigkeit vorherzusagen.

This work applies a graph neural network to the prediction of the potential energy surface of ionic liquids and organic molecules. Representing atomistic systems as molecular graphs enables the use of message passing neural networks to propagate information using the graph connectivity structure, which has the potential to replace the need for computationally expensive atomistic descriptors. Graph Attention Networks are adapted to operate on sparse molecular graphs and perform message passing to propagate information using attention as an update rule. After multiple message passing layers, a hierarchical regression readout head transforms the obtained features into atomic contributions. The sum of these contributions is the model's prediction of the potential energy, while the atomic forces are calculated as its derivatives. An efficient JAX implementation leveraging a GPU and just-in-time compilation enables molecular dynamics simulations using the forces, while sparsity avoids the quadratic computational complexity of self-attention and enables the model to efficiently scale to larger systems. This work adapts the implementation from NeuralIL, a descriptor-based architecture using spherical Bessel descriptors and a neural network, which serves as a baseline against which to evaluate the model's predictive accuracy and speed. The model achieves chemical accuracy when predicting the forces of an ionic liquid, while achieving a significantly faster runtime with a 4x speedup during training. Evaluation on the very large ANI-1 dataset shows the model's ability to efficiently predict the energy of small organic molecules with chemical accuracy and describe displaced molecular states such as during dissociation or vibration.