modded-nanoTabPFN

Disclaimer: The training code has gotten a bit crowded. For a simpler starting point, check out commit b0f29b7 (git log --oneline -- train_nano.py)

This repository hosts the nanoTabPFN speedrun, in which we (collaboratively|competitively) search for the fastest way to use 1 NVIDIA L40S to train a tabular foundation model (nanoTabPFN) that beats Random Forest on TabArena datasets.

The code is derived from nanoTabPFN with the inspiration of modded-nanogpt.

This repo now contains a training algorithm which attains the target performance in:

• 0.79m minutes on 1xL40S (baseline needed 74.32)
• 3648 synthetic datasets (baseline needed 80576)

This improvement in training speed has been brought about by the following techniques:

• Muon optimizer
• Batched Muon zeropower update for grouped QKV matrices
• Scaled Dot-Product Attention rewrite with explicit QKV
• Pre-norm transformer blocks
• Compile TransformerEncoderLayer forward
• bfloat16 autocast in training and inference
• Set float32 matmul precision to high
• Increase learning rate from 1e-4 to 1e-3
• Increase embedding size from 192 to 256
• Reduce attention heads from 6 to 4
• Exponential decay of residual stream across layers
• Lower precision RMSNorm
• Prepend 24 learnable Thinking Rows
• Latest Weight Averaging of 10 checkpoints
• Set AdamWScheduleFree weight decay to 0.01
• Repeated feature grouping (group size 5)
• Reduce transformer layers from 6 to 5
• Increase batch size from 1 to 2
• Add decoupled weight decay 0.1 to Muon
• Increase Muon momentum from 0.95 to 0.96
• Increase gradient clip from 1.0 to 2.0
• Feed mean of test feature embeddings into output decoder
• Shape-grouped zeropower update across all Muon matrices
• Multi-tensor momentum and weight update in Muon
• Producer-thread dataloader with pinned memory
• Move the NaN check to the CPU
• Single SDPA call over all datapoint queries

Running the current record

To run the current record, run the following commands.

git clone https://github.com/borawhocodess/modded-nanotabpfn.git
cd modded-nanotabpfn
uv sync
./run.sh

Record history

The following is the historical progression of speed records for the following competitive task:

Pretrain a neural network to ≤0.8068462330697953 validation average ROC AUC on subsampled TabArena using 1 NVIDIA L40S.

Note: The 0.8068462330697953 target was selected to match the performance of Random Forest on the same subsampled TabArena evaluation.

Rules

New records must:

• Not modify the evaluation pipeline.
• Not load any pretrained weights.
• Run faster than prior record when baselined on the same hardware with the same seed.

Other than that, anything and everything is fair game!

Evaluation details

Evaluation is on all of 38 TabArena classification tasks.

• Subsampling:
  • if >100 features, randomly select 100
  • if >1000 rows, randomly select 1000 (stratified by class labels)
• Cross-validation:
  • 5-fold StratifiedKFold with shuffling
  • class labels are encoded with integers per fold
• Preprocessing (per fold, fit on train only):
  • constant columns: dropped
  • numeric columns: numeric coercion + mean imputation
  • categorical columns: ordinal encoding + most-frequent imputation
• Metric:
  • binary or one-vs-rest ROC AUC
  • average over all tasks

Citation

This repo now has an accompanying paper, Speedrunning Tabular Foundation Model Pretraining (arXiv), accepted (OpenReview) at the FM4SD workshop at ICML 2026.

@misc{ozturk2026speedrunningtabularfoundationmodel,
      title={Speedrunning Tabular Foundation Model Pretraining},
      author={Salih Bora Ozturk and Alexander Pfefferle and Frank Hutter},
      year={2026},
      eprint={2606.03681},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2606.03681},
}

References

1. S. Müller et al. Transformers Can Do Bayesian Inference. arXiv preprint (2021).
2. N. Hollmann et al. TabPFN: “A Transformer That Solves Small Tabular Classification Problems in a Second”. Presented at ICLR (2022)
3. N. Hollmann et al. Accurate predictions on small data with a tabular foundation model. Nature 637, 319–326 (2025).
4. N. Erickson et al. TabArena: A Living Benchmark for Machine Learning on Tabular Data. arXiv preprint (2025).
5. A. Pfefferle et al. nanoTabPFN: A Lightweight and Educational Reimplementation of TabPFN. arXiv preprint (2025).
6. K. Jordan et al. Muon: An optimizer for hidden layers in neural networks (2024).
7. R. Xiong et al. On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745.
8. PyTorch docs: torch.nn.functional.scaled_dot_product_attention.
9. PyTorch docs: Automatic Mixed Precision (torch.autocast).
10. PyTorch docs: torch.set_float32_matmul_precision.
11. PyTorch docs: torch.compile.
12. PyTorch docs: Dealing with Recompilations.
13. B. Zhang & R. Sennrich. Root Mean Square Layer Normalization. NeurIPS 2019. arXiv:1910.07467.
14. L. Grinsztajn et al. TabPFN-2.5: Advancing the State of the Art in Tabular Foundation Models. arXiv:2511.08667 (2026).
15. J. Kaddour. Stop Wasting My Time! Saving Days of ImageNet and BERT Training with Latest Weight Averaging. arXiv:2209.14981 (2022).
16. S. Sanyal et al. Early Weight Averaging meets High Learning Rates for LLM Pre-training. arXiv:2306.03241 (2023).
17. I. Loshchilov & F. Hutter. Decoupled Weight Decay Regularization. ICLR 2019. arXiv:1711.05101.
18. A. Defazio et al. Schedule-Free Learning: A New Way to Train. arXiv:2405.15682 (2024).
19. J. Qu et al. TabICLv2: A better, faster, scalable, and open tabular foundation model. arXiv:2602.11139 (2026).
20. A. Karpathy. autoresearch: AI agents running research on single-GPU nanochat training automatically (2026).