Technische Samenvatting: AugMask

Probleemstelling

Score-gebaseerde diffusiemodellen zijn uitgegroeid tot krachtige diepe generatieve modellen, maar kampen met aanzienlijke uitdagingen bij de toepassing op tabeldata met ontbrekende waarden. Standaard diffusiearchitecturen gaan uit van volledig gespecificeerde inputs, terwijl real-world tabeldata vaak ontbrekende vermeldingen bevatten (NaNs).

Bestaande benaderingen om deze kloof te overbruggen vallen doorgaans uiteen in twee categorieën, die beide beperkingen kennen:

1. Constante Placeholder Imputatie: Het vervangen van ontbrekende waarden door vaste constanten (bijv. nul of kolomgemiddelden) vóór de training. Dit dwingt het model om ontbrekende vermeldingen identiek te behandelen als geldige nulwaarden, wat een spaarsheidsbias induceert en voorspellingen afhankelijk maakt van het niveau van ontbreken. Bovendien, zelfs als de loss wordt gemaskeerd naar geobserveerde coördinaten, blijft de deterministische placeholder in de input aanwezig, wat fungeert als een spurieuze aanwijzing die correleert met de ontbrekende status.
2. Zelf-imputatie (bijv. DiffPuter): Het iteratief updaten van ontbrekende vermeldingen tijdens de training via een Expectation-Maximization (EM) benadering. Hoewel dit vaste placeholders vermijdt, creëert het een "moving target" optimalisatieprobleem waarbij het model wordt getraind op zijn eigen evoluerende voorspellingen. Dit kan leiden tot ruisende gradiënten, niet-stationaire dynamiek en overmatig verspreide completions, aangezien diffusie-objectieven prioriteit geven aan diversiteit boven punt-gewijze nauwkeurigheid.

De kernuitdaging is het bieden van stabiele numerieke inputs aan de denoiser zonder de ingevulde ontbrekende waarden als grondwaarheid-doelwitten (ground-truth targets) te behandelen, om zo zowel de bias van constante placeholders als de instabiliteit van zelf-imputatie te vermijden.

Methodologie: AugMask

De auteurs stellen AugMask voor, een plug-and-play trainingsframework dat "missing-unaware" backbones aanpast aan incomplete data door conditionering te scheiden van supervisie. De methode werkt in twee afzonderlijke stappen:

1. Stochastische Conditionele Augmentatie (Input Constructie)

In plaats van een deterministische placeholder te gebruiken, construeert AugMask een volledig numerieke input door ontbrekende vermeldingen te samplen uit een context-afhankelijke distributie $T(\cdot | x_{obs}, m)$.

• Mechanisme: Voor elke ontbrekende feature $j$ wordt een waarde $Z_j$ getrokken uit een distributie die geconditioneerd is op de geobserveerde features en de missingness mask.
• Implementatie: De auteurs maken gebruik van een lichtgewicht, gefactoriseerde aanpak met per-feature LightGBM-modellen. Voor continue features fitten zij een conditionele gemiddelde regressor en een log-variantie regressor om te samplen uit een heteroskedastische Gaussische distributie. Voor categorische features gebruiken zij een conditionele classifier met label smoothing.
• Rol: Deze gesamplede waarden dienen als onzekere conditionele context voor de denoiser, niet als grondwaarheid-doelwitten. De stochastiek codeert de inherente onzekerheid van de ontbrekende status.

2. Alleen Geobserveerde Supervisie (Loss Masking)

De denoising loss wordt alleen toegepast op de coördinaten die oorspronkelijk geobserveerd waren.

• Doelstelling: Het model minimaliseert de reconstructiefout voor geobserveerde coördinaten $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$, waarbij $x_{A,0}$ de geaugmenteerde input is.
• Effect: De geaugmenteerde (voorheen ontbrekende) vermeldingen komen via de ruisende input $x_{A,t}$ in het model terecht om context te bieden, maar het model wordt nooit gestraft voor het falen bij het reconstrueren ervan. Dit voorkomt dat het model te veel vertrouwt op potentieel onnauwkeurige completions.

Theoretische Rechtvaardiging

Het paper verbindt deze trainingsregel met een Rao–Blackwellized objectief. Door te marginaliseren over plausibele completions, tonen de auteurs aan dat het objectief een variantie-gewogen sensitiviteitsstraf oplevert.

