Skip to main content
TabPFN can be applied to regression tasks using the TabPFNRegressor class. Instead of a single point estimate, the regressor predicts an output distribution, enabling you to generate different types of forecasts and accurately measure uncertainty in your predictions.

​
Getting Started

The following shows a full example of using the TabPFNRegressor: 
from sklearn.datasets import load_diabetes
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from sklearn.model_selection import train_test_split
from tabpfn import TabPFNRegressor
from tabpfn.constants import ModelVersion

# Load example dataset
X, y = load_diabetes(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

# Initialize and fit
model = TabPFNRegressor(device="auto")
# To use TabPFNv2:
# model = TabPFNRegressor.create_default_for_version(ModelVersion.V2)
model.fit(X_train, y_train)

# Predict
preds = model.predict(X_test)

# Evaluate
print("MSE:", mean_squared_error(y_test, preds))
print("MAE:", mean_absolute_error(y_test, preds))
print("R²:", r2_score(y_test, preds))

​
Understanding Predictions

The TabPFNRegressor.predict() method can return different types of predictions based on the output_type parameter.

​
Point Predictions (Mean, Median, Mode)

By default, predict() returns the mean of the predicted distribution. You can also request the median or mode for different point estimates. 
# Get the mean prediction (default)
mean_preds = model.predict(X_test, output_type="mean")

# Get the median prediction
median_preds = model.predict(X_test, output_type="median")

# Get the mode prediction
mode_preds = model.predict(X_test, output_type="mode")

​
Quantile Regression

To perform quantile regression and get predictions for specific quantiles, set output_type="quantiles" and pass a list of desired quantiles. This is ideal for estimating prediction intervals. 
# Get predictions for the 10th, 50th (median), and 90th quantiles
quantile_preds = model.predict(
    X_test,
    output_type="quantiles",
    quantiles=[0.1, 0.5, 0.9],
)

# quantile_preds is a list of arrays
# q10_preds = quantile_preds[0]
# q50_preds = quantile_preds[1]
# q90_preds = quantile_preds[2]

​
Full Distribution

For advanced use cases, you can retrieve the full predictive distribution.
  • output_type="main" returns a dictionary containing the mean, median, mode, and default quantiles.
  • output_type="full" returns everything in "main" plus the raw logits and criterion object, allowing for custom loss calculations (see the FAQ).
# Get a dictionary of all main prediction types
main_outputs = model.predict(X_test, output_type="main")
# print(main_outputs.keys()) 
# > dict_keys(['mean', 'median', 'mode', 'quantiles'])

# Get the full distribution objects for custom analysis
full_outputs = model.predict(X_test, output_type="full")
# print(full_outputs.keys())
# > dict_keys(['mean', 'median', 'mode', 'quantiles', 'criterion', 'logits'])
TabPFNRegressor provides a full predictive distribution, enabling loss-aware predictions without retraining. You can compute the Bayes-optimal point prediction that minimizes the expected custom loss. This method gives flexible custom-loss predictions without modifying the model.
out = reg.predict(X, output_type="full")
criterion, logits = out["criterion"], out["logits"]

q = np.linspace(0.01, 0.99, 101)
samples = np.stack(
    [criterion.icdf(logits, float(x)).cpu().numpy() for x in q], axis=1
)

denom = np.maximum(np.abs(samples), 1e-6)
diffs = np.abs(samples[:, :, None] - samples[:, None, :])
exp_mape = (diffs / denom[:, :, None]).mean(axis=1)
best_idx = exp_mape.argmin(axis=1)
y_hat = samples[np.arange(samples.shape[0]), best_idx]

AutoTabPFN Ensembles

Automated ensembling of TabPFN models for top accuracy.

Interpretability

Explain TabPFN predictions with SHAP utilities.

Hyperparameter Optimization

Automated hyper-parameter tuning for TabPFN models.