📊 statistics

AugMask: Training Diffusion Models on Incomplete Tabular Data via Stochastic Augmentation and Masking

Das Paper stellt AugMask vor, ein Plug-and-Play-Trainingsframework, das es Standard-Score-basierten Diffusionsmodellen ermöglicht, effektiv unvollständige tabellarische Daten zu generieren, indem es stochastische Augmentierung nutzt, um fehlende Werte als unsicheren Konditionierungskontext aufzufüllen, während die Denoising-Überwachung auf nur die beobachteten Koordinaten beschränkt wird.

Ursprüngliche Autoren: Jungkyu Kim, Taeyoung Park, Kibok Lee

Veröffentlicht 2026-06-03

📖 5 Min. Lesezeit🧠 Tiefgang

CC BY 4.0

Ursprüngliche Autoren: Jungkyu Kim, Taeyoung Park, Kibok Lee

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Das Problem: Einem Koch beibringen, mit einem unvollständigen Rezept zu arbeiten

Stellen Sie sich vor, Sie versuchen, einem Meisterkoch (einem Diffusionsmodell) beizubringen, ein bestimmtes Gericht zu kochen (realistische Tabellendaten zu generieren). Der Koch ist brillant darin, Rezepten zu folgen, aber es gibt einen Haken: Die Rezepte, die Sie ihm geben, sind unvollständig. Einige Zutaten fehlen, und die Stellen, an denen sie sein sollten, sind einfach nur leere weiße Flächen oder mit „NaN“ (Not a Number) markiert.

Standard-Köche (KI-Modelle) werden durch leere Stellen verwirrt. Sie können nicht kochen, wenn im Rezept steht: „Fügen Sie 2 Tassen [LEER] hinzu.“

Der alte Weg (Zero-Filling): Früher haben die Leute die Lücken einfach mit einem Platzhalter gefüllt, wie zum Beispiel „0“ oder „Wasser“. Aber das ist so, als würde man dem Koch sagen: „Füge 0 Tassen Zucker hinzu.“ Der Koch lernt dann, dass „0“ eine echte Zutat ist, was den Geschmack (die Datenverteilung) verfälscht.
Der „Rate-Weg“ (Imputation): Eine andere Methode versucht, den Koch raten zu lassen, welche Zutat fehlt, sie aufzuschreiben und den Koch dann darauf zu trainieren, diese Schätzung perfekt zu machen. Aber wenn der Koch falsch rät, versucht er ständig, eine falsche Schätzung zu perfektionieren, was zu einer verrauschten, verwirrten Küche führt.

Die Lösung: AugMask

Die Autoren schlagen AugMask vor, eine neue Trainingsstrategie, die wie ein kluger Sous-Chef fungiert. Sie löst das Problem, indem sie zwei Aufgaben trennt: Das Szenario setzen und die Arbeit bewerten.

1. Das Szenario setzen (Stochastische Augmentierung)

Anstatt die leeren Stellen leer zu lassen oder mit einer langweiligen „0“ zu füllen, füllt der Sous-Chef die Lücken mit plausiblen Vermutungen.

Der Clou: Der Sous-Chef gibt nicht nur eine einzige Vermutung ab. Er gibt eine Spanne von Vermutungen an.
- Analogie: Wenn im Rezept „Salz“ fehlt, sagt der Sous-Chef nicht einfach „1 Teelöffel“. Er sagt: „Es könnte irgendwo zwischen 0,5 und 1,5 Teelöffeln liegen.“
Warum? Dies lehrt den Hauptkoch, dass die fehlende Zutat unsicher ist. Der Hauptkoch lernt, das Gesamtbild (die anderen Zutaten) zu betrachten und zu verstehen, dass das fehlende Teil ein „Vielleicht“ ist, kein „Definitiv“.

2. Die Arbeit bewerten (Maskierung des Verlusts/Loss-Masking)

Dies ist der wichtigste Teil. Wenn der Hauptkoch versucht, das Gericht nachzubauen, bewertet der Lehrer (der Trainingsalgorithmus) nur die Zutaten, die ursprünglich vorhanden waren.

