Resumo Técnico: AugMask

Declaração do Problema

Modelos de difusão baseados em escore emergiram como poderosos modelos generativos profundos, mas enfrentam desafios significativos quando aplicados a dados tabulares contendo valores ausentes. Arquiteturas de difusão padrão assumem entradas totalmente especificadas, enquanto dados tabulares do mundo real frequentemente contêm entradas ausentes (NaNs).

Abordagens existentes para preencher essa lacuna tipicamente recaem em duas categorias, ambas com limitações:

1. Imputação por Placeholder Constante: Substituir valores ausentes por constantes fixas (ex: zeros, médias das colunas) antes do treinamento. Isso força o modelo a tratar entradas ausentes identicamente a zeros válidos, induzindo um viés de esparsidade e tornando as previsões dependentes do nível de ausência. Além disso, mesmo que o erro seja mascarado para as coordenadas observadas, o placeholder determinístico permanece na entrada, atuando como uma pista espúria correlacionada com a ausência.
2. Autoimputação (ex: DiffPuter): Atualizar iterativamente as entradas ausentes durante o treinamento via uma abordagem de Expectativa-Maximização (EM). Embora isso evite placeholders fixos, cria um problema de otimização de "alvo móvel" onde o modelo é treinado em suas próprias previsões em evolução. Isso pode levar a gradientes ruidosos, dinâmicas não estacionárias e complementações excessivamente dispersas, já que os objetivos de difusão priorizam a diversidade em detrimento da precisão pontual.

O desafio central é fornecer entradas numéricas estáveis para o denoiser sem tratar as entradas preenchidas como alvos de verdade fundamental (ground-truth), evitando assim tanto o viés de placeholders constantes quanto a instabilidade da autoimputação.

Metodologia: AugMask

Os autores propõem o AugMask, um framework de treinamento plug-and-play que adapta backbones que desconhecem a ausência de dados para dados incompletos, separando o condicionamento da supervisão. O método opera em dois passos distintos:

1. Aumentação Condicional Estocástica (Construção da Entrada)

Em vez de usar um placeholder determinístico, o AugMask constró de uma entrada totalmente numérica amostrando entradas ausentes de uma distribuição dependente do contexto $T(\cdot | x_{obs}, m)$.

• Mecanismo: Para cada característica ausente $j$, um valor $Z_j$ é extraído de uma distribuição condicionada às características observadas e à máscara de ausência.
• Implementação: Os autores utilizam uma abordagem leve e fatorizada usando modelos LightGBM por característica. Para características contínuas, eles ajustam um regressor de média condicional e um regressor de log-variância para amostrar de uma Gaussiana heterocedástica. Para características categóricas, utilizam um classificador condicional com suavização de rótulos (label smoothing).
• Papel: Esses valores amostrados servem como contexto de condicionamento com incerteza para o denoiser, não como alvos de verdade fundamental. A estocasticidade codifica a incerteza inerente da ausência.

2. Supervisão Apenas de Observados (Mascaramento de Perda)

A perda de denoising é aplicada apenas às coordenadas que foram originalmente observadas.

• Objetivo: O modelo minimiza o erro de reconstrução para coordenadas observadas $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$, onde $x_{A,0}$ é a entrada aumentada.
• Efeito: As entradas aumentadas (previamente ausentes) entram no modelo através da entrada ruidosa $x_{A,t}$ para fornecer contexto, mas o modelo nunca é penalizado por falhar em reconstruí-las. Isso evita que o modelo dependa excessivamente de complementações potencialmente imprecisas.

Justificativa Teórica

O artigo conecta esta regra de treinamento a um objetivo de Rao–Blackwellização. Ao marginalizar sobre complementações plausíveis, os autores mostram que o objetivo produz uma penalidade de sensibilidade ponderada pela variância.

• Interpretação: O termo de segunda ordem na aproximação local do objetivo é proporcional à variância condicional da aumentação. Uma alta variância condicional penaliza a sensibilidade do modelo àquela complementação específica, desencorajando efetivamente a dependência de preenchimentos incertos. Uma baixa variância permite que a complementação seja usada mais diretamente como contexto.
• Contraste com AugFull: Os autores comparam o AugMask ao "AugFull" (que aplica a perda a todas as coordenadas, incluindo as aumentadas). Eles demonstram que o AugFull introduz viés se a regra de aumentação for mal especificada, pois o modelo é forçado a ajustar pseudo-alvos potencialmente incorretos. O AugMask evita isso tratando as complementações estritamente como contexto.

Principais Contribuições

1. Um Princípio de Treinamento para Difusão Tabular Incompleta: A introdução do AugMask, que desacopla o condicionamento (via aumentação estocástica) da supervisão (via mascaramento de perda apenas de observados).
2. Uma Visão de Marginalização: Uma análise teórica mostrando que a aumentação estocástica induz uma penalidade de sensibilidade ponderada pela variância, regularizando o modelo contra a dependência excessiva de complementações incertas.
3. Uma Receita Plug-and-Play: Um framework leve que permite que modelos de difusão baseados em escore padrão, que desconhecem a ausência de dados, treinem em dados incompletos sem mudanças arquiteturais.

Resultados Experimentais

Os autores avaliaram o AugMask em diversos conjuntos de dados (Adult, Default, Shoppers, Magic, Beijing, News) e regimes de ausência (MCAR e MAR) com razões de ausência variando de 0,1 a 0,9.

• Desempenho: O AugMask permite que modelos de difusão padrão que desconhecem a ausência (ex: CDTD, TabDiff) superem baselines especializados em detecção de ausência (ex: MIWAE, MissDiff, DiffPuter, ForestDiff).
• Qualidade Generativa: Modelos treinados com AugMask alcançaram fidelidade superior ao nível de amostra (medida por $\alpha$-Precision, $\beta$-Recall) e estatísticas ao nível de característica (Trend, Shape).
• Utilidade de Machine Learning: Tarefas de downstream (classificação e regressão) treinadas em dados sintéticos gerados por modelos AugMask mostraram desempenho melhorado em comparação aos baselines.
• Robustez: Os benefícios do AugMask tornaram-se mais pronunciados conforme a razão de ausência aumentava. Em altas taxas de ausência (ex: 90%), o AugMask evitou a degradação severa observada em outros métodos.
• Análise de Sensibilidade: O método provou ser robusto à má especificação da aumentação. Mesmo quando a regra de aumentação foi perturbada (ex: reduzindo a estocasticidade ou introduzindo viés), o AugMask manteve sua vantagem de desempenho sobre o AugFull, confirmando que tratar as complementações como contexto, em vez de alvos, é o fator crítico.

Significância e Alegações

O artigo alega que o AugMask fornece uma solução fundamentada para o descompasso entre backbones de difusão totalmente especificados e dados reais incompletos. Ao tratar as complementações estocásticas como contexto incerto em vez de alvos de treinamento, o método:

• Evita o viés de esparsidade e padrões artificiais introduzidos por placeholders constantes.
• Evita a instabilidade de otimização e a dinâmica de alvo móvel de métodos de autoimputação.
• Utiliza a variância da aumentação para regularizar o modelo, tornando-o menos sensível a entradas incertas.

Os autores posicionam o AugMask como uma estratégia geral que permite que modelos de difusão existentes e de alto desempenho sejam implantados em cenários com dados ausentes sem a necessidade de modificações arquiteturais complexas ou funções de perda especializadas para imputação. O trabalho foca em geração; os autores observam que estender isso para pipelines de difusão latente ou objetivos específicos de imputação permanece como uma direção para trabalhos futuros.