Context is the new training
Tabular foundation models replace retraining with editable inference-time data
Tabular foundation models such as TabPFN and TabICL don’t need to be trained to perform regression or classification. What they do is called in-context learning. What used to be the training data now becomes the context data at prediction time.
This post explores the idea of context data and contrasts it with “classic” training data. Does moving from training data to context data change how we model? Does it enable something new?
Let’s dive in.
Training versus context
For traditional machine learning (linear regression, XGBoost, SVM), the training data shapes the model. Especially for trees, this is vivid: change the training data, and you might get out a differently shaped tree. Fit a linear regression model, and the weights (coefficients) become a function of the training data.
Not with tabular foundation models. Pre-trained. No classic training step. Prediction happens via in-context learning. When predicting with a TFM, you have to provide both the “training” (aka context) and the test data. Through multiple steps, the table cells are embedded, and through attention mechanisms and fully connected layers, the TFM enriches the cell representations by attending to other cells and computing stuff. All based on what the model learned in pre-training. Training and test data points can attend to the training data points, but not to other test data. The cell embeddings for y_test are used to predict or classify. If you want to learn more about how tabular foundation models work, check out my TabPFN architecture post.
So what happens when we change the context data? The model weights don’t change. The only things that change are the predictions that come out of the model: because the context changed, the embeddings will change, and the set of points that can be attended to will change.
But besides not affecting the model, is it really worth thinking about the “training” data as context data?
What I tell you in the next two sections is not new: You could also do traditional ML and call it context instead of training. Wrap it all in a function where the .predict() actually does fit+predict. Except, there are things that make it worthwhile to think of context data instead of training data: the ability and necessity to work with smaller datasets (which TFMs work well on) and the inverted cost in which training cost effectively disappears (ignoring pre-training here) and prediction becomes the expensive step.
Smaller and smarter contexts
Making predictions with tabular foundation models is relatively expensive. Especially if you always use all your training data as context. The inference runtime complexity of TabICL is O(n^2 + nm^2), where n is the number of rows of context+test and m is the number of columns. Since it’s quadratic in the number of rows, it will absolutely explode when we move up the orders of magnitude: a 10x increase in training data rows is a 100x increase in runtime (assuming that training data size is much larger than test data size and much larger than the number of features). The quadratic scaling strongly incentivizes using a smaller context dataset.
There is a more positive outlook on small contexts: Tabular foundation models work especially well for smaller data. That’s where they typically outperform most other approaches. Perhaps it’s because, for smaller data, the inductive biases learned in pre-training can shine.
Anyways, we can think of downsizing our data. When I experimented with tabular foundation models I sometimes downsampled the data for improved runtime, but with the drawback of reduced predictive performance.
For example, the following figure shows the mean absolute error for predicting wine quality, comparing TabICL and a random forest. The TabICL model has the same performance with 700 context data points as the random forest with 1200 training data points. However, increasing the context data for TabICL to the full 1200 data points gives us a big boost in MAE. 1200 data points is easily handled by a TFM nowadays. The question is whether, for example, downsampling from 1 million to 10k data points shows a relevant loss in predictive performance or not.
We can also be a bit smarter about downsizing the context:
This paper suggests using a kNN model for each test data point to decide on the context. An interesting quote from the paper: “We thus believe that using nearby points as context is a good inductive bias for tabular data classification.”
Another paper suggests clustering the training data. For a data point in the test data, check which cluster it belongs to, and then use the respective training cluster as context.
You can reduce the training data to representative data points.
The following figure shows the cluster approach: Using 2 or 3 clusters is actually slightly more performant than using all data.
Again, all of these approaches are possible with the classic training-test scheme, but with tabular foundation models, we are more incentivized to reduce the context size while at the same time the TFM models work quite well on smaller data.
Context change is the new re-training
With in-context learning, the cost inverts between training and inference: training costs close to nothing while inference becomes more expensive.
This has implications for topics such as interpretability, feature selection, and robustness analysis. For example, in model-agnostic model interpretability, we have methods that rely on shuffling a feature column, but there is also a version of that which relies on removing that column and re-training the model. I’ve covered this in the following post:
In general, methods that require re-training are now much cheaper, relative to methods that require predicting twice with the same model. Usually, re-training was expensive. But for TFMs, re-training boils down to adapting the context data. The cost of re-training is close to zero, we only pay for making the predictions with TFMs.
Context engineering for TFMs?
So far, we have covered needs of smaller contexts for better scaling and changed costs for some post-hoc methods. Doesn’t feel like a big shift away from classic training+predict paradigm. I’m still trying to wrap my head around the context-data-framing. We gain a lot of flexibility without having to think about hyperparameter tuning, model selection, cross-validation, and so on. Just some random thoughts on what a context-paradigm might enable or simplify:
Imagine a user requests their data to be fully deleted. Data your model was trained on. Removing them from a traditional ML model would require re-training, but for TFMs, it’s just a matter of removing their data from the context.
Imagine a forecasting task for which only the last 30 days are relevant, because data distribution is drifting. No problem with TFMs, we can just have a sliding context window.
Context data makes it easier to have per-request contexts. For example, contexts only with data from a certain region or customer segment, or even per user.
If you have to remove poisoned/wrong/problematic data, just delete the rows from the context.
Specialization is much easier. Imagine you have a global model that classifies transactions in your banking app. Then the user classifies some items on their own. These can be simply added to the context.
Customization is also simplified. Continuing with the transactions example, the user could say: Don’t learn from old transactions (= remove rows from context), or don’t use the merchant name as a feature (= remove column from context).
None of this is truly new. But still, I found it very refreshing to think of context data rather than training data. And I feel like I have yet to fully grasp the context-mindset.
Retrieval-augmented tabular prediction?