Frequentist Consistency of Prior-Data Fitted Networks for Causal Inference

Diese Arbeit zeigt, dass Prior-Data-Fitted-Networks (PFNs) für kausale Inferenz ohne Korrektur zu einer durch den Prior verursachten Verzerrung führen, und schlägt eine Kalibrierung mittels einer einstufigen Posterior-Korrektur (OSPC) auf Basis von Martingal-Posterior-Verteilungen vor, um die frequentistische Konsistenz und eine korrekte Unsicherheitsquantifizierung für den durchschnittlichen Behandlungseffekt (ATE) wiederherzustellen.

Das Wesentliche

Das große Problem: Der "Vorurteils-Filter"

Stell dir vor, du möchtest herausfinden, ob ein neues Medikament wirklich hilft. Du hast keine kontrollierte Studie (wo man die Hälfte der Leute zufällig behandelt und die andere Hälfte nicht), sondern nur alte Daten aus der Praxis. Das ist wie ein Detektiv, der nur verworrene Tatorte hat, aber keine Zeugen.

In der modernen KI gibt es eine neue Art von "Super-Detektiven", die PFNs (Prior-Data Fitted Networks) genannt werden.

• Wie sie funktionieren: Diese Detektive wurden nicht auf echten Tatorten trainiert, sondern in einer riesigen, künstlichen Simulation. Sie haben Millionen von "fiktiven Verbrechen" gesehen, die von einem bestimmten "Regisseur" (dem Prior) inszeniert wurden.
• Das Problem: Wenn diese Detektive nun echte Fälle analysieren, schauen sie durch eine Brille, die der Regisseur ihnen aufgesetzt hat. Wenn der Regisseur in seiner Simulation angenommen hat, dass Verbrechen nie durch bestimmte Umstände verursacht werden, dann wird der Detektiv auch in der echten Welt blind für diese Zusammenhänge sein.

Die Forscher haben entdeckt: Diese KI-Detektive neigen dazu, Verzerrungen (Bias) zu unterschätzen. Sie glauben fälschlicherweise, die Welt sei "sauberer" und weniger verworren, als sie wirklich ist. Das führt dazu, dass ihre Unsicherheitsangaben (z. B. "Wir sind zu 95 % sicher") oft falsch sind. Sie sind zu selbstsicher, obwohl sie sich irren könnten.

Die Lösung: Der "Korrektur-Schritt" (OSPC)

Die Autoren sagen: "Wir müssen diese KI-Detektive nicht neu erfinden, wir müssen ihnen nur eine Korrektur-Brille aufsetzen."

Sie haben eine Methode namens OSPC (One-Step Posterior Correction) entwickelt.

• Die Analogie: Stell dir vor, der Detektiv macht eine erste Schätzung. Dann kommt ein erfahrener, skeptischer Statistiker (der klassische frequentistische Ansatz) und sagt: "Moment mal, deine Schätzung ist gut, aber du hast einen systematischen Fehler gemacht, weil du zu sehr auf deine alten Trainingsdaten vertraust. Hier ist eine kleine Korrekturformel, die den Fehler ausgleicht."
• Das Ergebnis: Nach dieser Korrektur stimmt die Unsicherheit der KI plötzlich perfekt mit der eines klassischen, mathematisch bewährten Statistikers überein. Die KI wird "frequentistisch konsistent". Das bedeutet: Wenn man immer mehr Daten hat, wird ihre Antwort immer genauer und verlässlicher, genau wie bei den alten, bewährten Methoden.

Der Trick: Wie man die "Innere Welt" der KI sieht (Martingale Posteriors)

Es gibt noch ein technisches Hindernis. Die KI-Detektive geben normalerweise nur eine Antwort für einen einzelnen Fall aus (z. B. "Patient A hat 80 % Heilungschance"). Um die Korrektur anzuwenden, brauchen wir aber die Antwort für alle Patienten gleichzeitig als ein zusammenhängendes Bild (eine "funktionale Posterior").

• Das Problem: Die KI gibt uns nur einzelne Puzzleteile, aber kein Gesamtbild.
• Die Lösung (MP-OSPC): Die Forscher nutzen eine Technik namens Martingale Posteriors.
  • Die Analogie: Stell dir vor, du hast einen Zauberwürfel, der dir nur sagt, welche Farbe eine bestimmte Seite hat. Du willst aber wissen, wie der ganze Würfel aussieht. Die neue Methode nimmt diesen Würfel, dreht ihn ein paar Mal (simuliert neue Daten), schaut sich die Farben an und nutzt eine spezielle "Klebe-Technik" (Copulas), um aus den einzelnen Farben ein konsistentes Gesamtbild zu rekonstruieren.
  • So können sie aus den einzelnen KI-Antworten ein vollständiges, glattes Bild der Unsicherheit für alle Patienten gleichzeitig erstellen.

Warum ist das wichtig?

