Log In to continue learning
Tabular Foundation Models for Data Science Agents
Intermediate
1h32m
You've used TabPFN throughout the course. Now you'll look under the hood. We'll actually train a tiny version of TabPFN from scratch. After that, we'll wrap up by looking at where the field is heading. Let's go. Let's quickly recall what I showed in the first video, on the slides about TabPFN pre-training and inference. In the inference phase, you give TabPFN an entire new dataset in context, the training rows with their labels and the test rows with a question mark. And one forward pass gives you predictions for the test rows. Now, let's recall the pretraining. is what allows TabPFN to do this. During pretraining, it sees datasets that look exactly the same as during inference, except that they're synthetic, so we know the answers for the test rows. We score the predictions on those answers, take a gradient step, throw the dataset away, and generate a new one, over and over again. That pretraining usually takes days on many GPUs. Today, we'll do a tiny version ourselves right here in the browser in less than a minute. Here, we implement this entire pre-training pipeline in three parts: the prior, the model, and the training loop. The prior allows us to sample new classification datasets. The model is a transformer-like architecture that alternates attention across rows and columns, and the training loop is fairly standard and implements in-context learning. The model follows nanoTabPFN, which is a small version of TabPFN for teaching, created by my student Alexander Pfefferle at the University of Freiburg. Let's get started with some setup. The most interesting part here is that the plotting and evaluation code lives in a helper file, and we seed the random number generator so that you should get exactly the same results as I'm showing here. The first ingredient is the prior, the program that generates the synthetic datasets. Every call creates a batch of brand-new tables, each with 50 rows and 2 features, and the rows are split randomly into training and test rows. We start with a really simple prior. We draw random lines and have two classes, 1 are the red dots above the line, and class 0 are the blue crosses below the line. And we have in each of the samples a different line, so that we actually learn to generalize across different lines. In code, this here defines a line. This defines the features of the data points, and the class is defined here, depending on whether a data point is above the line or below the line, just with a little bit of noise. So there are, for example, some data points that are on the wrong side of the line, just barely. And then we split the rows into training rows, the first n_train, and test rows, the remainder of the 50 rows. The second ingredient is the model. It takes a table, turns every cell into a vector, mixes these vectors with attention, and then reads out a prediction for each test row. The heart of it is this attention here, so let's start there. Each of these layers applies two different types of attentions. First, across a row, the cells of one data point attend to each other, so the features and the label can exchange information. Then, down a column, every row attends to the training rows. The test rows look at the training rows, but not the other way around. That's where the in-context learning happens. A test row looks for similar training rows and uses their labels. Then we have a small MLP on every cell, and all the MLPs share parameters. And then that's one layer, and we stack a few of them. Here is exactly that layer in code. We have two standard PyTorch attention modules, one between features, one between data points, and then we have the MLP and the usual LayerNorms. The input h has four dimensions: a batch of tables, each with rows and columns, and in every cell, there's an embedding vector. Attention only takes a three-dimensional input, so in order to run self-attention over the features of each row, we fold the row dimension into the batch. And then we run attention across it. Queries, keys, and values all come from the same row, and then a residual connection and a LayerNorm. For the second attention, we flip the table. Now the column dimension is folded into the batch dimension. And this is the key line. The queries are all rows, but the keys and values come only from the first n_train rows. Those are the labeled ones. And that's the arrow we saw in the figure. All of the rows attend to the training rows. We then run the attention, plus again the residual connection and the LayerNorm, flip back and apply the MLP, and that's one layer. Now let's have a look at the rest of the model. The input is a table: labeled training rows, and test rows whose labels are hidden from the model. We use those hidden labels only to score the predictions, and use zeros going into the model. Then every cell becomes a vector, using a linear layer for the features and another one for the labels. The label is simply another column here. After the last TransformerLayer, we take the label cell of each of the test rows and use a small decoder MLP to read out the logits, to then get the classification probabilities. And then we look how likely are the actual labels under the predictive distribution. So how likely is this one under this predictive distribution and this zero under this predictive distribution? And that is going to be our cross-entropy loss. Now let's have a look at how all of these parts come together in the NanoTabPFN class. In the constructor, there are four pieces: the two encoders from panel two, the stack of transformer layers from panel three, and the decoder MLP from panel 4, with two outputs, one per class. The forward pass takes exactly what TabPFN takes: the training rows with their labels and the test rows. that was panel one. We first standardize each feature using only the labeled rows. And here are the labels: 2 * y - 1 turns class 0 and class 1 into a -1 and a +1, and the padding adds a 0 for every test row that signals that the class is unknown. We then embed every cell and put the label column next to the feature columns, so it looks like just another column. Then we run the attention layers and finally decode the label column of the test rows into two logits per test row. That's the whole model. The third ingredient is the training loop, and it's completely standard. We use the AdamW optimizer and a cosine decay on the learning rate. In every step, we generate fresh datasets from the prior. Notice that there is no real data here anywhere in this loop. We then split each dataset into the training rows with their labels and the test rows without their labels, and pass them through the model to get out the logits. And then we compute the usual cross-entropy only on the test rows. This is exactly the loss from the initial TabPFN pre-training slide that quantifies how likely the true test labels are under the model's predictions. Then we do a gradient step and on to the next batch of datasets. That's all of TabPFN pre-training. Let's train a really tiny model: two layers, 48-dimensional embeddings, and 750 training steps. Now the model is training, and every few steps we check how well it does on three test datasets that the model never actually trained on. And we output both the loss, cross-entropy loss, and then test accuracy across these datasets. And after about 22 seconds, it is actually done training. So I said we test on three datasets, and these are the three datasets. There's a line here, 20 red data points and 20 blue data points on top and below. Then there is a sine curve again with 40 training data points, and a sinusoidal on top of a line. And the model has only seen lines during pretraining, so it's only going to do well here. We also see the predictive surface here in this area is all red, so it's very certain that this is red. It's blue here. And in between there's this white areas where it's very uncertain. And we see that there is larger areas here where it's uncertain. We can also evaluate this quantitatively. Here, we compare the small TabPFN trained on this prior of lines against logistic regression and a random forest. And we see that it does well on the lines, but not well on the other two datasets. And this is not surprising because the model only has seen lines during its pre-training. So naturally, let's make our prior richer. Now, we actually add sinusoidals on top of our lines. So it's the same function as before, except there's a sine wave on top of it, and this is what it looks like in code. Now, let's train again with this richer prior. We train the same model with the same learning rate, same optimizer, same everything. The only thing that changed is the prior. We even have the same embedding size, the same number of layers and the same number of steps. We again trained for roughly 22 seconds and output the cross-entropy and the test accuracies, and you already see the test accuracies are looking dramatically better than before. And when we look at the visualizations of our predictive surfaces, we do see that we can now actually capture sine waves like this or also lines with sinusoidals on top. And quantitatively, we see that actually on these types of data sets, this tiny TabPFN that trained for 22 seconds is actually already better than logistic regression and random forest trained on the same data sets. But of course, the predictions aren't perfect. There's lots of mismatches if you look at it visually. And so let's actually now scale up and train a bigger model for a bit longer and see if it can get better. And in particular, we scale up by turning three different knobs. We use a wider embedding, more layers, and more training steps. And of course, this takes a bit longer now to train, so let's fast-forward. Okay, and now we have converged, still after only about four and a half minutes. and the test accuracy is much better now. And when we look at the predicted boundary, we see that it follows the waves much more closely. And also in terms of baseline comparison, we see the large model has much better accuracy. Of course, this example was very toy, though. only two features, 50 rows per table, and a prior that only knows wavy lines. And I actually modified the original nanoTabPFN to simplify this for teaching. But now let's have a look at the original nanoTabPFN. is a simplified and scaled down variant of TabPFN v2 for teaching. It has exactly the same architecture that we just trained. The main thing I simplified a bit further for the code part of this lecture was the prior. There, the original NanoTabPFN used a much more diverse prior to capture all kinds of patterns. And it also used tables of different sizes, in order to learn to length generalize to more rows and to more columns. Finally, the evaluation of NanoTabPFN was also different. In contrast to our toy functions, NanoTabPFN was actually evaluated on real tabular prediction tasks from the literature, specifically a subset of the TabArena benchmark, with datasets up to 10 columns and subsampled to 200 rows. And here is our result on these datasets. The x-axis is pre-training time on a single consumer GPU, and within 1 minute of pre-training, nanoTabPFN, the blue curve here, climbs past the Decision Tree, past K-Nearest Neighbors, and reaches the Random Forest and a plain TabPFN-v2 without preprocessing. And this is on real datasets that are entirely different from the datasets that the prior generates. So this is a general tabular ML algorithm, and it was learned from scratch in a minute. I still find it remarkable that that is possible. The NanoTabPFN code is on GitHub, so if the notebook we saw in this lecture made you curious, and if you want to go further, there is a TFM playground, where you can try your own ideas for pre-training tabular foundation models. So where is the field heading? Let me pick three directions to talk about. The first one we have already seen ourselves: scaling up. In January 2025, TabPFN v2 was published in Nature and scaled to 10,000 rows. TabPFN 2.5 then scaled to 100,000 rows, TabPFN 3 to a million, and TabPFN 4 to 10 million. that is a 1,000-fold increase in row limit in less than two years. The second direction is the messiness of real data. The example tables we used were clean: two numeric columns, independent rows, no outliers, balanced classes. Real data is rarely like that. Tabular data is often grouped, from several hospitals or several stores. And predicting for a new group is much harder than predicting for a known one. Tabular data is also often temporal, meaning it has timestamps. You need to watch out for these to avoid leakage, but they can also provide really helpful features if used right. Tabular data can also be really high-dimensional in some applications, with far more features than data points. This often happens in biomedicine, for example, in omics data. And finding the signal can then really be like looking for the needle in the haystack. Tabular data also often lives in a relational database with many linked tables that you need to connect in the right way in order to set up your machine learning problem. Tabular data is also often very imbalanced. For example, in fraud detection, maybe one row in 1,000 is actually fraud. Finally, tabular data often contains columns with text, and beyond that, you often have images and whole documents alongside a table. And treating these modalities right makes a big difference to performance. There are many more types of data complexities, and tabular foundation models are quickly getting better at all of these. The third direction I want to talk about goes beyond making models stronger on the current tasks. It rather changes the questions you can answer. And this is the one that I'm most excited about: it's causality. Everything we talked about in this course was prediction. Given the features, what is the label? But many of the most important questions are about interventions. The probability of an outcome given that we observe a treatment is not the same as when we choose to apply that treatment. So, for example, patients who get a drug differ from patients who don't get that drug. Today, the reliable way to make interventional predictions is to obtain interventional data through a randomized controlled trial or an A/B test, It's expensive, slow, and sometimes even impossible. And what we would really like to have are what-if answers for each patient or each customer, and to really understand all of the causal relationships between all of these variables. And all of that with a minimum of expensive interventional trials. Causal extensions of TabPFN could be a potential solution to this. The prior of TabPFN actually already samples causal graphs to create synthetic data, and we can also train the PFN to solve causal problems. For example, it is possible to train PFNs to not only output predictions for y_test, but also output probabilities for connections in the causal graph. And with Do-PFN, we can also simulate interventions in the prior and then predict the effect of interventions based purely on observational data. This line of work still needs to mature, but I'm super excited about the potential of these causal PFNs to take off the same way that TabPFN has taken off for purely predictive tasks. And that's it for this lecture. I hope you had fun building a tiny TabPFN and seeing where the field is going.
