Technische Samenvatting: Het Leren van Causale Ordeningen voor In-Context Voorspelling op Tabulaire Data

1. Probleemstelling

In-context learning-architecturen, zoals Tabulaire Fundamentmodellen (TFM's), hebben sterke voorspellende prestaties op tabulaire data bereikt door gebruik te maken van grootschalige pre-training. Deze modellen zijn echter voornamelijk associatief; ze vertrouwen op correlatiestructuren in plaats van expliciete causale richting. Bijgevolg kunnen hun voorspellingen onbetrouwbaar zijn onder distributieveranderingen of interventies, aangezien ze mogelijk gebruikmaken van niet-causale statistische associaties die niet generaliseren.

Hoewel er gevestigde methoden voor causale ontdekking bestaan om Gerichte Acielische Grafen (DAG's) te herstellen, zijn deze vaak specifiek voor een dataset, rekenintensief en losgekoppeld van moderne voorspellende architecturen. Daarentegen vertrouwen recente geamortiseerde methoden voor causale ontdekking op toezichttraining over synthetische grafen, wat grondware-causale labels vereist. Er bestaat een gat in het systematisch integreren van causale structuurontdekking in voorspellende architecturen binnen een ontoezicht, in-context learning-kader dat om kan gaan met ontbrekende data en interventies.

2. Methodologie: TABORDER

De auteurs stellen TABORDER voor, een transformer-gebaseerde architectuur die is ontworpen om simultaan causale variabele ordeningen af te leiden en robuuste voorspellingen te verrichten. De kerninnovatie is het gebruik van causaal orde-gedwongen attention.

Architectuurontwerp

TABORDER functioneert als een dual-transformer-systeem dat een dataset $D$ verwerkt met ontbrekende invoer aangegeven door een masker $M$:

1. Embedding: Elke cel $x_i^{(r)}$ en de bijbehorende indicator voor ontbrekende waarden worden geëmbed in een vectorrepresentatie.
2. Inferentie-tak voor Ordening: Een standaard multi-head self-attention transformer verwerkt de embeddings om een scalair score $s_i$ af te leiden voor elke variabele $X_i$. Lagere scores duiden op een eerdere positie in de causale orde. Deze scores worden gemiddeld over de rijen om een globale ordening $\hat{\pi}$ te produceren.
3. Voorspellings-tak met Orde-beperking: Een tweede transformer modelleert de conditionele verdeling van elke cel. Cruciaal is dat de kolomsgewijze attention wordt beperkt door een masker dat is afgeleid van de geleerde scores $s$. Specifiek kan variabele $X_i$ alleen attention geven aan variabelen $X_j$ waarvoor $s_j \leq s_i$. Dit dwingt een geordende factorisatie af van de gezamenlijke verdeling:
   $$p_\theta(x^{(r)} | D, \pi) = \prod_{k=1}^d p_\theta(x_{\pi(k)}^{(r)} | x_{\pi(1:k-1)}^{(r)}, D)$$
   Dit voorkomt dat het model "anti-causale" kenmerken (toekomstige oorzaken) gebruikt om een doel te voorspellen, waardoor de gevoeligheid voor schijnbare correlaties wordt verminderd.

Omgaan met Ontbrekende Data en Onzekerheid

Het model houdt expliciet rekening met ontbrekende waarden in causale ouders. Onder de aanname van een Additief Ruismodel (ANM) ($X_i = f(X_{pa(i)}) + N_i$), decodeert het model een conditioneel gemiddelde $\hat{\mu}$ en variantie $\hat{\sigma}^2$.

• Variantie-increment: Als een causale ouder ontbreekt, verhoogt het model de voorspellende variantie voor de doelvariabele met een leerbare hoeveelheid $\delta_{i,j}$. Dit vangt de toegenomen onzekerheid op wanneer relevante invoer niet wordt waargenomen.
• Trainingsdoel: Het model wordt end-to-end getraind door de log-likelihood van gemaskeerde invoer te maximaliseren. Het doel koppelt het leren van ordeningen en voorspelling: ordeningen die de gemaskeerde data beter verklaren (door de residuvariantie te minimaliseren) ontvangen een hogere likelihood.

Theoretische Rechtvaardiging

De auteurs rechtvaardigen de likelihood-gebaseerde aanpak binnen het kader van het Additief Ruismodel (ANM). Zij tonen aan dat voor causale additieve modellen de likelihood-optimale orde consistent is met de topologische orde van de onderliggende ware DAG. Bovendien bewijzen zij een stelling met betrekking tot ontbrekingsbewuste likelihood-asymmetrie: rekening houden met ontbrekende waarden via variantie-incrementen versterkt het likelihood-gat tussen de juiste causale richting en onjuiste richtingen, wat de identificatie van causale ordeningen faciliteert, zelfs in aanwezigheid van ontbrekende data.

3. Belangrijkste Bijdragen

1. Integratie van Causale Ontdekking en Voorspelling: TABORDER overbrugt het gat tussen het leren van causale structuren en in-context voorspelling op tabulaire data door causale ordeningen ontoezicht af te leiden binnen een voorspellend transformer-kader.
2. Orde-gedwongen Attention: De introductie van een differentieerbaar, geleerd attention-maas dat de informatiestroom beperkt tot variabelen die voorafgaan in een geleerde causale orde, waardoor wordt gegarandeerd dat voorspellingen alleen steunen op potentiële oorzaken.
3. Ontbrekingsbewuste Causale Leer: Het ontwerpen van een mechanisme waarbij ontbrekende waarden in causale ouders de voorspellende variantie expliciet verhogen, wat theoretisch en empirisch helpt bij het onderscheiden van causale richtingen.
4. Ontoezicht Leren: In tegenstelling tot geamortiseerde methoden voor causale ontdekking vereist TABORDER geen grondware-causale grafen tijdens de training; het leert optimale ordeningen uitsluitend via likelihood-maximalisatie op gemaskeerde data.

4. Experimentele Resultaten

De auteurs evalueren TABORDER op synthetische data met bekende causale structuren en real-world benchmarks.

• Inferentie van Causale Ordening: Op synthetische data gegenereerd uit niet-lineaire Gaussian process ANM's bereikt TABORDER een topologische divergentie ($d_{TOP}$) die vergelijkbaar is met gespecialiseerde algoritmen voor causale ontdekking (zoals CAM, SCORE, NOTEARS), maar leidt ordeningen af in bijna onmiddellijke tijd via een enkele forward pass. Op de real-world Sachs-proteïne-signaleringsdataset ontdekt TABORDER causale ordeningen met een lage divergentie ($d_{TOP} \approx 0.21$) in vergelijking met het consensusgraf, en presteert beter dan andere baselines.
• Voorspellende Prestaties & Imputatie: Bij taken voor het imputeren van ontbrekende waarden (OpenML CC18 en CTR23 benchmarks) bereikt TABORDER concurrerende prestaties, met name bij hoge percentages ontbrekende waarden ($\geq 40\%$), en presteert beter dan standaardmethoden zoals KNN, MICE en GAIN.
• Robuustheid tegen Interventies: In een keten van drie variabelen ($X \to Y \to Z$) met een interventie op het $Y \to Z$-mechanisme, behoudt TABORDER een accurate voorspelling van de mediator $Y$ gegeven $X$ en $Z$. Daarentegen verslechteren standaard TFM's (zoals TABPFN) en XGBoost aanzienlijk, aangezien ze vertrouwen op de verbroken niet-causale associatie tussen $Y$ en $Z$.
• Biologische Inzichten: Op de Sachs-dataset onder verschillende interventievoorwaarden leidt TABORDER ordeningen af die biologisch plausibele herorganisaties weerspiegelen. Bijvoorbeeld, upstream-interventies (zoals PKC-remming) veroorzaken globale veranderingen in de afgeleide orde, terwijl downstream-interventies lokale veranderingen veroorzaken, wat overeenkomt met bekende dynamiek van signaaltransductiepaden.

5. Betekenis en Beweringen

Het artikel beweert dat TABORDER een principiële aanpak biedt om in-context learning voor tabulaire data robuuster te maken tegen distributieveranderingen en interventies door expliciet causale richting te modelleren.

• Robuustheid: Door voorspellingen te dwingen om uitsluitend te steunen op voorgangers in een geleerde causale orde, vermijdt het model het exploiteren van schijnbare correlaties die falen onder interventie.
• Efficiëntie: In tegenstelling tot klassieke causale ontdekking, die optimalisatieproblemen per dataset oplost, leidt TABORDER structuren efficiënt af binnen een enkele forward pass, waardoor het schaalbaar is voor downstream-taken.
• Theoretische Grondslag: Het werk stelt vast dat likelihood-maximalisatie in aanwezigheid van ontbrekende data kan dienen als een proxy voor causale ontdekking onder standaard ANM-aannames, en biedt zo een theoretische basis voor ontoezicht ordeningsleren.

De auteurs erkennen beperkingen, waarbij zij opmerken dat de huidige implementatie zich richt op topologische ordeningen in plaats van volledige DAG's en uitgaat van reële kenmerken. Zij suggereren dat hoewel de methode robuust is, de afgeleide structuren hypotheses zijn die domeinvvalidatie vereisen, en dat toekomstig werk rijkere trainingspriors en data van gemengde typen zou kunnen verkennen.