Embeddings

The Embeddings extension extracts latent feature representations (embeddings) from TabPFN models. These dense vectors capture the representations learned by TabPFN’s transformer and can be reused for downstream tasks such as clustering, search, visualization, meta-learning, or as features for a simpler model. TabPFNEmbedding is a scikit-learn style transformer with the familiar fit / fit_transform / transform API. It supports two extraction modes:

Out-of-fold embeddings (n_fold >= 2, recommended) — robust, leakage-free training-set embeddings extracted via K-fold cross-validation. These generalize better and give stronger downstream performance.
Vanilla embeddings (n_fold=0) — a single model is trained on the full dataset and used for everything; cheaper, but the training embeddings leak label information.

Embeddings require the full local tabpfn package — they expose internal model representations that the tabpfn-client cloud backend does not provide. Passing a client model raises a TypeError.

Getting Started

The embedding module ships in the base tabpfn-extensions package (no extra needed). Install it alongside the local tabpfn engine:

pip install tabpfn tabpfn-extensions

The Interface

TabPFNEmbedding follows the scikit-learn transformer contract, and the method you call depends on whose embeddings you want:

Method	Use for	Returns
`fit_transform(X_train, y_train)`	the training data	Out-of-fold embeddings when `n_fold >= 2` (no label leakage); full-model embeddings when `n_fold == 0`
`transform(X)`	unseen / held-out data	Embeddings from the model trained on the full training set

transform always runs through the final full-data model. It never returns cached training embeddings, even if X happens to equal the training set — so for leakage-free training embeddings, always use fit_transform (or read the train_embeddings_ attribute after fit).

Pass a configured TabPFN model via the model= parameter. Use a classifier or regressor depending on your task — the examples below show both.

1. Out-of-fold (robust) embeddings — recommended

With n_fold >= 2, the training data is split into K folds; a fresh model is trained on each fold and used to embed its held-out partition. The out-of-fold (OOF) embeddings are reassembled into the original sample order, and a final model is refit on the full training set to embed unseen data.

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from tabpfn import TabPFNClassifier
from tabpfn_extensions.embedding import TabPFNEmbedding

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)

embedding = TabPFNEmbedding(
    n_fold=10,
    model=TabPFNClassifier(n_estimators=1, random_state=42),
)

train_embeddings = embedding.fit_transform(X_train, y_train)  # out-of-fold
test_embeddings = embedding.transform(X_test)                 # final model

Why prefer out-of-fold embeddings? Vanilla embeddings use a single model to embed the training and test rows. The problem with this approach is that the training rows contain target information, the test rows do not. This introduces the risk of information leakage. OOF embeddings break this leakage: every training point is embedded by a model that never saw it, so the training embeddings match the statistics of the held-out embeddings produced by transform. This is the robust variant introduced in “A Closer Look at TabPFN v2: Strength, Limitation, and Extension” (arXiv:2502.17361), and larger n_fold values yield more robust embeddings. In practice this lifts downstream performance: the get_embeddings.py example compares a baseline linear model, vanilla TabPFN embeddings, and K-fold embeddings on the same data — the K-fold embeddings come out ahead for both classification accuracy and regression R². Classifiers use StratifiedKFold and regressors use KFold. Set shuffle=True (with an optional random_state) to shuffle the split. n_fold=1 is invalid — use 0 for vanilla or >= 2 for cross-validation.

2. Vanilla embeddings

With n_fold=0, a single model is trained on the entire training set and reused for both training and unseen data. This is cheaper (one fit instead of K + 1) and fine when you only need embeddings for unseen data via transform, but avoid it for training-set embeddings you plan to feed into a downstream model — see the leakage caveat above.

embedding = TabPFNEmbedding(
    n_fold=0,
    model=TabPFNClassifier(n_estimators=1, random_state=42),
)

train_embeddings = embedding.fit_transform(X_train, y_train)  # training data
test_embeddings = embedding.transform(X_test)                 # unseen data

Output shape. Both fit_transform and transform return a 3D array of shape (n_estimators, n_samples, embed_dim) — one embedding matrix per ensemble member. This is not a drop-in 2D input for an sklearn Pipeline. Select a single member (embeddings[0]) or aggregate across axis=0 before passing the result to a downstream estimator.

What the Embeddings Capture

To see what the transformer actually adds, hold the projection fixed and change only the representation. PCA can only rotate and rescale the space it is handed, so whatever separation shows up in two components was already present in the input. Below are scikit-learn’s 8x8 handwritten digits (load_digits: 1,797 samples, 64 pixel features, 10 classes), projected onto their first two principal components twice. On the left, straight from the raw pixels. On the right, from TabPFN’s 512-dimensional out-of-fold embeddings of the same rows.

PCA of raw digit pixels compared with PCA of TabPFN embeddings

Both panels are the same linear projection of the same 1,797 digits; only the input representation differs. Axis units are arbitrary and each panel is standardized, so the two are directly comparable.

from sklearn.datasets import load_digits
from sklearn.decomposition import PCA
from tabpfn import TabPFNClassifier
from tabpfn_extensions.embedding import TabPFNEmbedding

