Réseaux de neurones sur graphe par transport optimal

Les architectures actuelles de réseaux neuronaux sur graphes (GNN) moyennent ou somment naïvement les embeddings des nœuds pour obtenir une représentation agrégée du graphe — ce qui peut entraîner une perte d’informations structurelles ou sémantiques. Nous introduisons ici OT-GNN, un modèle qui calcule des embeddings de graphes à l’aide de prototypes paramétriques mettant en évidence les aspects clés de différentes facettes d’un graphe. À cette fin, nous combinons avec succès le transport optimal (OT) avec des modèles graphiques paramétriques. Les représentations de graphes sont obtenues à partir des distances de Wasserstein entre l’ensemble des embeddings de nœuds fournis par le GNN et des « prototypes » sous forme de nuages de points considérés comme des paramètres libres. Nous prouvons théoriquement que, contrairement à l’agrégation par somme classique, notre classe de fonctions définie sur les nuages de points satisfait un théorème fondamental d’approximation universelle. Expérimentalement, nous résolvons un problème d’optimisation inhérent de type « effondrement » en proposant un régulariseur basé sur la contrastive bruit, qui guide le modèle vers une exploitation réelle de la géométrie du transport optimal. Enfin, nous surpassons les méthodes populaires sur plusieurs tâches de prédiction de propriétés moléculaires, tout en exhibant des représentations de graphes plus lisses.