Résumé Technique : CauScale – Découverte Causale Neurale à Grande Échelle

1. Énoncé du Problème

La découverte causale, la tâche consistant à inférer des graphes acycliques dirigés (DAG) représentant des relations causales à partir de données observationnelles, est cruciale pour des domaines allant de la bioinformatique à l'économie. Cependant, les méthodes existantes font face à de graves goulots d'étranglement de scalabilité lorsqu'elles sont appliquées à de grands graphes (par exemple, des centaines ou des milliers de nœuds).

• Les méthodes basées sur les contraintes (ex. PC, FCI) reposent sur des tests d'indépendance conditionnelle, ce qui entraîne une complexité temporelle exponentielle à mesure que le nombre de variables augmente.
• Les méthodes basées sur le score (ex. NOTEARS, SDCD) reformulent le problème en tant qu'optimisation continue mais nécessitent de résoudre un nouveau problème d'optimisation pour chaque nouveau jeu de données, ce qui les rend coûteuses en calcul à grande échelle.
• Les méthodes amorties (zero-shot) (ex. AVICI) préentraînent des modèles pour prédire directement les graphes, offrant des avantages de vitesse. Cependant, elles souffrent d'une faible efficacité en termes d'espace et de temps sur les grands graphes en raison de mécanismes d'attention qui croissent de manière quadratique avec le nombre de variables et d'une demande mémoire élevée pour maintenir des cartes d'attention spécifiques à chaque axe.

Le défi central est de développer une architecture neurale capable de réaliser une découverte causale amortie avec une grande précision tout en évoluant efficacement vers des graphes comprenant jusqu'à 1 000 nœuds sans coûts prohibitifs de mémoire ou de temps.

2. Méthodologie : L'Architecture CauScale

CauScale est une architecture neurale conçue pour une découverte causale amortie et efficace. Elle utilise une conception à deux flux (Flux de Données et Flux de Graphe) intégrée à des mécanismes d'efficacité spécifiques.

2.1. Entrée et Aperçu

Le modèle prend deux entrées :

1. Données ($\mathcal{D}$) : Un tenseur concaténant les variables observationnelles et les indicateurs d'intervention binaires.
2. Prior de Graphe ($\rho$) : Un prior statistique calculé comme l'inverse de la matrice de covariance des données.

L'architecture traite ces entrations à travers des piles alternées de Blocs Données-Graphe et d'Unités de Réduction, alimentant finalement la tête de prédiction du flux de graphe pour générer une matrice d'adjacence probabiliste.

2.2. Composants Clés

A. Conception à Deux Flux

• Flux de Données : Extrait les preuves relationnelles à partir d'observations de haute dimension. Il traite initialement la dimension d'échantillonnage complète.
• Flux de Graphe : Intègre les priors statistiques de graphes et préserve les signaux structurels clés. Il opère sur l'espace de relation nœud-nœud.
• Bloc Données-Graphe : Ce module fait le pont entre les deux flux.
  • Couche de Données : Met à jour l'embedding du flux de données.
  • Couche Data2Graph : Compresse le flux de données en une matrice de relation compressée par observation ($\omega$) en appliquant une attention axiale et un pooling. Cela distille l'évidence relationnelle en un message de graphe.
  • Couche de Graphe : Injecte le message $\omega$ dans le flux de graphe, mettant à jour l'embedding du graphe. Cela garantit que le flux de graphe bénéficie des signaux dérivés des données tout en maintenant l'intégrité structurelle.

B. Mécanismes d'Efficacité

• Unité de Réduction : Pour traiter l'efficacité temporelle, le modèle compresse l'embedding du flux de données le long de la dimension d'observation (échantillon). Tous les $k$ blocs, l'unité regroupe les échantillons en blocs de taille $r$ et applique un average pooling. Cela réduit la charge de calcul de $O(m)$ à $O(m/r)$ dans les couches suivantes, où $m$ est le nombre d'échantillons. Ceci est motivé par le fait que les signaux causaux sont principalement des dépendances entre les nœuds au sein d'un échantillon, et l'agrégation à travers des échantillons i.i.d. est moins destructrice que l'effondrement de la dimension des variables.
• Poids d'Attention Liés (Tied Attention Weights) : Pour traiter l'efficacité de l'espace, CauScale adopte des poids d'attention liés (suivant Rao et al., 2021) dans les deux flux. Au lieu de maintenir des cartes d'attention spécifiques à l'axe (qui nécessitent une mémoire de $O(R \cdot C^2)$ pour l'attention par ligne), le modèle partage les poids d'attention à travers l'axe cible. Cela réduit l'empreinte mémoire des cartes d'attention à $O(C^2)$ (ou $O(R^2)$ pour l'attention par colonne), diminuant considérablement l'utilisation de la mémoire GPU.