Die Regel: Wenn im Rezept ursprünglich „2 Tassen Mehl“ stand, prüft der Lehrer, ob der Koch das Mehl richtig gemacht hat.
Die Schlupflücke: Wenn im Rezept ursprünglich eine Lücke für „Salz“ war, ignoriert der Lehrer die Schätzung des Kochs für das Salz. Der Lehrer sagt: „Es ist mir egal, was du für das Salz geraten hast; stelle nur sicher, dass das Mehl stimmt.“
Das Ergebnis: Der Koch lernt, die „Vermutungen“ als Kontext (Hinweise) zu nutzen, um die echten Zutaten richtig zu bekommen, aber er wird nicht dafür bestraft, die Vermutungen falsch zu machen. Er hört auf, die Vermutungen auswendig zu lernen, und beginnt stattdan die Beziehungen zwischen den echten Daten zu verstehen.

Warum das funktioniert: Die „Unsicherheitspenalty“

Das Paper verwendet einige komplexe mathematische Erklärungen, warum dies funktioniert, was durch die einfache Metapher „Die zittrige Hand“ verstanden werden kann.

Deterministisches Raten (Schlecht): Wenn der Sous-Chef eine einzelne, sichere Vermutung abgibt (z. B. „Es ist exakt 1,0 Teelöffel“), denkt der Hauptkoch: „Okay, ich muss das exakt treffen.“ Wenn die Schätzung leicht falsch ist, wird der Koch verwirrt.
Stochastisches Raten (Gut): Wenn der Sous-Chef sagt: „Es liegt zwischen 0,5 und 1,5“, erkennt der Hauptkoch: „Das ist ein zittriger, unsicherer Hinweis.“
Die Strafe: Die Mathematik zeigt, dass der Hauptkoch bei einem unsicheren Hinweis (hohe Unsicherheit) ganz natürlich aufhört, diesen zu stark zu gewichten. Es ist wie ein Schüler, der einen unscharfen Hinweis in einer Prüfung ignoriert, weil er weiß, dass dieser falsch sein könnte. Dies verhindert, dass das Modell an schlechten Schätzungen „hängen bleibt“.

Das Ergebnis: Ein besserer Koch

Die Autoren haben AugMask auf vielen verschiedenen Datensätzen getestet (wie z. B. Daten zum Einkommen von Erwachsenen, Einkaufsgewohnheiten usw.) mit unterschiedlichen Mengen an fehlenden Informationen (von 10 % bis 90 % fehlend).

Das Ergebnis: Durch die Verwendung dieser „Szenario setzen, Schätzungen ignorieren“-Strategie wurden Standard-KI-Modelle (die normalerweise schlecht mit fehlenden Daten umgehen können) besser als spezialisierte Modelle, die explizit für fehlende Daten entwickelt wurden.
Die Erkenntnis: Man braucht keine spezielle, komplexe Küche für unvollständige Rezepte. Man braucht nur eine kluge Art, die Lücken mit „unsicheren Vermutungen“ zu füllen und dem Koch dann zu sagen, dass er diese Vermutungen ignorieren soll, wenn er das fertige Gericht bewertet.

Zusammenfassung in einem Satz

AugMask lehrt KI-Modelle, mit fehlenden Daten umzugehen, indem es die Lücken mit „unsicheren Vermutungen“ füllt, um Kontext zu bieten, aber das Modell streng nur anhand der Daten bewertet, die tatsächlich vorhanden waren, um zu verhindern, dass es durch seine eigenen Vermutungen verwirrt wird.

Technisches Resümee: AugMask

Problemstellung

Score-basierte Diffusionsmodelle haben sich als leistungsstarke generative Tiefenmodelle etabliert, stehen jedoch vor erheblichen Herausforderungen, wenn sie auf tabellarische Daten mit fehlenden Werten angewendet werden. Standardmäßige Diffusionsarchitekturen setzen vollständig spezifizierte Eingaben voraus, während reale tabellarische Daten oft fehlende Einträge (NaNs) enthalten.

Bestehende Ansätze zur Überbrückung dieser Lücke fallen typischerweise in zwei Kategorien, die beide Limitationen aufweisen:

