Finetune regression with high-dimensional features

[ivandon15]

Hi team,

Thank you very much for the excellent work on TabPFN! It has been extremely helpful in my research.

I have a question regarding fine-tuning TabPFN for regression with high-dimensional pretrained embeddings.
I am working with protein-related data where each sample is represented by embeddings from pretrained models, typically 1000–2000 features per sample.

Using standard TabPFN in-context inference works well and is memory-efficient.
However, when I try to use FinetunedTabPFNRegressor, GPU memory usage explodes during training (e.g. in the KV projection step such as
torch.einsum("... s, j h d s -> ... j h d", x_kv, w_kv)), exceeding 48 GB VRAM (RTX Ada 6000).

My questions are:
Are there recommended best practices for fine-tuning TabPFN with high-dimensional pretrained embeddings?
Is dimensionality reduction generally expected before fine-tuning, and are there suggested target dimensions?
Or is fine-tuning simply not recommended for such high-dimensional inputs, with in-context inference being the intended use case?

Any guidance would be greatly appreciated.
Thanks again for your great work!

[bejaeger]

Hi @ivandon15 ,

Great to see that you check out our new finetuning pipelines. I'll try to go through your questions one by one below:

Our current model is quite memory hungry and scales linearly with increasing feature size and sample size. So it's quite normal that a forward/backward pass can take multiple GBs. Things you can do to reduce memory usage:

• Use activation checkpointing: Set use_activation_checkpointing to True if you haven't yet done so.
• Reduce the size of the meta-datasets by setting n_finetune_ctx_plus_query_samples to lower values. This will make the initial performance of the model worse, but with fine-tuning you might recover what the reduced context size might have lost.
• If that doesn't work (or you have to reduce the context size by too much) you unfortuntaly would need to find a GPU with more VRAM.

Regarding your other questions:

• Dimensionality reduction can help in some cases, but not necessarily. Largely depends on your embeddings and how dense the representations are.
• For these kind of high-dimensional inputs we recommend finetuning as an option to improve performance. The model has seen very few of such samples during pre-training, so tuning it to that data domain could give big gains.

One additional pointer I can give you is this PR: #711
We introduce a new loss which is expected to give big improvements.

I hope this info helps, closing the PR for now but feel free to re-open if there are more things you would like to know.