Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Das Papier stellt CausalLongPFN vor, ein eingefrorenes, auf synthetischen temporalen Kausalmodellen vortrainiertes Prior-Fitted Network, das kontrafaktische longitudinale Ergebnisse ohne domänenspezifisches Training vorhersagt und dabei eine wettbewerbsfähige Leistung bei Benchmarks wie Krebs, HIV und MIMIC-III zeigt, während es gleichzeitig eine kosteneffiziente Alternative zur traditionellen kohortenspezifischen Modellierung bietet.

Das Wesentliche

Stellen Sie sich vor, Sie sind ein Arzt, der versucht vorherzusagen, wie ein Patient in den nächsten Wochen auf einen bestimmten Behandlungsplan reagieren wird. Das Problem ist, dass Patienten komplex sind: Ihr Gesundheitszustand verändert sich von Tag zu Tag, vergangene Behandlungen beeinflussen ihren aktuellen Zustand und andere Faktoren (wie Ernährung oder Stress) greifen in die Ergebnisse ein. Normalerweise müssen Sie, um solche Vorhersagen zu treffen, für jede neue Gruppe von Patienten, mit der Sie konfrontiert werden, ein brandneues, hochspezialisiertes Computermodell erstellen müssen. Das ist so, als würde man jedes Mal einen neuen Architekten engagieren, um ein Haus zu entwerfen, wenn man in eine neue Nachbarschaft zieht. Es ist langsam, teuer und erfordert viel Datenmaterial.

Dieses Paper stellt ein neues Werkzeug namens CAUSALLONGPFN vor. Betrachten Sie es als eine „Universelle Medizinische Intuitions-Maschine“, die die Regeln des Spiels bereits gelernt hat, bevor sie jemals einen echten Patienten sieht.

So funktioniert es, unterteilt in einfache Konzepte:

1. Das „Trainingslager“ (Synthetisches Pretraining)

Anstatt darauf zu warten, dass echte Patienten eintreffen, haben die Schöpfer ein riesiges, virtuelles „Trainingslager“ gebaut. In diesem Lager wurden Millionen von fiktiven Patienten mit Millionen von verschiedenen Körpertypen, Krankheiten und Reaktionsweisen auf Behandlungen simuliert. Sie haben diese fiktiven Patienten so programmiert, dass sie komplexes Verhalten zeigen:

• Einige werden langsam gesund; andere stürzen schnell ab.
• Einige Behandlungen wirken sofort; andere haben einen verzögerten Effekt.
• Einige Patienten reagieren unterschiedlich basierend auf ihrer verborgenen Genetik.

Das KI-Modell hat seine gesamte Zeit in diesem Lager verbracht und gelernt, die Ergebnisse für diese fiktiven Patienten vorherzusagen. Es hat nicht einfach nur Antworten auswendig gelernt; es hat die zugrunde liegende Logik davon gelernt, wie Behandlungen, Zeit und Biologie interagieren.

2. Der „Eingefrorene Experte“ (Kein neues Training erforderlich)

Hier liegt der magische Trick: Sobald das Modell sein Trainingslager abgeschlossen hatte, haben sie es eingefroren. Sie haben sein Gehirn versiegelt. Es kann nichts Neues lernen oder seine internen Einstellungen ändern.

Wenn eine neue Gruppe echter Patienten eintrifft (wie etwa Krebspatienten in einem Krankenhaus), beginnt das Modell nicht bei Null. Es muss nicht neu trainiert werden. Stattdessen agiert es wie ein Super-Praktikant, der die Krankenakte liest.

• Die Support-Trajektorien: Sie zeigen dem Modell einige Beispiele von Patienten aus dem aktuellen Krankenhaus (den „Support“). Dies ist so, als würde man dem Praktikanten ein paar Fallakten zum Lesen geben, damit er den spezifischen Stil dieses Krankenhauses versteht.
• Die Abfrage (Query): Sie fragen: „Was wird mit diesem spezifischen Patienten passieren, wenn wir ihm 5 Tage lang Behandlung A geben?“
• Die Antwort: Das Modell nutzt sofort das, was es im Trainingslager gelernt hat, kombiniert mit dem Kontext der Fallakten, die Sie ihm gerade gezeigt haben, um das Ergebnis vorherzusagen. Es tut dies, ohne auch nur einen einzigen Schritt des „Gradientenabstiegs“ (ein technischer Begriff für den üblichen Trainingsprozess) durchzuführen.

3. Der „Zeitreise-Simulator“

Das Modell ist darauf ausgelegt, mit longitudinalen Daten umzugehen, was bedeutet, dass es die Zeit versteht. Es rät nicht nur das Ergebnis von morgen; es simuliert die Zukunft Schritt für Schritt.

