Regression should predict full distributions

A "hidden" feature of tabular foundation models

When I took on a side quest to forecast water supply, I had to predict the 10%, 50%, and 90% quantiles. I solved this by training three separate quantile models (ensembles of xgboost, actually). I would have preferred to only train a single model that could predict all three quantiles at once. While approaches like linear regression can output full predictive distributions, these often come with (too) strong distributional assumptions.

What if we always worked with machine learning models that produce the full predictive distribution? With classification, we are already at this point: Modern machine learning approaches output not just the majority class, but a probability for each class. Whether this probability is calibrated is another question.

With regression, we are a bit stuck with a point-based mindset. However, this could change with tabular foundation models. At least in theory: While these models produce the full predictive distribution (or at least a discretized approximation over a fixed support) it’s not the default and the output is a bit hidden.

Tabular foundation models are pre-trained to output the full predictive distribution

TabICL and TabPFN predict the full predictive distribution (discretized approximation). But by default, they only give you the conditional mean. Consider the following code snippet:

reg = TabPFNRegressor()
reg.fit(X=X_train, y=y_train)
predictions = reg.predict(X_test)

The predictions object in the code above contains just the means of the full predictive distributions. But in the background, the foundation model actually produces the full distribution (or at least a discretized version over a fixed support) and later aggregates it.

If you want the predictive distribution, you need to change a paramter in .predict():

quantiles_range = np.arange(0.05, 0.95, 0.05)
quantiles = reg.predict(X_test, output_type='quantiles', quantiles=quantiles_range)

Both output type options (mean, quantile) take the same amount of compute time, since the tabular foundation models predict the distribution anyways.

Why prefer distributions over points?

Starting from P(Y|X = x) may seem more tedious, but it gives us more modeling freedom. There are many situations where the mean is not appropriate, or where we want more informative predictions:

• You might prefer the predictive median over the mean to get more robust predictions.

• If you are interested in the tails of the distribution, you can extract the 10% and the 90% quantiles (or any other quantiles for that matter).

• Computing quantile intervals or the variance of the predictive distribution can help quantify uncertainty.

• Maybe you need the probability that the prediction exceeds a certain threshold, P(Y>threshold).

• Maybe you want to extract the modalities.

• Or you can also visualize the entire distribution.

We have so many more options when we work with the predictive distribution. And we get it without any additional computations when using tabular foundation models.

Are the predictive distributions calibrated?

While TFMs are directly pre-trained to predict the full distribution, there is no guarantee that this distribution is calibrated, meaning, e.g., that the predicted 10% quantile matches the actual 10% quantile. What distinguishes TFMs from approaches like linear regression or Bayesian regression models is the lack of explicit distributional assumptions.¹ But to get a better idea whether to trust the full distributions, we need benchmarks. Unfortunately, popular benchmarks like TabArena focus on point estimates, not the entire distribution. This paper criticizes the field’s focus on the conditional mean and proposes reporting metrics in benchmarks that reflect calibration. Keep in mind that even when a model scores well in a benchmark, you have no guarantees for your own project. It’s just a rough pointer.

In the end, you need to validate calibration yourself, using proper scoring rules, or even calibrate using conformal prediction. But still, I find it exciting that tabular foundation models give us the option of predicting other aspects of the predictive distribution.

1. TFMs do have some implicitly encoded assumptions through their pre-training on synthetic data.