Extracting Feature and Outcome Embeddings from TabPFN

[lincj1994]

Hi.
I'm wondering how to extract feature embeddings and outcome embeddings from the TabPFN? I’m exploring ways to infer feature importance using embeddings from TabPFN using cosine similarities between these embeddings.
Thanks.
Caijin

[iivalchev]

For the feature embeddings, you can get them on a fitted model via

TabPFN/src/tabpfn/regressor.py

Line 711 in 05ab7da

def get_embeddings(

[lincj1994]

Hi. I'm still confused of how to get embeddings from a fitted model, such as the clf in the code example below, where I fit the model using AutoTabPFNClassifier. Could you please provide some additional guidance?
Thanks.

import pandas as pd
import numpy as np
from tabpfn import TabPFNClassifier
from tabpfn_extensions.post_hoc_ensembles.sklearn_interface import AutoTabPFNClassifier
from sklearn.model_selection import train_test_split


df = pd.read_csv('df2.csv')
X_features = list(df.drop(['resp_crpr', 'Tumor_Sample_Barcode'], axis=1))
y_features = ['resp_crpr']
clf = AutoTabPFNClassifier(device='auto')


args = Namespace(
    seed = 444444,
    test_size = 0.3
)
X_train, X_test, y_train, y_test = train_test_split(df[X_features].values.reshape(-1, 63), df[y_features], 
                                                test_size=args.test_size, random_state=args.seed)
clf.fit(X_train, y_train)

[iivalchev]

clf.get_embeddings(X_test, data_source='train')
or
clf.get_embeddings(X_test, data_source='test')

X_train is retained after the fit, so data_source dictates which part of X in embedding space you would like to get, if you need both simply concat the result from the calls above and you would get all of X as embeddings.

[lincj1994]

It seems that only classifier.py has attribute "get_embeddings" while not the AutoTabPFNClassifier.

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 1
----> 1 em_train = clf.get_embeddings(X_test, data_source='train')
      2 em_test = clf.get_embeddings(X_test, data_source='test')

AttributeError: 'AutoTabPFNClassifier' object has no attribute 'get_embeddings'

Package                 Version     Editable project location
----------------------- ----------- -------------------------------------------------------------
asttokens               3.0.0
certifi                 2025.1.31
charset-normalizer      3.4.1
cloudpickle             3.1.1
colorama                0.4.6
comm                    0.2.2
contourpy               1.3.1
cycler                  0.12.1
debugpy                 1.8.11
decorator               5.2.1
einops                  0.8.1
exceptiongroup          1.2.2
executing               2.1.0
filelock                3.17.0
fonttools               4.56.0
fsspec                  2025.3.0
future                  1.0.0
huggingface-hub         0.29.2
hyperopt                0.2.7
idna                    3.10
importlib_metadata      8.6.1
ipykernel               6.29.5
ipython                 9.0.2
ipython_pygments_lexers 1.1.1
jedi                    0.19.2
Jinja2                  3.1.6
joblib                  1.4.2
jupyter_client          8.6.3
jupyter_core            5.7.2
kiwisolver              1.4.8
MarkupSafe              3.0.2
matplotlib              3.10.1
matplotlib-inline       0.1.7
mpmath                  1.3.0
nest_asyncio            1.6.0
networkx                3.4.2
numpy                   2.2.3
packaging               24.2
pandas                  2.2.3
parso                   0.8.4
pickleshare             0.7.5
pillow                  11.1.0
pip                     25.0
platformdirs            4.3.6
prompt_toolkit          3.0.50
psutil                  5.9.0
pure_eval               0.2.3
py4j                    0.10.9.9
Pygments                2.19.1
pyparsing               3.2.1
python-dateutil         2.9.0.post0
pytz                    2025.1
pywin32                 308
PyYAML                  6.0.2
pyzmq                   26.2.0
requests                2.32.3
scikit-learn            1.6.1
scipy                   1.15.2
seaborn                 0.12.2
setuptools              75.8.0
six                     1.17.0
stack_data              0.6.3
sympy                   1.13.1
tabpfn                  2.0.6
tabpfn-extensions       0.0.4       D:\miniconda3\envs\tabpfn\Lib\site-packages\tabpfn-extensions
threadpoolctl           3.5.0
torch                   2.6.0
tornado                 6.4.2
tqdm                    4.67.1
traitlets               5.14.3
typing_extensions       4.12.2
tzdata                  2025.1
urllib3                 2.3.0
wcwidth                 0.2.13
wheel                   0.45.1
zipp                    3.21.0

[iivalchev]

Oh, I totally missed that you are using the AutoTabPFNClassifier, well in this case getting hold on the underlying TabPFNClassifiers is somewhat more challenging.

For the default config you can inspect clf.predictor_.estimators_ and see if you can call `get_embeddings' on them.

Not quite familiar with the phes, perhaps there is better way.

[LeoGrin]

Thanks @iivalchev ! Yeah I don't think there's a better way for now :)

[lincj1994]

Thanks for the reply. I cannot call get_embeddings from clf.predictor_._estimators either.
Appreciate the effort though!

[lincj1994]

@iivalchev Hi! Could you clarify whether clf.get_embeddings(X_test, data_source='train') returns embeddings for both features and outcomes? When I ran this code, the output was a NumPy array with shape (n_estimators, number_of_samples, 3×(features + outcomes)). I’m trying to confirm if this includes both feature and outcome embeddings in the final dimension. Additionally, does the column order in the embeddings correspond to the input data? For instance, do embeddings 0, 1, and 2 represent the outcome, while embeddings 3, 4, and 5 correspond to the first feature?

[iivalchev]

@lincj1994 I would assume so. However I don't think this is the case by looking at the code what you get from get_embeddings is just the embedding for y as just before it being passed to the decoder head.

        test_encoder_out = encoder_out[:, single_eval_pos_:, -1].transpose(0, 1)

where encoder_out is of shape (batch_size, seq_len, n_features + 1 (for y), embedding_size)

Also keep in mind that you would not get 1:1 mapping of features of X to embeddings as two features are mapped into a single embeddings vector. Not to forget that after default pre-processing the number of features doubles.

Important to understand is what do you mean by embeddings? Is it just the linear transform that maps scalars (in this case two) to a high dimensional representation (192)? Or do you want those representations after they passed through the transformer? In both case a bit of work would be needed to get them out but is doable.

[lincj1994]

@iivalchev I'm a bit confused. Are you saying that the get_embeddings only returns embeddings for the target y (the value to predict) even though the annotation for get_embeddings states Get the embeddings for the input data XType?

[iivalchev]

@lincj1994 yeah me too :) but double checked it just now. What you get is just the embedding of y after passing through the transformer either y_trains or y_tests.

@LeoGrin can someone else confirm that, maybe I am missing something.

[noahho]

@lincj1994

When I ran this code, the output was a NumPy array with shape (n_estimators, number_of_samples, 3×(features + outcomes)). I’m trying to confirm if this includes both feature and outcome embeddings in the final dimension.
Additionally, does the column order in the embeddings correspond to the input data? For instance, do embeddings 0, 1, and 2 represent the outcome, while embeddings 3, 4, and 5 correspond to the first feature?

Embeddings include both X and y but there is no ordering in the embddings. Embeddings contain mixed information and while the shape depends on the features and outcomes there is no clear separation for where the information is stored. To reduce the embedding size to something smaller you can use PCA.

For a more convenient way to retrieve embeddings see the PR PriorLabs/tabpfn-extensions#39 in process. I hope we merge this today! It depends on the work done by @iivalchev and builds a.) a convenience wrapper b.) improves embeddings using a K-fold technique.

[iivalchev]

@noahho

Embeddings include both X and y but there is no ordering in the embddings. Embeddings contain mixed information and while the shape depends on the features and outcomes there is no clear separation for where the information is stored.

I find this even more confusing.
What is returned by

TabPFN/src/tabpfn/utils.py

Line 47 in 05ab7da

def _get_embeddings(

for train_embeddings is

TabPFN/src/tabpfn/model/transformer.py

Line 676 in 05ab7da

train_encoder_out = encoder_out[:, :single_eval_pos_, -1].transpose(0, 1)

encoder_out is the transformed embedded_input which is of shape b s f+1 e and embedded_input is

TabPFN/src/tabpfn/model/transformer.py

Line 617 in 05ab7da

embedded_input = torch.cat((embedded_x, embedded_y.unsqueeze(2)), dim=2)

so encoder_out[:, :single_eval_pos_, -1] is the transformed representation of y_train, right? I don't think it's shape depends on the shape of the features. You can argue that it holds information passed from X_train and from the MLP layers of the transformer. However as far as I understand @lincj1994 is also interested in encoder_out[:, :single_eval_pos_, :-1] values which are not returned.