Causal Longitudinal Prior-Fitted Networks for Counterfactual Outcome Prediction

Het artikel introduceert CausalLongPFN, een bevroren prior-fitted netwerk dat vooraf is getraind op synthetische temporele causale modellen en die contrafectuele longitudinale uitkomsten voorspelt zonder domeinspecifieke training, waarbij het competitieve prestaties demonstreert op benchmarks zoals kanker, HIV en MIMIC-III, terwijl het een kosteneffectief alternatief biedt voor traditionele cohortspecifieke modellering.

De kern

Stel je voor dat je een arts bent die probeert te voorspellen hoe een patiënt zal reageren op een specifiek behandelplan in de komende weken. Het probleem is dat patiënten complex zijn: hun gezondheid verandert van dag tot dag, hun eerdere behandelingen beïnvloeden hun huidige staat, en andere factoren (zoals dieet of stress) verstoren de resultaten. Normaal gesproken moet je, om deze voorspellingen te doen, voor elke nieuwe groep patiënten die je tegenkomt een volledig nieuw, hooggespecialiseerd computermodel bouwen. Dit is als het inhuren van een nieuwe architect om een huis te ontwerpen telkens wanneer je naar een nieuwe buurt verhuist. Het is traag, duur en vereist veel data.

Dit artikel introduceert een nieuwe tool genaamd CAUSALLONGPFN. Denk aan dit als een "Universele Medische Intuïtie-motor" die de regels van het spel al heeft geleerd voordat hij ooit een echte patiënt heeft gezien.

Hier is hoe het werkt, onderverdeeld in eenvoudige concepten:

1. Het "Trainingskamp" (Synthetische Pretraining)

In plaats van te wachten tot er echte patiënten arriveren, hebben de makers een enorme, virtuele "trainingskamp" gebouwd. In dit kamp hebben ze miljo's nep-patiënten gesimuleerd met miljo's verschillende lichaamstypen, ziekten en reacties op behandelingen. Ze hebben deze nep-patiënten geprogrammeerd met complexe gedragingen:

• Sommigen worden langzaam beter; anderen gaan snel achteruit.
• Sommige behandelingen werken onmiddellijk; andere hebben een vertraagd effect.
• Sommige patiënten reageren anders op basis van hun verborgen genetica.

Het AI-model heeft al zijn tijd in dit kamp doorgebracht om de uitkomsten voor deze nep-patiënten te voorspellen. Het heeft niet alleen antwoorden uit het hoofd geleerd; het heeft de onderliggende logica geleerd van hoe behandelingen, tijd en biologie met elkaar interageren.

2. De "Bevroren Expert" (Geen Nieuwe Training Nodig)

Hier is de magische truc: nadat het model zijn trainingskamp had voltooid, hebben ze het bevroren. Ze hebben het brein op slot gezet. Het kan niets nieuws leren of zijn interne instellingen veranderen.

Wanneer er een nieuwe groep echte patiënten arriveert (zoals kankerpatiënten in een ziekenhuis), begint het model niet vanaf nul. Het hoeft niet opnieuw getraind te worden. In plaats daarvan fungeert het als een super-stagiair die het dossier leest.

• De Support Trajectories: Je laat het model een paar voorbeelden zien van patiënten uit het huidige ziekenhuis (de "support"). Dit zijn als het geven van een paar dossiers aan de stagiair, zodat hij de specifieelijke stijl van dit ziekenhuis begrijpt.
• De Query: Je vraagt: "Wat zal er met deze specifieke patiënt gebeuren als we hen de komende 5 dagen Behandeling A geven?"
• Het Antwoord: Het model gebruikt direct wat het heeft geleerd in het trainingskamp, gecombineerd met de context van de dossiers die je er net aan hebt getoond, om de uitkomst te voorspellen. Dit doet het zonder ook maar één stap van "gradient descent" te nemen (een technische term voor het gebruikelijke hertrainingsproces).

3. De "Tijdreizende Simulator"

Het model is ontworpen om met longitudinale data om te gaan, wat betekent dat het tijd begrijpt. Het raadt niet alleen de uitslag van morgen; het simuleert de toekomst stap voor stap.

