CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

Technische Samenvatting: CRUMB – Efficiënte Inferentie van Prior Fitted Networks via Distributionally Matched Context Batching

Probleemstelling
Prior-fitted networks (PFNs), zoals TabPFN en TabICL, vormen een klasse van tabulaire fundamentmodellen die in staat zijn tot in-context leren. Ze accepteren een volledig gelabelde trainingsset als context en genereren voorspellingen voor test-queries in een enkele forward pass, waarbij ze vaak de prestaties van traditionele gradient-boosting methoden op kleine tot middelgrote datasets overtreffen. Echter, een kritieke bottleneck beperkt hun schaalbaarheid: het self-attention mechanisme in transformer-gebaseerde PFN's schaalt kwadratisch ($O(N^2)$) met het aantal trainingssamples ($N$). Dit maakt inferentie computationeel onhaalbaar voor grote datasets ($N > 50.000$).

Bestaande oplossingen staan voor een afweging tussen nauwkeurigheid en efficiëntie. Uniforme subsampling verbetert de snelheid door willekeurig trainingsdata weg te gooien, maar offert informatie op. Query-afhankelijke methoden zoals $k$-nearest neighbors ($k$NN) verbeteren de nauwkeurigheid door de context aan te passen aan elk testpunt, maar voorkomen batching, waardoor $T$ aparte forward passes nodig zijn voor $T$ testpunten. Andere benaderingen, zoals de Mixture of In-Context Prompters (MICP), proberen queries te batchen door trainingsdata te clusteren, maar kunnen er niet in slagen om de trainingscontext af te stemmen op de specifieke distributie van de testdata, met name onder covariate drift.

Methodologie
De auteurs stellen CRUMB (Clustered Retrieval Using Minimised-MMD Batching) voor, een drie-fasen, architectuur-agnostische inferentie-wrapper die geen hertraining van het model vereist. De methode lost de spanning tussen contextkwaliteit en batching-efficiëntie op via de volgende fasen:

1. Test-zijde Clustering: De test-queries ($D_{test}$) worden gepartitioneerd in $K$ clusters ($C_1, \dots, C_K$) met behulp van $k$-means op gestandaardiseerde inputkenmerken. Dit stelt het systeem in staat om vergelijkbare testpunten te groeperen die een gemeenschappelijke trainingscontext kunnen delen.
2. Distributionally Matched Context Selectie: Voor elk testcluster $C_k$ wordt een kleine trainingssubset $S_k$ (grootte $n \ll N$) geselecteerd uit de volledige trainingspool ($D_{train}$). In tegenstelling tot geometrische nabijheidsmethoden (bijv. centroid-nearest-neighbor), selecteert CRUMB $S_k$ door de Maximum Mean Discrepancy (MMD) tussen de empirische distributie van het testcluster en de kandidaat-trainingssubset greedily te minimaliseren.
   • Dit wordt bereikt via kernel herding, een greedy algoritme dat iteratief trainingspunten selecteert om de MMD-doelfunctie te minimaliseren. De doelfunctie balanceert twee termen: het belonen van nabijheid tot het testcluster (relevantie) en het bestraffen van redundantie tussen geselecteerde trainingspunten (diversiteit).
   • Om dit proces te versnellen, gebruiken de auteurs Random Fourier Features (RFF) om de Gaussian RBF-kernel te benaderen en gebruiken zij gebatched greedy selectie.
3. Gebatched Inferentie: Het PFN-model voert $K$ onafhankelijke forward passes uit. In elke pass ontvangt het model de geselecteerde trainingssubset $S_k$ als context en de gehele cluster $C_k$ als queries. Dit vermindert het totale aantal forward passes van $T$ (in per-query $k$NN) naar $K$, wat efficiënt GPU-gebruik mogelijk maakt terwijl de distributionele uitlijning behouden blijft.

Belangrijkste Bijdragen

