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

Résumé Technique : CRUMB – Inférence de Réseaux Pré-ajustés Efficace via le Batching de Contexte par Appariement de Distribution

Énoncé du Problème
Les réseaux pré-ajustés (PFN - Prior-Fitted Networks), tels que TabPFN et TabICL, représentent une classe de modèles de fondation tabulaires capables d'apprentissage en contexte (in-context learning). Ils acceptent un ensemble d'entraînement entièrement étiqueté comme contexte et génèrent des prédictions pour des requêtes de test en une seule passe directe (forward pass), surpassant souvent les méthodes traditionnelles de boosting de gradient sur de petits à moyens ensembles de données. Cependant, un goulot d'étranglement critique limite leur extensibilité : le mécanisme d'auto-attention des transformeurs utilisés dans les PFN scale de manière quadratique ($O(N^2)$) avec le nombre d'échantillons d'entraînement ($N$). Cela rend l'inférence computationnellement prohibitive pour les grands ensembles de données ($N > 50\,000$).

Les solutions existantes font face à un compromis entre précision et efficacité. Le sous-échantillonnage uniforme améliore la vitesse en écartant aléatoirement des données d'entraînement, mais sacrifie l'information. Les méthodes dépendantes des requêtes, comme les $k$-plus proches voisins ($k$NN), améliorent la précision en adaptant le contexte à chaque point de test, mais empêchent le traitement par lots (batching), nécessitant $T$ passes directes séparées pour $T$ points de test. D'autres approches, telles que le Mixture of In-Context Prompters (MICP), tentent de traiter les requêtes par lots en regroupant les données d'entraînement, mais elles peuvent échouer à aligner le contexte d'entraînement avec la distribution spécifique des données de test, particulièrement en cas de dérive de covariable (covariate drift).

Méthodologie
Les auteurs proposent CRUMB (Clustered Retrieval Using Minimised-MMD Batching), un wrapper d'inférence agnostique à l'architecture, en trois étapes, qui ne nécessite aucun ré-entraînement du modèle. La méthode résout la tension entre la qualité du contexte et l'efficacité du batching grâce aux étapes suivantes :

1. Clustering côté Test : Les requêtes de test ($D_{test}$) sont partitionnées en $K$ clusters ($C_1, \dots, C_K$) en utilisant le $k$-means sur les caractéristiques d'entrée standardisées. Cela permet au système de regrouper des points de test similaires qui peuvent partager un contexte d'entraînement commun.
2. Sélection de Contexte par Appariement de Distribution : Pour chaque cluster de test $C_k$, un petit sous-ensemble d'entraînement $S_k$ (taille $n \ll N$) est sélectionné dans le pool complet d'entraînement ($D_{train}$). Contrairement aux méthodes de proximité géométrique (ex: voisin le plus proche du centroïde), CRUMB sélectionne $S_k$ en minimisant de manière gloutonne la Divergence de Maximum Mean Discrepancy (MMD) entre la distribution empirique du cluster de test et le sous-ensemble d'entraînement candidat.
   • Ceci est réalisé via le kernel herding, un algorithme glouton qui sélectionne itérativement des points d'entraînement pour minimiser l'objectif MMD. La fonction objectif équilibre deux termes : récompenser la proximité avec le cluster de test (pertinence) et pénaliser la redondance parmi les points d'entraînement sélectionnés (diversité).
   • Pour accélérer ce processus, les auteurs utilisent des caractéristiques de Fourier aléatoires (RFF - Random Fourier Features) pour approximer le noyau RBF gaussien et utilisent une sélection gloutonne par lots (batched greedy selection).
3. Inférence par Lots (Batched Inference) : Le modèle PFN exécute $K$ passes directes indépendantes. Dans chaque passe, le modèle reçoit le sous-ensemble d'entraînement sélectionné $S_k$ comme contexte et l'intégralité du cluster $C_k$ comme requêtes. Cela réduit le nombre total de passes directes de $T$ (dans le cas du $k$NN par requête) à $K$, permettant une utilisation efficace du GPU tout en maintenant l'alignement de la distribution.

Contributions Clés