• Es sagt Tag 1 voraus.
• Es nimmt diese Vorhersage und nutzt sie als Ausgangspunkt für Tag 2.
• Es wiederholt diesen Prozess, um zu sehen, was an Tag 5, 6 oder 7 passiert.

Dies ist wie ein Flugsimulator, der nicht nur den Startflug zeigt, sondern den gesamten Flugpfad basierend auf den Entscheidungen des Piloten simuliert, selbst wenn sich das Wetter während des Fluges ändert.

4. Warum das wichtig ist (Die Ergebnisse)

Die Autoren haben diesen „eingefrorenen Experten“ gegen die alte Methode getestet (für jedes Krankenhaus ein neues Modell zu bauen).

• Der Test: Sie verwendeten Daten aus den Bereichen Krebs, HIV, Blutverdünner (Warfarin) und echte Intensivstations-Aufzeichnungen.
• Das Ergebnis: Das eingefrorene Modell schnitt genauso gut ab wie – und manchmal sogar besser als – die Modelle, die speziell für jede einzelne Krankheit trainiert wurden.
• Der große Gewinn: Es funktionierte besonders gut bei echten Intensivstations-Daten, bei denen man keine „Was-wäre-wenn“-Szenarien testen konnte (da es ethisch nicht vertretbar ist, in einer Simulation verschiedene Behandlungen an echten Patienten auszuprobieren). Das Modell konnte allein basierend auf den beobachteten Daten vorhersagen, was wahrscheinlich als Nächstes passieren wird.

Das Wesentliche

Die Autoren behaupten, dass man nicht immer ein neues, maßgeschneidertes Modell für jeden neuen medizinischen Datensatz bauen muss. Stattdessen kann man ein einziges, massives Modell auf einer breiten Vielfalt von „Was-wäre-wenn“-Szenarien (synthetischen Daten) trainieren und es dann als ein eingefrorenes, einsatzbereites Werkzeug nutzen.

Es ist, als hätte man einen Meisterkoch, der in einer virtuellen Küche jede Art von Küche geübt hat. Wenn man ihm nun einen neuen Satz Zutaten bringt (eine neue Patientengruppe), muss er nicht erst lernen, wie man kocht; er sieht sich einfach Ihre Zutaten an und weiß sofort, wie er das Gericht zubereitet, indem er seine umfassende, vorab gelernte Intuition nutzt.

Wichtiger Hinweis: Die Autoren betonen vorsichtig, dass dieses Werkzeug für die Vorhersage und Forschung gedacht ist, nicht für die Treffung endgültiger lebensentscheidender Entscheidungen in einer Klinik. Es hilft Ärzten, potenzielle Ergebnisse zu verstehen, stützt sich aber dennoch auf dieselben medizinischen Standardannahmen, auf denen jedes andere kausale Modell basiert. Es ist ein leistungsstarkes Forschungswerkzeug, kein Ersatz für das Urteilsvermögen eines Arztes.

Für Leser, die die technische Umsetzung oder das Modell selbst untersuchen möchten, stehen die Ressourcen hier zur Verfügung:

• Code on GitHub: https://github.com/Amirhossein-Zare/causal-long-pfn
• Pretrained model weights on Hugging Face: https://huggingface.co/Amirhossein-Zare/causal-long-pfn

Technische Zusammenfassung

Technisches Resümee: Causal Longitudinal Prior-Fitted Networks für die kontrafaktische Ergebnisschätzung

Problemformulierung

Die Arbeit befasst sich mit der Herausforderung, potenzielle Ergebnisse unter zukünftigen Behandlungssequenzen in longitudinalen Beobachtungsdaten vorherzusagen. Die Kernaufgabe besteht darin, das geschichtungsbedingte potenzielle Ergebnis $E[Y(\bar{a})_{t+\tau} | H_t]$ zu schätzen, wobei $H_t$ die beobachtete Historie bis zum Zeitpunkt $t$ darstellt und $\bar{a}$ eine geplante zukünftige Behandlungssequenz ist.

Dieses Problem wird durch drei primäre Faktoren erschwert:

1. Zeitvariante Konfundierung: Behandlungszuweisungen zu jedem Schritt hängen von Kovariaten ab, die selbst Ergebnisse vorangegangener Behandlungen sind.
2. Heterogene Patientendynamik: Individuelle Trajektorien weisen komplexe, nichtlineare Zustandsentwicklungen und latente Heterogenität auf.
3. Datenbeschränkungen: Beobachtungskohorten sind oft zu klein, um zuverlässige Deep-Sequence-Modelle für jede neue Domäne oder jeden Simulator von Grund auf zu trainieren.

