Résumé technique : Apprentissage d'ordres causaux pour la prédiction tabulaire en contexte

1. Énoncé du problème

Les architectures d'apprentissage en contexte, telles que les modèles de base tabulaires (TFM), ont atteint de fortes performances prédictives sur des données tabulaires en exploitant un pré-entraînement à grande échelle. Cependant, ces modèles sont principalement associatifs ; ils reposent sur des structures corrélatoires plutôt que sur une directionnalité causale explicite. Par conséquent, leurs prédictions peuvent être peu fiables en cas de décalage de distribution ou d'interventions, car ils peuvent exploiter des associations statistiques non causales qui ne se généralisent pas.

Bien que des méthodes établies de découverte causale existent pour retrouver des graphes acycliques orientés (DAG), elles sont souvent spécifiques à un jeu de données, intensives en calcul et découplées des architectures prédictives modernes. À l'inverse, les approches récentes de découverte causale amortie reposent sur un entraînement supervisé sur des graphes synthétiques, nécessitant des étiquettes causales de vérité terrain. Il existe un vide dans l'intégration systématique de la découverte de structures causales au sein d'architectures prédictives, dans un cadre d'apprentissage en contexte non supervisé capable de gérer les données manquantes et les interventions.

2. Méthodologie : TABORDER

Les auteurs proposent TABORDER, une architecture basée sur les transformateurs conçue pour inférer simultanément des ordres causaux de variables et effectuer des prédictions robustes. L'innovation centrale réside dans l'utilisation d'une attention contrainte par l'ordre causal.

Conception de l'architecture

TABORDER fonctionne comme un système à double transformateur qui traite un jeu de données $D$ avec des entrées manquantes indiquées par un masque $M$ :

1. Encodage : Chaque cellule $x_i^{(r)}$ et son indicateur de manquance sont encodés dans une représentation vectorielle.
2. Branche d'inférence d'ordre : Un transformateur standard à attention multi-têtes traite les encodages pour inférer un score scalaire $s_i$ pour chaque variable $X_i$. Des scores plus faibles indiquent une position antérieure dans l'ordre causal. Ces scores sont moyennés sur les lignes pour produire un ordre global $\hat{\pi}$.
3. Branche de prédiction contrainte par l'ordre : Un second transformateur modélise la distribution conditionnelle de chaque cellule. Crucialement, son attention par colonne est contrainte par un masque dérivé des scores appris $s$. Plus précisément, la variable $X_i$ ne peut porter attention qu'aux variables $X_j$ où $s_j \leq s_i$. Cela impose une factorisation ordonnée de la distribution conjointe :
   $$p_\theta(x^{(r)} | D, \pi) = \prod_{k=1}^d p_\theta(x_{\pi(k)}^{(r)} | x_{\pi(1:k-1)}^{(r)}, D)$$
   Cela empêche le modèle d'utiliser des caractéristiques « anti-causales » (causes futures) pour prédire une cible, réduisant ainsi la sensibilité aux corrélations fallacieuses.

Gestion des données manquantes et de l'incertitude

Le modèle prend explicitement en compte la manquance chez les parents causaux. Sous l'hypothèse d'un modèle à bruit additif (ANM) ($X_i = f(X_{pa(i)}) + N_i$), le modèle décode une moyenne conditionnelle $\hat{\mu}$ et une variance $\hat{\sigma}^2$.

• Augmentation de la variance : Si un parent causal est manquant, le modèle augmente la variance prédictive pour la variable cible d'une quantité apprenable $\delta_{i,j}$. Cela capture l'incertitude accrue lorsque les entrées pertinentes ne sont pas observées.
• Objectif d'entraînement : Le modèle est entraîné de bout en bout en maximisant la vraisemblance des entrées masquées. L'objectif couple l'apprentissage de l'ordre et la prédiction : les ordres qui expliquent mieux les données masquées (en minimisant la variance résiduelle) reçoivent une vraisemblance plus élevée.

Justification théorique

Les auteurs justifient l'approche basée sur la vraisemblance dans le cadre du modèle à bruit additif (ANM). Ils démontrent que pour les modèles causaux additifs, l'ordre optimal de vraisemblance est cohérent avec l'ordre topologique du DAG sous-jacent réel. De plus, ils prouvent un théorème concernant l'asymétrie de vraisemblance consciente de la manquance : la prise en compte de la manquance via des augmentations de variance amplifie l'écart de vraisemblance entre la direction causale correcte et les directions incorrectes, facilitant l'identification des ordres causaux même en présence de données manquantes.

