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

Resumen Técnico: UD-DML: Muestreo por Diseño Uniforme para Aprendizaje Automático Doble sobre Datos Masivos

Enunciado del Problema
El Aprendizaje Automático Doble (DML, por sus siglas en inglés) ha surgido como un marco líder para la inferencia válida sobre parámetros causales de baja dimensión (por ejemplo, el Efecto Promedio del Tratamiento, ATE) en presencia de componentes de confusión de alta dimensión y flexibles. Sin embargo, el DML estándar se vuelve computacionalmente prohibitivo al aplicarse a conjuntos de datos observacionales masivos ($n \in [10^5, 10^7]$), ya que requiere el entrenamiento repetido de modelos complejos de aprendizaje automático a través de pliegues de ajuste cruzado. Un remedio natural es aplicar DML a una submuestra extraída uniformemente ($r \ll n$). No obstante, el muestreo uniforme ingenuo es estadísticamente defectuoso: ignora la geometría del espacio de covariables, a menudo fallando en preservar la distribución de las covariables o mantener un equilibrio y superposición adecuados entre tratados y controles. En contextos observacionales, tales deficiencias desestabilizan la estimación de los componentes de confusión y degradan el rendimiento de la inferencia basada en puntuaciones ortogonales, particularmente en regímenes de baja superposición.

Metodología: UD-DML
Los autores proponen Aprendizaje Automático Doble por Diseño Uniforme (UD-DML), una estrategia de muestreo basada en diseño que integra los principios del diseño uniforme (DU) con el paradigma DML. El método opera en tres fases:

1. Construcción del Espacio Rotado por PCA:

   • Las covariables se estandarizan y se proyectan en un espacio de menor dimensión utilizando el Análisis de Componentes Principales (PCA).
   • Se retiene la dimensión $q$ de tal manera que capture un umbral prescrito de varianza acumulada (por ejemplo, $\rho_0 = 0.85$).
   • Esta rotación aborda la multicolinealidad y centra el diseño en la estructura geométrica dominante de los datos.

2. Construcción del Esqueleto de Baja Discrepancia:

   • Se genera un conjunto de $r_p$ puntos "esqueleto" en el hipercubo unitario $[0, 1]^q$ utilizando el método de puntos de retícula buenos con salida de uno, con un generador de potencia.
   • El generador óptimo se selecciona minimizando la discrepancia de mezcla, un criterio generalizado de discrepancia $L_2$ que asegura que los puntos llenen el espacio y estén distribuidos uniformemente.
   • Estos puntos de diseño se mapean de vuelta a la distribución empírica de las covariables rotadas mediante la función de distribución acumulada (CDF) empírica inversa, creando puntos de anclaje $v_j$ en el espacio rotado por PCA.

3. Selección de la Submuestra Equilibrada:

   • Para cada punto esqueleto $v_j$, el algoritmo identifica la unidad tratada más cercana y la unidad control más cercana en el espacio rotado utilizando una búsqueda KD-tree.
   • Esto resulta en una submuestra emparejada de tamaño $r = 2r_p$ compuesta por observaciones originales.
   • Crucialmente, los puntos esqueleto sirven solo como anclajes de diseño; la estimación DML subsiguiente se realiza sobre los datos observados originales $(Y_i, W_i, X_i)$ correspondientes a los índices seleccionados.

4. Inferencia:

   • Se aplica el DML estándar con ajuste cruzado (utilizando puntuaciones AIPW) a la submuestra seleccionada.
   • La inferencia se realiza mediante estadísticos de Wald basados en la varianza empírica de los pseudo-resultados con ajuste cruzado.

Contribuciones Clave

1. Innovación Algorítmica: El desarrollo de un algoritmo de muestreo novedoso que acopla la cobertura geométrica basada en diseño uniforme con la inferencia causal basada en puntuaciones ortogonales, específicamente adaptado para la estimación del Efecto Promedio del Tratamiento (ATE).
2. Garantías Teóricas:
   • Representatividad y Equilibrio: Los autores establecen cotas basadas en discrepancia que muestran que la submuestra seleccionada hereda una representatividad favorable de la distribución completa de covariables y un equilibrio inherente entre tratados y controles dentro de un marco de kernels.
   • Normalidad Asintótica: Bajo condiciones de regularidad leves (incluyendo condiciones de tasa de producto para los estimadores de componentes de confusión y una estructura aproximada de baja dimensión de la señal óptima), se demuestra que el estimador UD-DML es $\sqrt{r}$-asintóticamente normal, apoyando una inferencia de Wald válida a pesar de que $r \ll n$.