• MMD-gebaseerde Context Retrieval: Het artikel introduceert greedy MMD-minimalisatie (kernel herding) als een mechanisme voor het selecteren van trainingssubsets. De auteurs demonstreren dat deze distributionele afstemmingsbenadering de prestaties van centroid-nearest-neighbor en Voronoi-uniforme retrieval-strategieën binnen hetzelfde clusteringframework overtreft.
• Gebatched Context Selectie via Test-zijde Clustering: Door test-queries te clusteren en een enkele, distributioneel passende trainingscontext toe te wijzen aan elk cluster, maakt CRUMB efficiënte batching mogelijk zonder verlies van query-relevantie.
• Robuustheid tegen Covariate Drift: De methode gaat op natuurlijke wijze om met covariate drift. Omdat de trainingscontext wordt geselecteerd om overeen te komen met de huidige test-batch distributie (in plaats van een vaste trainingsdistributie), past CRUMB zich aan bij verschuivingen in de testdata. In contrast hiermee presteren methoden zoals MICP, die vertrouwen op vaste trainingsclusters, significant slechter wanneer de testdistributie verschuift van de trainingsdistributie.

Experimentele Resultaten
De auteurs evalueerden CRUMB op de TabArena benchmark (51 diverse tabulaire datasets) over drie PFN-architecturen: TabPFNv2, TabICLv1 en TabICLv2.

• Prestaties vs. Baselines: Op de TabArena benchmark presteert CRUMB significant beter dan uniforme subsampling en de MICP-techniek bij identieke context-budgetten. Hoewel CRUMB statistisch gezien dicht bij de prestaties van per-query $k$NN ligt, bereikt het dit met een vast aantal $K$ forward passes (bijv. $K=20$) in plaats van $T$ passes, wat resulteert in een substantiële reductie in inferentietijd (bijv. 35s vs. 361s in gerapporteerde experimenten).
• Grote Datasets: Experimenten op datasets met tot 120.000 samples (bijv. GiveMeSomeCredit, Higgs) tonen aan dat CRUMB met early stopping (het beëindigen van contextselectie wanneer de MMD-verbetering een plateau bereikt) vaker wint van MICP, met name voor TabPFNv2 en TabICLv1.
• Covariate Drift: In gecontroleerde drift-experimenten waar de testdistributie verschuift langs de eerste hoofdcomponent, groeit het voordeel van CRUMB ten opzichte van MICP monotoon met de intensiteit van de drift. Bij nul drift leidt CRUMB met +4,9%; bij maximale drift (volledig out-of-sample) wordt het gat vergroot tot +17,1%. Dit wordt toegeschreven aan het vermogen van CRUMB om de trainingscontext opnieuw af te stemmen op de verschoven testdistributie, terwijl de vaste trainingsclusters van MICP uit balans raken.

Betekenis en Claims
Het artikel claimt dat CRUMB inferentie van prior-fitted networks praktisch maakt voor grootschalige tabulaire datasets zonder dat architecturale wijzigingen of model fine-tuning vereist zijn. Door test-zijde clustering te combineren met distributioneel passende context retrieval (MMD), bereikt CRUMB een "sweet spot" tussen de nauwkeurigheid van per-query retrieval en de efficiëntie van uniforme batching.

De auteurs benadrukken dat de kerninnovatie ligt in het test-adaptieve ontwerp: het clusteren van testpunten en het selecteren van contexten op basis van distributionele afstemming (MMD) in plaats van louter geometrische nabijheid. Dit ontwerp biedt een natuurlijk robuustheidsmechanisme tegen covariate drift, een regime waarin bestaande context-selectiemethoden vaak falen. Het artikel concludeert dat hoewel CRUMB de prestatiegrens van de onderliggende PFN-pretraining erft, het effectief de computationele barrière wegneemt die de toepassing van PFN's op grote datasets verhindert.

Beperkingen
De auteurs erkennen dat CRUMB de volledige set test-queries vooraf nodig heeft om clustering uit te voeren, wat het minder geschikt maakt voor echte online streaming-settings zonder modificatie (bijv. het cachen van centroids, wat effectief terugvalt op een MICP-stijl routing). Daarnaast kan de afhankelijkheid van Euclidische $k$-means en isotrope Gaussische kernels de prestaties in hoog-dimensionale of gemengde feature-ruimtes beperken. Tot slot kan CRumb als een pure inferentie-wrapper niet compenseren voor een slechte match tussen de PFN-pretraining priors en de specifieke dataset.