Graph Sampling Priors for nanoTabPFN

Dominik Scheuer, Kai Nipken, Eric Pfitzenmaier, University of Freiburg. Supervision: Magnus Bühler

Tabular Prior-Data Fitted Networks such as nanoTabPFN [1] are pretrained on synthetic data drawn from structural causal models (SCMs), and every SCM needs an underlying graph. Erdős–Rényi random DAGs are the de-facto default, yet their optimality has never been tested. This project asks:

How does the inductive bias of the sampled DAG family affect downstream generalization?

We pretrain nanoTabPFN across four DAG samplers, tune each one with its own random search (100 configs per sampler), and evaluate on real data (TabArena [2]) under a shared nested-CV protocol.

Results

Prior	mean ROC AUC (38 tasks)	Δ vs ER (32 tasks)	beats ER on	p
Erdős–Rényi (baseline)	0.716 ± 0.122	n/a	n/a	n/a
Barabási–Albert (hierarchical)	0.702 ± 0.127	−0.016	9/32	0.006
Barabási–Albert (random topology)	0.683 ± 0.132	−0.033	6/32	< 0.001
Stochastic Block Model	0.674 ± 0.133	−0.042	4/32	< 0.001
Uniform mixture (¼ each)	0.736 ± 0.121	+0.013	18/32	0.012
ER-heavy mixture (0.4 ER, 0.2 each)	0.685 ± 0.122	−0.039	3/32	< 0.001

• No single alternative prior beats ER. Every structured single-family prior loses to the Erdős–Rényi baseline, and the performance ordering of the alternative priors is stable: BA-hierarchical > BA-random > SBM.
• A uniform mixture of all four does beat ER. Diversity across graph families helps where any one family alone hurts, but the weights matter: making the mixture ER-heavy eliminates the performance benefit.

Absolute AUCs are the concatenated out-of-fold ROC AUC averaged over all 38 TabArena classification tasks (± std across datasets, not a confidence interval). Δ, win counts and p-values are per-dataset paired comparisons on the 32 tasks that survive the minority-class filter (the 6 tasks with fewer than 10 minority-class rows in a test fold are dropped).

1. Setup: sampler -> SCM -> synthetic table -> nanoTabPFN

The pretraining pipeline is fixed; the graph sampler is the only part we swap. Each batch draws one DAG, instantiates an SCM over it (per-node mechanisms + exogenous noise), and emits batch_size independent tables from that same graph, so structure is shared within a batch, noise is not.

The four graph samplers

Defined within graphtabpfn/prior/graph.py, every sampler follows the same interface. They must output a DAG featuring a target node relabelled as "y", along with per-node feature dimensions and a dropout fraction representing hidden, confounding columns.

Sampler	key	Structure it induces
Erdős–Rényi (baseline)	erdos_renyi_baseline	Random topological order + Bernoulli(edge_prob) edges. Homogeneous degrees, no hubs, no hierarchy: a broad, unbiased generalist.
Barabási–Albert, random topology	barabasi_albert_random_topo	Preferential attachment builds hubs, then a random topological order orients the edges. Hubs therefore mostly become colliders: a hub is an effect of many columns rather than a driver of many.
Barabási–Albert, hierarchical	barabasi_albert_age_constrained	Same hub-building, but edges always point old → young. Hubs become sources: a few early drivers fan out into many measured columns, and the target sits inside the hierarchy with both causes and effects.
Stochastic Block Model	stochastic_block_model	Nodes partitioned into communities with p_in ≫ p_out, edges oriented u < v for acyclicity. Produces semantic silos: tightly coupled feature blocks with sparse cross-talk, cascading from early to late blocks.

To ensure the model learns bidirectional causal inference, every sampler strictly requires the target to be an interior node (possessing a visible parent, a visible child, and a visible column). Because ER and BA-random randomize their topological order, their target can be reliably drawn from a fixed node index. Conversely, BA-hierarchical and SBM explicitly tie node indices to their causal hierarchy; to ensure the target varies in causal depth across batches, these samplers dynamically select their target from a uniformly sampled pool of eligible interior nodes.

What the priors actually look like

