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

Sintesi Tecnica: CRUMB – Inferenza di Reti con Fit del Prior Efficiente tramite Batching di Contesto Distribuito Corrispondente

Problematica
Le reti con fit del prior (PFN), come TabPFN e TabICL, rappresentano una classe di modelli fondazionali tabulari capaci di apprendimento in-context. Esse accettano un insieme di addestramento completamente etichettato come contesto e generano previsioni per query di test in un singolo passaggio forward, spesso superando i metodi tradizionali di gradient boosting su dataset di piccole o medie dimensioni. Tuttavia, un collo di bottiglia critico ne limita la scalabilità: il meccanismo di self-attention nelle reti transformer-based delle PFN scala quadraticamente ($O(N^2)$) rispetto al numero di campioni di addestramento ($N$). Ciò rende l'inferenza computazionalmente proibitiva per dataset di grandi dimensioni ($N > 50.000$).

Le soluzioni esistenti affrontano un compromesso tra accuratezza ed efficienza. Il sottocampionamento uniforme migliora la velocità scartando casualmente i dati di addestramento, ma sacrifica l'informazione. I metodi dipendenti dalla query, come i $k$-vicini più prossimi ($k$NN), migliorano l'accuratezza adattando il contesto a ogni punto di test, ma impediscono il batching, richiedendo $T$ passaggi forward separati per $T$ punti di test. Altri approcci, come il Mixture of In-Context Prompters (MICP), tentano di raggruppare le query tramite clustering dei dati di addestramento, ma possono fallire nell'allineare il contesto di addestramento alla specifica distribuzione dei dati di test, in particolare in presenza di drift covariata.

Metodologia
Gli autori propongono CRUMB (Clustered Retrieval Using Minimised-MMD Batching), un wrapper di inferenza a tre stadi, agnostico rispetto all'architettura, che non richiede il riaddestramento del modello. Il metodo risolve la tensione tra qualità del contesto ed efficienza del batching attraverso i seguenti stadi:

1. Clustering lato Test: Le query di test ($D_{test}$) sono partizionate in $K$ cluster ($C_1, \dots, C_K$) utilizzando il $k$-means sulle feature di input standardizzate. Ciò consente al sistema di raggruppare punti di test simili che possono condividere un comune contesto di addestramento.
2. Selezione del Contesto con Corrispondenza Distribuionale: Per ogni cluster di test $C_k$, viene selezionato un piccolo sottoinsieme di addestramento $S_k$ (dimensione $n \ll N$) dal pool completo di addestramento ($D_{train}$). A differenza dei metodi di prossimità geometrica (es. centroid-nearest-neighbor), CRUMB seleziona $S_k$ minimizzando greedymente la Massima Discrepanza Media (MMD) tra la distribuzione empirica del cluster di test e il sottoinsieme di addestramento candidato.
   • Ciò viene ottenuto tramite il kernel herding, un algoritmo greedy che seleziona iterativamente punti di addestramento per minimizzare l'obiettivo MMD. La funzione obiettivo bilancia due termini: premiare la prossimità al cluster di test (rilevanza) e penalizzare la ridondanza tra i punti di addestramento selezionati (diversità).
   • Per accelerare questo processo, gli autori impiegano le Random Fourier Features (RFF) per approssimare il kernel RBF Gaussiano e utilizzano una selezione greedy in batch.
3. Inferenza in Batch: Il modello PFN esegue $K$ passaggi forward indipendenti. In ogni passaggio, il modello riceve il sottoinsieme di addestramento selezionato $S_k$ come contesto e l'intero cluster $C_k$ come query. Ciò riduce il numero totale di passaggi forward da $T$ (nel caso di $k$NN per query) a $K$, consentendo un utilizzo efficiente della GPU pur mantenendo l'allineamento distributivo.

Contributi Chiave

