Resumen Técnico: AugMask

Planteamiento del Problema

Los modelos de difusión basados en puntuación (score-based) han emergido como potentes modelos generativos profundos, pero enfrentan desafíos significativos cuando se aplican a datos tabulares que contienen valores faltantes. Las arquitecturas de difusión estándar asumen entradas completamente especificadas, mientras que los datos tabulares del mundo real suelen contener entradas faltantes (NaNs).

Los enfoques existentes para cerrar esta brecha suelen caer en dos categorías, ambas con limitaciones:

1. Imputación por Marcadores de Posición Constantes: Reemplazar los valores faltantes con constantes fijas (por ejemplo, ceros o la media de la columna) antes del entrenamiento. Esto obliga al modelo a tratar las entradas faltantes de la misma manera que los ceros válidos, induciendo un sesgo de dispersión y haciendo que las predicciones dependan del nivel de ausencia de datos. Además, incluso si la pérdida se enmascara para las coordenadas observadas, el marcador de posición determinista permanece en la entrada, actuando como una pista espuria correlacionada con la ausencia de datos.
2. Autoimputación (ej. DiffPuter): Actualizar iterativamente las entradas faltantes durante el entrenamiento mediante un enfoque de Esperanza-Maximización (EM). Aunque esto evita los marcadores de posición fijos, crea un problema de optimización de "objetivo móvil" donde el modelo se entrena sobre sus propias predicciones en evolución. Esto puede conducir a gradientes ruidosos, dinámicas no estacionarias y completaciones excesivamente dispersas, ya que los objetivos de difusión priorizan la diversidad sobre la precisión puntual.

El desafío central es proporcionar entradas numéricas estables para el denoiser sin tratar los valores completados como objetivos de verdad fundamental (ground-truth), evitando así tanto el sesgo de los marcadores de posición constantes como la inestabilidad de la autoimputación.

Metodología: AugMask

Los autores proponen AugMask, un marco de entrenamiento plug-and-play que adapta los núcleos (backbones) que ignoran la ausencia de datos para trabajar con datos incompletos, separando el condicionamiento de la supervisión. El método opera en dos pasos distintos:

1. Aumento Condicional Estocástico (Construcción de la Entrada)

En lugar de utilizar un marcador de posición determinista, AugMask construye una entrada totalmente numérica muestreando las entradas faltantes de una distribución dependiente del contexto $T(\cdot | x_{obs}, m)$.

• Mecanismo: Para cada característica faltante $j$, se extrae un valor $Z_j$ de una distribución condicionada a las características observadas y a la máscara de ausencia.
• Implementación: Los autores utilizan un enfoque ligero y factorizado empleando modelos LightGBM por característica. Para características continuas, ajustan un regresor de media condicional y un regresor de log-varianza para muestrear de una Gaussiana heterocedástica. Para características categóricas, utilizan un clasificador condicional con suavizado de etiquetas (label smoothing).
• Rol: Estos valores muestreados sirven como contexto de condicionamiento con incertidumbre para el denoiser, no como objetivos de verdad fundamental. La estocasticidad codifica la incertidumbre inherente de la ausencia de datos.

2. Supervisión Solo de Observados (Enmascaramiento de la Pérdida)

La pérdida de eliminación de ruido (denoising loss) se aplica únicamente a las coordenadas que originalmente fueron observadas.

• Objetivo: El modelo minimiza el error de reconstrucción para las coordenadas observadas $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$, donde $x_{A,0}$ es la entrada aumentada.
• Efecto: Las entradas aumentadas (previamente faltantes) entran al modelo a través de la entrada ruidosa $x_{A,t}$ para proporcionar contexto, pero el modelo nunca es penalizado por fallar en su reconstrucción. Esto evita que el modelo dependa excesivamente de completaciones potencialmente inexactas.

Justificación Teórica

El artículo conecta esta regla de entrenamiento con un objetivo de Rao–Blackwellized. Al marginalizar sobre las completaciones plausibles, los autores demuestven que el objetivo produce una penalización de sensibilidad ponderada por la varianza.

