Graphiquement des réseaux de neurones dans TensorFlow

Les objets et leurs relations sont omniprésents dans le monde qui nous entoure, et les relations peuvent être aussi importantes pour comprendre un objet que ses propres attributs considérés de manière isolée – prendre par exemple les réseaux de transport, les réseaux de production, les graphiques de connaissances ou les réseaux sociaux. Les mathématiques discrètes et l’informatique ont une longue histoire de formalisation de réseaux tels que graphiquescomposé de nœuds connecté par bords de diverses manières irrégulières. Pourtant, la plupart des algorithmes d’apprentissage automatique (ML) ne permettent que des relations régulières et uniformes entre les objets d’entrée, tels qu’une grille de pixels, une séquence de mots ou aucune relation.

Graphique des réseaux de neuronesou GNNS pour faire court, sont devenus une technique puissante pour tirer parti de la connectivité du graphique (comme dans les algorithmes plus anciens Deepwalk et Node2vec) et les fonctionnalités d’entrée sur les différents nœuds et bords. Les GNN peuvent faire des prédictions des graphiques dans son ensemble (cette molécule réagit-elle d’une certaine manière?), Pour les nœuds individuels (quel est le sujet de ce document, compte tenu de ses citations?) Ou pour les bords potentiels (ce produit est-il susceptible d’être acheté avec ce produit?). En plus de faire des prédictions sur les graphiques, les GNN sont un outil puissant utilisé pour combler le gouffre vers des cas d’utilisation du réseau neuronal plus typique. Ils codent pour un graphique discret, relationnel informations dans un continu moyen pour qu’il puisse être inclus naturellement dans un autre système d’apprentissage en profondeur.

Nous sommes ravis d’annoncer la sortie de TensorFlow GNN 1.0 (TF-GNN), une bibliothèque testée en production pour construire des GNN à grande échelle. Il prend en charge la modélisation et la formation à TensorFlow ainsi que l’extraction de graphiques d’entrée à partir de magasins de données énormes. TF-GNN est construit à partir de zéro pour les graphiques hétérogènes, où les types d’objets et les relations sont représentés par des ensembles distincts de nœuds et de bords. Les objets du monde réel et leurs relations se produisent dans des types distincts, et l’attention hétérogène de TF-GNN rend naturel de les représenter.

À l’intérieur de TensorFlow, ces graphiques sont représentés par des objets de type tfgnn.GraphTensor. Il s’agit d’un type de tenseur composite (une collection de tenseurs dans une classe Python) accepté comme un citoyen de première classe dans tf.data.Dataset, tf.functionetc. Il stocke à la fois la structure du graphique et ses fonctionnalités attachées aux nœuds, aux bords et au graphique dans son ensemble. Les transformations formables des graphtenseurs peuvent être définies comme des objets de couches au niveau élevé API Kerasou en utilisant directement le tfgnn.GraphTensor primitif.

GNNS: faire des prédictions pour un objet en contexte

Pour l’illustration, examinons une application typique de TF-GNN: prédire une propriété d’un certain type de nœud dans un graphique défini par des tables de référencement d’une énorme base de données. Par exemple, une base de données de citation des articles arxiv informatiques (CS) avec des citations un-à-plusieurs et des relations citées par plusieurs à un où nous aimerions prédire le domaine de chaque article.

Comme la plupart des réseaux de neurones, un GNN est formé sur un ensemble de données de nombreux exemples étiquetés (~ millions), mais chaque étape de formation ne se compose que d’un lot beaucoup plus petit d’exemples de formation (disons, des centaines). Pour évoluer à des millions, le GNN est formé sur un flux de sous-graphes raisonnablement petits du graphique sous-jacent. Chaque sous-graphe contient suffisamment de données d’origine pour calculer le résultat GNN pour le nœud étiqueté en son centre et former le modèle. Ce processus – généralement appelé échantillonnage de sous-graphes – est extrêmement consécutif pour la formation GNN. La plupart des outils existants accomplissent l’échantillonnage d’une manière par lots, produisant des sous-graphiques statiques pour la formation. TF-GNN fournit des outils pour améliorer cela en échantillonnant dynamiquement et interactivement.

Sur la photo, le processus d’échantillonnage de sous-graphes où de petits sous-graphes traitables sont échantillonnés à partir d’un graphique plus grand pour créer des exemples d’entrée pour la formation GNN.

TF-GNN 1.0 lance une API Python flexible pour configurer l’échantillonnage de sous-graphe dynamique ou par lots à toutes les échelles pertinentes: interactivement dans un cahier Colab (comme celui-ci), pour un échantillon efficace d’un petit ensemble de données stocké dans la mémoire principale d’un seul hôte d’entraînement, ou distribué par Faisceau Apache Pour d’énormes ensembles de données stockés sur un système de fichiers réseau (jusqu’à des centaines de millions de nœuds et de milliards de bords). Pour plus de détails, veuillez consulter nos guides d’utilisateur pour en mémoire et à base de faisceau Échantillonnage, respectivement.

Sur ces mêmes sous-graphes échantillonnés, la tâche de GNN est de calculer un état caché (ou latent) au nœud racine; L’état caché agrége et code les informations pertinentes du quartier du nœud racine. Une approche classique est réseaux de neurones qui transfèrent des messages. Dans chaque cycle de message, les nœuds reçoivent des messages de leurs voisins le long des bords entrants et mettent à jour leur propre état caché chez eux. Après n Rounds, l’état caché du nœud racine reflète les informations agrégées de tous les nœuds à l’intérieur n bords (photo ci-dessous pour n = 2). Les messages et les nouveaux états cachés sont calculés par des couches cachées du réseau neuronal. Dans un graphique hétérogène, il est souvent logique d’utiliser des couches cachées formées séparément pour les différents types de nœuds et de bords