• Interpretatie: De tweede-orde term in de lokale benadering van het objectief is proportioneel aan de conditionele variantie van de augmentatie. Een hoge conditionele variantie straft de sensitiviteit van het model voor die specifieke completion af, wat effectief het vertrouwen op onzekere invullingen ontmoedigt. Een lage variantie staat toe dat de completion directer als context wordt gebruikt.
• Contrast met AugFull: De auteurs vergelijken AugMask met "AugFull" (die de loss toepast op alle coördinaten, inclusief geaugmenteerde coördinaten). Zij demonstreren dat AugFull bias introduceert als de augmentatie-regel fout gespecificeerd is, omdat het model gedwongen wordt om potentieel onjuiste pseudo-doelwitten te fitten. AugMask vermijdt dit door completions strikt als context te behandelen.

Belangrijkste Bijdragen

1. Een Trainingsprincipe voor Incomplete Tabulaire Diffusie: De introductie van AugMask, dat conditionering (via stochastische augmentatie) ontkoppelt van supervisie (via alleen geobserveerde loss masking).
2. Een Marginalisatie-perspectief: Een theoretische analyse die aantoont dat stochastische augmentatie een variantie-gewogen sensitiviteitsstraf induceert, wat het model reguleert tegen een te grote afhankelijkheid van onzekere invullingen.
3. Een Plug-and-Play Recept: Een lichtgewicht framework dat standaard, "missing-unaware" score-gebaseerde diffusiemodellen in staat stelt om op incomplete data te trainen zonder architecturale wijzigingen.

Experimentele Resultaten

De auteurs evalueerden AugMask over diverse datasets (Adult, Default, Shoppers, Magic, Beijing, News) en missingness-regimes (MCAR en MAR) met ontbrekende ratio's variërend van 0,1 tot 0,9.

• Prestaties: AugMask stelt standaard "missing-unaware" diffusiemodellen (bijv. CDTD, TabDiff) in staat om gespecialiseerde "missing-aware" baselines (bijv. MIWAE, MissDiff, DiffPuter, ForestDiff) te overtreffen.
• Generatieve Kwaliteit: Modellen getraind met AugMask bereikten superieure sample-niveau getrouwheid (gemeten via $\alpha$-Precision, $\beta$-Recall) en feature-niveau statistieken (Trend, Shape).
• Machine Learning Nut: Downstream taken (classificatie en regressie) getraind op synthetische data gegenereerd door AugMask-modellen vertoonden verbeterde prestaties vergeleken met baselines.
• Robuustheid: De voordelen van AugMask werden prominenter naarmate de ontbrekende ratio toenam. Bij hoge ontbrekende ratio's (bijv. 90%) voorkwam AugMask de ernstige degradatie die bij andere methoden werd waargenomen.
• Sensitiviteitsanalyse: De methode bleek robuust tegen misspecificatie van de augmentatie. Zelfs wanneer de augmentatie-regel werd verstoord (bijv. door de stochastiek te verminderen of bias te introduceren), behield AugMask zijn prestatievoordeel ten opzichte van AugFull, wat bevestigt dat het behandelen van completions als context in plaats van doelwitten de cruciale factor is.

Betekenis en Claims

Het paper claimt dat AugMask een principieel oplossing biedt voor de mismatch tussen volledig gespecificeerde diffusie-backbones en incomplete real-world data. Door stochastische completions te behandelen als onzekere context in plaats van trainingsdoelwitten, vermindert de methode:

• De spaarsheidsbias en kunstmatige patronen die door constante placeholders worden geïntroduceerd.
• De optimalisatie-instabiliteit en "moving-target" dynamiek van zelf-imputatiemethoden.
• Het benut van de variantie van de augmentatie om het model te reguleren, waardoor het minder gevoelig wordt voor onzekere inputs.

De auteurs positioneren AugMask als een algemene strategie die het mogelijk maakt om bestaande, hoogwaardige "missing-unaware" diffusiemodellen in te zetten in omgevingen met ontbrekende data zonder complexe architecturale wijzigingen of gespecialiseerde loss-functies voor imputatie te vereisen. De focus ligt op generatie; de auteurs merken op dat het uitbreiden hiervan naar latent diffusion pipelines of imputatie-specifieke objectieven een richting is voor toekomstig werk.