Chunk-TabPFN

This repository contains an extension to TabPFN with chunked attention implementation, benchmarking pipelines on TabArena, and supporting scripts/notebooks.

---

🚀 How to use

We use uv for environment management.

uv sync

This installs all dependencies specified in pyproject.toml.

Make sure your PYTHONPATH env var is set to the project root:

export PYTHONPATH=$PYTHONPATH:$(pwd)

---

📂 Main contributions

Our contributions live in src/tabpfn/extensions:

• For math kernel implementation — pure PyTorch.
• For efficient attention implementation — memory- and compute-efficient attention.

---

🧪 Running experiments on TabArena

We provide scripts to run different scales of experiments:

• run_single_tabarena_large_task_manual.py — run on a large TabArena dataset with a given grid of context lengths, recording all the metrics and performance.
• run_single_tabarena_task.py — run on a single TabArena task.
• tabarena_leaderboard.py — aggregate results into leaderboard format.

Results for large datasets are saved under results/.

Results for TabArena are stored in ./tabarena_benchmarking_examples/tabarena_minimal_example/custom_tabpfn.

---

📓 Notebooks

For debugging and experimentation:

• demo_notebook.ipynb — quick demo of pipeline.
• notebook_analyze_parity.ipynb — parity analysis and long-sequence plotting.

Figures generated by notebooks are stored under figures/.

---

📖 Citing

This work builds on the following excellent repositories:

• TabRepo
• TabArena
• TabPFN
• ChunkLlama

Please cite them if you use this codebase in your research.