Konstante Platzhalter-Imputation: Ersetzen fehlender Werte durch feste Konstanten (z. B. Nullen, Spaltenmittelwerte) vor dem Training. Dies zwingt das Modell dazu, fehlende Einträge identisch mit gültigen Nullwerten zu behandeln, was einen Sparsitäts-Bias induziert und Vorhersagen abhängig vom Grad der Fehlstellen macht. Selbst wenn der Verlust auf beobachtete Koordinaten maskiert wird, bleibt der deterministische Platzhalter in der Eingabe bestehen und wirkt als spurendes Signal, das mit der Fehlstellenrate korreliert.
Selbst-Imputation (z. B. DiffPuter): Iterative Aktualisierung fehlender Einträge während des Trainings über einen Expectation-Maximization (EM)-Ansatz. Obwohl dies feste Platzhalter vermeidet, erzeugt es ein Optimierungsproblem mit einem „beweglichen Ziel“ (moving target), bei dem das Modell auf seinen eigenen, sich entwickelnden Vorhersagen trainiert wird. Dies kann zu verrauschten Gradienten, nicht-stationären Dynamiken und übermäßig dispersen Vervollständigungen führen, da Diffusions-Objectives Diversität gegenüber punktweiser Genauigkeit priorisieren.

Die zentrale Herausforderung besteht darin, dem Denoiser stabile numerische Eingaben bereitzustellen, ohne die aufgefüllten fehlenden Werte als Ground-Truth-Ziele zu behandend, um sowohl den Bias konstanter Platzhalter als auch die Instabilität der Selbst-Imputation zu vermeiden.

Methodik: AugMask

Die Autoren schlagen AugMask vor, ein Plug-and-Play-Trainingsframework, das fehlstellen-unbewusste Backbones an unvollständige Daten anpasst, indem es Conditioning (Bedingung) von der Supervision (Überwachung) trennt. Die Methode opert in zwei distinkten Schritten:

1. Stochastische konditionale Augmentierung (Input-Konstruktion)

Anstatt einen deterministischen Platzhalter zu verwenden, konstruiert AugMask eine vollständig numerische Eingabe, indem fehlende Einträge aus einer kontextabhängigen Verteilung $T(\cdot | x_{obs}, m)$ gezogen werden.

Mechanismus: Für jedes fehlende Merkmal $j$ wird ein Wert $Z_j$ aus einer Verteilung gezogen, die auf den beobachteten Merkmalen und der Maske der Fehlstellen basiert.
Implementierung: Die Autoren nutzen einen leichtgewichtigen, faktorisierten Ansatz unter Verwendung von pro-Feature LightGBM-Modellen. Für kontinuierliche Merkmale passen sie einen bedingten Mittelwert-Regressor und einen Log-Varianz-Regressor an, um aus einer heteroskedastischen Gaußverteilung zu sampeln. Für kategoriale Merkmale verwenden sie einen bedingten Klassifikator mit Label Smoothing.
Rolle: Diese gesampelten Werte dienen als unsicheres Konditionierungs-Kontext für den Denoiser, nicht als Ground-Truth-Ziele. Die Stochastizität kodiert die inhärente Unsicherheit der Fehlstellen.

2. Nur beobachtete Supervision (Loss-Maskierung)

Der Denoising-Loss wird nur auf die Koordinaten angewendet, die ursprünglich beobachtet wurden.

Zielsetzung: Das Modell minimiert den Rekonstruktionsfehler für beobachtete Koordinaten $m \odot (D_\theta(x_{A,t}, t) - x_{A,0})^2$ , wobei $x_{A,0}$ die augmentierte Eingabe ist.
Effekt: Die augmentierten (zuvor fehlenden) Einträge treten durch die verrauschte Eingabe $x_{A,t}$ in das Modell ein, um Kontext zu liefern, aber das Modell wird niemals dafür bestraft, sie zu rekonstruieren. Dies verhindert, dass das Modell zu stark auf potenziell ungenaue Vervollständigungen vertraut.

Theoretische Begründung

Das Paper verbindet diese Trainingsregel mit einem Rao-Blackwellized Objective. Durch die Marginalisierung über plausible Vervollständigungen zeigen die Autoren, dass das Objective eine varianzgewichtete Sensitivitätsstrafe liefert.

Interpretation: Der Term zweiter Ordnung in der lokalen Approximation des Objectives ist proportional zur bedingten Varianz der Augmentierung. Eine hohe bedingte Varianz bestraft die Sensitivität des Modells gegenüber dieser spezifischen Vervollständigung und verhindert so effektiv die Abhängigkeit von unsicheren Füllungen. Eine niedrige Varianz erlaubt es, die Vervollständigung direkter als Kontext zu nutzen.
Kontrast zu AugFull: Die Autoren vergleichen AugMask mit „AugFull“ (welches den Loss auf alle Koordinaten anwendet, einschließlich der augmentierten). Sie demonstrieren, dass AugFull eine Verzerrung (Bias) einführt, wenn die Augmentierungsregel falsch spezifiziert ist, da das Modell gezwungen wird, potenziell falsche Pseudo-Ziele anzupassen. AugMask vermeidet dies, indem es Vervollständigungen strikt als Kontext behandelt.