Bestehende longitudinale kausale Schätzer (z. B. RMSN, CRN, G-Net, Causal Transformer) erfordern typischerweise einen separaten überwachten Trainingslauf, einschließlich Hyperparameter-Auswahl und Propensity-Modellierung, für jede neue Kohorte. Diese Pipeline ist kostspielig und unpraktisch, wenn wiederholtes domänenspezifisches Training erforderlich ist.

Methodik: CAUSALLONGPFN

Die Autoren führen Causal Longitudinal Prior-Fitted Networks (CAUSALLONGPFN) ein, einen Prior-Fitted In-Context-Prädiktor, der die longitudinale kausale Vorhersage über einen breiten Prior über temporale strukturelle Kausalmodelle (TSCMs) amortisieren soll.

1. Synthetisches Pretraining auf einem breiten Prior

Das Modell wird vollständig auf synthetischen Episoden vortrainiert, die aus einem breiten Prior über TSCMs gesampelt wurden. Dieser Prior ist darauf ausgelegt, eine weite Klasse longitudinaler kausaler Dynamiken abzudecken, anstatt einen einzelnen Simulator zu replizieren. Zu den Hauptmerkmalen des synthetischen Priors gehören:

• Kausale temporale Graphen: Dünnbesetzte simultane und verzögerte Abhängigkeiten mit azyklischen instantanen Graphen.
• Nichtlineare Mechanismen: Zustands-Koordinaten folgen dünnbesetzten nichtlinearen autoregressiven Updates unter Verwendung diverser elementarer Nichtlinearitäten (Identität, Tanh, Sinusoidal, ReLU, etc.) und verschiedener Rauschverteilungen.
• Dynamische Motive: Strukturierte Motive wie Action-Memory, Sättigung, Homöostase, Feedback-Kontrolle und geglättete Readout-Kanäle sind überlagert, um qualitative Mechanismen wie verzögerte Effekte und regulatorisches Feedback zu erfassen.
• Konfundiertes Verhalten der Policies: Behandlungen werden aus zustandsabhängigen stochastischen Policies gesampelt, die durch latente Einheit-Heterogenität ($Z_i$) beeinflusst werden, was ein zeitvariantes Feedback zwischen Behandlung und Konfounder erzeugt.
• Outcome-Modelle: Skalare Ergebnisse werden über autoregressive Readouts mit direkten und kumulativen Behandlungseffekten generiert.

2. Architektur

CAUSALLONGPFN verwendet eine Dual-Encoder-Architektur:

• Causal History Encoder: Ein Trajektorien-basierter kausaler Transformer (unter Verwendung von Masked Self-Attention), der longitudinale Sequenzen auf Historien-Repräsentationen abbildet und sicherstellt, dass die Repräsentation zum Zeitpunkt $r$ nur von Informationen abhängt, die bis zu diesem Zeitpunkt verfügbar sind.
• PFN Context Encoder: Ein Transformer-basierter Kontext-Encoder, der eine In-Context-Adaption durchführt. Er verarbeitet Support-Trajektorien (die als ungeordnete Menge behandelt werden) und ein Query-Token gemeinsam mittels Self-Attention. Es wird keine Positionskodierung für die Ordnung der Support-Trajektorien zugewiesen.
• Gaussian-Mixture Prediction Head: Der finale Query-Repräsentations-Head parametrisiert eine 5-komponentige Gaußsche Mischverteilung für das normalisierte nächste Ergebnis, was sowohl Punktprognosen als auch Unsicherheitsschätzungen ermöglicht.

3. In-Context-Vorhersage und Rollout

Zum Testzeitpunkt ist das Modell eingefroren (frozen). Es erhält:

• Support-Trajektorien: Beispiele aus der neuen Domäne/Aufgabe.
• Query-Historie: Beobachtet bis zum Zeitpunkt $t_{obs}$.
• Vorgeschlagene zukünftige Behandlungssequenz: Die geplante Intervention.

Das Modell liefert eine prädiktive Verteilung ohne Gradienten-Updates, Propensity-Modell-Fitting oder adversarielles Balancing. Für die Mehrschritt-Vorhersage ($\tau > 1$) führt das Modell einen autoregressiven Plug-in Rollout durch: Es sagt die Ein-Schritt-Ergebnisverteilung voraus, fügt den Mittelwert der Mischverteilung in die Query-Historie ein und wiederholt den Prozess unter der spezifizierten Behandlungssequenz.

Zentrale Beiträge

