Optimal Transport Graph Neural Networks

Aktuelle Architekturen von Graph Neural Networks (GNNs) aggregieren Knoten-Embeddings naiv durch Mittelwertbildung oder Summierung zu einer globalen Graphrepräsentation – was strukturelle oder semantische Informationen potenziell verlieren kann. In diesem Beitrag stellen wir OT-GNN vor, ein Modell, das Graph-Embeddings mithilfe parametrischer Prototypen berechnet, die wesentliche Aspekte verschiedener Graph-Aspekte hervorheben. Hierfür kombinieren wir erfolgreich die Theorie des optimalen Transports (Optimal Transport, OT) mit parametrischen Graphmodellen. Die Graphrepräsentationen werden durch Wasserstein-Distanzen zwischen der Menge der GNN-Knoten-Embeddings und „Prototyp“-Punktwolken ermittelt, die als freie Parameter fungieren. Theoretisch zeigen wir, dass unsere Funktionsklasse auf Punktwolken im Gegensatz zur herkömmlichen Summenaggregation einen fundamentalen universellen Approximationssatz erfüllt. Empirisch lösen wir ein inhärentes Optimierungsproblem der Degenerierung, indem wir einen Noise-Contrastive-Regularisierer vorschlagen, der das Modell dazu anleitet, die Geometrie des optimalen Transports tatsächlich auszunutzen. Schließlich übertrifft unser Ansatz etablierte Methoden bei mehreren Aufgaben zur Vorhersage molekularer Eigenschaften und zeichnet sich zudem durch glattere Graphrepräsentationen aus.