Par Asic Q. Chen, Ruian Shi, Xiang Gao, Ricardo Baptista, Rahul G. Krishnan

Dans ce billet de blog, nous présentons StrNN, un moyen efficace d'injecter une structure variable précédemment supposée dans des réseaux neuronaux arbitraires via le masquage de poids. La nature modulaire de StrNN conduit à des applications dans l'estimation de la densité, la modélisation générative, l'inférence causale, et bien plus encore.

Basé sur l'article de NeurIPS 2023 : Réseaux neuronaux structurés pour l'estimation de la densité et l'inférence causale.

Pour exécuter votre propre version de StrNN, trouvez notre paquetage Python sur GitHub.

Figure 1: StrNN injecte une structure en masquant les poids d'un réseau neuronal. En haut : connexions StrNN (vertes) comparées à un réseau entièrement connecté (gris). En bas : La factorisation binaire d'une matrice d'adjacence produit des masques de poids. Les poids masqués sont représentés en gris.

Bien que les réseaux neuronaux soient des estimateurs de fonctions universels, il peut souvent être utile de limiter la classe de fonctions qu'un réseau neuronal peut modéliser. Par exemple, dans l'étude largement citée Deep Sets¹ Zaheer et al. se sont concentrés sur les fonctions invariantes par permutation, ce qui signifie que la sortie de la fonction reste inchangée quel que soit l'ordre des entrées, afin de mieux traiter les données d'apprentissage qui se présentent sous forme d'ensembles non ordonnés. Dans notre article, nous démontrons comment la structure des réseaux neuronaux peut conduire à des invariances de fonction bénéfiques pour d'autres applications. 

L'un des principaux cas d'utilisation motivants est l'estimation de la densité de probabilité, où nous utilisons des réseaux neuronaux pour modéliser la densité de probabilité conjointe de variables aléatoires. Dans ce contexte, nous connaissons souvent déjà certains faits sur le processus de génération des données lui-même, à savoir les relations d'indépendance entre les variables aléatoires. Nous obtenons généralement ces déclarations d'indépendance grâce à notre expertise dans le domaine ou à des algorithmes de découverte de structures. La structure de ces indépendances est généralement décrite à l'aide de réseaux bayésiens.

Plus généralement, nous utilisons une matrice d'adjacence pour modéliser des déclarations d'indépendance arbitraires pour les entrées et les sorties. Si nous disposons de données à d dimensions, notre matrice d'adjacence est une matrice binaire A∈{0, 1}^{d × d} où A_ij=0 si et seulement si _xi⊥ x_j | x{1_{, ..., i \ pa(i)}} et _Aij=1 sinon. Considérons un autoencodeur avec d entrées et sorties. Lorsque les sorties sont autorégressives par rapport aux entrées, A comporte tous les uns sous la diagonale et tous les zéros ailleurs. Les travaux existants, tels que l'autoencodeur masqué pour l'estimation de la densité (MADE)² ont exploité cette structure pour utiliser chaque nœud de sortie d'un autoencodeur afin de modéliser un facteur de probabilité marginale dans la règle de la chaîne de probabilité. Nous nous intéressons plutôt à des structures d'indépendance connues plus complexes, ce qui signifie que A n'est pas seulement triangulaire inférieur, mais aussi clairsemé sous la diagonale, comme dans l'exemple illustré à la figure 1.

En nous inspirant de l'approche utilisée dans MADE, nous renforçons la structure donnée par A en multipliant les matrices de poids du réseau neuronal par des masques binaires. De cette manière, nous annulons certaines connexions dans le réseau de sorte qu'il n'y ait pas de chemins entre des entrées et des sorties indépendantes selon A. Plus concrètement, pour un réseau neuronal jouet y=f(x) avec une seule couche cachée, nous multiplions par élément (noté ʘ) les matrices de poids W et V avec des masques binaires ^MW et ^MV.

h(x) = g((W ʘ^MW)x+b), y = f((V ʘ^MV)h(x)+c)

Pour injecter la structure prescrite par A, il suffit de trouver les matrices de masque appropriées ^MW et ^MV.

L'idée clé est que l'injection de structure dans un réseau neuronal se résume à un problème de factorisation de matrice binaire. Nous formulons le problème de la manière suivante : 

Given an adjacency matrix A∈{0,1 }^{d × d} and a neural network with L hidden layers, each with h₁, h₂, …, h_L (≥ d) hidden units, we want to factor A into mask matrices M¹∈ {0, 1}^h₁×d, M²∈ {0, 1}^h₂xh₁, …, M^L ∈ {0, 1}^{d × h_L} such that A’ ~ A where A’ = M^L × … × M² × M¹. We use the notation A’ ~ A to denote that matrices A’ and A share the same sparsity pattern, i.e.: exact same locations of zeros and non-zeros. (Overloading notation here – it doesn’t mean matrix similarity like in linear algebra!) Note that here A is a binary matrix and A’ is an integer matrix. We then mask the neural network’s hidden layers using M¹, M², …, M^L through element-wise multiplication like in the above equation to obtain a Structured Neural Network (StrNN), which respects the independence constraints prescribed by A. The value of each entry A’_ij thus corresponds to the number of connections flowing from input x_j to output x̂ _i in the StrNN.