Sur la photo, un réseau neuronal de transfert de messages simple où, à chaque étape, l’état de nœud est propagé des nœuds extérieurs aux nœuds intérieurs où il est regroupé pour calculer de nouveaux états de nœud. Une fois le nœud racine atteint, une prédiction finale peut être faite.

La configuration de la formation est terminée en plaçant une couche de sortie sur l’état caché de GNN pour les nœuds étiquetés, en calculant le perte (pour mesurer l’erreur de prédiction) et mettre à jour les poids des modèles par rétro-propagation, comme d’habitude dans toute formation de réseau neuronal.

Au-delà de la formation supervisée (c’est-à-dire, minimisant une perte définie par les étiquettes), les GNN peuvent également être formés de manière non supervisée (c’est-à-dire sans étiquettes). Cela nous permet de calculer un continu représentation (ou intégration) de la discret Structure graphique des nœuds et de leurs caractéristiques. Ces représentations sont ensuite généralement utilisées dans d’autres systèmes ML. De cette façon, les informations relationnelles discrètes codées par un graphique peuvent être incluses dans des cas d’utilisation de réseau neuronal plus typiques. TF-GNN prend en charge une spécification à grain fin des objectifs non supervisés pour les graphiques hétérogènes.

Bâtiment des architectures GNN

La bibliothèque TF-GNN soutient la construction et la formation des GNN à différents niveaux d’abstraction.

Au plus haut niveau, les utilisateurs peuvent prendre l’un des modèles prédéfinis regroupés avec la bibliothèque exprimée en couches Keras. Outre une petite collection de modèles de la littérature de recherche, TF-GNN est livré avec un modèle de modèle hautement configurable qui fournit une sélection organisée de choix de modélisation que nous avons trouvés pour fournir de fortes lignes de base sur bon nombre de nos problèmes internes. Les modèles implémentent les couches GNN; Les utilisateurs n’ont qu’à initialiser les couches Keras.

Au niveau le plus bas, les utilisateurs peuvent écrire un modèle GNN à partir de zéro en termes de primitives pour passer des données autour du graphique, tels que la diffusion des données d’un nœud à tous ses bords sortants ou la mise en commun des données dans un nœud de tous ses bords entrants (par exemple, calculant la somme des messages entrants). Le modèle de données graphiques de TF-GNN traite également les nœuds, les bords et les graphiques d’entrée entiers en ce qui concerne les fonctionnalités ou les états cachés, ce qui rend simple pour exprimer non seulement des modèles centrés sur le nœud comme le MPNN discuté ci-dessus mais aussi des formes plus générales de Graphiques. Cela peut, mais pas à faire avec Keras comme cadre de modélisation en haut de Core TensorFlow. Pour plus de détails et des niveaux de modélisation intermédiaires, voir le TF-GNN guide de l’utilisateur et collection de modèles.

Orchestration de formation

Alors que les utilisateurs avancés sont gratuits pour effectuer une formation sur des modèles personnalisés, le Runner TF-GNN fournit également un moyen succinct d’orchestrer la formation des modèles Keras dans les cas communs. Une simple invocation peut ressembler à ceci:

Le coureur fournit des solutions prêtes à l’emploi pour les douleurs ML comme l’entraînement distribué et tfgnn.GraphTensor rembourrage pour formes fixes sur les TPU cloud. Au-delà de la formation sur une seule tâche (comme indiqué ci-dessus), il soutient une formation conjointe sur plusieurs (deux ou plus) tâches de concert. Par exemple, les tâches non supervisées peuvent être mélangées avec des tâches supervisées pour éclairer une représentation continue (ou intégration) finale avec des biais inductifs spécifiques d’application. Les appelants ont seulement besoin de remplacer l’argument de la tâche par une cartographie des tâches:

De plus, le coureur TF-GNN comprend également une implémentation de gradients intégrés pour une utilisation dans l’attribution du modèle. La sortie des gradients intégrés est un graphtenseur avec la même connectivité que le graphtenseur observé, mais ses caractéristiques remplacées par des valeurs de gradient où des valeurs plus grandes contribuent plus que des valeurs plus petites dans la prédiction GNN. Les utilisateurs peuvent inspecter les valeurs de gradient pour voir quelles fonctionnalités leur GNN utilise le plus.

Conclusion

En bref, nous espérons que TF-GNN sera utile pour faire avancer l’application des GNN dans Tensorflow à grande échelle et alimenter l’innovation supplémentaire dans le domaine. Si vous êtes curieux d’en savoir plus, essayez notre Démo colab Avec la référence populaire OGBN-MAG (dans votre navigateur, aucune installation requise), parcourez le reste de notre Guides et colabs utilisateurou jetez un œil à notre papier.

Remerciements

La version 1.0 de TF-GNN a été développée par une collaboration entre Google Research: Sami Abu-El-Haija, Neslihan Bulut, Bahar Fatemi, Johannes Gasteiger, Pedro Gonnet, Jonathan Halcrow, Liangze Jiang, Silvio Lattanzi, Brandon Mayer, Vahab Mirrokni, Bryan Perozi, ANTONSUL Zelle, Google Core ML: Arno Eigenwillig, Oleksandr Ferludin, Parth Kothari, Mihir Paradkar, Jan Pfeifer, Rachael Tamakloe et Google Deepmind: Alvaro Sanchez-Gonzalez et Lisa Wang.