3. Contributions clés

1. Intégration de la découverte causale et de la prédiction : TABORDER comble le fossé entre l'apprentissage de structures causales et la prédiction tabulaire en contexte en inférant des ordres causaux de manière non supervisée au sein d'un cadre de transformateur prédictif.
2. Attention contrainte par l'ordre : L'introduction d'un masque d'attention différentiable et appris qui restreint le flux d'information aux variables précédant dans un ordre causal appris, garantissant que les prédictions reposent uniquement sur des causes potentielles.
3. Apprentissage causal conscient de la manquance : La conception d'un mécanisme où la manquance chez les parents causaux augmente explicitement la variance prédictive, ce qui aide théoriquement et empiriquement à distinguer les directions causales.
4. Apprentissage non supervisé : Contrairement aux méthodes de découverte causale amortie, TABORDER ne nécessite pas de graphes causaux de vérité terrain lors de l'entraînement ; il apprend les ordres optimaux uniquement par maximisation de la vraisemblance sur des données masquées.

4. Résultats expérimentaux

Les auteurs évaluent TABORDER sur des données synthétiques avec des structures causales connues et des benchmarks du monde réel.

• Inférence d'ordre causal : Sur des données synthétiques générées à partir d'ANM non linéaires à processus gaussiens, TABORDER atteint une divergence topologique ($d_{TOP}$) comparable à celle d'algorithmes spécialisés de découverte causale (par exemple, CAM, SCORE, NOTEARS), mais infère les ordres en un temps quasi instantané via un seul passage avant. Sur le jeu de données réel de signalisation protéique Sachs, TABORDER découvre des ordres causaux avec une faible divergence ($d_{TOP} \approx 0,21$) par rapport au graphe de consensus, surpassant d'autres références.
• Performance prédictive et imputation : Dans les tâches d'imputation de valeurs manquantes (benchmarks OpenML CC18 et CTR23), TABORDER atteint des performances compétitives, en particulier à des taux de manquance élevés ($\geq 40\%$), surpassant les méthodes standard comme KNN, MICE et GAIN.
• Robustesse aux interventions : Dans une chaîne à trois variables ($X \to Y \to Z$) avec une intervention sur le mécanisme $Y \to Z$, TABORDER maintient une prédiction précise du médiateur $Y$ étant donné $X$ et $Z$. En revanche, les TFM standards (comme TABPFN) et XGBoost se dégradent considérablement, car ils reposent sur l'association non causale brisée entre $Y$ et $Z$.
• Insights biologiques : Sur le jeu de données Sachs sous diverses conditions d'intervention, TABORDER infère des ordres qui reflètent des réorganisations biologiquement plausibles. Par exemple, les interventions en amont (par exemple, l'inhibition de PKC) provoquent des changements globaux dans l'ordre inféré, tandis que les interventions en aval provoquent des changements localisés, s'alignant sur la dynamique connue des voies de signalisation.

5. Importance et affirmations

L'article affirme que TABORDER fournit une approche rationnelle pour rendre l'apprentissage en contexte pour les données tabulaires plus robuste aux décalages de distribution et aux interventions en modélisant explicitement la directionnalité causale.

• Robustesse : En imposant que les prédictions reposent uniquement sur les prédécesseurs dans un ordre causal appris, le modèle évite d'exploiter des corrélations fallacieuses qui échouent sous intervention.
• Efficacité : Contrairement à la découverte causale classique qui résout des problèmes d'optimisation par jeu de données, TABORDER infère la structure efficacement au sein d'un seul passage avant, la rendant évolutive pour les tâches en aval.
• Fondement théorique : Ce travail établit que la maximisation de la vraisemblance en présence de données manquantes peut servir d'approximation pour la découverte causale sous des hypothèses ANM standard, fournissant une base théorique pour l'apprentissage non supervisé d'ordres.

Les auteurs reconnaissent des limites, notant que l'implémentation actuelle se concentre sur les ordres topologiques plutôt que sur des DAG complets et suppose des caractéristiques à valeurs réelles. Ils suggèrent que bien que la méthode soit robuste, les structures inférées sont des hypothèses nécessitant une validation par domaine, et que des travaux futurs pourraient explorer des priors d'entraînement plus riches et des données de types mixtes.