Access attention weights/scores

[ruid7181]

Describe the workflow you want to enable

Thank you for the wonderful work of TabPFN!

I would like to check the attention pattern learned by TabPFN from a tabular data, and explore the potential sparsity in the attention towards different tokens. Can you provide guides how I can access the attention weights (in each attention layer of TabPFN)?

Thanks,
Rui

Describe your proposed solution

I noticed the method get_embeddings, which returns the feature vectors, can I use this to calculate the softmaxed cosine similarity as attention weights?

Describe alternatives you've considered, if relevant

No response

Additional context

I'm using the model tabpfn.regressor.TabPFNRegressor locally.

Impact

Medium (Significant enhancement)

[LeoGrin]

Hello @ruid7181 !

Looking at the cosine similarity between embedding will give you a related metric but not exactly the attention weights.
To get the attention weights, you can look at the .model_ attribute (created after fit) and add PyTorch hooks to it (you can see the model architecture in tabpfn.model.transformer).