Wasserstein Distances Made Explainable: Insights Into Dataset Shifts and Transport Phenomena

Diese Arbeit stellt eine neue Methode vor, die mithilfe von Explainable AI die Beiträge verschiedener Datenkomponenten zu Wasserstein-Abständen erklärt, um so Ursachen für Verteilungsverschiebungen und Transportphänomene präzise zu identifizieren.

Das Wesentliche

🌊 Der große Umzug: Warum WaX wie ein Detektiv funktioniert

Stellen Sie sich vor, Sie haben zwei riesige Lagerhallen voller verschiedener Gegenstände.

• Lager A ist das „Quell-Lager" (z. B. alte Fotos von Menschen).
• Lager B ist das „Ziel-Lager" (z. B. neue Fotos von Menschen).

Ihre Aufgabe ist es, die Unterschiede zwischen diesen beiden Lagern zu verstehen. Wie viel Arbeit ist nötig, um alles aus Lager A so umzuräumen, dass es genau wie Lager B aussieht?

In der Welt der Datenwissenschaft nennt man diese „Arbeitsmenge" den Wasserstein-Abstand. Er ist wie ein Maß für den Aufwand eines riesigen Umzugs.

Das Problem: Der Umzugsplan ist verwirrend

Bisher konnten Mathematiker und KI-Experten den Gesamtaufwand berechnen. Sie wusten: „Oh, es kostet 1000 Energieeinheiten, Lager A in Lager B zu verwandeln."
Aber sie wusten nicht, warum es so teuer war.

• War es, weil die Möbel zu schwer waren?
• War es, weil die Treppe zu steil war?
• Oder weil jemand versehentlich 500 Stühle in den falschen Raum geschoben hat?

Bisherige Methoden schauten sich nur den Umzugsplan an (welcher Stuhl geht zu welchem Regal). Das ist wie ein langer, unübersichtlicher Excel-Sheet. Man sieht die Bewegungen, aber man versteht nicht, welche einzelnen Faktoren den Umzug so schwierig gemacht haben.

Die Lösung: WaX (Wasserstein made Explainable)

Die Autoren dieses Papers haben eine neue Methode namens WaX entwickelt. Man kann sich WaX wie einen super-scharfsinnigen Detektiv vorstellen, der den Umzugsplan nicht nur liest, sondern ihn zerlegt und analysiert.

WaX beantwortet drei wichtige Fragen:

1. Welche Gegenstände (Datenpunkte) machen den Umzug am schwersten?
   • Beispiel: WaX könnte sagen: „Schau mal, die 500 Stühle in der Mitte sind das Problem. Wenn wir die weglassen, wird der Umzug viel einfacher."
2. Welche Eigenschaften (Merkmale) sind schuld?
   • Beispiel: „Es liegt nicht am Gewicht der Möbel, sondern daran, dass alle Möbel in Lager A blau sind und in Lager B rot. Die Farbe ist der Hauptgrund für den Unterschied."
3. Gibt es versteckte Gruppen?
   • Beispiel: „Eigentlich gibt es zwei verschiedene Umzüge: Einen für die Küche und einen für das Schlafzimmer. Diese laufen völlig unterschiedlich ab."

Wie funktioniert das magisch? (Die Analogie)

Stellen Sie sich vor, der Umzugsplan ist ein riesiges, komplexes Netzwerk aus Seilen, die von Lager A nach Lager B führen.

• Früher: Man sah nur, wie straff die Seile waren (der Gesamtabstand).
• Mit WaX: Man nimmt eine Lupe und verfolgt jedes Seil zurück. WaX nutzt eine Technik namens „Layer-wise Relevance Propagation" (LRP). Das ist wie ein Rückwärts-Fluss von Wasser.
  • Man füllt den Umzug mit Wasser (die Gesamtarbeit).
  • Dann lässt man das Wasser rückwärts durch die Rohre fließen.
  • Wo das Wasser am stärksten herauskommt, dort liegt die Ursache.
  • So sieht man genau: „Aha! 80 % der Arbeit wurden durch die Stühle in der Mitte verursacht, und 20 % durch die Farbe der Wände."

Wofür ist das gut? (Echte Anwendungen)

Die Autoren zeigen in ihrem Paper drei coole Beispiele, wie WaX hilft:

