Résumé Technique : AugMask

Énoncé du Problème

Les modèles de diffusion basés sur le score sont apparus comme de puissants modèles génératifs profonds, mais ils font face à des défis importants lorsqu'ils sont appliqués à des données tabulaires contenant des valeurs manquantes. Les architectures de diffusion standard supposent des entrées entièrement spécifiées, alors que les données tabulaires du monde réel contiennent souvent des entrées manquantes (NaN).

Les approches existantes pour combler cette lacune tombent généralement dans deux catégories, qui présentent toutes deux des limites :

1. Imputation par Valeur Constante (Placeholder) : Remplacer les valeurs manquantes par des constantes fixes (ex: zéros, moyennes de colonnes) avant l'entraînement. Cela force le modèle à traiter les entrées manquantes de la même manière que des zéros valides, induisant un biais de parcimonie et rendant les prédictions dépendantes du niveau de manque de données. De plus, même si la perte est masquée pour les coordonnées observées, le placeholder déterministe reste dans l'entrée, agissant comme un indice fallacieux corrélé à l'absence de données.
2. Auto-imputation (ex: DiffPuter) : Mettre à jour de manière itérative les entrées manquantes pendant l'entraînement via une approche d'Espérance-Maximisation (EM). Bien que cela évite les placeholders fixes, cela crée un problème d'optimisation à "cible mouvante" où le modèle est entraîné sur ses propres prédictions évolutives. Cela peut conduire à des gradients bruités, des dynamiques non stationnaires et des complétions trop dispersées, car les objectifs de diffusion privilégient la diversité plutôt que la précision ponctuelle.

Le défi central consiste à fournir des entrées numériques stables au dénoiseur sans traiter les valeurs complétées comme des cibles de vérité terrain (ground-truth), évitant ainsi à la fois le biais des placeholders constants et l'instabilité de l'auto-imputation.

Méthodologie : AugMask

Les auteurs proposent AugMask, un cadre d'entraînement "plug-and-play" qui adapte les backbones ignorant les données manquantes aux données incomplètes en séparant le conditionnement de la supervision. La méthode opère en deux étapes distinctes :

1. Augmentation Conditionnelle Stochastique (Construction de l'Entrée)

Au lieu d'utiliser un placeholder déterministe, AugMask construit une entrée entièrement numérique en échantillonnant les entrées manquantes à partir d'une distribution dépendante du contexte $T(\cdot | x_{obs}, m)$.

• Mécanisme : Pour chaque caractéristique manquante $j$, une valeur $Z_j$ est tirée d'une distribution conditionnée sur les caractéristiques observées et le masque de présence.
• Implémentation : Les auteurs utilisent une approche factorisée légère utilisant des modèles LightGBM par caractéristique. Pour les caractéristiques continues, ils ajustent un régresseur de moyenne conditionnelle et un régresseur de log-variance pour échantillonner à partir d'une Gaussienne hétéroscédastique. Pour les caractéristiques catégorielles, ils utilisent un classificateur conditionnel avec lissage deラベル (label smoothing).
• Rôle : Ces valeurs échantillonnées servent de contexte de conditionnement incertain pour le dénoiseur, et non de cibles de vérité terrain. La stochasticité encode l'incertitude inhérente à l'absence de données.

2. Supervision Uniquement sur les Données Observées (Masquage de la Perte)

La perte de débruitage est appliquée uniquement aux coordonnées qui étaient initialement observées.

• Objectif : Le modèle minimise l'erreur de reconstruction pour les coordonnées observées $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$, où $x_{A,0}$ est l'entrée augmentée.
• Effet : Les entrées augmentées (précédemment manquantes) entrent dans le modèle via l'entrée bruitée $x_{A,t}$ pour fournir du contexte, mais le modèle n'est jamais pénalisé pour échouer à les reconstruire. Cela empêche le modèle de trop s'appuyer sur des complétions potentiellement inexactes.

Justification Théorique

Le papier relie cette règle d'entraînement à un objectif de Rao–Blackwellisé. En effectuant la marginalisation sur les complétions plausibles, les auteurs montrent que l'objectif produit une pénalité de sensibilité pondérée par la variance.

• Interprétation : Le terme de second ordre dans l'approximation locale de l'objectif est proportionnel à la variance conditionnelle de l'augmentation. Une variance conditionnelle élevée pénalise la sensibilité du modèle à cette complétion spécifique, décourageant efficacement l'reliance sur des remplissages incertains. Une faible variance permet d'utiliser la complétion plus directement comme contexte.
• Contraste avec AugFull : Les auteurs comparent AugMask à "AugFull" (qui applique la perte à toutes les coordonnées, y compris les augmentées). Ils démontrent qu'AugFull introduit un biais si la règle d'augmentation est mal spécifiée, car le modèle est forcé de s'ajuster à des pseudo-cibles potentiellement incorrectes. AugMask évite cela en traitant les complétions strictement comme du contexte.

Contributions Clés

1. Un Principe d'Entraînement pour la Diffusion Tabulaire Incomplète : L'introduction d'AugMask, qui découple le conditionnement (via l'augmentation stochastique) de la supervision (via le masquage de perte sur les données observées uniquement).
2. Une Vue de Marginalisation : Une analyse théorique montrant que l'augmentation stochastique induit une pénalité de sensibilité pondérée par la variance, régularisant le modèle contre une dépendance excessive aux complétions incertaines.
3. Une Recette "Plug-and-Play" : Un cadre léger qui permet aux modèles de diffusion de score standards, ignorant les données manquantes, de s'entraîner sur des données incomplètes sans changements architecturaux.