1. Ein Prior-Fitted Modell für die longitudinale kausale Vorhersage: CAUSALLONGPFN ist der erste PFN-Stil-Modell für geschichtungsbedingte potenzielle Ergebnisvorhersagen unter geplanten longitudinalen Behandlungssequenzen. Es operiert als eingefrorenes Modell, das keine Testzeit-Adaption erfordert.
2. Ein synthetischer Prior über longitudinale kausale Aufgaben: Die Autoren entwerfen einen TSCM-Prior, der diverse Aufgaben mit Treatment-Confounder-Feedback, latenter Heterogenität, nichtlinearen Dynamiken, verzögerten/kumulativen Effekten und Regimewechseln erzeugt.
3. Architektur für longitudinale In-Context-Inferenz: Ein neuartiger Dual-Encoder, der einen kausalen Transformer-History-Encoder mit einem PFN-Context-Encoder und einem Gaussian-Mixture-Head kombiniert.
4. Autoregressiver Counterfactual Rollout: Eine Erweiterung des gelernten Ein-Schritt-Prädiktors auf die Mehrschritt-Vorhersage mittels deterministischem Plug-in Rollout.
5. Zero-Shot Evaluierung: Eine umfassende Evaluierung eines einzelnen eingefrorenen Modells gegen domänenspezifisch trainierte Baselines (MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer) auf verzweigbaren kontrafaktischen Benchmarks und faktischen Realdaten.

Ergebnisse

Das Modell wurde auf vier Benchmarks evaluiert: Krebsgeschwür-Wachstum, Warfarin PK/PD, HIV-Behandlungsdynamik und MIMIC-III ICU-Trajektorien.

• Domänen-balancierte Performance: CAUSALLONGPFN erreichte den besten domänen-balancierten One-Step-normalisierten RMSE (0,222) und übertraf damit knapp MSM und RMSN. Bei der Fünf-Schritt-Vorhersage belegte es den dritten Platz insgesamt, hinter RMSN und G-Net, übertraf jedoch MSM, CRN und Transformer-basierte Baselines.
• Pro-Domäne Performance:
  • MIMIC-III (Faktisch): Das Modell belegte den ersten Platz sowohl in der Ein-Schritt- als auch in der Fünf-Schritt-Vorhersage, was eine starke Übertragung auf reale klinische Trajektorien ohne domänenspezifisches Training zeigt.
  • Kontrafaktische Benchmarks (Krebs, HIV, Warfarin): Das Modell blieb wettbewerbsfähig und belegte bei Ein-Schritt-Aufgaben den zweiten oder dritten Platz. Bei längerfristigen kontrafaktischen Aufgaben (z. B. 5-Schritt-Krebsvorhersage) erreichten spezialisierte, domänenspezifisch trainierte rekursive Modelle (RMSN, CRN) jedoch geringere Fehler, was auf einen Vorteil hindeutet, wenn ausreichend Daten der Zieldomäne für ein spezifisches Fitting vorhanden sind.
• Unsicherheit-Kalibrierung: Der Gaussian-Mixture-Head lieferte nützliche Verteilungsinformationen. Die Kalibrierung variierte je nach Domäne, wobei Warfarin die beste Kalibrierung zeigte und MIMIC-III aufgrund höherer Heterogenität weitere Intervalle aufwies.

Bedeutung, Ansprüche und Verfügbarkeit

Das Paper behauptet, dass breites synthetisches Causal-Pretraining eine nützliche, eingefrorene Alternative zum wiederholten domänenspezifischen Training bieten kann, wenn:

• Das erneute Training spezialisierter Modelle kostspielig oder unpraktisch ist.
• Eine schnelle Anpassung an eine neue Kohorte erforderlich ist.
• Kontrafaktische Supervision nicht verfügbar ist (wie bei realen faktischen Vorhersageaufgaben wie MIMIC-III).

Die Autoren betonen, dass CAUSALLONGPFN die Standard-Kausalannahmen (Konsistenz, Positivität, sequentielle Austauschbarkeit), die für die Interpretation von Beobachtungsdaten erforderlich sind, nicht aufhebt. Stattdessen amortisiert es das Schätzproblem. Die Ergebnisse legen nahe, dass ein ausreichend breiter synthetischer Prior reusable Strukturen über Behandlungs-Reaktions-Aufgaben hinweg erfassen kann, was das Modell zu einem starken, allgemeinen In-Context-Prädiktor macht. Dennoch merkt das Paper bescheiden an, dass domänenspezifisches Training wertvoll bleibt, wenn ausreichend Daten der Zieldomäne und Validierungssignale für spezifische, langfristige kontrafaktische Vorhersagen verfügbar sind.

Die Arbeit positioniert sich als Forschungswerkzeug für kausale Sequenzmodellierung und Hypothesengenerierung und warnt davor, die Vorhersagen als eigenständiges klinisches Entscheidungssystem zu nutzen, insbesondere wenn kausale Annahmen oder der Prior-Support unzureichend sind.

Code & Model Availability:
Die Implementierung ist auf GitHub verfügbar unter https://github.com/Amirhossein-Zare/causal-long-pfn. Die vortrainierten Modellgewichte sind auf Hugging Face verfügbar unter https://huggingface.co/Amirhossein-Zare/causal-long-pfn.