Riassunto Tecnico: AugMask

Definizione del Problema

I modelli di diffusione basati sul punteggio (score-based diffusion models) sono emersi come potenti modelli generativi profondi, ma affrontano sfide significative quando applicati a dati tabulari contenenti valori mancanti. Le architetture di diffusione standard assumono input completamente specificati, mentre i dati tabulari del mondo reale contengono spesso voci mancanti (NaN).

Gli approcci esistenti per colmare questo divario rientrano tipicamente in due categorie, entrambe con limitazioni:

1. Imputazione con Placeholder Costanti: Sostituzione dei valori mancanti con costanti fisse (ad es. zeri, medie di colonna) prima dell'addestramento. Ciò costringe il modello a trattare le voci mancanti in modo identico a zeri validi, inducendo un bias di sparsità e rendendo le previsioni dipendenti dal livello di mancanza dei dati. Inoltre, anche se la perdita (loss) viene mascherata alle coordinate osservate, il placeholder deterministico rimane nell'input, agendo come un indizio spuria correlato alla mancanza dei dati.
2. Auto-imputazione (es. DiffPuter): Aggiornamento iterativo delle voci mancanti durante l'addestramento tramite un approccio Expectation-Maximization (EM). Sebbene questo eviti i placeholder fissi, crea un problema di ottimizzazione a "bersaglio mobile" (moving target) in cui il modello viene addestrato sulle proprie previsioni in evoluzione. Ciò può portare a gradienti rumorosi, dinamiche non stazionarie e completamenti eccessivamente dispersi, poiché gli obiettivi di diffusione danno priorità alla diversità rispetto all'accuratezza puntuale.

La sfida centrale è fornire input numerici stabili al denoiser senza trattare i valori completati come target di ground-truth, evitando così sia il bias dei placeholder costanti che l'instabilità dell'auto-imputazione.

Metodologia: AugMask

Gli autori propongono AugMask, un framework di addestramento plug-and-play che adatta backbone ignari della mancanza di dati (missing-unaware) a dati incompleti separando il condizionamento dalla supervisione. Il metodo opera in due fasi distinte:

