UD-DML: Uniform Design Subsampling for Double Machine Learning over Massive Data

Résumé technique : UD-DML : Sous-échantillonnage par plan uniforme pour l'apprentissage automatique double sur des données massives

Énoncé du problème
L'apprentissage automatique double (Double Machine Learning, DML) est devenu un cadre de référence pour l'inférence valide sur des paramètres causaux de faible dimension (par exemple, l'effet moyen du traitement, ATE) en présence de composantes de nuisance flexibles et de haute dimension. Cependant, le DML standard devient prohibitif sur le plan computationnel lorsqu'il est appliqué à des ensembles de données observationnelles massives ($n \in [10^5, 10^7]$), car il nécessite un entraînement répété de modèles d'apprentissage automatique complexes à travers des plis de couplage croisé (cross-fitting). Un remède naturel consiste à appliquer le DML à un sous-échantillon tiré uniformément ($r \ll n$). Pourtant, le sous-échantillonnage uniforme naïf est statistiquement défectueux : il néglige la géométrie de l'espace des covariables, échouant souvent à préserver la distribution des covariables ou à maintenir un équilibre et un recouvrement adéquats entre les groupes traités et témoins. Dans les contextes observationnels, ces déficiences déstabilisent l'estimation des composantes de nuisance et dégradent les performances de l'inférence basée sur les scores orthogonaux, en particulier dans les régimes de faible recouvrement.

Méthodologie : UD-DML
Les auteurs proposent le Double Apprentissage Automatique par Plan Uniforme (UD-DML), une stratégie de sous-échantillonnage basée sur la conception qui intègre les principes des plans uniformes (UD) au paradigme du DML. La méthode opère en trois phases :

1. Construction d'un espace tourné par ACP :

   • Les covariables sont standardisées et projetées dans un espace de dimension inférieure en utilisant l'Analyse en Composantes Principales (ACP).
   • La dimension $q$ est retenue de manière à capturer un seuil de variance cumulative prescrit (par exemple, $\rho_0 = 0,85$).
   • Cette rotation traite la multicolinéarité et concentre la conception sur la structure géométrique dominante des données.

2. Construction d'un squelette à faible discrépance :

   • Un ensemble de $r_p$ points « squelette » est généré dans l'hypercube unité $[0, 1]^q$ en utilisant la méthode des points de réseau optimaux avec un générateur de puissance, en laissant un point de côté (leave-one-out).
   • Le générateur optimal est sélectionné en minimisant la discrépance de mélange, un critère généralisé de discrépance $L_2$ qui garantit que les points remplissent l'espace et sont uniformément distribués.
   • Ces points de conception sont mappés vers la distribution empirique des covariables tournées via la fonction de répartition cumulative (CDF) empirique inverse, créant des points d'ancrage $v_j$ dans l'espace tourné par ACP.

3. Sélection d'un sous-échantillon équilibré :

   • Pour chaque point squelette $v_j$, l'algorithme identifie l'unité traitée la plus proche et l'unité témoin la plus proche dans l'espace tourné en utilisant une recherche par arbre KD.
   • Cela aboutit à un sous-échantillon apparié de taille $r = 2r_p$ composé d'observations originales.
   • Crucialement, les points squelette servent uniquement d'ancres de conception ; l'estimation DML subséquente est effectuée sur les données observées originales $(Y_i, W_i, X_i)$ correspondant aux indices sélectionnés.

4. Inférence :

   • Le DML couplé croisé standard (utilisant les scores AIPW) est appliqué au sous-échantillon sélectionné.
   • L'inférence est réalisée via des statistiques de Wald basées sur la variance empirique des pseudo-résultats couplés croisés.

Contributions clés

1. Innovation algorithmique : Le développement d'un nouvel algorithme de sous-échantillonnage qui couple la couverture géométrique basée sur les plans uniformes avec l'inférence causale basée sur les scores orthogonaux, spécifiquement conçu pour l'estimation de l'effet moyen du traitement (ATE).
2. Garanties théoriques :
   • Représentativité et équilibre : Les auteurs établissent des bornes basées sur la discrépance montrant que le sous-échantillon sélectionné hérite d'une représentativité favorable de la distribution complète des covariables et d'un équilibre inhérent entre les groupes traités et témoins dans un cadre à noyau.
   • Normalité asymptotique : Sous des conditions de régularité légères (incluant des conditions de taux de produit pour les estimateurs de nuisance et une structure de signal de faible dimension approximative du signal oracle), l'estimateur UD-DML est prouvé être $\sqrt{r}$-asymptotiquement normal, soutenant une inférence de Wald valide malgré $r \ll n$.
3. Efficacité computationnelle : La méthode réduit le coût dominant de l'ajustement des nuisances de $O(n)$ à $O(r)$, où $r \ll n$, tout en ajoutant seulement une surcharge de prétraitement quasi-linéaire.
4. Validation empirique : Des simulations extensives et une application sur des données réelles démontrent que UD-DML surpasse le sous-échantillonnage uniforme naïf (UNIF-DML) en termes d'efficacité statistique (RMSE plus faible, intervalles de confiance plus étroits) et de stabilité inférentielle (couverture fiable), en particulier dans les régimes de faible recouvrement et de mauvaise spécification des nuisances.

Résultats

• Études de simulation : À travers trois processus de génération de données (variabilité d'hétérogénéité et de recouvrement), UD-DML a constamment atteint un Erreur Quadratique Moyenne (RMSE) plus faible et des intervalles de confiance plus serrés que UNIF-DML. L'écart de performance s'est considérablement élargi à mesure que le recouvrement se détériorait.
• Évolutivité : Alors que le DML sur données complètes (FULL-DML) évolue linéairement (ou pire) avec $n$, le temps d'exécution de UD-DML est resté stable à mesure que $n$ augmentait, à condition que $r$ soit fixe.
• Robustesse : Sous une mauvaise spécification conjointe des modèles de résultat et de score de propension, UD-DML a maintenu une couverture nominale dans les scénarios de faible recouvrement où UNIF-DML a échoué de manière catastrophique. Cela a été attribué à l'équilibre induit par la conception atténuant la confusion résiduelle.
• Application sur données réelles : Appliqué à l'ensemble de données de natalité américaines de 2021 ($N \approx 3,67$ millions), UD-DML a estimé l'effet du tabagisme maternel sur le poids de naissance. Il a obtenu une accélération de 12 fois par rapport au DML sur données complètes tout en maintenant une précision statistique proche de la référence sur données complètes, tandis que UNIF-DML a montré une variance et un biais significativement plus élevés.

Signification et affirmations
L'article affirme que UD-DML résout la tension entre l'évolutivité computationnelle et la validité inférentielle dans l'inférence causale massive. En traitant le sous-échantillonnage comme un problème de conception statistique plutôt que comme une réduction probabiliste, UD-DML construit un échantillon de travail qui est simultanément :

• Représentatif : Il capture la géométrie dominante de la distribution complète des covariables.
• Équilibré : Il assure que les groupes traités et témoins sont alignés le long des directions latentes principales.
• Efficace : Il permet l'utilisation d'apprenants flexibles en apprentissage automatique sur une fraction des données sans sacrifier les propriétés de consistance $\sqrt{n}$ du cadre DML sous-jacent.

Les auteurs soulignent que la méthode est particulièrement précieuse dans les régimes de « faible recouvrement » courants dans les études observationnelles, où le sous-échantillonnage aléatoire naïf est le plus susceptible de produire des estimations instables. L'approche offre un compromis fondé sur des principes : statistiquement proche du DML sur données complètes et computationnellement proche du sous-échantillonnage uniforme, mais avec des compromis biais-variance supérieurs.

Limites et orientations futures
Les auteurs reconnaissent que l'implémentation actuelle repose sur l'appariement par arbre KD dans un espace tourné par ACP, ce qui peut souffrir du fléau de la dimensionnalité si le signal pertinent est de haute dimension ($p \gg r$). Un travail futur pourrait explorer des plongements de faible dimension appris. De plus, le cadre est actuellement limité aux traitements binaires via un appariement 1:1 ; la généralisation aux traitements continus ou multi-valeurs et les stratégies hybrides avec un sous-échantillonnage optimal basé sur le modèle sont notées comme des axes de recherche ouverts.