Zentrale Beiträge

Ein Trainingsprinzip für unvollständige tabellarische Diffusion: Die Einführung von AugMask, das die Konditionierung (via stochastischer Augmentierung) von der Supervision (via beobachtungs-exklusiver Loss-Maskierung) entkoppelt.
Eine Marginalisierungs-Perspektive: Eine theoretische Analyse, die zeigt, dass stochastische Augmentierung eine varianzgewichtete Sensitivitätsstrafe induziert, was das Modell gegen eine übermäßige Abhängigkeit von unsicheren Vervollständigungen regularisiert.
Ein Plug-and-Play-Rezept: Ein leichtgewichtiges Framework, das es Standard-Modellen für fehlstellen-unbewusste score-basierte Diffusion ermöglicht, auf unvollständigen Daten zu trainieren, ohne architektonische Änderungen zu benötigen.

Experimentelle Ergebnisse

Die Autoren evaluierten AugMask über diverse Datensätze (Adult, Default, Shoppers, Magic, Beijing, News) und Fehlstellen-Regimes (MCAR und MAR) mit Fehlstellenraten von 0,1 bis 0,9.

Performance: AugMask ermöglicht es standardmäßigen, fehlstellen-unbewussten Diffusionsmodellen (z. B. CDTD, TabDiff), spezialisierte fehlstellen-bewusste Baselines (z. B. MIWAE, MissDiff, DiffPuter, ForestDiff) zu übertreffen.
Generative Qualität: Mit AugMask trainierte Modelle erreichten eine überlegene Sample-Fidelity (gemessen an $\alpha$ -Precision, $\beta$ -Recall) und Feature-Level-Statistiken (Trend, Shape).
Maschinelle Lernnutzen: Downstream-Aufgaben (Klassifikation und Regression), die auf synthetischen Daten trainiert wurden, die von AugMask-Modellen generiert wurden, zeigten eine verbesserte Performance im Vergleich zu den Baselines.
Robustheit: Die Vorteile von AugMask wurden mit zunehmendem Anteil an Fehlstellen deutlicher. Bei hoher Fehlstellenrate (z. B. 90 %) verhinderte AugMask die starke Degradation, die bei anderen Methoden beobachtet wurde.
Sensitivitätsanalyse: Die Methode erwies sich als robust gegenüber einer Fehlspezifikation der Augmentierung. Selbst wenn die Augmentierungsregel gestört wurde (z. B. durch Reduzierung der Stochastizität oder Einführung von Bias), behielt AugMask seinen Leistungsvorteil gegenüber AugFull bei, was bestätigt, dass die Behandlung von Vervollständigungen als Kontext der entscheidende Faktor ist.

Bedeutung und Ansprüche

Das Paper behauptet, dass AugMask eine fundierte Lösung für die Diskrepanz zwischen vollständig spezifizierten Diffusions-Backbones und unvollständigen realen Daten bietet. Indem es stochastische Vervollständigungen als unsicheren Kontext statt als Trainingsziele behandelt, bewirkt die Methode:

Vermeidung des Sparsitäts-Bias und künstlicher Muster, die durch konstante Platzhalter entstehen.
Vermeidung der Optimierungsinstabilität und der Moving-Target-Dynamik von Selbst-Imputationsmethoden.
Nutzung der Varianz der Augmentierung zur Regularisierung des Modells, wodurch es weniger sensitiv gegenüber unsicheren Eingaben wird.

Die Autoren positionieren AugMask als eine allgemeine Strategie, die es existierenden, leistungsstarken fehlstellen-unbewussten Modellen ermöglicht, in Umgebungen mit fehlenden Daten eingesetzt zu werden, ohne komplexe architektonische Modifikationen oder spezialisierte Loss-Funktionen für die Imputation zu benötigen. Die Arbeit konzentriert sich auf die Generierung; die Autoren merken an, dass die Erweiterung auf Latent-Diffusion-Pipelines oder imputationsspezifische Objectives eine Richtung für zukünftige Arbeiten bleibt.