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

Resumo Técnico: UD-DML: Subamostragem de Design Uniforme para Duplo Aprendizado de Máquina sobre Dados Massivos

Declaração do Problema
O Duplo Aprendizado de Máquina (DML) emergiu como uma estrutura líder para inferência válida sobre parâmetros causais de baixa dimensionalidade (por exemplo, Efeito Médio do Tratamento, ATE) na presença de componentes de incômodo de alta dimensionalidade e flexíveis. No entanto, o DML padrão torna-se computacionalmente proibitivo quando aplicado a conjuntos de dados observacionais massivos ($n \in [10^5, 10^7]$), pois requer o treinamento repetido de modelos complexos de aprendizado de máquina através de dobras de ajuste cruzado. Um remédio natural é aplicar o DML a uma subamostra retirada uniformemente ($r \ll n$). Contudo, a subamostragem uniforme ingênua é estatisticamente falha: ela desconsidera a geometria do espaço de covariáveis, frequentemente falhando em preservar a distribuição das covariáveis ou manter um equilíbrio e sobreposição adequados entre tratados e controles. Em configurações observacionais, tais deficiências desestabilizam a estimação de incômodo e degradam o desempenho da inferência baseada em escores ortogonais, particularmente em regimes de baixa sobreposição.

Metodologia: UD-DML
Os autores propõem o Duplo Aprendizado de Máquina de Design Uniforme (UD-DML), uma estratégia de subamostragem baseada em design que integra princípios de design uniforme (UD) com o paradigma DML. O método opera em três fases:

1. Construção do Espaço Rotacionado por PCA:

   • As covariáveis são padronizadas e projetadas em um espaço de dimensão inferior usando Análise de Componentes Principais (PCA).
   • A dimensão $q$ é retida de modo que capture um limiar prescrito de variância cumulativa (por exemplo, $\rho_0 = 0.85$).
   • Esta rotação aborda a multicolinearidade e foca o design na estrutura geométrica dominante dos dados.

2. Construção do Esqueleto de Baixa Discrepância:

   • Um conjunto de $r_p$ pontos "esqueleto" é gerado no hipercubo unitário $[0, 1]^q$ usando o método de bons pontos de rede com saída (leave-one-out) com um gerador de potência.
   • O gerador ótimo é selecionado minimizando a discrepância de mistura, um critério de discrepância $L_2$ generalizado que garante que os pontos preencham o espaço e estejam uniformemente distribuídos.
   • Esses pontos de design são mapeados de volta para a distribuição empírica das covariáveis rotacionadas via a função de distribuição acumulada (CDF) empírica inversa, criando pontos de ancoragem $v_j$ no espaço rotacionado por PCA.

3. Seleção de Subamostra Balanceada:

   • Para cada ponto esqueleto $v_j$, o algoritmo identifica a unidade tratada mais próxima e a unidade controle mais próxima no espaço rotacionado usando uma busca em árvore KD.
   • Isso resulta em uma subamostra pareada de tamanho $r = 2r_p$ consistindo de observações originais.
   • Crucialmente, os pontos esqueleto servem apenas como âncoras de design; a estimação subsequente do DML é realizada sobre os dados observados originais $(Y_i, W_i, X_i)$ correspondentes aos índices selecionados.

4. Inferência:

   • O DML padrão com ajuste cruzado (usando escores AIPW) é aplicado à subamostra selecionada.
   • A inferência é conduzida via estatísticas de Wald baseadas na variância empírica dos pseudo-resultados com ajuste cruzado.

Contribuições Principais

1. Inovação Algorítmica: O desenvolvimento de um algoritmo de subamostragem novel que acopra a cobertura geométrica baseada em design uniforme com a inferência causal baseada em escores ortogonais, especificamente adaptado para a estimação do Efeito Médio do Tratamento (ATE).
2. Garantias Teóricas:
   • Representatividade e Equilíbrio: Os autores estabelecem limites baseados em discrepância mostrando que a subamostra selecionada herda uma representatividade favorável da distribuição completa de covariáveis e um equilíbrio inerente entre tratados e controles dentro de um framework de kernel.
   • Normalidade Assintótica: Sob condições de regularidade leves (incluindo condições de taxa de produto para estimadores de incômodo e estrutura aproximada de baixa dimensionalidade do sinal do oráculo), o estimador UD-DML é provado ser assintoticamente normal em $\sqrt{r}$, apoiando inferência Wald válida apesar de $r \ll n$.