3. Eficiencia Computacional: El método reduce el costo dominante de ajuste de componentes de confusión de $O(n)$ a $O(r)$, donde $r \ll n$, añadiendo solo una sobrecarga de preprocesamiento cuasi-lineal.
4. Validación Empírica: Simulaciones extensas y una aplicación con datos reales demuestran que UD-DML supera al muestreo uniforme ingenuo (UNIF-DML) en eficiencia estadística (menor Error Cuadrático Medio, intervalos de confianza más estrechos) y estabilidad inferencial (cobertura confiable), particularmente en regímenes de baja superposición y de mala especificación de componentes de confusión.

Resultados

• Estudios de Simulación: A través de tres procesos generadores de datos (variando heterogeneidad y superposición), UD-DML logró consistentemente un menor Error Cuadrático Medio (RMSE) e intervalos de confianza más ajustados que UNIF-DML. La brecha de rendimiento se amplió significativamente a medida que empeoraba la superposición.
• Escalabilidad: Mientras que el DML con datos completos (FULL-DML) escala linealmente (o peor) con $n$, el tiempo de ejecución de UD-DML se mantuvo estable a medida que aumentaba $n$, siempre que $r$ se mantuviera fijo.
• Robustez: Bajo una mala especificación conjunta de los modelos de resultado y de puntuación de propensión, UD-DML mantuvo la cobertura nominal en escenarios de baja superposición donde UNIF-DML falló catastróficamente. Esto se atribuyó al equilibrio inducido por el diseño que mitiga la confusión residual.
• Aplicación con Datos Reales: Aplicado al conjunto de datos de Natalidad de EE. UU. de 2021 ($N \approx 3.67$ millones), UD-DML estimó el efecto del tabaquismo materno en el peso al nacer. Logró una aceleración de 12 veces sobre el DML con datos completos mientras mantenía una precisión estadística cercana a la referencia de datos completos, whereas UNIF-DML mostró una varianza y un sesgo significativamente mayores.

Significado y Afirmaciones
El artículo afirma que UD-DML resuelve la tensión entre la escalabilidad computacional y la validez inferencial en la inferencia causal masiva. Al tratar el muestreo como un problema de diseño estadístico en lugar de una reducción probabilística, UD-DML construye una muestra de trabajo que es simultáneamente:

• Representativa: Captura la geometría dominante de la distribución completa de covariables.
• Equilibrada: Asegura que los grupos tratados y control estén alineados a lo largo de las direcciones latentes principales.
• Eficiente: Permite el uso de aprendices flexibles de aprendizaje automático en una fracción de los datos sin sacrificar las propiedades de consistencia $\sqrt{n}$ del marco DML subyacente.

Los autores enfatizan que el método es particularmente valioso en regímenes de "baja superposición" comunes en estudios observacionales, donde el muestreo aleatorio ingenuo es más propenso a producir estimaciones inestables. El enfoque ofrece un punto medio principista: estadísticamente cercano al DML con datos completos y computacionalmente cercano al muestreo uniforme, pero con compensaciones de sesgo-varianza superiores.

Limitaciones y Direcciones Futuras
Los autores reconocen que la implementación actual depende del emparejamiento KD-tree en un espacio rotado por PCA, lo cual puede sufrir la maldición de la dimensionalidad si la señal relevante es de alta dimensión ($p \gg r$). El trabajo futuro podría explorar incrustaciones de baja dimensión aprendidas. Además, el marco se limita actualmente a tratamientos binarios mediante emparejamiento 1:1; la generalización a tratamientos continuos o de múltiples valores y estrategias híbridas con muestreo óptimo basado en modelos se señalan como direcciones de investigación abiertas.