• Het voorspelt Dag 1.
• Het gebruikt die voorspelling vervolgens als startpunt voor Dag 2.
• Het herhaalt dit proces om te zien wat er gebeurt op Dag 5, 6 of 7.

Dit is als een vluchtsimulator die niet alleen de startvlucht laat zien, maar de gehele vluchtroute simuleert op basis van de beslissingen van de piloot, zelfs als het weer halverwege verandert.

4. Waarom Dit Er Toe Doet (De Resultaten)

De auteurs hebben deze "bevroren expert" getest tegen de oude methode (het bouwen van een nieuwe model voor elk ziekenhuis).

• De Test: Ze gebruikten data van kanker, HIV, bloedverdunners (warfarin) en echte IC-gegevens.
• De Uitkomst: Het bevroren model presteerde net zo goed als, en soms zelfs beter dan, de modellen die speciaal waren getraind voor elke specifieke ziekte.
• De Grote Winst: Het werkte vooral goed op echte IC-data waar ze geen "wat als"-scenario's konden testen (omdat het onethisch is om in een simulatie verschillende behandelingen uit te proberen op echte patiënten). Het model kon de volgende stap voorspellen op basis van de geobserveerde gegevens alleen.

De Kernboodschap

De auteurs beweren dat je niet altijd een nieuwe, aangepaste model nodig hebt voor elke nieuwe medische dataset. In plaats daarvan kun je één massaal model trainen op een brede variëteit aan "wat als"-scenario's (synthetische data) en het vervolgens gebruiken als een bevroren, direct inzetbare tool.

Het is als het hebben van een meesterkok die in een virtuele keuken elk type keuken heeft geoefend. Wanneer je hem een nieuwe set ingrediënten brengt (een nieuwe patiëntengroep), hoeft hij niet opnieuw te leren hoe hij moet koken; hij kijkt gewoon naar je ingrediënten en weet onmiddellijk hoe hij het gerecht moet bereiden, gebruikmakend van zijn uitgebreide, vooraf geleerde intuïtie.

Belangrijke Opmerking: De auteurs benadrukken dat deze tool bedoeld is voor voorspelling en onderzoek, en niet voor het maken van definitieve beslissingen over leven of dood in een klinische setting. Het heldt artsen om potentiële uitkomsten te begrijpen, maar het vertrouwt nog steeds op dezelfde standaard medische aannames waar elk ander causaal model ook op vertrouwt. Het is een krachtig onderzoeksinstrument, geen vervanging voor het oordeel van een arts.

---

Wil je zelf aan de slag met dit model?
De code en de vooraf getrainde modelgewichten zijn openbaar beschikbaar:

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

Technische samenvatting

Technische Samenvatting: Causal Longitudinal Prior-Fitted Networks voor Counterfactual Outcome Voorspelling

Probleemformulering

Het artikel behandelt de uitdaging van het voorspellen van potentiële uitkomsten onder toekomstige behandelingssequenties in longitudinale observationele data. De kerntaak is het schatten van de geschiedenis-conditionele potentiële uitkomst $E[Y(\bar{a})_{t+\tau} | H_t]$, waarbij $H_t$ de geobserveerde geschiedenis tot tijd $t$ vertegenwoordigt, en $\bar{a}$ een geplande toekomstige behandelingssequentie is.

Dit probleem wordt bemoeilijkt door drie primaire factoren:

1. Tijdvariërende confounding: Toewijzingen van behandelingen op elk stap zijn afhankelijk van covariaten die op zichzelf uitkomsten zijn van eerdere behandelingen.
2. Heterogene patiëntdynamiek: Individuele trajecten vertonen complexe, niet-lineaire staatsevolutie en latente heterogeniteit.
3. Databeperkingen: Observationele cohorten zijn vaak te klein om betrouwbare deep sequence-modellen vanaf nul te trainen voor elke nieuwe domein of simulator.

Bestaande longitudinale causale estimators (bijv. RMSN, CRN, G-Net, Causal Transformer) vereisen doorgaans een aparte supervised training run, inclusief hyperparameterselectie en propensity-modellering, voor elk nieuw cohort. Deze pipeline is kostbaar en onpraktisch wanneer herhaalde domeinspecifieke training vereist is.