Structural statistics of the graphs as the trained models see them, with parameters drawn per graph from each prior's tuned HPO config (logs/hpo/<prior>.best.json) on top of the fixed data-shape ranges. The four panels separate the families cleanly: in/out-degree asymmetry, largest hub per graph, community modularity, and causal depth (longest path). Regenerate with uv run python poster/graph_stats_hpo.py; poster/graph_stats.py is the untuned counterpart.

2. Methodology

Architecture & training

• nanoTabPFN: 4 layers, 128-d embedding, 4 heads, 512-d MLP (~1.13 M parameters).
• 300 epochs for the HPO proxy budget, 1000 epochs for the final models (100 steps/epoch, batch 4 -> 100,000 graphs, 400,000 tables).
• Training pipeline adapted from modded-nanotabpfn [3], TFM-Playground [4] and the Do-PFN prior [5].
• Trained on a single NVIDIA A100 (40 GB).

Evaluation protocol: nested CV

• TabArena classification tasks, subsampled to ≤ 1000 rows / ≤ 100 features (protocol reused from modded-nanotabpfn [3]).
• Outer: stratified 5-fold per dataset (80/20 train/test) -> unbiased test estimate.
• Inner: each outer-train split 80/20 into train/validation -> config selection.
• Splits are precomputed once and stored on disk (logs/hpo/splits_seed42_k5_h02.json), so every prior and every config is scored on byte-for-byte identical folds.
• Metric: concatenated out-of-fold ROC AUC per dataset, then averaged over datasets.
• Pretraining never touches TabArena rows: the model is trained purely on the synthetic prior and sees dataset rows only in-context at fit() time. That is why search models can be shared across outer folds without leakage, avoiding the need to multiply the hyperparameter search budget by K.

Hyperparameter search

scripts/hpo/run_hpo.py runs 100 random-search configs per sampler in two phases:

1. Search: every config is pretrained once at the cheap proxy budget (--epochs-search 300) with a shared seed and scored on each fold's inner-validation split.
2. Final: the best config is re-selected per outer fold, retrained at the full budget (--epochs-final 1000) with a per-fold seed (set_seed(outer_fold)), and scored on the held-out outer test.

Fixed (pinned to the TabArena eval budget, graphtabpfn/hpo.py): num_nodes ~ U(5,35), features_per_node ~ U(1,3), train/test sample counts, max_classes = 10, and the four SCM noise parameters. Tuned: the density parameters per family (ER edge_prob range, BA num_edges range, SBM num_blocks / p_in / p_out ranges), the feature-dropout upper bound, and the class-binning upper bound.

Runs are resumable: search trials append to logs/hpo/<sampler>.search.jsonl, outer folds to logs/hpo/<sampler>.cv.json, and the frozen config plus aggregates land in logs/hpo/<sampler>.best.json.

3. Single graph-sampler evaluation

Per-dataset Δ ROC AUC of each contender prior against the Erdős–Rényi baseline, sorted by ER-relative performance. Error bars are ±1 std from a paired bootstrap of the pooled Δ (1000 resamples).

Finding: The ER baseline beats every other single graph type. BA-hierarchical is the closest contender (mean Δ −0.016, 9/32 wins), consistent with the idea that hubs-as-sources is the least unrealistic of the structured biases, while random-topology BA (hubs as colliders) and SBM (modular silos) lose substantially. The datasets where the structured priors do win (Website Phishing, Student Dropout, Amazon Employee Access) are the same across all three, hinting that the gain is dataset-driven rather than family-specific.

4. Graph-sampler mixture prior

A single nanoTabPFN is trained on a mixture of all four graph types: each batch draws one family by weighted random sampling, so mixing happens at the dataset level, not inside a single graph. Each family contributes its own frozen HPO-selected config, and everything else (splits, budget, per-fold seeds) is identical to the single-prior final phase.

• Uniform mixture (¼ each): mean Δ +0.013, better than ER on 18/32 datasets, p = 0.012.
• ER-heavy mixture (0.4 ER, 0.2 each other): mean Δ −0.039, better on only 3/32, p < 0.001.

Finding: A uniform mixture outperforms all individual priors. The performance improvement stems from the overall diversity of the graph families rather than the inclusion of one specific topology. The failure of the ER-heavy mixture highlights that skewing the distribution toward the strongest baseline collapses this benefit, resulting in worse performance than using the pure ER prior alone.

