priorforge 0.1.0

Reusable data-generating-process (DGP) priors for Prior-Fitted Network (PFN) training

PriorForge

Reusable data-generating process (DGP) priors for Prior-Fitted Network (PFN) training.

Every PFN paper hand-crafts its own DGP from scratch — the structural generator that turns sampled latent structure into a dataset (X, y). That generator is the hard, paper-defining part: an SCM (TabPFN), a Gaussian process (PFNs4BO), an Error-Trend-Seasonality process (ForecastPFN), or a covariate→treatment→outcome chain (CausalPFN). PriorForge ships these generators as composable components, so you import and train instead of rebuilding the prior each time.

Installation

pip install priorforge

For GPU support with a specific CUDA version, install PyTorch separately first:

# CUDA 13.0
pip install torch --index-url https://download.pytorch.org/whl/cu130
pip install priorforge

Quick Start

A DGP emits synthetic datasets; PriorDataLoader turns them into the (context, query) splits a PFN trains on.

import torch
from priorforge import SCMPrior, PriorDataLoader, BarDistribution

dgp    = SCMPrior(n_features=5)                 # TabPFN-style structural prior
loader = PriorDataLoader(dgp, n=128, batch=32)  # context/query batches
bar    = BarDistribution(n_bins=1000)           # output head

batch = loader.sample_batch()
ctx_x, ctx_y, qry_x, qry_y = batch              # raw (X, y) context — no summary stats
# logits     = model(ctx_x, ctx_y, qry_x)
# boundaries = bar.fit_boundaries(ctx_y[0])
# loss       = bar.nll(logits, qry_y, boundaries)

Generators

Each generator targets a distinct research community. All share the DGP interface: sample_dataset(n) -> (X, y), plus latents() and prior_predictive() for diagnostics.

Generator	Paper lineage	Use case
SCMPrior	TabPFN	tabular classification / regression (sparse layered DAG)
MLPSCMPrior	TabPFN-v2 / GraphPFN	tabular (dense random-MLP neural SCM)
GPPrior	PFNs4BO	Bayesian optimization, regression w/ UQ
ETSPrior	ForecastPFN	zero-shot time-series forecasting
CausalDGP	CausalPFN	treatment-effect estimation (back-door)
FrontDoorDGP	CausalFM	effect estimation via the front-door criterion
IVDGP	CausalFM	effect estimation via an instrumental variable
MixtureDGP	Mitra	mixed synthetic priors (a prior over the priors above)

from priorforge import GPPrior, ETSPrior, CausalDGP, MixtureDGP, SCMPrior, MLPSCMPrior

GPPrior(n_features=2, input_warping=True)       # ARD Matérn-3/2 + linear + warping
ETSPrior(periods=(7.0, 365.25))                 # trend × seasonality × Weibull noise
CausalDGP(n_covariates=5, confounding=1.0)      # propensity → treatment → outcome

# Mitra finding: a curated mixture generalises better than any single prior
MixtureDGP([SCMPrior(n_features=8), MLPSCMPrior(n_features=8)], weights=[0.5, 0.5])

For the causal generators, latents() carries the ground truth needed to train and score an effect-estimation PFN: CausalDGP exposes y0, y1, cate, ate, propensity; FrontDoorDGP adds the (hidden) confounder/mediator and a closed-form ate with the biased naive_diff; IVDGP exposes the structural tau, the biased naive_ols, and the first_stage_r2 instrument strength.

Output heads

import torch
from priorforge import BarDistribution, MartingalePosterior

bar  = BarDistribution(n_bins=1000)             # discretized predictive head
post = MartingalePosterior(bar)                 # posterior over functionals
draws  = post.quantity(logits, boundaries, fn=torch.mean)   # predictive resampling
lo, hi = post.credible_interval(draws, alpha=0.9)

Training a PFN

A reference in-context regressor ships with the library, so you can go from a generator to a trained model in one call — and see the prior actually teach a transformer to predict held-out points.

import torch
from priorforge import SCMPrior, train_pfn

dgp    = SCMPrior(n_features=5)
result = train_pfn(dgp, steps=300, batch=16, n_bins=100, generator=torch.Generator().manual_seed(0))
print(result.first_loss, result.last_loss)      # BarDistribution NLL drops over training

PFNTransformer consumes raw (context_x, context_y, query_x) with the standard PFN attention mask (queries never see other queries) and predicts over a BarDistribution head. Swap SCMPrior for any regression generator (GPPrior, MLPSCMPrior, MixtureDGP, …). See examples/train_pfn.py. Training auto-selects a GPU when one is available (device=None); pass device="cpu" to force CPU.