3. Eficiência Computacional: O método reduz o custo dominante de ajuste de incômodo de $O(n)$ para $O(r)$, onde $r \ll n$, adicionando apenas uma sobrecarga de pré-processamento quase linear.
4. Validação Empírica: Simulações extensas e uma aplicação com dados reais demonstram que o UD-DML supera a subamostragem uniforme ingênua (UNIF-DML) em eficiência estatística (menor Erro Quadrático Médio, RMSE, e intervalos de confiança mais estreitos) e estabilidade inferencial (cobertura confiável), particularmente em regimes de baixa sobreposição e de má especificação de incômodo.

Resultados

• Estudos de Simulação: Em três processos de geração de dados (variando heterogeneidade e sobreposição), o UD-DML consistentemente alcançou menor Erro Quadrático Médio (RMSE) e intervalos de confiança mais apertados do que o UNIF-DML. A lacuna de desempenho ampliou-se significativamente à medida que a sobreposição deteriorava.
• Escalabilidade: Enquanto o DML com dados completos (FULL-DML) escala linearmente (ou pior) com $n$, o tempo de execução do UD-DML permaneceu estável à medida que $n$ aumentava, desde que $r$ fosse fixo.
• Robustez: Sob má especificação conjunta dos modelos de resultado e de escore de propensão, o UD-DML manteve a cobertura nominal em cenários de baixa sobreposição onde o UNIF-DML falhou catastróficamente. Isso foi atribuído ao equilíbrio induzido pelo design mitigando a confusão residual.
• Aplicação com Dados Reais: Aplicado ao conjunto de dados de Natalidade dos EUA de 2021 ($N \approx 3,67$ milhões), o UD-DML estimou o efeito do tabagismo materno no peso ao nascer. Alcançou uma aceleração de 12 vezes sobre o DML com dados completos, mantendo precisão estatística próxima ao benchmark de dados completos, enquanto o UNIF-DML mostrou variância e viés significativamente maiores.

Significado e Alegações
O artigo alega que o UD-DML resolve a tensão entre escalabilidade computacional e validade inferencial na inferência causal massiva. Ao tratar a subamostragem como um problema de design estatístico em vez de uma redução probabilística, o UD-DML constrói uma amostra de trabalho que é simultaneamente:

• Representativa: Captura a geometria dominante da distribuição completa de covariáveis.
• Balanceada: Garante que os grupos tratados e controle estejam alinhados ao longo das direções latentes principais.
• Eficiente: Permite o uso de aprendizes flexíveis de aprendizado de máquina em uma fração dos dados sem sacrificar as propriedades de consistência em $\sqrt{n}$ do framework DML subjacente.

Os autores enfatizam que o método é particularmente valioso em regimes de "baixa sobreposição" comuns em estudos observacionais, onde a subamostragem aleatória ingênua é mais propensa a produzir estimativas instáveis. A abordagem oferece um meio-termo principiado: estatisticamente próximo do DML com dados completos e computacionalmente próximo da subamostragem uniforme, mas com trade-offs viés-variância superiores.

Limitações e Direções Futuras
Os autores reconhecem que a implementação atual depende de pareamento em árvore KD em um espaço rotacionado por PCA, o que pode sofrer da maldição da dimensionalidade se o sinal relevante for de alta dimensionalidade ($p \gg r$). Trabalhos futuros poderiam explorar embeddings de baixa dimensionalidade aprendidos. Além disso, o framework atualmente limita-se a tratamentos binários via pareamento 1:1; a generalização para tratamentos contínuos ou multivalorados e estratégias híbridas com subamostragem ótima baseada em modelo são notadas como direções de pesquisa em aberto.