5. Future work

• Non-acyclic graphs: drop the DAG constraint and incorporate cyclical feedback loops into the prior.
• Curriculum learning over graph types: schedule the mixture weights during pretraining instead of holding them fixed.

Repository layout

graphtabpfn/            library
├── prior/graph.py      the four DAG samplers
├── prior/scm.py        SCM + mechanisms hung on the DAG
├── prior/dataloader.py graph -> synthetic classification batches
├── prior/config.py     PriorConfig (all sampler/SCM parameters)
├── model.py, train.py  nanoTabPFN + training loop
├── evaluation.py       TabArena tasks, subsampling, nested splits, scoring
├── hpo.py              fixed vs tuned params, random-search config sampler
└── interface.py        sklearn-style fit/predict wrapper

scripts/
├── hpo/                precompute_splits.py, run_hpo.py, hpo_on_cluster.sh
├── ablation/           mixture-prior ablation (see its README)
└── pretrain/           standalone pretrain.py / evaluate.py

poster/                 every poster figure + the script that generates it
logs/                   hpo/, ablation/, poster_significance.json

Setup

uv venv
uv pip install -e .

Reproducing the results

# 1. shared nested-CV splits (once; every prior is compared on these exact folds)
uv run python scripts/hpo/precompute_splits.py --k-outer 5 --split-seed 42

# 2. per-prior random search + final phase (one run per sampler)
uv run python scripts/hpo/run_hpo.py --graph-sampler erdos_renyi_baseline            --num-configs 100
uv run python scripts/hpo/run_hpo.py --graph-sampler barabasi_albert_age_constrained --num-configs 100
uv run python scripts/hpo/run_hpo.py --graph-sampler barabasi_albert_random_topo     --num-configs 100
uv run python scripts/hpo/run_hpo.py --graph-sampler stochastic_block_model          --num-configs 100

# 3. mixture ablation (reuses the frozen per-prior configs; --check runs offline)
uv run python scripts/ablation/run_mixture_ablation.py --tag mixture_uniform
uv run python scripts/ablation/run_mixture_ablation.py --tag mixture_er_heavy --weights 0.4,0.2,0.2,0.2

# on the cluster: one job per prior / per mixture
mkdir -p logs/hpo
GRAPH_SAMPLER=barabasi_albert_random_topo sbatch scripts/hpo/hpo_on_cluster.sh
sbatch scripts/ablation/mixture_on_cluster.sh

Standalone pretraining without the HPO harness:

uv run python scripts/pretrain/pretrain.py --epochs 100 --steps 100 --eval tabarena --eval-every 10
uv run python scripts/pretrain/evaluate.py --checkpoint workdir/nanotabpfn.pth

Poster figures

Each figure is self-contained and reads straight from logs/:

uv run python poster/pipeline_priors.py   # pipeline + the four priors (intro banner)
uv run python poster/graph_stats_hpo.py   # structural fingerprints at tuned configs
uv run python poster/hero.py              # single-prior Δ vs ER
uv run python poster/mixture.py           # mixture Δ vs ER
uv run python poster/significance.py      # -> logs/poster_significance.json

Lint / format

uv run ruff check graphtabpfn scripts
uv run ruff format graphtabpfn scripts

Acknowledgement

The authors acknowledge the support by the state of Baden-Württemberg through bwHPC and by GCP through education credits.

References

1. A. Pfefferle, J. Hog, L. Purucker, and F. Hutter. nanoTabPFN: A Lightweight and Educational Reimplementation of TabPFN. arXiv:2511.03634, Code

2. N. Erickson, L. Purucker, A. Tschalzev, D. Holzmüller, P. Mutalik Desai, D. Salinas, and F. Hutter. TabArena: A Living Benchmark for Machine Learning on Tabular Data. NeurIPS 2025 (Datasets & Benchmarks, Spotlight). arXiv:2506.16791, Code

3. S. B. Ozturk, A. Pfefferle, and F. Hutter. Speedrunning Tabular Foundation Model Pretraining. arXiv:2606.03681, Code

4. AutoML Freiburg. TFM-Playground. Code

5. Ole Ossen. Do-PFN Prior. Code