La factorisation des matrices binaires est un problème NP-hard. Bien que la littérature existante soit abondante, elle se concentre principalement sur la recherche de facteurs de faible rang tout en minimisant la perte de reconstruction. Dans notre cas, nous avons besoin d'une perte de reconstruction nulle si nous voulons respecter la structure donnée par A, nous devons donc trouver nos propres méthodes de factorisation.

L'identifiabilité est également un problème. Lorsque la taille des couches cachées est importante, il existe de nombreux masques possibles pour une même matrice d'adjacence. Nous devons donc spécifier des objectifs d'optimisation, qui sont directement liés à l'architecture du réseau neuronal. Par exemple, dans le présent document, nous testons principalement l'hypothèse selon laquelle la maximisation du nombre de chemins restants dans notre réseau neuronal masqué permet d'améliorer l'expressivité et la généralisation. En partant de cette idée, nous étudions deux objectifs. Le premier est illustré par l'équation 2 ci-dessous :

Il maximise la somme de toutes les entrées de A'. Comme nous l'avons mentionné précédemment, cela équivaut à maximiser les chemins totaux entre toutes les entrées et sorties du StrNN. L'équation 3 constitue le deuxième objectif que nous considérons :

Nous ajoutons ici un terme de pénalité de variance, afin que les chemins restants ne soient pas trop concentrés sur une sortie. Grâce à des évaluations empiriques sur divers ensembles de données synthétiques, nous concluons que la pénalité de variance ajoutée ne fait pas de différence significative dans les performances d'estimation de densité de StrNN, et nous adoptons donc l'équation 2 comme objectif pour le reste du projet.

Nous allons maintenant nous pencher sur les méthodes de résolution du problème de factorisation de la matrice binaire. Nous pouvons trouver des solutions exactes pour maximiser les équations 2 et 3 via la programmation en nombres entiers, mais empiriquement, cette méthode est excessivement lente pour des dimensions d'entrée plus importantes (par exemple : d ≥ 20). Par conséquent, nous proposons un algorithme simple et efficace qui se rapproche de l'objectif d'optimisation (équation 2) tout en conservant une perte de reconstruction nulle lorsqu'il s'agit du modèle de rareté de la matrice d'adjacence A. Le pseudocode de notre algorithme est décrit dans l'algorithme 1 de notre article, et nous fournissons ci-dessous une visualisation d'un exemple de matrice d'adjacence.

Figure 2 : Illustration visuelle de l'algorithme 1 : l'objectif est de trouver des masques binaires appropriés à partir de la matrice d'adjacence A. À l'étape 1, nous remplissons le premier masque en copiant les lignes non nulles de A - en répétant l'opération, si nécessaire, jusqu'à ce que nous ayons rempli toutes les lignes. À l'étape 2, nous remplissons le deuxième masque en conservant autant de 1 que possible. Prenons la dernière ligne par exemple : nous initialisons d'abord les lignes à tous les 1. Nous observons que dans la dernière ligne de A, il y a des 0 dans la première et la dernière colonne. Nous trouvons alors les emplacements uniques des 1 dans la première et la dernière colonne du premier masque, en les mettant à zéro dans la dernière ligne du deuxième masque, de sorte que le produit des deux masques respecte toujours la densité donnée par la dernière ligne de A. Nous répétons cette procédure pour toutes les lignes du deuxième masque.

Voyons maintenant quelques applications de StrNN.

La normalisation des flux est un cadre populaire pour l'estimation de la densité de probabilité et les tâches de modélisation générative. Il appartient à un sous-ensemble de modèles d'estimation de la densité qui encadrent le problème comme l'apprentissage de fonctions complexes entre des espaces de haute dimension. En particulier, le cadre des flux de normalisation apprend des cartes inversibles entre une distribution de base simple et une distribution cible complexe. Cela permet une estimation simple et efficace de la vraisemblance ainsi que la génération d'échantillons. Il n'est donc pas difficile de comprendre pourquoi nous pourrions vouloir utiliser les indépendances conditionnelles connues entre les variables d'entrée lors de l'apprentissage d'un flux. Pour ce faire, nous utilisons StrNN pour renforcer les invariances de fonction dans les réseaux de flux.

La classe la plus naturelle de flux de normalisation à étendre est celle des flux autorégressifs (ARF³), qui impose une structure autorégressive entre les entrées et les sorties, de sorte que la formule de changement de variables de la probabilité, essentielle au calcul des densités de probabilité, est facile à calculer. En remplaçant les conditionneurs autorégressifs des ARF par des StrNN, nous pouvons coder des déclarations d'indépendance supplémentaires afin d'améliorer à la fois l'estimation de la vraisemblance et la qualité de la génération d'échantillons. Nous appelons cela le flux autorégressif structuré (StrAF), comme le montre la figure 3. Sur la base d'un raisonnement similaire, nous introduisons une structure dans un flux de normalisation continu appelé FFJORD⁴ en utilisant StrNN pour paramétrer l'équation différentielle qui décrit la dynamique de génération des données continues. La figure 4 montre une comparaison des échantillons générés par StrAF, StrCNF et les lignes de base.