Methodologie: CAUSALLONGPFN

De auteurs introduceren Causal Longitudinal Prior-Fitted Networks (CAUSALLONGPFN), een prior-fitted in-context voorspeller ontworpen om longitudinale causale voorspelling te ammortiseren over een brede prior over temporele structurele causale modellen (TSCM's).

1. Synthetische Pretraining op een Brede Prior

Het model wordt volledig voorgetraind op synthetische episodes gesampled uit een brede prior over TSCM's. Deze prior is ontworpen om een brede klasse van longitudinale causale dynamieken te bestrijken in plaats van één enkele simulator te repliceren. Belangrijke kenmerken van de synthetische prior zijn:

• Causale Temporele Grafen: Spaarzame gelijktijdige en vertraagde afhankelijkheden met acyclische instantane grafen.
• Niet-lineaire Mechanismen: Staatscoördinaten volgen spaarzame niet-lineaire autoregressieve updates met diverse elementaire niet-lineariteiten (identity, tanh, sinusoidale, ReLU, etc.) en diverse ruisverdelingen.
• Dynamische Motieven: Gestructureerde motieven zoals action-memory, saturating, homeostatic, feedback-control en smoothed-readout kanalen worden overlapt om kwalitatieve mechanismen zoals vertraagde effecten en regulerende feedback te vangen.
• Confounded Behavior Policies: Behandelingen worden gesampled uit staat-afhankelijke stochastische policies beïnvloed door latente eenheid-heterogeniteit ($Z_i$), wat zorgt voor een tijdvariërende behandeling-confounder feedback.
• Outcome Modellen: Scalaire uitkomsten worden gegenereerd via autoregressieve readouts met directe en cumulatieve behandelings-effecten.

2. Architectuur

CAUSALLONGPFN maakt gebruik van een dual-encoder architectuur:

• Causal History Encoder: Een traject-niveau causale Transformer (met behulp van masked self-attention) die longitudinale sequenties naar geschiedenis-representaties mapt, waardoor de representatie op tijdstip $r$ alleen afhankelijk is van informatie beschikbaar tot en met dat tijdstip.
• PFN Context Encoder: Een Transformer-gebaseerde context encoder die in-context adaptatie uitvoert. Het verwerkt support-trajecten (behandeld als een ongeordende verzameling) en een query-token gezamenlijk via self-attention. Er wordt geen positional encoding toegewezen aan de ordening van de support-trajecten.
• Gaussian-Mixture Prediction Head: De uiteindelijke query-representatie parametriseert een 5-componenten Gaussian mixture distributie voor de genormaliseerde volgende uitkomst, wat zowel puntvoorspellingen als onzekerheidsschattingen biedt.

3. In-Context Voorspelling en Rollout

Op testtijd is het model bevroren (frozen). Het ontvangt:

• Support trajecten: Voorbeelden uit het nieuwe domein/taak.
• Query geschiedenis: Geobserveerd tot tijd $t_{obs}$.
• Voorgestelde toekomstige behandelingssequentie: De geplande interventies.

Het model geeft een voorspellende distributie terug zonder gradiënt-updates, propensity-model fitting of adversarial balancing. Voor multi-step voorspelling ($\tau > 1$) voert het model een autoregressieve plug-in rollout uit: het voorspelt de one-step uitkomstdistributie, voegt de mixture mean toe aan de query-geschiedenis, en herhaalt het proces onder de gespecificeerde behandelingssequentie.

Belangrijkste Bijdragen

1. Een Prior-Fitted Model voor Longitudinale Causale Voorspelling: CAUSALLONGPFN is de eerste PFN-stijl model voor geschiedenis-conditionele potentiële-uitkomst voorspelling onder geplande longitudinale behandelingssequenties. Het werkt als een bevroren model dat geen test-tijd adaptatie vereist.
2. Een Synthetische Prior over Longitudinale Causale Taken: De auteurs ontwerpen een TSCM-prior die diverse taken genereert met behandeling-confounder feedback, latente heterogeniteit, niet-lineaire dynamiek, vertraagde/cumulatieve effecten en regime-veranderingen.
3. Architectuur voor Longitudinale In-Context Inferentie: Een nieuwe dual-encoder die een causale Transformer history encoder combineert met een PFN context encoder en een Gaussian-mixture head.
4. Autoregressieve Counterfactual Rollout: Een uitbreiding van de geleerde one-step predictor naar multi-step voorspelling via deterministische plug-in rollout.
5. Zero-Shot Evaluatie: Een uitgebreide evaluatie van een enkel bevroren model tegen domein-getrainde baselines (MSM, RMSN, G-Net, CRN, Causal Transformer, G-Transformer) op branchable counterfactual benchmarks en feitelijke real-world data.

Resultaten

Het model werd geëvalueerd op vier benchmarks: Kanker tumor groei, Warfarin PK/PD, HIV behandelingsdynamiek en MIMIC-III ICU trajecten.

• Domein-Gebalanceerde Prestaties: CAUSALLONGPFN behaalde de beste domein-gebalanceerde one-step genormaliseerde RMSE (0.222), waarbij het MSM en RMSN nauwelijks versloeg. Voor vijf-staps voorspelling rangschikte het derde overall, achter RMSN en G-Net, maar presteerde het beter dan MSM, CRN en transformer-gebaseerde baselines.
• Per-Domein Prestaties:
  • MIMIC-III (Factual): Het model rangschikte eerste in zowel one-step als vijf-staps voorspelling, wat duidt op een sterke transfer naar real-world klinische trajecten zonder domeinspecifieke training.
  • Counterfactual Benchmarks (Kanker, HIV, Warfarin): Het model bleef competitief, met een tweede of derde plaats op one-step taken. Echter, op langere-horizon counterfactual taken (bijv. 5-staps kanker voorspelling), bereikten gespecialiseerde domein-getrainde recurrente modellen (RMSN, CRN) lagere fouten, wat suggereert dat er een voordeel is wanneer voldoende target-domein data aanwezig is voor specifieke fitting.
• Onzekerheid Calibratie: De Gaussian-mixture head bood nuttige distributionele informatie. Calibratie varieerde per domein, waarbij Warfarin de beste calibratie vertoonde en MIMIC-III bredere intervallen had vanwege hogere heterogeniteit.

Betekenis en Claims

Het artikel claimt dat brede synthetische causale pretraining een nuttige bevroren alternatief kan bieden voor herhaalde domeinspecifieke training wanneer:

• Het hertrainen van gespecialiseerde modellen kostbaar of onpraktisch is.
• Snelle adaptatie aan een nieuw cohort vereist is.
• Counterfactual supervisie niet beschikbaar is (zoals bij real-world factual voorspellingstaken zoals MIMIC-III).

De auteurs benadrukken dat CAUSALLONGPFN de standaard causale aannames (consistentie, positiviteit, sequentiële exchangeability) die nodig zijn voor het interpreteren van observationele data niet wegneemt. In plaats daarvan ammortiseert het het schattingsprobleem. De resultaten suggereren dat een voldoende brede synthetische prior een herbruikbare structuur over behandelings-respons taken kan vangen, waardoor het model een sterke general-purpose in-context voorspeller is. De paper merkt echter bescheiden op dat domeinspecifieke training waardevol blijft wanneer voldoende target-domein data en validatiesignalen beschikbaar zijn, met name voor lang-horizon counterfactual voorspellingen in specifieke domeinen.

Het werk positioneert zich als een onderzoeksinstrument voor causale sequentiemodellering en hypothesegeneratie in plaats van een zelfstandig klinisch beslissingssysteem, en waarschuwt tegen overmatig vertrouwen in voorspellingen wanneer causale aannames of de ondersteuning van de prior ontoereikend zijn.

Code & Model Beschikbaarheid

De implementatie is beschikbaar op GitHub op https://github.com/Amirhossein-Zare/causal-long-pfn en de voorgetrainde modelgewichten zijn beschikbaar op Hugging Face op https://huggingface.co/Amirhossein-Zare/causal-long-pfn .