GEAR: Generative Expansion and Real Anchoring for Two-Stage Distillation of Tabular Foundation Models
Authors: Qi Qin, Jiajie Zhu, Dali Chen, Yuzhao Zhang, Jia-Xing Han, Peng Zhang, Ying Yan, Yifan Sun, Yu Su
Organizations: Renmin University of China, Center for Applied Statistics and School of Statistics, Beijing, China · Ant Digital Technologies, Ant Group, Hangzhou, China · Nanjing University, Nanjing, China
Tabular foundation models (TFMs) achieve strong performance through in-context learning, but context-dependent inference imposes substantial latency and memory costs, hindering large-scale deployment. We propose GEAR (Generative Expansion and Real Anchoring), a modular two-stage framework that distills TFMs into lightweight MLP or tree-based predictors that can be deployed on commodity CPUs. Stage 1 uses synthetic covariates solely as teacher-query locations and trains the student on soft TFM targets, expanding coverage beyond observed rows. Stage 2 re-anchors the student to the target distribution using real labels and out-of-fold teacher predictions, whitch avoids self-labeling leakage. We further derive a risk certificate characterizing the trade-off between generated-query volume and generator fidelity. Experiments on TALENT and TabArena demonstrate the broad applicability of GEAR. Two-stage MLPs outperform supervised MLPs by 1.81--2.00 AUC points on binary tasks and 1.19--1.35 points on multiclass tasks, with additional gains over real-data-only distillation of 1.76--2.19 and 2.09--2.40 points, respectively. On binary tasks, the gains also transfer to LightGBM and XGBoost, and all three student families outperform CatBoost, the strongest non-TFM baseline, in mean AUC. Ablations show gains beyond longer training or alternative warm starts, greater stability from staged than mixed optimization, and generator-dependent diminishing returns as query volume increases. Finally, GEAR reduces median inference time by 57--2866 times and peak prediction memory by 1.9--3.3 times, while retaining higher AUC than matched supervised baselines.
Figures & tables
Figure 1: Accuracy–efficiency trade-off on TALENT. Mean AUC is plotted against median prediction time; bubble size denotes peak RSS. Shapes denote TFM families, colors denote model types, crosses mark supervised baselines, light fills indicate real-only distillation, and dark fills indicate GEAR. TabM real-only distillation reaches only ≈0.76 AUC and is omitted; GEAR students retain teacher accuracy at much lower cost.
Figure 2: Overview of the proposed two-stage generative distillation protocol.
Figure 3: Generator trajectories on TALENT. Columns correspond to TabICL, TabPFN, and TabDPT, and rows to binary and multiclass tasks. Curves show mean gains in AUC points over the supervised MLP across generated-to-real ratios K . Colors denote generators; dashed circles and solid squares indicate Stage 1 and Stage 2, respectively. Horizontal lines mark the teacher, supervised-MLP, and real-only-distillation references.
Task
TabPFGen
Copula
CTGAN
TabDiff
Bin.
+0.87/0.04
+0.81/0.19
+0.71/0.10
+1.18/<0.01
Multi.
+0.26/0.05
−0.05/0.08
−0.17/0.28
+0.18/0.02
Table 1: GEAR MLP versus CatBoost at K=40 . Cells report mean ΔGEAR−CB / two-sided Wilcoxon p -value.
Figure 4: Stage 1 replacement ablations using TabPFN as the teacher. Boxplots report paired gains in AUC points of GEAR over five alternatives, ordered from real-only mixed distillation to real-only soft initialization.
Teacher
Task
Student
S1 Δ
Real Δ
S2 Δ
TabICL
Bin.
LGBM
+0.79(0.25)
+0.92(0.24)
+1.53(0.20)
XGB
+0.96(0.21)
+1.11(0.15)
+1.37(0.18)
Multi.
LGBM
−6.34(1.14)
−8.69(1.04)
−5.93(1.14)
XGB
+0.22(0.22)
+0.40(0.13)
+0.51(0.12)
TabPFN
Bin.
LGBM
−2.57(0.76)
+0.91(0.24)
+1.32(0.20)
XGB
−2.71(0.90)
+1.14(0.16)
+1.38(0.18)
Table 2: Tree-student distillation with TabPFGen. Values are paired gains in AUC points over supervised training; subscripts denote standard errors.
Figure 5: Two-stage versus mixed real–generated distillation on TALENT binary tasks. Rows denote TFM teachers, and columns denote TabPFGen and TabDiff. Solid curves show paired median gains in AUC points over Stage 1 with interquartile bands; dashed hollow-marker curves show non-degradation rates relative to Stage 1 (right axis).
Tabular foundation models (TFMs) achieve strong predictive performance through in-context learning, yet repeatedly conditioning on labeled data makes inference expensive. Knowledge distillation can reduce this cost by transferring their predictive ability to lightweight, dataset-specific students. However, the dependence of TFM predictions on both a labeled context and a query introduces two design questions: how to construct teacher supervision and whether expanding query coverage improves distillation. We examine these questions across two TFMs and both neural and tree-based students, and derive an effective distillation recipe. The recipe uses the full labeled training set as teacher context and trains students solely on teacher predictions for observed and synthetic queries. On TabArena, the resulting students outperform their supervised trained tuned-and-ensembled counterparts by 57-98 Elo points. Applied unchanged to TALENT, the same recipe improves matched default students on 236-258 of 300 datasets and reduces median primary error by 4.0-6.4%. The distilled students also achieve median inference speedups of 3.0-21.6 times over their teachers, offering a practical trade-off between predictive performance and repeated inference cost. Code is available at https://github.com/nums-ai/TFM_Distillation .
A fraud scorer needs to answer in under 2 ms. The best tabular foundation models (TFMs) take 151-1,275 ms on GPU. We close this gap by distilling the TFM offline into an XGBoost or CatBoost student that runs natively on CPU. The central obstacle is specific to in-context learning (ICL) teachers: they leak labels when scoring their own training set, so the soft targets collapse to near-one-hot vectors with no inter-class structure left to distill. Stratified out-of-fold (OOF) teacher labeling prevents this. Across 153 classification datasets drawn from TALENT, OpenML-CC18, TabZilla, and TabArena, distilling TabICLv2 into XGBoost gives 0.882 macro-mean AUC (96.5% of teacher AUC) at 1.9 ms on CPU, a 38x to 860x speedup across teacher-student pairs with a statistically significant edge over a tuned CatBoost baseline (Wilcoxon p = 0.0008; 51% win rate). Four further findings: teacher rank transfers exactly to student rank; gains concentrate on low-dimensional data (< 21 features: +0.011 over CatBoost vs. >21 features: +0.001); multi-teacher averaging helps MLP students (+0.006, p = 0.003) but adds less than 0.001 for tree students; and on high-dimensional tasks where the teacher itself trails CatBoost, distillation makes things worse rather than better. The full pipeline is open-sourced as part of the TabTune library.
Tabular foundation models (TFMs) based on prior-data fitted networks typically use sampled hard query labels, leaving avoidable sampling variability in pretraining. A key observation is that intermediate generator states can reveal label uncertainty beyond a sampled outcome. We propose Generator-state-Informed Supervision for TFMs (GIST), which uses these states to average query losses over conditional label distributions. GIST jointly designs task generation and query supervision, using entropy-calibrated encodings to obtain tractable, input-consistent targets on supported label-dependent branches. We show that this averaging preserves the population objective and yields no larger gradient covariance at fixed parameters within each generating distribution. Under matched pretraining schedules, GIST improves mean rank across three architectures on both TabArena and TALENT without changing architecture or inference. Across these benchmarks and architectures, GIST wins 62.1% of paired dataset-level comparisons against the corresponding hard-label controls.
Zhangyang Peng, Jingze Shi, Tongyi Zhang +1
The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China.