Measuring Variable Importance in Heterogeneous Treatment Effects with Confidence

Resumen Técnico: Medición de la Importancia de Variables en Efectos de Tratamiento Heterogéneos con Confianza

Enunciado del Problema
El artículo aborda el desafío de estimar la importancia de las variables (covariables) en el contexto de Efectos de Tratamiento Heterogéneos (HTE), específicamente el Efecto Promedio de Tratamiento Condicional (CATE). Si bien el Aprendizaje Automático (ML) ofrece herramientas poderosas para predecir respuestas individuales al tratamiento a partir de datos biomédicos de alta dimensión, su falta de interpretabilidad dificulta su adopción científica y clínica. Los métodos existentes para la importancia de variables en la predicción estándar (por ejemplo, Leave-One-Covariate-Out, LOCO) o en la inferencia causal a menudo enfrentan compensaciones entre el costo computacional, las garantías estadísticas (control del error Tipo I) y la robustez frente a errores de estimación en muestras finitas. En el contexto causal, la indisponibilidad de resultados potenciales de verdad fundamental complica la puntuación de los estimadores, y la necesidad de volver a ajustar modelos complejos para cada variable (como en LOCO) puede conducir a una alta varianza y a una potencia estadística reducida, particularmente en los regímenes de muestras pequeñas a medianas comunes en aplicaciones biomédicas.

Metodología: PermuCATE
Los autores proponen PermuCATE, un algoritmo para la evaluación estadísticamente rigurosa de la importancia global de variables en la estimación de CATE. El método se basa en la Importancia de Permutación Condicional (CPI) adaptada para la inferencia causal.

• Mecanismo Central: A diferencia de LOCO, que vuelve a ajustar el estimador de CATE en un subconjunto de covariables (excluyendo la variable de interés), PermuCATE reutiliza el estimador de CATE entrenado. Perturba los datos de entrada permutando los residuos de la variable de interés ($X_j$) después de regresarla sobre las covariables restantes ($X_{-j}$).
• Muestreo Condicional: El algoritmo asume que la relación entre las covariables es más simple que la relación entre las covariables y el resultado (Suposición 3.1). Estima un predictor condicional $\hat{\nu}_j(X_{-j}) \approx E[X_j | X_{-j}]$ y construye una covariable perturbada $X'_j = \hat{\nu}_j + \text{shuffle}(X_j - \hat{\nu}_j)$.
• Puntuación: La importancia de la variable se mide como el aumento en un riesgo causal factible (específicamente el riesgo PO) cuando se utiliza la covariable perturbada frente a la covariable original. Este riesgo se calcula utilizando pseudo-resultados derivados de funciones de molestia (funciones de respuesta y puntuaciones de propensión).
• Base Teórica: Los autores proporcionan un análisis teórico de muestras finitas que muestra que la varianza de PermuCATE está impulsada principalmente por el error en la estimación de la distribución condicional ($\hat{\nu}_j$), mientras que la varianza de LOCO está impulsada por el error en la estimación del propio modelo CATE ($\hat{\tau}$) y sus submodelos. Dado que estimar la distribución condicional de las covariables es generalmente más fácil que estimar la función CATE compleja, se espera que PermuCATE tenga una varianza menor.

Contribuciones Clave

1. Propuesta de Algoritmo: Introducción de PermuCATE, un procedimiento agnóstico al modelo para la importancia de variables en CATE que generaliza CPI a contextos causales.
2. Análisis Teórico: Desarrollo de un análisis de varianza de muestras finitas que compara PermuCATE y LOCO. El análisis demuestra que PermuCATE es más robusto frente a errores de estimación en el modelo CATE, un factor crítico en tamaños de muestra pequeños o al utilizar estimadores complejos de aprendizaje profundo.
3. Validación Empírica: Experimentos exhaustivos en datos simulados (Baja Dimensión, Lineal de Alta Dimensión, Polinomial de Alta Dimensión) y en un benchmark semisintético del mundo real (IHDP). El estudio evalúa el rendimiento en diversos estimadores de CATE, incluidos Bosques Aleatorios (Causal Forest), Redes Neuronales Profundas (CateNet) y modelos base basados en Transformers (TabPFN).

Resultados

• Reducción de Varianza: Tanto en escenarios lineales como no lineales, PermuCATE demostró consistentemente una varianza menor que LOCO. Esta diferencia fue más pronunciada en tamaños de muestra pequeños y al utilizar modelos estocásticos complejos (por ejemplo, redes neuronales profundas) para la estimación de CATE.
• Potencia Estadística: Debido a la menor varianza, PermuCATE logró una mayor potencia estadística (la capacidad de identificar correctamente variables verdaderamente positivas) en comparación con LOCO. En configuraciones de alta dimensión ($d=50, 100$), PermuCATE mantuvo su potencia a medida que aumentaba la dimensionalidad, mientras que la potencia de LOCO disminuyó significativamente.
• Robustez frente a Especificación Incorrecta del Modelo: En el benchmark IHDP, PermuCATE proporcionó mejores clasificaciones de importancia de variables (mayor AUC) y potencia estadística en todos los aprendices probados (Causal Forest, CateNet, TabPFN). Cabe destacar que, al utilizar CateNet, LOCO sufrió de ruido de optimización inherente al entrenamiento estocástico, lo que llevó a tasas más altas de error Tipo I en algunos casos, mientras que PermuCATE mantuvo el control.
• Eficiencia Computacional: PermuCATE evita la necesidad de volver a ajustar el modelo CATE completo para cada variable, ofreciendo ventajas computacionales significativas sobre LOCO, especialmente cuando el estimador de CATE es computacionalmente costoso.

Significado y Afirmaciones
El artículo afirma que PermuCATE ofrece una alternativa superior a LOCO para la importancia de variables en la inferencia causal, particularmente en el "régimen de datos limitados" típico de la investigación biomédica. Los autores enfatizan que, mientras que el trabajo anterior se centró intensamente en el control del error Tipo I, su análisis destaca el impacto crítico del error Tipo II (falsos negativos) impulsado por la varianza de muestras finitas.

El significado de este trabajo radica en permitir el uso de modelos de ML complejos de alta capacidad (como aprendizaje profundo y modelos base) para la estimación de CATE, manteniendo al mismo tiempo una interpretabilidad fiable y estadísticamente rigurosa. Esto cierra la brecha entre el modelado predictivo avanzado y la necesidad de conocimientos científicos accionables en estudios intervencionistas, como los ensayos clínicos. Los autores afirman que su método facilita el descubrimiento de biomarcadores predictivos al proporcionar explicaciones robustas para la heterogeneidad del tratamiento, apoyando así la toma de decisiones médicas y acelerando el desarrollo de tratamientos. El enfoque se presenta como aplicable más allá de la biomedicina a cualquier dominio que requiera inferencia causal con ML interpretable.