🤖 machine learning

Pocket Foundation Models: Distilling TFMs into CPU-Ready Gradient-Boosted Trees

This paper introduces a method to distill slow, GPU-dependent tabular foundation models into fast, CPU-native gradient-boosted trees using stratified out-of-fold labeling to prevent label leakage, achieving near-teacher accuracy with up to an 860x speedup for real-time applications like fraud detection.

Original authors: Aditya Tanna, Nassim Bouarour, Mohamed Bouadi, Vinay kumar Sankarapu, Pratinav Seth

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Aditya Tanna, Nassim Bouarour, Mohamed Bouadi, Vinay kumar Sankarapu, Pratinav Seth

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The Genius Who Can't Run

Imagine you have a Genius Professor (the "Foundation Model") who is incredibly smart at solving puzzles. This professor can look at a massive library of past puzzles and instantly figure out the answer to a new one.

However, there's a catch:

They are slow: To solve a puzzle, the Professor has to read through their entire library of notes every single time. This takes a long time (about 150 milliseconds to over a second).
They are expensive: The Professor needs a super-computer (a powerful GPU) to work.
The real world is impatient: In situations like stopping credit card fraud or triaging patients, you need an answer in under 2 milliseconds. The Professor is too slow and too heavy to be useful here.

The Solution: The "Pocket" Student

The authors wanted to create a Student who is as smart as the Professor but runs on a regular laptop (CPU) and answers in a blink of an eye.

They used a technique called Knowledge Distillation. Think of this as the Professor giving the Student a cheat sheet. Instead of just telling the Student the right answer (e.g., "Yes, this is fraud"), the Professor explains how they thought about it (e.g., "I'm 90% sure it's fraud, but there's a 10% chance it's a false alarm"). This nuance helps the Student learn better than just memorizing the answers.

The Trap: The "Mirror" Problem

Here is the tricky part the paper discovered.

If you ask the Professor to grade the exact same puzzles they just used to study, they get confused. Because the Professor has the answers right in front of them in their "context," they don't actually think; they just recall. They say, "I know this one! It's definitely Fraud!" with 100% certainty.

If the Student learns from this, they get a bad cheat sheet. The Professor isn't sharing any "wisdom" or "uncertainty"; they are just repeating the answer key. The Student learns nothing new and ends up dumber than if they had just studied the answer key directly.

The Fix: The "Blind" Test
To fix this, the authors used a method called Stratified Out-of-Fold (OOF) labeling.

Imagine splitting the library of puzzles into 5 separate piles.
The Professor studies 4 piles and is then tested only on the 5th pile they haven't seen yet.
Then, they switch piles. They study the new 4 piles and are tested on the next unseen pile.
By doing this, the Professor is forced to actually think and guess, rather than just recall. This creates a high-quality, honest cheat sheet for the Student.

The Results: A Fast, Smart Student

The team tested this on 153 different datasets (ranging from small medical records to large financial data). Here is what they found:

Speed: The new Student (a model called XGBoost) runs 38 to 860 times faster than the Professor. It answers in about 1.9 milliseconds on a standard computer, meeting the strict "under 2ms" requirement.
Smarts: The Student keeps 96.5% of the Professor's intelligence. In fact, on 51% of the tests, the Student actually beat a standard, well-tuned model (CatBoost) that is usually the gold standard for this type of work.
The "Low-Dimension" Sweet Spot: The Student shines brightest on problems with fewer features (like a puzzle with fewer than 21 pieces). On these, the Student is significantly better than the competition. However, on massive, complex puzzles with thousands of features, the Student doesn't gain much advantage over standard models.
Teamwork (Multi-Teacher):
- For Tree-based Students (like XGBoost), having multiple Professors give advice didn't help much. One strong Professor was enough.
- For Neural Network Students (MLPs), having multiple Professors average their advice did help, acting like a "smoothing" effect that made the Student more accurate.

The Bottom Line

The paper proves that you don't need a super-computer to use the world's smartest AI models for real-time tasks. By using a clever "blind test" method to generate a cheat sheet, you can train a lightweight, fast "Student" model that runs on a regular CPU.

