Résumé technique : Compression de bout en bout pour les modèles de fondation tabulaires

Énoncé du problème

Les modèles de fondation tabulaires (TFM), particulièrement ceux utilisant l'apprentissage en contexte (ICL) avec des architectures de transformeurs de type décodeur uniquement, ont récemment défié la dominance des arbres de décision boostés par gradient. Cependant, ces modèles font face à un goulot d'étranglement fondamental : leur complexité d'inférence évolue de manière quadratique ($O(N^2)$) ou linéaire ($O(N)$ avec le cache KV) par rapport au nombre d'instances d'entraînement $N$. Cette dépendance quadratique de la taille du contexte conditionnel crée des surcoûts de mémoire et de calcul prohibitifs, limitant les TFM aux jeux de données de petite ou moyenne taille. Par conséquent, les performances se dégradent souvent sur les jeux de données volumineux et riches en informations, et les modèles peinent à gérer les scénarios du monde réel nécessitant des fenêtres de contexte massives.

Méthodologie : TACO

Pour remédier à ce goulot d'étranglement, les auteurs proposent TACO (Tabular Compression), une architecture modulaire de bout en bout conçue pour compresser le jeu de données d'entraînement dans un espace latent avant la prédiction. Le cadre se compose de deux composants entraînés conjointement :

1. Compresseur ($g$) : Un module basé sur un transformeur qui projette le jeu de données d'entraînement original $D_{train}$ (avec $N$ lignes) vers une représentation compacte de rang inférieur $D_{mini-train}$ contenant $K$ lignes, où $K \ll N$. Le compresseur est entraîné pour apprendre une incrustation condensée qui capture les motifs prototypiques des données tabulaires. Il utilise une table « fictive » initialisée à partir d'un sous-ensemble de données, qui est affinée de manière itérative pour représenter l'ensemble du jeu de données.
2. Prédicteur ($f$) : Un transformeur tabulaire standard (basé sur l'architecture TabPFNv2) qui effectue la prédiction en contexte. Au lieu d'être attentif à l'ensemble du jeu de données de $N$ lignes, le prédicteur se conditionne sur le contexte compressé $D_{mini-train}$ et les caractéristiques du lot de test.

Stratégie d'entraînement :
Le compresseur et le prédicteur sont méta-appris conjointement de manière de bout en compte. L'objectif est de minimiser la perte de prédiction sur un ensemble de jeux de données, incitant le compresseur à produire une représentation qui maximise la performance du prédicteur. Le pipeline d'entraînement comprend :

• Pré-entraînement synthétique : Utilisation de modèles causaux structurels (SCM) pour générer des jeux de données synthétiques diversifiés.
• Apprentissage curriculaire : Augmentation progressive des longueurs de séquence pendant l'entraînement (de 1k à 60k lignes).
• Adaptation aux données réelles : Ajustement fin (fine-tuning) sur des jeux de données tabulaires réels via des tâches auto-supervisées (bootstrapping de lignes/caractéristiques).
• Entraînement multi-taux : Le modèle est entraîné pour gérer simultanément diverses taux de compression ($r \in \{1, 2, 4, 8, 16\}\%$), permettant à un modèle unique d'opérer sous différentes contraintes de ressources.

Mécanisme de scalabilité :
Pour les jeux de données dépassant la capacité de mémoire d'un seul passage direct (ex: >1 million de lignes), TACO emploie une stratégie de découpage et d'assemblage (chunking-and-stitching). Le jeu de données est partitionné en blocs disjoints, chacun étant compressé indépendamment, et les résumés résultants sont concaténés pour former un contexte global pour le prédicteur.

Contributions clés

1. Architecture novatrice : Introduction de TACO, la première méthode de bout en bout pour compresser le contexte d'entraînement des modèles de fondation tabulaires, découplant la qualité prédictive de la taille brute du contexte conditionnel.
2. Gains d'efficacité : Démonstration empirique que TACO atteint jusqu'à 94× de gain de vitesse en temps d'inférence et 98 % d'économies de mémoire par rapport aux transformeurs de prédicteurs seuls (POT) de pointe, tout en maintenant la performance à un taux de compression de 1 %.
3. Scalabilité : Une stratégie validée pour ingérer des jeux de données de plus d'un million de lignes (ex: le jeu de données MetroPT-3 avec ~1,5M de lignes) en les compressant en quelques milliers de lignes latentes, une tâche irréalisable pour les TFM standards en raison des contraintes de mémoire.
4. Optimisation conjointe : Preuve que l'entraînement conjoint du compresseur et du prédicteur est supérieur au gel du prédicteur, car il permet au prédicteur d'adapter son espace de représentation aux entrées compressées.

Résultats expérimentaux

Les auteurs ont évalué TACO sur le benchmark TabArena et sur des jeux de données supplémentaires à grande échelle (TabFSBench, TableShift, MetroPT-3) :

• Performance vs Vitesse : À un taux de compression de 1 %, TACO obtient une accélération de 94× pour les prédictions ultérieures et réduit la consommation de mémoire de 97,6 % par rapport à la ligne de base POT non compressée. La dégradation de la performance est statistiquement insignifiante (différence d'AUC-ROC proche de zéro) par rapport à la ligne de base à contexte complet.
• Efficacité de la mémoire : Dans les scénarios avec cache KV, POT rencontre fréquemment des erreurs de dépassement de mémoire (OOM) sur des jeux de données de plus de 32 000 lignes, tandis que TACO reste faisable. TACO réduit l'utilisation de la VRAM jusqu'à 92 % dans les configurations de cache KV.
• Comparaisons avec les bases de référence : TACO surpasse significativement les simples méthodes de compression comme l'échantillonnage aléatoire et le sous-échantillonnage k-plus proches voisins (kNN) à tous les taux de compression.
• Gestion de grande échelle : Sur le jeu de données MetroPT-3 (~1,5M de lignes), TACO a compressé le jeu d'entraînement à environ 1 200 lignes et a obtenu une précision et une AUPRC supérieures aux bases de référence et à XGBoost, alors que POT et TabPFNv2 ont échoué à s'exécuter.
• Robustesse : Sur le benchmark TableShift, TACO a fonctionné avec succès sur 12 des 15 jeux de données où les TFM standards ont échoué en raison des limites de mémoire, tout en maintenant une précision compétitive sous des changements de distribution.

Signification et affirmations

L'article affirme que TACO recadre fondamentalement l'inférence en contexte tabulaire en apprenant des représentations de données qui restent prédictives sous des budgets stricts lors de l'inférence. Les auteurs soutiennent que leur méthode permet un compromis précision-latence principled, permettant aux prédicteurs tabulaires en contexte de passer à l'échelle sur des jeux de données réels et massifs qui étaient auparavant inaccessibles aux approches basées sur les transformeurs.

Ce travail souligne que la complexité quadratique des mécanismes d'attention n'est pas une barrière insurmontable si le contexte est intelligemment compressé. En réduisant la taille effective du contexte par un facteur de $N/K$, TACO linéarise l'évolution de la mémoire et des coûts de calcul, rendant les modèles de fondation de haute performance viables pour des applications industrielles impliquant des données tabulaires massives. Les auteurs concluent que, bien que TACO fasse progresser la scalabilité, les travaux futurs devront aborder les divers priors pour les valeurs manquantes, les changements de distribution et les extensions vers la régression et les tâches de séries temporelles.