CWhy/tabpfn-jax-ckpts

tabpfn-jax checkpoints

Model weights from a from-scratch replication of TabPFN v1 (Hollmann et al., ICLR 2023) in JAX + optax (no flax). Prior-data-fitted networks that do in-context learning over rows of a tabular dataset — no gradient updates at predict time.

These are params-only checkpoints (optimizer state stripped). Each is a Python pickle of {"params": <nested dict>, "cfg": <ModelConfig kwargs>, "step": int}.

The training/inference code is not public. This repo is weights + metrics only. Enough is documented here and in manifest.json to load and run a forward pass given a compatible model definition.

Results

Evaluated under the TabPFN v1 paper's protocol — 18 numerical OpenML-CC18 datasets (≤1,000 training points), 5× 50/50 splits, ROC AUC one-vs-one macro.

Checkpoint	Training prior	CC18 AUC	TabArena-Lite AUC
v1_d	TabICL mix_scm	0.9332	0.8401
v1_g	TabPFN v1's own (open) prior	0.9308	0.8341
v2_a	TabICLv2 graph SCM	0.9322	0.8433
m6_v2	TabICL mix_scm (3-stage arch)	—	0.835
Official TabPFN v1	—	0.9340	—
Published (ICLR'23)	—	0.934 ± 0.0086	—

v1_d matches the published v1 result to within the paper's own ±0.0086 std.

Scope

Parity is on small, clean, numerical classification (v1's home turf, where v1 and v2 are within 0.002 of each other). On the larger, messier TabArena-Lite benchmark these models trail TabPFN v2 (0.875) by ~0.03 — a gap that scales with dataset size and traces to training context length (T=1024), not the prior.

Load

import pickle, jax
ck = pickle.load(open("v1_d.params.pkl", "rb"))
params, cfg, step = ck["params"], ck["cfg"], ck["step"]
# feed to a compatible forward(params, ModelConfig(**cfg), X, y, feat_mask, train_size)

See manifest.json for each model's prior, metrics, and parameter count.