X, y = load_digits(return_X_y=True)

embedding = TabPFNEmbedding(
    n_fold=5,
    shuffle=True,
    random_state=42,
    model=TabPFNClassifier(n_estimators=1, random_state=42),
)
E = embedding.fit_transform(X, y)[0]        # (1797, 512), out-of-fold

Z_pixels = PCA(n_components=2, random_state=42).fit_transform(X)   # left panel
Z_tabpfn = PCA(n_components=2, random_state=42).fit_transform(E)   # right panel

Reading the panels. In pixel space the ten digits collapse into one cloud. Only 0, 4 and 6 drift to the edges; 1, 5, 7 and 8 are interleaved in the middle with no boundary between them, which is why a linear model on raw pixels struggles. On the embeddings the same two components pull 3, 2, 0, 6, 4 and 9 apart into distinct regions with clear whitespace between them, and the residual crowding is confined to the 1/7/5/8 band. Two linear dimensions out of 512 are enough to expose most of the class structure, because the transformer has already done the work that PCA cannot do on its own.

This is a supervised-versus-unsupervised comparison. PCA on pixels sees only pixels; TabPFN’s encoder is conditioned on labeled rows, so its embedding space is organized around the target rather than around pixel variance. That is the whole point of using it as a feature extractor, but it does mean the figure is not a like-for-like dimensionality-reduction benchmark.

This separability is the practical payoff. It is the same property that lets a plain linear model on top of the embeddings do well, as in the next section.

Using Embeddings as Features

A common pattern is to use TabPFN embeddings as features for a lightweight downstream model. Because the embeddings are 3D, select an ensemble member (embeddings[0]) to get a 2D feature matrix.

Classification
Regression

from sklearn.datasets import load_breast_cancer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
from tabpfn import TabPFNClassifier
from tabpfn_extensions.embedding import TabPFNEmbedding

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)

embedding = TabPFNEmbedding(
    n_fold=10,
    model=TabPFNClassifier(n_estimators=1, random_state=42),
)
train_embeddings = embedding.fit_transform(X_train, y_train)
test_embeddings = embedding.transform(X_test)

clf = LogisticRegression(max_iter=5000)
clf.fit(train_embeddings[0], y_train)          # pick ensemble member 0 -> 2D
y_pred = clf.predict(test_embeddings[0])
print(f"Accuracy: {accuracy_score(y_test, y_pred):.4f}")

from sklearn.datasets import fetch_openml
from sklearn.linear_model import Ridge
from sklearn.metrics import r2_score
from sklearn.model_selection import train_test_split
from tabpfn import TabPFNRegressor
from tabpfn_extensions.embedding import TabPFNEmbedding

dataset = fetch_openml("space_ga", version=1, as_frame=False)
X, y = dataset["data"], dataset["target"].astype(float)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)

embedding = TabPFNEmbedding(
    n_fold=10,
    model=TabPFNRegressor(n_estimators=1, random_state=42),
)
train_embeddings = embedding.fit_transform(X_train, y_train)
test_embeddings = embedding.transform(X_test)

reg = Ridge()
reg.fit(train_embeddings[0], y_train)          # pick ensemble member 0 -> 2D
y_pred = reg.predict(test_embeddings[0])
print(f"R² Score: {r2_score(y_test, y_pred):.4f}")

Parameters

Parameter	Type	Default	Description
`n_fold`	`int`	`0`	`0` disables CV (vanilla). `>= 2` enables K-fold out-of-fold embeddings. `1` is invalid.
`model`	`TabPFNClassifier \| TabPFNRegressor \| None`	`None`	Pre-configured TabPFN estimator. When `None`, the task is inferred from `y` at `fit` time (with a warning). Passing it explicitly is recommended.
`shuffle`	`bool`	`False`	Whether to shuffle the K-fold split.
`random_state`	`int \| None`	`None`	Seed used by the K-fold split when `shuffle=True`.

After fitting, two attributes are available: model_ (the fitted full-data model) and train_embeddings_ (the training-set embeddings, OOF when n_fold >= 2).

Migration. The old get_embeddings(X_train, y_train, X, data_source=...) method and the tabpfn_clf / tabpfn_reg constructor arguments are deprecated. Use model= together with fit_transform (training, OOF) and transform (unseen data) instead.

Example Script

Full runnable example for classification and regression.

Google Colab Example

Check out our Google Colab for a demo.

Getting Started

API

Capabilities

Agentic

Integrations

Troubleshooting

Use Cases

Changelog

Getting Started

The Interface

1. Out-of-fold (robust) embeddings — recommended

2. Vanilla embeddings

What the Embeddings Capture

Using Embeddings as Features

Parameters

Example Script

Google Colab Example

​Getting Started

​The Interface

​1. Out-of-fold (robust) embeddings — recommended

​2. Vanilla embeddings

​What the Embeddings Capture

​Using Embeddings as Features

​Parameters

Example Script

Google Colab Example

Getting Started

The Interface

1. Out-of-fold (robust) embeddings — recommended

2. Vanilla embeddings

What the Embeddings Capture

Using Embeddings as Features

Parameters