If the Professor is good: The Student becomes a fast, accurate replacement.
If the Professor is bad: The Student will also be bad (the method doesn't magically fix a weak Professor).
The Golden Rule: You must ensure the Professor is "blind" to the test data, or the Student learns nothing.

The authors have made all the tools to do this open-source, so anyone can build their own "Pocket Foundation Model."

Technical Summary: Pocket Foundation Models

Problem Statement

Tabular Foundation Models (TFMs) such as TabICLv2, TabPFNv2.6, and LimiX have demonstrated state-of-the-art accuracy on small structured datasets by utilizing in-context learning (ICL) via large transformers. However, these models face a critical deployment barrier: they require attending to the entire training set at query time, resulting in inference latencies between 151 ms and 1,275 ms on GPUs. This latency is incompatible with real-time applications like fraud detection, credit scoring, or patient triage, which often require sub-2 ms responses. Furthermore, GPU costs and CPU-only production constraints render direct TFM deployment impractical.

While knowledge distillation offers a pathway to compress these models into faster Gradient-Boosted Decision Tree (GBDT) students, a specific technical obstacle arises with ICL-based teachers. When an ICL teacher scores examples that are already present in its context (i.e., its own training set), its predictions collapse to near-one-hot vectors with near-zero entropy. This "leakage" destroys the inter-class structure (soft labels) necessary for effective distillation, causing the student to learn no more than it would from hard labels, often performing worse than standard training.

Methodology

The paper proposes a leakage-aware distillation pipeline that converts TFM teachers into CPU-ready GBDT or MLP students. The core innovation is the use of Stratified Out-of-Fold (OOF) labeling to prevent context leakage.

OOF Labeling: The training dataset $D$ is partitioned into $K=5$ stratified folds. For each fold $k$ , the teacher model is fitted on $D \setminus D_k$ and used to predict probabilities only for the held-out set $D_k$ . This ensures the teacher never sees the query examples in its context, preserving the soft-label signal (inter-class structure) required for distillation.
Multi-Teacher Averaging: In multi-teacher settings, OOF predictions from different TFM families (e.g., TabPFN, TabICL, LimiX) are averaged with equal weights before being used as targets.
Student Objective: The student (XGBoost, CatBoost, LightGBM, or MLP) is trained using a Hinton-style mixed loss function:
$\mathcal{L} = \alpha \sum_i w_i T_i^2 \text{KL}(\hat{p}_i^{T_i} \parallel q_i^{T_i}) + (1-\alpha) \sum_i w_i \ell_{\text{CE}}(y_i, q_i)$
Where $\alpha=0.7$ , $T_i$ is a per-sample temperature scaling with teacher entropy, and $w_i$ is a confidence weight that down-weights overconfident or near-random predictions. For tree-based students, the KL term simplifies to per-class MSE regression on soft-label logits.
Teacher Selection: The paper finds that the ranking of teachers by their solo AUC on a held-out sample transfers exactly to the ranking of the resulting students. Thus, selecting the best teacher requires no complex architecture search; one simply picks the teacher with the highest solo AUC.

Key Contributions

Mandatory OOF Labeling: The authors identify that OOF labeling is not optional for ICL-based TFMs. Without it, the soft-label signal collapses, and distillation yields students inferior to hard-label baselines.
Performance vs. Latency: Distilling TabICLv2 into XGBoost achieves a macro-mean AUC of 0.882 (96.5% of the teacher's AUC) with a latency of 1.9 ms on a single CPU core. This represents a 38× to 860× speedup compared to the teacher, while statistically outperforming a tuned CatBoost baseline on 51% of 153 datasets (Wilcoxon $p=0.0008$ ).
Predictable Gains: The benefits of distillation are highly predictable based on data dimensionality:
- Low-dimensional data ( $\le$ 21 features): Distillation provides a significant gain (+0.011 AUC over CatBoost).
- High-dimensional data ( $>$ 21 features): Gains are negligible (+0.001 AUC), and if the teacher itself underperforms CatBoost, distillation can degrade performance further.
Multi-Teacher Dynamics:
- Tree Students: Multi-teacher averaging provides a statistically detectable but practically negligible gain (+0.0006 AUC). A single strong teacher is sufficient.
- MLP Students: Multi-teacher averaging acts as effective label smoothing, yielding a significant gain (+0.006 AUC, $p=0.003$ ) over single-teacher training.
Open-Source Pipeline: The full pipeline is implemented and released as part of the TabTune library.

Results

The evaluation covers 153 classification datasets from TALENT, OpenML-CC18, TabZilla, and TabArena.

Accuracy: The best distilled configuration (TabICLv2 $\to$ XGBoost) achieves 0.882 AUC, retaining 96.5% of the teacher's performance.
Latency: Distilled tree students run in 1.9–4.0 ms on CPU, compared to 151–1,275 ms for teachers on GPU.
Baselines: The distilled models beat a tuned CatBoost baseline on 51% of datasets. In contrast, OrionMSP, a weaker teacher, produces students that win only 22.6% of the time.
Ablation: For MLP students on low-dimensional data, using hard labels with OOF ( $\alpha=0$ ) actually outperformed the full Hinton soft-label pipeline by 0.034 AUC, suggesting that for clean, low-dimensional data, the complex soft-label machinery is optional, but OOF labeling remains critical.

Significance

The paper demonstrates that the "deployment problem" of tabular foundation models can be solved without sacrificing their accuracy advantages. By addressing the specific leakage issue inherent to ICL teachers, the authors show that it is possible to distill the "inter-class structure" of a TFM into a lightweight, CPU-native GBDT.

The significance lies in the predictability and graceful failure of the approach:

If a TFM is strong, distillation transfers that strength to a fast student.
If a TFM is weak (or the data is high-dimensional), distillation does not artificially inflate performance; it fails cleanly, often performing worse than a standard GBDT, signaling that the teacher was not a good candidate.
The method requires no complex hyperparameter tuning for the student; the primary decision is simply selecting the best-performing teacher.

The authors conclude that for low-dimensional structured data, practitioners can achieve near-teacher accuracy at 1% of the inference cost, provided they strictly adhere to the OOF labeling protocol to avoid signal collapse.