C. Tête de Prédiction

L'embedding de graphe final est traité par une tête de prédiction décomposée. Pour chaque paire de nœuds non ordonnée $\{i, j\}$, le modèle prédit les logits pour trois états : absence d'arête, $i \to j$, ou $j \to i$. Contrareusement à certains travaux antérieurs, CauScale n'impose pas explicitement de contraintes d'acyclicité pendant le passage direct (forward pass) pour éviter une optimisation contrainte coûteuse, s'appuyant plutôt sur les représentations apprises et une stratégie de rupture de cycle post-hoc si nécessaire.

3. Principales Contributions

1. Pré-entraînement Scalable : L'une des premières études démontrant que le pré-entraînement de réseaux neuraux pour la découverte causale peut passer à l'échelle pour des graphes allant jusqu'à 500 nœuds lors de l'entraînement et 1 000 nœuds lors de l'inférence.
2. Architecture CauScale : Une nouvelle architecture neurale qui optimise conjointement l'efficacité temporelle et la mémoire sans sacrifier la précision de la découverte, utilisant une conception à deux flux, des unités de réduction et une attention liée.
3. Validation Exhaustive : Expérimentations étendues sur des jeux de données synthétiques et semi-synthétiques (expression génique de cellule unique) démontrant une performance supérieure à travers diverses échelles de graphes, mécanismes causaux et contextes hors distribution (OOD).

de Résultats Expérimentaux

Les auteurs ont évalué CauScale par rapport aux bases de référence par contraintes (FCI), par score (NOTEARS, SDCD) et amorties (AVICI, SEA).

• Précision :

  • Sur les données synthétiques en distribution ($n=100$), CauScale a atteint 99,6 % de mAP (Mean Average Precision).
  • Sur les données hors distribution (OOD) avec des mécanismes complexes (ex. fonctions polynomiales), il a atteint 84,4 % de mAP, surpassant nettement les bases comme SEA (36,2 %) et SDCD (41,9 %).
  • Sur le jeu de données de cellule unique SERGIO-GRN, il a obtenu la meilleure performance sur tous les indicateurs par rapport à toutes les bases de référence de découverte causale.
  • Il a maintenu une haute performance (96,6 % mAP) sur des graphes de 1 000 nœuds, une échelle jamais vue lors de l'entraînement (taille max d'entraînement de 500 nœuds).

• Efficacité :

  • Temps : CauScale a atteint des vitesses d'inférence 4× à 13 000× plus rapides que les méthodes précédentes. Par exemple, pour des graphes de $n=1000$, il a pris 0,83 seconde, contre 10 896 secondes pour NOTEARS et 3,34 secondes pour AVICI.
  • Espace : CauScale a réussi à s'entraîner sur des graphes de 500 nœuds et à inférer sur 1 000 nœuds. En revanche, AVICI a échoué avec une erreur "Out of Memory" sur le jeu de données SERGIO-GRN ($n=100$) en raison des contraintes de mémoire.

• Études d'Ablation :

  • La suppression de l'Unité de Réduction a entraîné des erreurs OOM sur des ensembles d'entraînement plus larges et a considérablement dégradé l'efficacité du temps et de la mémoire.
  • La suppression du Flux de Graphe a causé la chute la plus importante de précision, soulignant l'importance de la conception à deux flux.
  • L'Attention Liée a fourni une accélération de 6× par rapport à l'attention vanilla tout en maintenant ou en améliorant la précision.

5. Signification et Revendications

L'article affirme que CauScale représente une étape concrète vers la rendre la découverte de structures causales à grande échelle accessible aux domaines scientifiques intensifs en données.

• Scalabilité : Il comble avec succès le fossé entre le besoin théorique de découverte causale à grande échelle et les limitations pratiques de mémoire et de calcul, permettant l'inférence sur des graphes de 1 000 nœuds là où les méthodes neurales précédentes échouent.
• Généralisation : Le modèle démontre une robuste généralisation aux structures de graphes OOD, bien qu'il montre une certaine sensibilité aux distributions de bruit OOD et aux fonctions de mécanisme, suggérant que les travaux futurs devraient se concentrer sur des générations de bruit plus diverses.
• Limites : Les auteurs notent modestement que la dépendance aux données d'entraînement simulées (SCM) peut limiter la généralisation aux contextes réels où la génération de données dévie de l'architecture supposée. Ils insistent sur le fait que les graphes prédits doivent être traités comme des hypothèses nécessitant une validation par des experts du domaine et des expériences en aval avant toute prise de décision à enjeux élevés.

En résumé, CauScale offre un cadre scalable, efficace et précis pour la découverte causale neurale, surmontant les goulots d'étranglement de temps et d'espace qui ont précédemment entravé l'application de l'apprentissage profond aux grands graphes causaux.