Code for a study on predicting the runtime distributions (RTDs) of randomised algorithms. We assess Prior-Data Fitted Networks, specifically TabPFN [4], on this task and compare it against established baselines: Random Forests and Gaussian Processes [1], DistNet [2], and Bayes DistNet [3]. Our TabPFN approach sets a new state of the art for this problem.
This repository accompanies the AutoML 2026 paper “From Tables to Runtime: Predicting Algorithm Runtime Distributions with TabPFN.”
If you use this code in your research, please cite:
@inproceedings{ibrahimli2026tables,
title = {From Tables to Runtime: Predicting Algorithm Runtime Distributions with {TabPFN}},
author = {Ibrahimli, Hagverdi and Eggensperger, Katharina and Adriaensen, Steven},
booktitle = {Proceedings of the Fifth International Conference on Automated Machine Learning},
year = {2026},
series = {Proceedings of Machine Learning Research},
publisher = {PMLR}
}Carried out at the AutoML Lab, University of Freiburg, with the codebase implemented primarily by Hagverdi Ibrahimli as part of his MSc. research project.
- Motivation
- Task
- Data
- Models
- Metrics
- Experimental protocol
- Installation
- Downloading the data
- Repository structure
- Usage
- Command-line reference
- Reproducing the results
- Figures
- References
Runtime distributions are typically multi-modal and heavy-tailed, so standard regression is not enough.
An empirical runtime distribution, showing its multi-modal and heavy-tailed shape.
We use TabPFN to fit a posterior predictive distribution over runtime, which captures this structure.
TabPFN vs. Oracle Normal on a bimodal log runtime distribution (from spear_qcp). The unimodal Oracle Normal (red) misses the second mode; TabPFN's posterior predictive distribution (green) captures both, against the observed runtimes (blue crosses).
A randomised solver does not have one runtime per instance, it has a distribution over runtimes. Given the feature vector of an unseen instance, the goal is to predict that full distribution.
Every instance in the benchmark was attempted 100 times with independent random seeds, each attempt bounded by the scenario's cutoff. Those 100 measurements are its empirical RTD. A model trains on the instances of a training split, predicts a distribution for each held-out test instance, and is scored against that instance's 100 observed runtimes rather than their mean.
We use the benchmark released with DistNet [2], fetched by the script in Downloading the data.
A scenario pairs one randomised algorithm with one family of problem instances. It provides a numeric feature vector per instance (the inputs), 100 measured runtimes per instance (the targets), and the captime used when the data was generated. Scenario names and their on-disk paths are registered in data_source_release.py.
load_data.py reproduces the original DistNet cleaning. An instance is dropped if any of its runs crashed or timed out, if any runtime reached the cutoff (a censored measurement), if its feature vector is constant, or if its solver status falls outside the scenario's class. Every surviving instance therefore has 100 uncensored measurements, which is why the counts below are lower than in the raw archive. The -512 missing-value sentinel is replaced by the column median.
| Scenario | Domain | # Instances | # Features | Cutoff (s) |
|---|---|---|---|---|
| clasp_factoring | SAT | 2000 | 102 | 5000 |
| saps-CVVAR | SAT | 10011 | 46 | 60 |
| spear_qcp | SAT | 8072 | 91 | 5000 |
| yalsat_qcp | SAT | 11743 | 91 | 5000 |
| spear_swgcp | SAT | 11182 | 76 | 5000 |
| yalsat_swgcp | SAT | 11182 | 76 | 5000 |
| lpg-zeno | Planning | 3999 | 165 | 300 |
Counts are after cleaning. One cross-validation fold trains on about 90% of them, which matters when choosing --context_size under --subsample_unflattened.
All models solve the same task and are scored by the same code. They differ in which distributions they can express and how they are fitted.
| Model | --model |
Predicted distribution | Fitted by | Multi-modal |
|---|---|---|---|---|
| TabPFN [4] | tabpfn |
Learned piecewise density over bucket borders | In-context, one forward pass, no training | Yes |
| DistNet [2] | distnet |
Log-normal, or normal over log-runtime | SGD, maximum likelihood | No |
| Bayes DistNet [3] | bayesian_distnet |
Log-normal fitted to samples from a weight posterior | Bayes-by-Backprop | No |
| Random Forest [1] | random_forest |
Normal over log-runtime, from mean and predictive variance | Forest fit, optional SMAC tuning | No |
| Gaussian Process [1] | gp |
Normal over log-runtime | Exact GP, marginal likelihood | No |
| Log-normal | lognormal |
Single log-normal, features ignored | Closed-form moments | No |
- TabPFN is the method under study. It is never trained on this data: the training rows are passed as context and one forward pass gives a predictive distribution per test instance. That distribution is a learned piecewise density, not a member of a two-parameter family, so it more reliably captures the bimodal RTDs that the baselines struggle with. See tabpfn_helpers.py.
- DistNet is the reference neural baseline: a two-hidden-layer MLP (16 units, tanh, batch norm) predicting the two parameters of a log-normal. See distnet_lognormal.py.
- Bayes DistNet replaces DistNet's dense layers with Bayes-by-Backprop layers. Each forward pass samples one runtime;
n_enspasses define the predictive distribution. It targets better calibration when few observations per instance are available. See bayesian_distnet.py. - Random Forest is made distributional by the Law of Total Variance: the mean of the per-leaf variances plus the variance across tree means gives a predictive variance alongside the mean. It needs the patched scikit-learn that draws exact split points uniformly at random. See random_forest.py.
- Gaussian Process uses a scaled Matern-3/2 kernel plus a white-noise term, which absorbs the run-to-run variability of identical feature rows. Exact inference is cubic in the number of training rows, making it the bottleneck of the context-size sweep.
- Log-normal ignores the features and fits one distribution to the pooled training runtimes. It is the floor a feature-based model must beat, and is not part of the main comparison.
--oracle ignores the instance features and fits each test instance's own 100 observed runtimes. This upper-bounds what a model of that family can score on this data.
--model tabpfn --oraclefits TabPFN on the test instance's own runtimes, using constant dummy features.--model lognormal --oracleuses the per-instance mean and standard deviation of the observed log-runtimes.
Comparing the two oracles isolates the cost of the unimodal parametric form, since both are given the same information.
Runtimes span orders of magnitude, so every model is fitted on a transformed target chosen with --target_scale. The metric code inverts the transform, including the density's Jacobian, so scores from different target scales are directly comparable. This is what makes the scaler ablations meaningful.
--target_scale |
Transform | Supported by |
|---|---|---|
log |
log1p(y) |
tabpfn, distnet, random_forest (required), gp (required) |
max |
y / max(y_train), min-max with an implicit minimum of 0 |
tabpfn, distnet, bayesian_distnet (required) |
original |
none | tabpfn |
lognormal always works on natural log-runtime, but the CLI still requires the flag.
All metrics are lower-is-better, computed per test instance and then aggregated. See calculate_metrics.py.
Each prediction is first mapped into one shared space, z = log1p(y), so models fitted under different target scales are compared on equal terms. The distributional metrics are then integrated numerically on a per-instance grid of 15,000 points covering the union of the empirical support and the model's own (0.0001, 0.9999) quantile range, so neither an over-dispersed nor an over-confident prediction is truncated.
| Metric | Definition |
|---|---|
| NLLH | Primary metric. Mean negative log-density of the 100 observed runtimes, normalised per instance by that instance's largest observed log-runtime. Rescaling the log-runtime axis shifts a log-density by a constant, and this constant makes NLLH comparable across instances and scenarios whose runtimes differ by orders of magnitude. Densities are floored at 1e-10. |
| CRPS | Continuous ranked probability score against the 100 observations, using the exact decomposition ∫(F - F_emp)² + ∫F_emp(1 - F_emp). The second term is the intrinsic spread of the sample, a floor no model can beat. |
| Wasserstein | 1-Wasserstein distance, the integrated absolute CDF gap. |
| KS | Kolmogorov-Smirnov statistic, the largest absolute CDF gap. |
| MAE | Absolute difference between predicted and empirical median, in seconds. Reported for reference; it rewards central tendency and ignores distributional shape. |
Cross-validation. A single KFold(n_splits=10, shuffle=True, random_state=0) splits instances, never individual runs, so all 100 runtimes of an instance stay on one side and a test instance is genuinely unseen. The seed is fixed project-wide, so fold k holds the same instances for every model and scenario. Per-fold scores are therefore paired across models, which is what the signed-rank test used for context-size selection relies on.
Repetitions. Any run involving random subsampling (of context, features, or per-instance samples) is repeated with 5 seeds.
Aggregation. Every plot reduces scores in this order:
- average the per-test-instance values within a run;
- average over the repetition seeds, collapsing replicates that share a fold;
- average over the 10 folds for the plotted mean, and take the standard deviation across folds for the shaded band.
Averaging seeds before the fold statistics keeps the error bands a measure of cross-fold variability rather than of resampling noise.
Compute budgets. DistNet training and Random Forest tuning are each capped at one hour of wall-clock time per run, so the baselines get comparable compute. Epoch and trial counts are only safety caps. A curve ending in an x marker means the model hit its time limit or ran out of memory there.
One CLI invocation runs one (model, scenario, fold) point. Sweeping the grids is the caller's job, typically a cluster array job.
The environment is defined in pyproject.toml and pinned by the committed uv.lock.
Linux is the supported platform and the one the reported experiments ran in. Some dependencies are compiled from source at install time, so a C++ toolchain must be available:
| Platform | Requirement |
|---|---|
| Linux | gcc and g++. |
| Windows | Visual Studio Build Tools with the Desktop development with C++ workload, or WSL2. |
| macOS | Xcode command line tools (xcode-select --install). |
-
Install
uv:pip install uv
-
From the repository root:
uv sync
This installs the locked dependency versions, including the patched forks of TabPFN and scikit-learn pinned to exact commits, and PyTorch for CUDA 13.0. uv fetches CPython 3.13.0 automatically if needed.
Two notes:
- The TabPFN fork provides the predictive distributions this project consumes:
predict(..., output_type="full")must return the bucket logits together with the bar-distribution criterion. The scikit-learn fork draws exact split points uniformly at random, as the Random Forest of [1] requires. uv syncalso installs thedevgroup (pytest, black, ipykernel).ipykernelis what lets the notebooks run. Useuv sync --no-devto skip it.
A GPU is strongly recommended for TabPFN, and a cluster or batch environment for the full experiment grid.
The benchmark archive is not redistributed here. From the repository root:
python download_distnet_data.pyThis extracts the data into data/distnet_data/.
project_root_dir/
├── download_distnet_data.py # Downloads and extracts DistNet benchmark datasets
├── pyproject.toml # Package and dependency configuration (uv)
├── uv.lock # Fully resolved dependency versions
├── data/distnet_data/ # Benchmark data, created by the download script
├── notebooks/ # Analysis notebooks producing the figures
│ ├── feature_robustness.ipynb
│ ├── heatmap_factorial_design.ipynb
│ ├── predictive_performance_scaling.ipynb
│ └── wilcoxon_signed_rank_test.ipynb
├── results/ # Per-run metadata and compiled .pkl/.csv results
└── src/
└── tabpfn_project/
├── experiment_config.py # Typed description of a run; documents every knob
├── globals.py # Scenarios, metrics, scales, and budgets
├── paths.py # Project root and key directories
├── helper/
│ ├── bayesian_distnet.py # Bayes DistNet [3]
│ ├── calculate_metrics.py # NLLH, CRPS, Wasserstein, KS, MAE
│ ├── data_source_release.py # Scenario registry
│ ├── distnet_lognormal.py # DistNet [2]
│ ├── load_data.py # Loading, cleaning, and KFold splitting
│ ├── preprocess.py # Instance filters, imputation, standardisation
│ ├── random_forest.py # Random Forest baseline [1]
│ ├── tabpfn_helpers.py # Batched and oracle TabPFN inference
│ ├── utils.py # Subsampling, aggregation, and plotting
│ └── y_scalers.py # Target transforms
└── scripts/
├── main.py # CLI entrypoint
├── model_handler.py # Per-model training and evaluation
└── prepare_data.py # Builds train/test arrays for a run
Each run writes one pickle to results/<save_dir>/metadata/<experiment_id>_metadata.pkl holding its full configuration and metrics. TabPFN also writes its raw predictive distributions to results/<save_dir>/tabpfn_preds_full/, which are too large for the metadata. helper.utils.fetch_save_dict compiles a metadata directory into one flat .pkl list, which is what the notebooks load.
generate_experiment_id does not encode --do_hpo or --rf_new_default. Two Random Forest runs differing only in those flags produce the same filename and overwrite each other, so every model variant needs its own --save_dir.
Run a single fold from the repository root:
python -m tabpfn_project.scripts.main \
--scenario lpg-zeno \
--model tabpfn \
--fold 0 \
--context_size 128 \
--seed_context_size 100 \
--target_scale log \
--save_dir test_run_dirEach flag maps to one attribute of ExperimentConfig, whose docstring gives the full semantics.
| Flag | Type | Description |
|---|---|---|
--scenario |
choice | One of the seven scenarios above. |
--model |
choice | distnet, tabpfn, bayesian_distnet, random_forest, lognormal, gp. |
--fold |
int 0-9 | Cross-validation fold index. |
--target_scale |
choice | log, max, or original. See the support matrix above. |
--save_dir |
str | Sub-directory of results/ for this run's output. |
| Flag | Type | Default | Description |
|---|---|---|---|
--context_size |
int | all | Size of the training set given to the model. Counts rows of the flattened table, or instances with --subsample_unflattened. Requires --seed_context_size. |
--seed_context_size |
int | Seed for the context subsampling. | |
--subsample_unflattened |
flag | off | Subsample whole instances before flattening, without replacement, giving exactly context_size * num_samples_per_instance rows. Without it, rows are drawn after flattening with replacement, which allows context sizes beyond the natural training-set size. |
--num_samples_per_instance |
int 1-100 | 100 | Runtimes kept per training instance. Requires --seed_samples_per_instance when below 100. Test targets are never subsampled. |
--seed_samples_per_instance |
int | Seed for the per-instance runtime subsampling. | |
--remove_duplicates |
flag | off | TabPFN only. Keep one runtime per training instance, collapsing the duplicate feature rows. |
| Flag | Type | Default | Description |
|---|---|---|---|
--n_features_keep |
int | all | Number of randomly chosen features to keep. Takes precedence over --feature_drop_rate. 0 is meaningful: features become a single zero column, giving a featureless baseline that can only learn the scenario's marginal RTD. |
--feature_drop_rate |
float 0-1 | Fraction of features to drop. Ignored if --n_features_keep is given. |
|
--seed_feature_drop_rate |
int | Seed for feature subsampling. Required with either flag above. | |
--jitter_x |
flag | off | TabPFN diagnostic. Adds zero-mean Gaussian noise to the training features. Requires --jitter_val. Skipped in oracle mode. |
--jitter_val |
float | Relative noise intensity: column j gets noise of std jitter_val * std(column j). |
|
--rand_extend_x |
flag | off | TabPFN diagnostic. Appends noise columns to train and test features, breaking exact row duplication without adding signal. Requires --n_rand_cols. Skipped in oracle mode. |
--n_rand_cols |
int | Number of random columns to append. |
| Flag | Type | Default | Description |
|---|---|---|---|
--early_stopping |
flag | off | DistNet and Bayes DistNet. Holds out 20% of the training data with a GroupShuffleSplit on instance ids and restores the best-validation checkpoint. |
--do_hpo |
flag | off | Random Forest. Tunes with SMAC over a 3-fold GroupKFold NLLH objective, budgeted at one hour. Supersedes --rf_new_default. |
--rf_new_default |
flag | off | Random Forest. Uses n_estimators=50, var_min=1e-6 instead of the values of [1] (10, 0.01). |
--oracle |
flag | off | TabPFN and lognormal only. Fits the diagnostic upper bound described above. |
--use_cpu |
flag | off | Forces CPU execution even when CUDA is available. |
python -m tabpfn_project.scripts.main [args]Sweeps the context size for every model, showing how predictive quality scales with the amount of training data.
Base arguments: --scenario, --model, --fold, --context_size, --seed_context_size, --save_dir. Results land in RESULTS_DIR/<save_dir>.
Grid:
--context_size:2**iforiinrange(5, 17); forgpthe maximumiis 12--seed_context_size:i * 100foriinrange(1, 6)--fold:range(10)
Run each variant below across that grid, giving each its own --save_dir:
- TabPFN
- Variant 0:
--target_scale log - Variant 1:
--target_scale max - Variant 2 (noDUPS):
--target_scale log --remove_duplicates - Variant 3 (naive):
--target_scale log --n_features_keep 0 --seed_feature_drop_rate -1
- Variant 0:
- DistNet
- Variant 0:
--early_stopping --use_cpu --target_scale max - Variant 1:
--early_stopping --use_cpu --target_scale log
- Variant 0:
- Random Forest
- Variant 0:
--use_cpu --rf_new_default --target_scale log - Variant 1:
--use_cpu --rf_new_default --target_scale log --do_hpo - Variant 2:
--use_cpu --target_scale log
- Variant 0:
- Bayes DistNet
- Variant 0:
--early_stopping --use_cpu --target_scale max
- Variant 0:
- GP
- Variant 0:
--use_cpu --target_scale log
- Variant 0:
Sweeps the number of retained features down to zero, measuring how much each model exploits instance features rather than the scenario's marginal runtime distribution. Requires Experiment 1 to be complete.
First extract the optimal context size per variant, so each model is compared at its own operating point. Taking the best mean score would overfit fold noise and favour the largest context tried, so instead each candidate is compared to the empirical best with a one-sided Wilcoxon signed-rank test over the 10 paired per-fold scores, and the smallest candidate that is not significantly worse is chosen.
from tabpfn_project.helper.utils import fetch_save_dict, analyze_optimal_context
from tabpfn_project.paths import RESULTS_DIR
from tabpfn_project.globals import DISTNET_SCENARIOS
# 1. Compile the per-run metadata into one list (repeat per variant)
fetch_save_dict(
results_dir=RESULTS_DIR / "model_save_dir_name",
metadata_dir=RESULTS_DIR / "model_save_dir_name" / "metadata",
model_name="tabpfn", # matches the run's --model
model_file_name="tabpfn_var0", # output filename stem, distinguishes variants
scenario=None,
)
# 2. Extract optimal context sizes to CSV (repeat per variant)
analyze_optimal_context(
model_results_path=RESULTS_DIR / "model_save_dir_name" / "tabpfn_var0.pkl",
scenarios=DISTNET_SCENARIOS,
metric="NLLH",
output_csv_path=RESULTS_DIR / "model_save_dir_name" / "tabpfn_var0.csv",
)Do this for four variants: TabPFN 0, DistNet 0, Random Forest 0, and GP 0.
Then run those four with the extracted --context_size.
Base arguments: --scenario, --model, --fold, --context_size (from the CSV), --seed_context_size, --save_dir, --n_features_keep, --seed_feature_drop_rate
Grid:
--seed_context_size:i * 100foriinrange(1, 6)--seed_feature_drop_rate:i * 1000foriinrange(1, 6)--n_features_keep: per scenario, start at 0, then 1, doubling up to the scenario's feature count without passing it, then append the exact feature count. For 12 features:[0, 1, 2, 4, 8, 12]. At0and at the full count, use the single seed--seed_feature_drop_rate -1.
Crosses the number of training instances with the number of runtime samples per instance, separating whether performance comes from seeing more instances or from seeing each instance more often. TabPFN variant 0 only.
Base arguments: --scenario, --model, --fold, --context_size, --seed_context_size, --save_dir, --subsample_unflattened, --num_samples_per_instance, --seed_samples_per_instance
Grid:
--seed_context_size:200, 400, 600--seed_samples_per_instance:2000, 4000, 6000--num_samples_per_instance:[2**i for i in range(0, 7)] + [100]--context_size: 1-2-5 magnitude scaling,10, 20, 50, 100, 200, 500, 1000, ..., up to the scenario's instance count
Under --subsample_unflattened, --context_size counts instances drawn without replacement, so it must not exceed the fold's training instances, about 90% of the totals in the table above. For the full-data point at the top of each heatmap, omit --context_size and --seed_context_size; plot_ml_heatmap places such runs at their effective instance count.
The notebooks in notebooks/ build the figures from the compiled .pkl files. Each figure-producing cell is marked with a # Figures ... comment on its first line.
Cell references give the cell index (0-based position in the .ipynb cells array) and the code-cell ordinal (what the execution counter shows after a clean top-to-bottom run).
| Figure | Notebook | Cell index | Code cell | Plotting function |
|---|---|---|---|---|
| 1, 4, 7 | predictive_performance_scaling.ipynb |
18 | 9 | plot_main_results |
| 2 | predictive_performance_scaling.ipynb |
21 | 11 | plot_instance_level_bar_chart |
| 3(a), 8 | predictive_performance_scaling.ipynb |
7 | 4 | plot_main_results |
| 3(b), 11 | heatmap_factorial_design.ipynb |
5 | 3 | plot_ml_heatmap |
| 5, 10 | feature_robustness.ipynb |
5 | 3 | plot_feat_dropping_results |
| 6, 9 | predictive_performance_scaling.ipynb |
25 | 13 | plot_combined_time_vram_results |
| 12 | predictive_performance_scaling.ipynb |
9 | 5 | plot_main_results |
| 13 | predictive_performance_scaling.ipynb |
5 | 3 | plot_main_results |
| 14 | predictive_performance_scaling.ipynb |
11 | 6 | plot_main_results |
| 15 | predictive_performance_scaling.ipynb |
13 | 7 | plot_main_results |
| 16 | predictive_performance_scaling.ipynb |
15 | 8 | plot_main_results |
The notebooks read fixed paths, so the model_file_name and results_dir passed to fetch_save_dict must match these.
predictive_performance_scaling.ipynb, from Experiment 1:
File under results/ |
Variant |
|---|---|
tabpfn/tabpfn_scalerLOG.pkl |
TabPFN 0 |
tabpfn/tabpfn_scalerMAX.pkl |
TabPFN 1 |
tabpfn/tabpfn_scalerLOG_noDUPS.pkl |
TabPFN 2 |
tabpfn/tabpfn_naive.pkl |
TabPFN 3 |
tabpfn/tabpfn_oracle.pkl |
TabPFN with --oracle |
distnet/distnet_scalerMAX.pkl |
DistNet 0 |
distnet/distnet_scalerLOG.pkl |
DistNet 1 |
bayesian_distnet/bayesian_distnet_scalerMAX.pkl |
Bayes DistNet 0 |
random_forest/random_forest_newDEFS.pkl |
Random Forest 0 |
random_forest/random_forest_tuned.pkl |
Random Forest 1 |
random_forest/random_forest_oldDEFS.pkl |
Random Forest 2 |
gaussian_process/gp_scalerLOG.pkl |
GP 0 |
feature_robustness.ipynb, from Experiment 2: tabpfn/tabpfn_fdrop.pkl, distnet/distnet_fdrop.pkl, random_forest/random_forest_fdrop.pkl, gaussian_process/gp_fdrop.pkl.
heatmap_factorial_design.ipynb, from Experiment 3: tabpfn/tabpfn_heatmap.pkl.
wilcoxon_signed_rank_test.ipynb runs the context-size selection for Experiment 2 and writes its CSV. The committed cell is a worked example pointing at results/GaussianProcess/gp_results_exp_random_subsample_log.pkl; adapt the paths to your own --save_dir names.
- Hutter, F., Xu, L., Hoos, H. H., & Leyton-Brown, K. (2014). Algorithm Runtime Prediction: Methods & Evaluation. Artificial Intelligence, 206, 79-111.
- Eggensperger, K., Lindauer, M., & Hutter, F. (2018). Neural Networks for Predicting Algorithm Runtime Distributions. IJCAI 2018.
- Tuero, J. E., & Buro, M. (2021). Bayes DistNet: A Robust Neural Network for Algorithm Runtime Distribution Predictions. AAAI 2021, 35(13).
- Hollmann, N., Müller, S., Eggensperger, K., & Hutter, F. (2023). TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second. ICLR 2023.