1. Vertrauen: Bisher war unklar, ob man den Unsicherheitsangaben von diesen neuen, super-schnellen KI-Modellen trauen konnte. Jetzt wissen wir: Mit der Korrektur sind sie verlässlich.
2. Best of Both Worlds: Man bekommt die Geschwindigkeit und Flexibilität der modernen KI (die in Sekunden lernt) und die mathematische Strenge der klassischen Statistik (die uns vor falschen Schlüssen schützt).
3. Anwendung: Das ist super wichtig für Bereiche wie Medizin (Welches Medikament hilft wirklich?), Politik (Hilft eine neue Steuer?) oder Wirtschaft. Man will keine falsche Sicherheit haben, wenn es um Menschenleben oder Geld geht.

Zusammenfassung in einem Satz

Die Forscher haben gezeigt, dass neue KI-Modelle für Kausalität zwar schnell sind, aber durch ihre Trainingsdaten verzerrt sein können; mit einer cleveren Korrektur-Methode (OSPC) und einem Trick zum Rekonstruieren von Gesamtbildern (Martingale Posteriors) machen sie diese Modelle so verlässlich wie die besten klassischen Statistiker – und das ohne sie neu zu trainieren.

Technische Zusammenfassung

1. Problemstellung und Motivation

Hintergrund:
Prior-Data Fitted Networks (PFNs) sind eine Klasse von Foundation-Modellen, die Bayessche Inferenz durch Vorab-Training auf synthetischen Datensätzen amortisieren. Sie haben sich in der kausalen Inferenz (z. B. zur Schätzung des durchschnittlichen Behandlungseffekts, ATE) als leistungsfähig erwiesen, indem sie die Aufgabe als „In-Context-Learning"-Problem formulieren. PFNs liefern eine Posterior Predictive Density (PPD), die eine „Out-of-the-Box"-Unsicherheitsquantifizierung ermöglicht.

Das Kernproblem:
Es ist unklar, ob PFN-basierte Schätzer für den ATE eine frequentistische Konsistenz aufweisen. Das bedeutet: Konvergieren die von PFNs generierten Posterior-Verteilungen asymptotisch gegen die Verteilung klassischer, effizienter frequentistischer Schätzer (wie den A-IPTW-Schätzer)?
Die Autoren identifizieren ein fundamentales Hindernis: den prior-induzierten Confounding-Bias.

• PFNs werden auf synthetischen Daten trainiert, die aus einer impliziten Prior-Verteilung stammen.
• Diese implizite Prior neigt dazu, beobachtete Confounding-Effekte systematisch in Richtung Null zu schrumpfen (Regularisierung-induced Confounding Bias).
• Folge: Selbst bei wachsendem Stichprobenumfang ($n \to \infty$) wird der Prior nicht vollständig durch die Daten „überschrieben". Die Posterior-Verteilung des ATE bleibt verzerrt, was die frequentistische Konsistenz und die Gültigkeit des Bernstein-von-Mises (BvM)-Theorems verhindert.

2. Methodik

Die Arbeit schlägt einen dreistufigen Ansatz vor, um dieses Problem zu lösen und PFNs konsistente kausale Schätzer zu machen:

A. Analyse des Bias (Prior-Induced Confounding Bias)

Die Autoren zeigen theoretisch und empirisch, dass naive PFN-basierte Schätzer (die PPDs direkt als funktionale Posteriors interpretieren) systematisch verzerrt sind. Der Bias entsteht, weil die im Training gelernte Prior-Verteilung hochgradige Confounding-Szenarien (hohe Werte von $\Delta = E[Y|A=1] - E[Y|A=0] - \text{ATE}$) unterrepräsentiert.

B. One-Step Posterior Correction (OSPC)

Um den Bias zu korrigieren, wird eine One-Step Posterior Correction (OSPC) eingeführt, basierend auf der effizienten Einflussfunktion (Efficient Influence Function, EIF) des ATE.

• Prinzip: Ähnlich wie der frequentistische A-IPTW-Schätzer eine Bias-Korrektur des Plug-in-Schätzers darstellt, wird die OSPC als Korrektur des Bayesschen Plug-in-Posteriors definiert.
• Formel: Der korrigierte Posterior für den ATE $\psi$ wird als Push-Forward der funktionalen Posteriors der Störvariablen (Nuisance Functions: $\mu_a$ und $\pi$) berechnet:
  $$\psi_{\text{OSPC}}(\tilde{\eta}) = \psi_{\text{PI}}(\tilde{\eta}) + \mathbb{E}_{\text{BB}}[\phi_\psi(Z; \tilde{\eta})]$$
  wobei $\phi_\psi$ die effiziente Einflussfunktion ist und $\mathbb{E}_{\text{BB}}$ eine Erwartung über das Bayesian Bootstrap darstellt.
• Theoretischer Gewinn: Unter milden Annahmen (insbesondere $L_2$-Konvergenz der Nuisance-Funktionen) stellt die OSPC sicher, dass der korrigierte Posterior asymptotisch normalverteilt ist und mit dem frequentistischen A-IPTW-Schätzer übereinstimmt (semi-parametrisches BvM-Theorem).