NanoTabPFN — a known-outcome reproduction

A faithful reimplementation of NanoTabPFN (Pfefferle et al., 2025) ships as a classification example. Its 150×5×2 SCM-style training prior is exactly SCMPrior(task="classification"), so the whole loop lives in the library:

from priorforge import SCMPrior, train_nanotabpfn

dgp    = SCMPrior(n_features=5, task="classification", n_classes=2)
result = train_nanotabpfn(dgp, steps=500)   # 356,066 params — matches upstream
# held-out query accuracy climbs ~0.5 -> ~0.8

The model uses TabPFN-v2's two-axis attention (feature attention across columns, datapoint attention across rows, with test rows attending to train rows only). Two documented deviations from upstream — AdamW instead of schedulefree.AdamWScheduleFree, plus an explicit LR warmup that ScheduleFree would otherwise provide (without it this post-norm transformer collapses to chance). See examples/train_nanotabpfn.py.

Zero-shot transfer to real datasets

The payoff demo. train_nanotabpfn_multi pretrains one model across varied SCM geometries (feature counts 2–12, binary tasks), then nanotabpfn_predict runs a single forward pass per dataset — no gradient steps, no per-dataset fitting — on real sklearn tables:

from priorforge import train_nanotabpfn_multi, nanotabpfn_predict

model = train_nanotabpfn_multi(steps=2000, feature_range=(2, 12),
                               class_range=(2, 2), n_classes=2).model
pred  = nanotabpfn_predict(model, train_x, train_y, test_x, n_classes=2)

Trained only on synthetic priors (it never sees real data), it lands ~2 pts behind the best fitted logistic-regression / KNN / decision-tree / random-forest baseline on breast-cancer, iris, wine and digits — competitive in-context with no fitting. See examples/eval_real_datasets.py.

Prior breadth matters. train_nanotabpfn_multi takes a dgp_factory (k_feat, k_cls) -> DGP hook, so you can broaden the pretraining prior in one line — e.g. mix the sparse-DAG SCMPrior, the dense MLPSCMPrior, and per-step MixtureDGP blends. Broadening the pretraining prior halves the zero-shot gap to the best fitted model — direct evidence that prior design, not the training loop, is the lever.

Faster training via prefetch. Profiling showed CPU data generation is ~66% of each step (the host→device copy is negligible), so the win is overlapping it with GPU compute, not porting the generator to CUDA. Pass num_workers=N to generate batches in background processes:

train_nanotabpfn_multi(steps=3000, feature_range=(2, 12), num_workers=2, amp=True)
# num_workers: ~3.2x (8.4 -> 26.8 steps/s), now GPU-bound. Default 0 keeps the
#              serial, reproducible path.
# amp=True:    bf16 autocast, a further ~15% (no accuracy change).

(torch.compile was tried and dropped — for this 356k model with per-step varying feature widths it gave no steady-state speedup over eager bf16 while adding heavy per-width compilation cost.)

Verification

Three levels of checks:

from priorforge import verify_dgp, prior_predictive_fidelity, verify_distribution

verify_dgp(SCMPrior(n_features=5), verbose=True)     # DGP contract (shapes, task, repro)
prior_predictive_fidelity(dgp, reference_y)          # does the prior cover real data?
verify_distribution(LogNormal(), params)             # leaf-level 5-level protocol

Distributions (observation/noise leaves)

The distribution library is a supporting layer — the noise/observation leaves used inside a generator.

Kind	Classes
Primitives	LogNormal, NegativeBinomial, BetaBin
Wrappers	ZeroInflated(base)
Combinators	Mixture([d1, d2], weights)
Aliases	ZILN, ZINB

Distributions inheriting CausalDistribution expose sample_with_effect for matched (control, treatment) pairs — a leaf used inside a causal DGP:

from priorforge import ZILN
ctrl, trt = ZILN().sample_with_effect(params, N_control=500, N_treatment=500, tau=0.2)

Design

• Structural generators are the product; the distribution catalog is a supporting layer. See REDESIGN.md.
• Raw-context path: PFNs consume raw (X, y) context, not summary statistics. Sufficient statistics are an opt-in leaf transform.
• PyTorch only: the Triton two-backend plan is deferred (REDESIGN §7).
• Reproducible: every generator threads a torch.Generator.

License

MIT