• Interpretación: El término de segundo orden en la aproximación local del objetivo es proporcional a la varianza condicional del aumento. Una alta varianza condicional penaliza la sensibilidad del modelo hacia esa completación específica, desalentando efectivamente la dependencia de llenados inciertos. Una baja varianza permite que la completación se utilice más directamente como contexto.
• Contraste con AugFull: Los autores comparan AugMask con "AugFull" (que aplica la pérdida a todas las coordenadas, incluyendo las aumentadas). Demuestran que AugFull introduce sesgo si la regla de aumento está mal especificada, ya que el modelo se ve obligado a ajustarse a posibles pseudo-objetivos incorrectos. AugMask evita esto al tratar las completaciones estrictamente como contexto.

Contribuciones Clave

1. Un Principio de Entrenamiento para la Difusión Tabular Incompleta: La introducción de AugMask, que desacopla el condicionamiento (mediante aumento estocástico) de la supervisión (mediante enmascaramiento de pérdida solo de observados).
2. Una Visión de Marginalización: Un análisis teórico que muestra cómo el aumento estocástico induce una penalización de sensibilidad ponderada por la varianza, regularizando al modelo contra la dependencia excesiva de completaciones inciertas.
3. Una Receta Plug-and-Play: Un marco ligero que permite a los modelos de difusión basados en puntuación estándar (que ignoran la ausencia de datos) entrenar con datos incompletos sin cambios arquitectónicos.

Resultados Experimentales

Los autores evaluaron AugMask en diversos conjuntos de datos (Adult, Default, Shoppers, Magic, Beijing, News) y regímenes de ausencia (MCAR y MAR) con proporciones de ausencia que van desde 0.1 hasta 0.9.

• Desempeño: AugMask permite que los modelos de difusión estándar que ignoran la ausencia de datos (ej. CDTD, TabDiff) superen a las líneas base especializadas en la detección de ausencias (ej. MIWAE, MissDiff, DiffPuter, ForestDiff).
• Calidad Generativa: Los modelos entrenados con AugMask lograron una fidelidad de muestra superior (medida por $\alpha$-Precision, $\beta$-Recall) y estadísticas de nivel de característica (Trend, Shape).
• Utilidad de Aprendizaje Automático: Las tareas de downstream (clasificación y regresión) entrenadas con datos sintéticos generados por modelos AugMask mostraron un rendimiento mejorado en comparación con las líneas base.
• Robustez: Los beneficios de AugMask se hicieron más pronunciados a medida que aumentaba la proporción de ausencia. Con una ausencia alta (ej. 90%), AugMask evitó la degradación severa observada en otros métodos.
• Análisis de Sensibilidad: El método demostró ser robusto a la especificación errónea del aumento. Incluso cuando la regla de aumento fue perturbada (ej. reduciendo la estocasticidad o introduciendo sesgo), AugMask mantuvo una ventaja de desempeño sobre AugFull, confirmando que tratar las completaciones como contexto en lugar de como objetivos es el factor crítico.

Significado y Reivindicaciones

El artículo afirma que AugMask proporciona una solución fundamentada al desajuste entre los backbones de difusión totalmente especificados y los datos reales incompletos. Al tratar las completaciones estocásticas como contexto incierto en lugar de objetivos de entrenamiento, el método:

• Evita el sesgo de dispersión y los patrones artificiales introducidos por los marcadores de posición constantes.
• Evita la inestabilidad de optimización y la dinámica de objetivo móvil de los métodos de autoimputación.
• Aprovecha la varianza del aumento para regularizar el modelo, haciéndolo menos sensible a las entradas inciertas.

Los autores posicionan a AugMask como una estrategia general que permite desplegar modelos de difusión de alto rendimiento existentes en entornos con datos faltantes sin requerir modificaciones arquitectónicas complejas o funciones de pérdida especializadas para la imputación. El trabajo se centra en la generación; los autores señalan que extender esto a pipelines de difusión latente o objetivos específicos de imputación sigue siendo una dirección para trabajos futuros.