1. Robustere KI-Modelle (Der „Batch-Effekt"-Detektiv):
   Manchmal lernen KI-Modelle falsche Dinge. Zum Beispiel lernt eine KI, Krebs zu erkennen, aber sie schaut sich eigentlich nur an, welches Krankenhaus das Bild gemacht hat (weil die Bilder dort einen anderen Kontrast haben).

   • WaX hilft: Es findet heraus: „Hey, die KI schaut nicht auf den Tumor, sondern auf den Hintergrund des Krankenhauses!" Man kann diese „falschen" Merkmale dann entfernen, damit die KI wirklich den Tumor erkennt.

2. Das Altern verstehen (Der Zeit-Reisende):
   Die Autoren haben Daten von Meeresschnecken (Abalone) genommen. Sie haben gesehen, wie sich eine Gruppe von Schnecken über ein Jahr verändert hat.

   • WaX hilft: Es hat aufgedeckt, dass nicht alle Schnecken gleich altern. Manche werden schwerer, während andere länger werden. WaX hat diese unterschiedlichen „Alterns-Typen" in Gruppen aufgeteilt und gezeigt, welche Körperteile bei welcher Gruppe wachsen. Ohne WaX hätte man nur gesagt: „Die Schnecken werden im Durchschnitt schwerer."

3. Daten-Check (Der Qualitätskontrolleur):
   Wenn man zwei große Foto-Datenbanken vergleicht (z. B. Prominente vs. normale Gesichter), will man wissen: „Was fehlt?"

   • WaX hilft: Es zeigt sofort: „In der Datenbank der Prominenten fehlen fast alle Bilder von Männern mit Brille" oder „In der anderen Datenbank gibt es zu viele Paare, die zusammen abgedrückt sind." Das hilft, Lücken in Datensätzen zu finden, bevor man eine KI darauf trainiert.

Fazit

WaX ist wie eine Übersetzungsmaschine für komplexe Daten.
Es nimmt eine trockene mathematische Zahl (den Abstand zwischen zwei Datenmengen) und verwandelt sie in eine verständliche Geschichte:

• „Hier ist das Problem."
• „Hier ist der Grund."
• „Hier ist die Lösung."

Anstatt nur zu sagen „Die Daten sind unterschiedlich", sagt WaX: „Die Daten sind unterschiedlich, weil diese spezifischen Merkmale diese spezifischen Gruppen betreffen." Das macht KI nicht nur stärker, sondern auch verständlicher und fairer.

Technische Zusammenfassung

1. Problemstellung

Wasserstein-Abstände (auch Earth Mover's Distance) sind ein etabliertes mathematisches Werkzeug aus dem Bereich des Optimalen Transports (Optimal Transport, OT), um Unterschiede zwischen Wahrscheinlichkeitsverteilungen zu quantifizieren. Sie werden häufig verwendet, um Datensatzverschiebungen (Dataset Shifts) zu analysieren oder Transportphänomene über die Zeit zu modellieren.

Das zentrale Problem, das die Autoren adressieren, ist die mangelnde Erklärbarkeit dieser Metriken:

• Die Berechnung des Wasserstein-Abstands liefert zwar einen skalaren Wert für die Distanz, gibt aber keine Auskunft darüber, welche spezifischen Datenkomponenten (z. B. einzelne Merkmale, Datenpunkte oder Untergruppen) zu dieser Distanz beitragen.
• Auch die Analyse des optimalen Transportplans (Kopplung $\gamma^\star$) reicht oft nicht aus, um zu verstehen, welche Eingabevariablen oder Subgruppen für hohe oder niedrige Distanzen verantwortlich sind. Der Transportplan zeigt zwar, wie Punkte verschoben werden, aber nicht, welche Merkmale den „Kostenaufwand" (die Distanz) dominieren.
• Bisherige Methoden der Explainable AI (XAI) konzentrieren sich meist auf die Erklärung von Modellvorhersagen (z. B. Klassifikation) oder auf Distanzen auf Instanzebene, nicht jedoch auf Distanzen auf Verteilungsebene.

2. Methodik: WaX (Wasserstein distances made explainable)

Die Autoren stellen WaX vor, ein neues Framework, das Prinzipien der Explainable AI (insbesondere Layer-wise Relevance Propagation, LRP) auf Wasserstein-Abstände anwendet.

Kernidee:
Der Ansatz basiert auf einer Neuralisierung-Propagierung-Strategie:

1. Neuralisierung: Der komplexe iterative Prozess der Berechnung des Wasserstein-Abstands $W_p$ wird in eine äquivalente, zweischichtige neuronale Netzwerk-Struktur umgewandelt.

   • Schicht 1: Berechnet die paarweisen Distanzen zwischen Quell- und Zielinstanzen ($z_{kl} = \|x_k - y_l\|_q$).
   • Schicht 2: Wendet eine normierte Aggregation über diese Distanzen an, gewichtet durch die optimale Kopplung $\gamma^\star$ ($W_p = (\sum \gamma^\star_{kl} z_{kl}^p)^{1/p}$).
   • Dabei wird die Kopplung $\gamma^\star$ als konstant betrachtet (bereits durch einen OT-Löser vorgegeben).

2. Propagierung (LRP): Sobald der Abstand als neuronales Netz modelliert ist, wird die Layer-wise Relevance Propagation rückwärts durch das Netzwerk durchgeführt, um die Relevanz des Gesamtabstands auf die Eingaben zurückzuführen.

   • Instanz-Attribution: Die Relevanz wird von $W_p$ auf die Paare von Instanzen $(x_k, y_l)$ verteilt ($R_{kl}$).
   • Feature-Attribution: Die Relevanz wird weiter auf die einzelnen Eingabe-Features ($R_i$) verteilt.

Hyperparameter und Heuristik:
Die Methode verwendet zwei Hyperparameter $\alpha$ und $\beta$, die steuern, wie stark die Relevanz auf Ausreißer (hohe Werte) oder gleichmäßig verteilt wird. Die Autoren schlagen eine Heuristik vor ($\alpha = p$, $\beta = \min(p+2, q)$), die sich als robust erwiesen hat.

Erweiterung: U-WaX (Subspace-based Explanations)
Für komplexere Szenarien wird U-WaX eingeführt. Dies ermöglicht die Attribution des Wasserstein-Abstands auf Unterräume (Subspaces) der Eingabedaten, die als abstrakte Konzepte oder Teilverschiebungen interpretiert werden können. Dies geschieht durch die Optimierung einer orthogonalen Matrix $U$, die die Daten in signifikante Unterräume zerlegt, basierend auf einer „Tailedness"-Statistik (Extremwertanalyse der Transportkosten).

3. Wichtige Beiträge

• Erstmalige systematische Erklärung von Verteilungs-Distanzen: WaX ist die erste Methode, die Wasserstein-Abstände systematisch auf Datenkomponenten (Features, Instanzen, Subgruppen) zurückführt, anstatt nur den Transportplan zu analysieren.
• Theoretische Fundierung: Die Methode erfüllt wichtige Axiome von Erklärungen, insbesondere die Erhaltungseigenschaft (Conservation), d. h., die Summe der attribuierten Relevanzen entspricht exakt dem berechneten Wasserstein-Abstand.
• Effizienz: Im Gegensatz zu Methoden, die Features entfernen und den Abstand neu berechnen müssen (Occlusion), benötigt WaX nur einen einzigen Vorwärts- und Rückwärtsdurchlauf (unter Nutzung von Gradienten-Äquivalenzen), was sie für große Datensätze skalierbar macht.
• Flexibilität: Das Framework ist unabhängig von der spezifischen Formulierung des Wasserstein-Abstands (z. B. $p$-Norm, $q$-Metrik, Sinkhorn-Regularisierung) anwendbar.

4. Ergebnisse und Evaluation

Die Autoren evaluieren WaX umfassend auf mehreren Datensätzen und Szenarien:

• Faithfulness (Treue der Erklärung): Gemessen durch die Symmetric Relevance Gain (SRG) Metrik. WaX übertrifft konsistent Baseline-Methoden wie MeanShift (nur Mittelwertvergleich), Occlusion (Feature-Entfernung) und Coupling (reine Analyse des Transportplans). Besonders bei nichtlinearen Modellen (hohe $p$ und $q$ Werte) und bei Ausreißer-Sensitivität zeigt WaX überlegene Leistung.
• Charakterisierung von Transportphänomenen: In Experimenten mit synthetischen Zeitreihen und histopathologischen Daten (PLISM-Datensatz) konnte WaX die „Ground Truth" der Transportrichtungen (welche Dimensionen sich ändern) genauer rekonstruieren als lineare Klassifikatoren oder Mittelwert-Vergleiche.
• Anwendungsfälle (Use Cases):
  1. Domain Adaptation: WaX hilft, domainspezifische (nicht robuste) Features zu identifizieren und zu entfernen, was die Robustheit von Klassifikatoren auf neuen Domänen (z. B. Office-Caltech10 Benchmark) signifikant verbessert.
  2. Alterungsprozess (Abalone-Datensatz): U-WaX konnte einen simulierten Alterungsprozess in eine heterogene Population zerlegen und zeigte, dass verschiedene Untergruppen (Subspaces) unterschiedliche Merkmale (z. B. Größe vs. Gewicht) unterschiedlich stark beeinflussen.
  3. Datensatz-Unterschiede (CelebA vs. LFW): U-WaX identifizierte semantische Unterschiede zwischen Gesichtsbild-Datensätzen, z. B. Unterschiede in der Darstellung von Geschlechtern, das Tragen von Brillen oder das Vorhandensein von Paaren, und verknüpfte diese mit textuellen Konzepten (via CLIP-Embeddings).

5. Bedeutung und Ausblick

Die Arbeit schließt eine Lücke in der Literatur zu Explainable AI, indem sie Distanzmetriken selbst erklärbar macht.

• Praktischer Nutzen: WaX ermöglicht es Praktikern, nicht nur zu wissen, dass sich zwei Datensätze unterscheiden, sondern zu verstehen, warum und wo diese Unterschiede liegen. Dies ist entscheidend für die Validierung von Modellen, die Auswahl robuster Features und das Verständnis komplexer physikalischer oder biologischer Prozesse.
• Interpretierbarkeit: Durch die Zerlegung in Subspaces (U-WaX) können komplexe Verschiebungen in verständliche, konzeptuelle Komponenten zerlegt werden.
• Zukunft: Die Autoren sehen Potenzial darin, die Methode auf fortgeschrittenere OT-Formulierungen (z. B. Gromov-Wasserstein, Sliced Wasserstein) und dynamische/kausale Transportmodelle zu erweitern.

Zusammenfassend bietet WaX einen effizienten, theoretisch fundierten und vielseitig anwendbaren Ansatz, um die „Black Box" der Verteilungsvergleiche zu öffnen und tiefgehende Einblicke in Datensatzverschiebungen zu gewinnen.