• Recupero del Contesto basato su MMD: Il documento introduce la minimizzazione greedy dell'MMD (kernel herding) come meccanismo per selezionare i sottoinsiemi di addestramento. Gli autori dimostrano che questo approccio di matching distributivo supera le strategie di recupero centroid-nearest-neighbor e Voronoi-uniform all'interno dello stesso framework di clustering.
• Batching del Contesto tramite Clustering lato Test: Raggruppando le query di test e assegnando un singolo contesto di addestramento, distribuito in modo corrispondente a ciascun cluster, CRUMB consente un efficiente batching senza sacrificare la rilevanza della query.
• Robustezza al Drift della Covariata: Il metodo gestisce naturalmente il drift della covariata. Poiché il contesto di addestramento è selezionato per corrispondere alla distribuzione dell'attuale batch di test (anziché a una distribuzione di addestramento fissa), CRUMB si adatta ai cambiamenti nella distribuzione dei dati di test. Al contrario, metodi come MICP, che si affidano a cluster di addestramento fissi, degradano significativamente quando la distribuzione di test subisce uno shift rispetto alla distribuzione di addenza.

Risultati Sperimentali
Gli autori hanno valutato CRUMB sul benchmark TabArena (51 dataset tabulari diversificati) su tre architetture PFN: TabPFNv2, TabICLv1 e TabICLv2.

• Performance rispetto ai Baseline: Sul benchmark TabArena, CRUMB supera significativamente il sottocampionamento uniforme e la tecnica MICP a parità di budget di contesto. Sebbene CRUMB sia statisticamente vicino alle performance del $k$NN per query, lo raggiunge con un numero fisso di $K$ passaggi forward (es. $K=20$) anziché $T$ passaggi, risultando in una riduzione sostanziale del tempo di inferenza (es. 35s rispetto a 361s negli esperimenti riportati).
• Dataset Grandi: Gli esperimenti su dataset con fino a 120.000 campioni (es. GiveMeSomeCredit, Higgs) mostrano che CRUMB con early stopping (terminando la selezione del contesto quando il miglioramento dell'MMD si stabilizza) vince più frequentemente contro MICP, particolarmente per TabPFNv2 e TabICLv1.
• Drift della Covariata: In esperimenti controllati di drift dove la distribuzione di test si sposta lungo la prima componente principale, il vantaggio di CRUMB rispetto a MICP cresce monotonicamente con l'intensità del drift. Con drift zero, CRUMB conduce di +4,9%; con drift massimo (completamente fuori campione), il divario si amplia a +17,1%. Ciò è attribuito alla capacità di CRUMB di ri-allineare il contesto di addestramento alla distribuzione di test spostata, mentre i cluster di addestramento fissi di MICP diventano disallineati.

Significato e Rivendicazioni
Il documento sostiene che CRUMB renda pratica l'inferenza delle reti con fit del prior per dataset tabulari su larga scala senza richiedere modifiche architettoniche o il fine-tuning del modello. Combinando il clustering lato test con il recupero del contesto con corrispondenza distributiva (MMD), CRUMB raggiunge un "punto ottimale" tra l'accuratezza del recupero per query e l'efficienza del batching uniforme.

Gli autori sottolineano che l'innovazione centrale risiede nel design adattivo al test: raggruppare i punti di test e selezionare i contesti basandosi sull'allineamento distributivo (MMD) piuttosto che sulla sola prossimità geometrica. Questo design fornisce un meccanismo naturale di robustezza contro il drift della covariata, un regime in cui i metodi di selezione del contesto esistenti spesso falliscono. Il documento conclude che, sebbene CRUMB erediti il limite superiore delle prestazioni del pretraining della PFN sottostante, esso rimuove efficacementamente la barriera computazionale che impedisce l'applicazione delle PFN a dataset di grandi dimensioni.

Limitazioni
Gli autori riconoscono che CRUMB richiede l'intero set di query di test in anticipo per eseguire il clustering, rendendolo inadatto a scenari di streaming online puro senza modifiche (es. il caching dei centroidi, che equivale a un routing in stile MICP). Inoltre, la dipendenza dal $k$-means euclideo e dai kernel gaussiani isotropici potrebbe limitarne le prestazioni in spazi di feature ad alta dimensionalità o di tipo misto. Infine, come puro wrapper di inferenza, CRUMB non può compensare un cattivo match tra i prior di pretraining della PFN e il dataset specifico.