1. Augmentation Condizionale Stocastica (Costruzione dell'Input)

Invece di utilizzare un placeholder deterministico, AugMask costruisce un input completamente numerico campionando le voci mancanti da una distribuzione dipendente dal contesto $T(\cdot | x_{obs}, m)$.

• Meccanismo: Per ogni caratteristica mancante $j$, un valore $Z_j$ viene estratto da una distribuzione condizionata sulle caratteristiche osservate e sulla maschera di mancanza.
• Implementazione: Gli autori utilizzano un approccio leggero e fattorizzato tramite modelli LightGBM per singola caratteristica. Per le caratteristiche continue, adattano un regressore della media condizionale e un regressore della log-varianza per campionare da una Gaussiana eteroschedastica. Per le caratteristiche categoriche, utilizzano un classificatore condizionale con label smoothing.
• Ruolo: Questi valori campionati servono come contesto di condizionamento incerto per il denoiser, non come target di ground-truth. La stocasticità codifica l'incertezza intrinseca della mancanza dei dati.

2. Supervisione Solo sulle Osservazioni (Mascheramento della Loss)

La perdita di denoising viene applicata solo alle coordinate che erano originariamente osservate.

• Obiettivo: Il modello minimizza l'errore di ricostruzione per le coordinate osservate $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$, dove $x_{A,0}$ è l'input aumentato.
• Effetto: Le voci aumentate (precedentemente mancanti) entrano nel modello attraverso l'input rumoroso $x_{A,t}$ per fornire contesto, ma il modello non viene mai penalizzato per il fallimento nella loro ricostruzione. Ciò impedisce al modello di fare troppo affidamento su completamenti potenzialmente inaccurati.

Giustificazione Teorica

Il documento collega questa regola di addestramento a un obiettivo di Rao–Blackwellized. Marginalizzando sulle completazioni plausibili, gli autori dimostrano che l'obiettivo produce una penalità di sensibilità pesata sulla varianza.

• Interpretazione: Il termine del secondo ordine nell'approssimazione locale dell'obiettivo è proporzionale alla varianza condizionale dell'augmentation. Un'alta varianza condizionale penalizza la sensibilità del modello a quella specifica completazione, scoraggiando efficacemente l'affidamento su riempimenti incerti. Una bassa varianza permette alla completazione di essere utilizzata più direttamente come contesto.
• Contrasto con AugFull: Gli autori confrontano AugMask con "AugFull" (che applica la loss a tutte le coordinate, incluse quelle aumentate). Dimostrano che AugFull introduce bias se la regola di augmentation è errata (misspecified), poiché il modello è costretto a adattarsi a potenziali pseudo-target errati. AugMask evita ciò trattando le completazioni strettamente come contesto.

Contributi Chiave

1. Un Principio di Addestramento per la Diffusione Tabulare Incompleta: L'introduzione di AugMask, che separa il condizionamento (tramite augmentation stocastica) dalla supervisione (tramite mascheramento della perdita solo sulle osservazioni).
2. Una Visione di Marginalizzazione: Un'analisi teorica che mostra come l'augmentation stocastica induca una penalità di sensibilità pesata sulla varianza, regolarizzando il modello contro l'eccessivo affidamento su completazioni incerte.
3. Una Ricetta Plug-and-Play: Un framework leggero che consente ai modelli di diffusione basati sul punteggio standard, ignari della mancanza di dati, di addestrarsi su dati incompleti senza modifiche architettoniche.

Risultati Sperimentali

Gli autori hanno valutato AugMask su diversi dataset (Adult, Default, Shoppers, Magic, Beijing, News) e regimi di mancanza (MCAR e MAR) con rapporti di mancanza compresi tra 0,1 e 0,9.

• Performance: AugMask consente ai modelli di diffusione standard (es. CDTD, TabDiff) di superare i baseline specializzati nella gestione dei dati mancanti (es. MIWAE, MissDiff, DiffPuter, ForestDiff).
• Qualità Generativa: I modelli addestrati con AugMask hanno ottenuto una precisione superiore a livello di campione (misurata tramite $\alpha$-Precision, $\beta$-Recall) e statistiche a livello di caratteristica (Trend, Shape).
• Utilità per il Machine Learning: I task a valle (classificazione e regressione) addestrati su dati sintetici generati dai modelli AugMask hanno mostrato prestazioni migliorate rispetto ai baseline.
• Robustezza: I benefici di AugMask sono diventati più pronunciati all'aumentare del rapporto di mancanza. In presenza di alta mancanza (es. 90%), AugMask ha evitato il grave degrado osservato in altri metodi.
• Analisi di Sensibilità: Il metodo si è dimostrato robusto alla misspecification dell'augmentation. Anche quando la regola di augmentation veniva perturbata (ad es. riducendo la stocasticità o introducendo bias), AugMask manteneva un vantaggio di performance rispetto ad AugFull, confermando che trattare le completazioni come contesto piuttosto che come target è il fattore critico.

Significato e Rivendicazioni

Il documento sostiene che AugMask fornisce una soluzione fondata al disallineamento tra i backbone di diffusione a specifica completa e i dati reali incompleti. Trattando le completazioni stocastiche come contesto incerto anziché come target di addestramento, il metodo:

• Evita il bias di sparsità e i pattern artificiali introdotti dai placeholder costanti.
• Evita l'instabilità di ottimizzazione e le dinamiche a bersaglio mobile dei metodi di auto-imputazione.
• Sfrutta la varianza dell'augmentation per regolarizzare il modello, rendendolo meno sensibile agli input incerti.

Gli autori posizionano AugMask come una strategia generale che consente ai modelli di diffusione esistenti e ad alte prestazioni di essere implementati in contesti con dati mancanti senza richiedere complesse modifiche architettoniche o funzioni di perdita specializzate. Il lavoro si concentra sulla generazione; gli autori osservano che estendere questo approccio a pipeline di diffusione latente o obiettivi specifici di imputazione rimane una direzione per il lavoro futuro.