• Récupération de Contexte basée sur la MMD : L'article introduit la minimisation gloutienne de la MMD (kernel herding) comme mécanisme de sélection des sous-ensembles d'entraînement. Les auteurs démontrent que cette approche d'appariement de distribution surpasse les stratégies de récupération par voisin du centroïde et par Voronoi-uniforme au sein du même cadre de clustering.
• Sélection de Contexte par Batching via le Clustering Côté Test : En regroupant les requêtes de test et en assignant un contexte d'entraînement unique et aligné sur la distribution à chaque cluster, CRUMB permet un batching efficace sans sacrifier la pertinence des requêtes.
• Robustesse à la Dérive de Covariable : La méthode gère naturellement la dérive de covariable. Parce que le contexte d'entraînement est sélectionné pour correspondre à la distribution du lot de test actuel (plutement qu'à une distribution d'entraînement fixe), CRUMB s'adapte aux changements de la distribution des données de test. En revanche, des méthodes comme MICP, qui reposent sur des clusters d'entraînement fixes, se dégradent considérablement lorsque la distribution de test s'éloigne de la distribution d'entraînement.

Résultats Expérimentaux
Les auteurs ont évalué CRUMB sur le benchmark TabArena (51 ensembles de données tabulaires divers) à travers trois architectures PFN : TabPFNv2, TabICLv1 et TabICLv2.

• Performance vs Baselines : Sur le benchmark TabArena, CRUMB surpasse significativement le sous-échantillonnage uniforme et la technique MICP à budgets de contexte identiques. Bien que CRUMB soit statistiquement proche de la performance du $k$NN par requête, il y parvient avec un nombre fixe de $K$ passes directes (ex: $K=20$) plutôt que $T$ passes, ce qui permet une réduction substantielle du temps d'inférence (ex: 35s contre 361s dans les expériences rapportées).
• Grands Ensembles de Données : Les expériences sur des jeux de données allant jusqu'à 120 000 échantillons (ex: GiveMeSomeCredit, Higgs) montrent que CRUMB avec un arrêt précoce (early stopping — terminant la sélection du contexte lorsque l'amélioration de la MMD plafonne) l'emporte plus fréquemment contre MICP, particulièrement pour TabPFNv2 et TabICLv1.
• Dérive de Covariable : Dans des expériences de dérive contrôlée où la distribution de test se déplace le long de la première composante principale, l'avantage de CRUMB sur MICP croît de manière monotone avec l'intensité de la dérive. À dérive nulle, CRUMB mène de +4,9 % ; à dérive maximale (complètement hors échantillon), l'écart s'élargit à +17,1 %. Cela est attribué à la capacité de CRUMB à ré-aligner le contexte d'entraînement sur la distribution de test décalée, là où les clusters d'entraînement fixes de MICP deviennent désalignés.

Signification et Revendications
L'article affirme que CRUMB rend l'inférence des réseaux pré-ajustés pratique pour les ensembles de données tabulaires à grande échelle sans nécessiter de changements architecturaux ou de fine-tuning du modèle. En combinant le clustering côté test avec la récupération de contexte par appariement de distribution (MMD), CRUMB atteint un "point d'équilibre" entre la précision de la récupération par requête et l'efficacité du batching uniforme.

Les auteurs soulignent que l'innovation centrale réside dans la conception adaptative au test : le clustering des points de test et la sélection de contextes basés sur l'alignement de distribution (MMD) plutôt que sur la simple proximité géométrique. Cette conception fournit un mécanisme naturel de robustesse contre la dérive de covariable, un régime où les méthodes de sélection de contexte existantes échouent souvent. L'article conclut que bien que CRUMB hérite du plafond de performance du pré-entraînement du PFN sous-jacent, il élimine efficacement la barrière computationnelle empêchant lient l'application des PFN à de grands ensembles de données.

Limites
Les auteurs reconnaissent que CRUMB nécessite l'ensemble complet des requêtes de test en amont pour effectuer le clustering, ce qui le rend inadapté aux environnements de streaming pur (online streaming) sans modification (ex: mettre en cache les centroïdes, ce qui revient à un routage de type MICP). De plus, le recours au $k$-means euclidien et aux noyaux gaussiens isotropes peut limiter les performances dans des espaces de caractéristiques de haute dimension ou de type mixte. Enfin, en tant que pur wrapper d'inférence, CRUMB ne peut pas compenser un mauvais appariement entre les priors de pré-entraînement du PFN et le jeu de données spécifique.