C. Martingale Posteriors (MPs) zur Gewinnung funktionaler Posteriors

Ein praktisches Problem bei PFNs ist, dass sie nur punktweise PPDs liefern, keine vollständigen funktionalen Posteriors über die Nuisance-Funktionen ($\tilde{\mu}, \tilde{\pi}$), die für die OSPC benötigt werden.

• Lösung: Die Autoren nutzen das Framework der Martingale Posteriors (MPs).
• Implementierung (MP-OSPC): Sie kombinieren PFNs mit Copulas (ein hybrider Ansatz).
  1. Der PFN liefert die initiale PPD (Schritt 1).
  2. Copulas werden verwendet, um die PPDs sequentiell zu aktualisieren und dabei die Abhängigkeiten zwischen den Datenpunkten zu modellieren (Schritt $N > 1$).
  3. Dies ermöglicht das Ziehen von Stichproben ganzer funktionaler Posteriors ($\tilde{\eta} \sim \Pi(\eta|D)$), die notwendig sind, um die OSPC durchzuführen.
  4. Der Ansatz erlaubt verschiedene Kopplungsarten (x-unabhängig, x-parallel, glatt/smooth), wobei die „smooth"-Variante als natürlichere Repräsentation der Unsicherheit identifiziert wird.

3. Wichtige Beiträge

1. Identifikation des Bias: Erstmals wird gezeigt, dass naive PFN-basierte Bayessche ATE-Schätzer aufgrund eines prior-induzierten Confounding-Bias nicht frequentistisch konsistent sind, da der Prior asymptotisch nicht verschwindet.
2. Theoretische Korrektur (OSPC): Entwicklung einer Kalibrierungsmethode (OSPC), die die frequentistische Konsistenz wiederherstellt und ein semi-parametrisches Bernstein-von-Mises-Theorem für PFNs beweist.
3. Praktische Implementierung (MP-OSPC): Einführung eines Frameworks, das Martingale Posteriors mit PFNs kombiniert, um die für die OSPC notwendigen funktionalen Posteriors der Störvariablen zu rekonstruieren.
4. Empirische Validierung: Umfassende Experimente zeigen, dass die MP-OSPC-Methodik sowohl asymptotisch als auch in endlichen Stichproben überlegene Ergebnisse liefert.

4. Ergebnisse

Die Autoren evaluieren ihre Methode an synthetischen Daten, dem IHDP-Datensatz und 77 Datensätzen aus dem ACIC 2016 Benchmark sowie in einer Fallstudie zu Lockdown-Maßnahmen während der COVID-19-Pandemie.

• Asymptotische Konsistenz: Die mit MP-OSPC kalibrierten PFN-Schätzer konvergieren in ihrer Verteilung gegen die des klassischen A-IPTW-Schätzers (gemessen durch die Total-Variation-Distanz $\hat{d}_{TV}$). Naive Plug-in-Methoden zeigen hier signifikante Abweichungen.
• Finite-Sample-Kalibrierung: In endlichen Stichproben liefern die MP-OSPC-Schätzer besser kalibrierte Konfidenzintervalle (gemessen durch die Kolmogorov-Smirnov-Distanz $\hat{d}_{KS}$) als andere Bayessche ATE-Schätzer und naive PFN-Methoden.
• Robustheit: Die Methode funktioniert robust über verschiedene Dimensionen der Kovariaten ($d_x$) und Stichprobengrößen hinweg. Sie verbessert insbesondere PFNs, deren implizite Prior-Verteilung nicht mit dem tatsächlichen Confounding-Niveau der Daten übereinstimmt.
• Fallstudie: In der Analyse von Lockdown-Effekten stimmen die Unsicherheitsverteilungen der MP-OSPC-Schätzer stark mit den frequentistischen A-IPTW-Ergebnissen überein, was die praktische Anwendbarkeit unterstreicht.

5. Bedeutung und Fazit

Diese Arbeit schließt eine kritische Lücke in der Forschung zu Foundation Models für die kausale Inferenz. Sie demonstriert, dass PFNs zwar mächtige Werkzeuge zur Unsicherheitsquantifizierung sind, aber ohne Korrektur systematische Fehler aufweisen können, die ihre theoretische Fundierung untergraben.

Durch die Einführung der MP-OSPC bieten die Autoren einen Weg, die Stärken beider Welten zu kombinieren:

1. Die frequentistische Konsistenz und asymptotische Effizienz klassischer Schätzer.
2. Die flexible, prior-gesteuerte Unsicherheitsquantifizierung und die Rechenleistung von Foundation Models in endlichen Stichproben.

Dies ermöglicht den zuverlässigen Einsatz von PFNs in hochriskanten Anwendungsbereichen (wie Medizin oder Politik), wo korrekte Unsicherheitsintervalle für fundierte Entscheidungen entscheidend sind.