Résultats Expérimentaux

Les auteurs ont évalué AugMask sur divers jeux de données (Adult, Default, Shoppers, Magic, Beijing, News) et différents régimes de données manquantes (MCAR et MAR) avec des taux de manque allant de 0,1 à 0,9.

• Performance : AugMask permet aux modèles de diffusion standards ignorant les données manquantes (ex: CDTD, TabDiff) de surpasser les bases spécialisées dans la gestion des données manquantes (ex: MIWAE, MissDiff, DiffPuter, ForestDiff).
• Qualité Générative : Les modèles entraînés avec AugMask ont atteint une fidélité au niveau de l'échantillon supérieure (mesurée par $\alpha$-Précision, $\beta$-Rappel) et des statistiques au niveau des caractéristiques supérieures (Trend, Shape).
• Utilité pour l'Apprentissage Automatique : Les tâches en aval (classification et régression) entraînées sur des données synthétiques générées par les modèles AugMask ont montré des performances améliorées par rapport aux bases de référence.
• Robustesse : Les bénéfices d'AugMask sont devenus plus prononcés à mesure que le ratio de données manquantes augmentait. À un taux de manque élevé (ex: 90 %), AugMask a évité la dégradation sévère observée avec les autres méthodes.
• Analyse de Sensibilité : La méthode s'est révélée robuste à la mauvaise spécification de l'augmentation. Même lorsque la règle d'augmentation était perturbée (ex: réduction de la stochasticité ou introduction de biais), AugMask maintenait un avantage de performance sur AugFull, confirmant que traiter les complétions comme du contexte plutôt que comme des cibles est le facteur critique.

Signification et Revendications

Le papier affirme qu'AugMask fournit une solution fondée sur des principes au décalage entre les backbones de diffusion entièrement spécifiés et les données réelles incomplètes. En traitant les complétions stochastiques comme un contexte incertain plutôt que comme des cibles d'entraînement, la méthode :

• Évite le biais de parcimonie et les motifs artificiels introduits par les placeholders constants.
• Évite l'instabilité d'optimisation et la dynamique de cible mouvante des méthodes d'auto-imputation.
• Utilise la variance de l'augmentation pour régulariser le modèle, le rendant moins sensible aux entrées incertaines.

Les auteurs positionnent AugMask comme une stratégie générale qui permet de déployer les modèles de diffusion de score performants existants dans des contextes de données manquantes sans nécessiter de modifications architecturales complexes ou de fonctions de perte spécialisées pour l'imputation. Le travail se concentre sur la génération ; les auteurs notent que l'extension de cela aux pipelines de diffusion latente ou aux objectifs spécifiques d'imputation reste une direction pour les travaux futurs.