Figure 3: Le StrAF injecte une adjacence prescrite dans chaque étape du flux à l'aide d'un conditionneur StrNN. Le StrAF ne permute pas les variables latentes, ce qui permet de respecter la matrice d'adjacence tout au long du flux.

Figure 4: Les échantillons générés par le modèle sont représentés par des points bleus pour des dimensions choisies au hasard. La densité de la vérité de terrain est représentée par les contours orange.

Tout en injectant une structure, StrAF et StrCNF héritent de l'efficacité de StrNN en raison de notre choix de masquage des poids. Plus précisément, la sortie du StrNN peut être calculée en un seul passage dans le réseau. En comparaison, les approches de masquage d'entrée telles que le Graphical Normalizing Flows⁵ doivent effectuer d'autres passages pour calculer la sortie d'une seule donnée. Cela empêche non seulement l'application efficace du masquage d'entrée à des données de grande dimension, mais constitue également un obstacle à l'intégration de la méthode dans certaines architectures. Par exemple, FFJORD nécessite déjà de nombreuses évaluations du réseau neuronal pour résoudre numériquement l'EDO définissant la carte des flux, de sorte qu'il est particulièrement inefficace d'effectuer d'autres passages par évaluation. Cela fait de notre StrNN le moyen le plus simple et le plus efficace d'injecter une structure dans ce type de flux continu.

Nous appliquons ensuite StrAF à l'estimation de l'effet causal. Nous nous appuyons sur des travaux antérieurs (Causal Autoregressive Flows⁶) qui modélisent les modèles d'équations structurelles comme des flux affines, ce qui conduit à des théorèmes d'identifiabilité favorables. Dans la figure 5, les résultats expérimentaux montrent que l'exploitation de la structure d'indépendance exacte améliore les performances sur les requêtes interventionnelles et contrefactuelles pour de nombreuses variables.

Figure 5 : Evaluations des prédictions causales (à gauche : interventions ; à droite : contrefactuels) sur des SEM à 5 et 10 variables réalisées par StrAF et CAREFL. La performance est mesurée par l'erreur quadratique moyenne totale correspondante avec l'écart-type sur dix exécutions. (a) mesure l'erreur de la valeur attendue d'une variable sous différentes interventions, tandis que (b) calcule l'erreur en dérivant des valeurs contrefactuelles sous différents échantillons et requêtes observés.

Conclusion

Nous avons présenté le réseau neuronal structuré, un approximateur de fonction qui nous permet d'injecter une structure variable arbitraire par le biais du masquage de poids. Nous avons présenté le masquage de poids comme un problème de factorisation de matrice binaire et proposé divers algorithmes pour le résoudre. Nous avons appliqué les réseaux neuronaux structurés à la normalisation des flux pour améliorer l'estimation de la densité et la modélisation générative, ce qui nous donne également un outil puissant pour modéliser les modèles d'équation de structure pour l'estimation de l'effet causal.

Travaux futurs

Dans notre travail, nous avons démontré l'avantage "plug-and-play" du StrNN en l'intégrant dans des architectures de flux pour effectuer l'estimation de la densité. De même, le StrNN peut être facilement incorporé dans d'autres architectures SOTA existantes afin d'appliquer une structure connue pour diverses tâches. Nous pensons que l'extension aux modèles de diffusion, à l'inférence variationnelle et même à l'apprentissage supervisé pourraient être des pistes prometteuses pour les travaux futurs.

Références

[1] Zaheer, Manzil, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R. Salakhutdinov et Alexander J. Smola. "Deep sets". In Advances in neural information processing systems 30 (2017).

[2] Germain, Mathieu, Karol Gregor, Iain Murray et Hugo Larochelle. "Made : Masked autoencoder for distribution estimation". In International conference on machine learning, pp. 881-889. PMLR, 2015.

[3] Huang, Chin-Wei, David Krueger, Alexandre Lacoste et Aaron Courville. "Neural autoregressive flows". In International Conference on Machine Learning, pp. 2078-2087. PMLR, 2018.

[4] Grathwohl, Will, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, et David Duvenaud. "Ffjord : Free-form continuous dynamics for scalable reversible generative models." arXiv preprint arXiv:1810.01367 (2018).

[5] Wehenkel, Antoine, et Gilles Louppe. "Graphical normalizing flows". In International Conference on Artificial Intelligence and Statistics, pp. 37-45. PMLR, 2021.

[6] Khemakhem, Ilyes, Ricardo Monti, Robert Leech et Aapo Hyvarinen. "Causal autoregressive flows". In International conference on artificial intelligence and statistics, pp. 3520-3528. PMLR, 2021.