PaperlayerView exact revision on arXiv

Source-linked AI summary

LinearPFN: Amortized Variable Selection for Linear Models with Interactions

Louis Schiekiera, Max Zimmer, Christophe Roux, Manuel Arnold, Sebastian Pokutta, Fritz Günther

arXiv:2610.06580v1cs.LGstat.ML

TL;DR

Spike-and-slab regression provides inclusion probabilities for candidate effects, but exact enumeration grows exponentially with their number. LinearPFN amortizes inference for linear models with main effects and pairwise interactions, and selection is ahead of every baseline in every reported slice.

  • Problem

    Exact spike-and-slab posterior enumeration is practical only for small predictor sets, so larger problems require approximations such as MCMC.

    Source: §1 · p. 2

  • Method

    LinearPFN is a transformer pretrained on synthetic draws from an explicit spike-and-slab prior, returning inclusion probabilities, posterior-mean coefficients, and predictive distributions in one forward pass.

    Source: §1 · p. 2

  • Results

    Selection is ahead of every baseline in every slice; against SUSIE, the paired difference in F1 per dataset is +0.118 [+0.114, +0.122].

    Source: §4 · p. 8

  • Takeaways & Limitations

    Amortization reduces variable selection to one forward pass, taking 0.03 seconds at the median on one GPU in the reported comparison.

    Source: §5 · p. 9

    With nearly collinear predictors, posterior mass can split so both PIPs fall below one half and the MPM rule selects neither; LINEARPFN reports no credible sets to group them.

    Source: §5 · p. 9

Abstract

from arXiv · show

Spike-and-slab regression is a standard Bayesian formulation of variable selection: it returns a posterior distribution over which candidate effects are active rather than a single selected subset, so that every candidate effect carries an inclusion probability. Its cost grows exponentially with the number of candidate effects, so the posterior can be enumerated exactly only when the number of predictors is small. Beyond that reach, the posterior has to be approximated, typically by Markov chain Monte Carlo over the model space, which requires a fresh run for every dataset and, within a fixed budget of steps, may fail to converge. We present LinearPFN, a prior-data fitted transformer network that amortizes spike-and-slab inference for linear models with main effects and pairwise interactions. The network is pretrained once on synthetic datasets, drawn from an explicitly specified prior, and a single forward pass over a new dataset returns posterior inclusion probabilities, posterior-mean coefficients and posterior predictive distributions with no per-dataset fitting. The prior is conjugate by design, so that the posterior for each fixed set of active effects has a closed form, and wherever the exact posterior is still computable by enumeration we verify the network's outputs against it. On real predictor matrices from published social-science datasets, with outcomes drawn from the prior so that the true active set is known, LinearPFN attains a higher per-dataset selection AUC and a higher F1 under the median probability model rule than five classical baselines. The lead holds when the coefficients, the interactions or the noise depart from the prior. Code: https://github.com/schiekiera/LinearPFN. Trained model: https://huggingface.co/schiekiera/LinearPFN.

1 INTRODUCTION

Spike-and-slab selection represents uncertainty over active effects, but exact posterior enumeration becomes impractical as the number of predictors grows. LinearPFN amortizes inference with a prior-trained transformer for main effects and pairwise interactions.

1 INTRODUCTION

  • Exact enumeration sums over 2d−1 models and is practical only at small p; larger problems rely on approximations such as MCMC.Other approaches mentioned include adaptive sampling and variational approximations.Source: §1 · p. 2
  • LinearPFN is pretrained on synthetic draws from an explicit spike-and-slab prior and returns inclusion probabilities, posterior-mean coefficients, and predictive distributions in one forward pass.Its selection head is trained against simulated active sets, while the regression head predicts query-row outcomes.Source: §1 · p. 2, Figure 1 · p. 2
  • The authors present LinearPFN as the first PFN to return inclusion probabilities over a regression model space.They state this novelty claim is to their knowledge.Source: §1 · p. 2
  • Because each candidate-model posterior has a closed form, exact posterior outputs can validate LinearPFN wherever model enumeration is feasible.The related description says the network is validated against enumeration for up to five predictors.Source: §1 · p. 2
  • The pretraining prior calibrates active main-effect sizes to meta-analytic norms for published psychology correlations.The effect-size calibration is described as resembling individual-level social-science data.Source: §1 · p. 3
  • Benchmarking uses 10,000 datasets with real social-science predictor matrices and prior-drawn outcomes, plus variants that violate outcome-model assumptions.On published tables with actual outcomes, LinearPFN is compared with an MCMC sampler targeting the same posterior.Source: §1 · p. 3

2 RELATED WORK

Prior-data fitted networks approximate posterior predictions by training on synthetic datasets drawn from a prior, while variable-selection approaches differ in whether they return a subset or uncertainty over subsets. LinearPFN combines PFN-style prediction with supervision of selection and coefficient outputs.

2 RELATED WORK

  • PFNs train on prior-sampled datasets and use query-row predictions to approximate the posterior predictive distribution without per-dataset fitting.TABPFN extends this approach to general tabular prediction; LINEARPFN instead uses a spike-and-slab prior over linear main effects and pairwise interactions.Source: §2 · p. 3
  • LINEARPFN pairs a PFN-trained regression head with selection and coefficient heads supervised by simulated active sets and coefficients.Neural posterior estimation instead targets parameters, usually through a normalizing flow over their joint posterior.Source: §2 · p. 3
  • Frequentist selection returns a subset, whereas spike-and-slab Bayesian selection assigns inclusion probabilities and averages coefficients across models.Lasso-based interaction methods can omit parent effects; GLINTERNET and HIERNET impose heredity constraints, while SUSIE uses a variational approximation.Source: §2 · p. 3

3 METHODS

LinearPFN uses a conjugate spike-and-slab prior over main effects and pairwise interactions to train a transformer that returns posterior quantities in one pass. The paper evaluates it against the exact posterior where enumeration is feasible and benchmarks selection on real predictor matrices paired with prior-drawn outcomes.

3.1 THE PRIOR

  • The predictor prior combines generic correlated columns and latent-factor columns, shuffles their order, and maps columns to Gaussian, Likert, binary, count or skewed marginals.The matrix is intended to mimic social-science data, with rows as cases and columns as measured characteristics.Source: §3.1 · p. 4, Figure 2 · p. 4
  • The active-effects prior enforces strong heredity and averages count-based and rate-based inclusion schemes, while the effect-size hyperprior calibrates standardized active main-effect magnitudes to psychology correlation norms.The prior quartiles are 0.078, 0.169, and 0.294, versus norms of 0.11, 0.19, and 0.29; ρ is fixed at 0.5, so interaction effects are smaller than main effects.Source: §3.1 · pp. 4–5

3.2 NETWORK AND TRAINING

  • The shared transformer processes table cells with feature attention, context-row attention and an MLP, while three heads estimate predictive distributions, PIPs and posterior-mean coefficients.The heads are trained toward the posterior predictive distribution, effect-inclusion probabilities and coefficient means.Source: §3.2 · pp. 5–6
  • Interaction inputs combine parent-predictor summaries with rowwise product summaries, preserving evidence that the product of the parent means cannot capture.The pair representation is symmetric in the two predictors and uses four parts.Source: §3.2 · pp. 5–6

3.3 EVALUATION

  • For validation, conjugacy gives each candidate model a closed-form marginal likelihood, enabling exact posterior enumeration for p ≤5 and comparison on six criteria.The evaluation integrates out the slab scale numerically and compares predictive distributions, coefficients and PIPs.Source: §3.3 · p. 6
  • 10,000 datasets pair real social-science predictor matrices with prior-drawn outcomes; five baselines are compared using per-dataset selection AUC and F1, plus coefficient RMSE.Sampling is stratified across density and R2 bins so sparse and dense active sets and weak and strong signals are represented equally.Source: §3.3 · p. 7
  • Three outcome-model variants test fixed-magnitude coefficients, more and larger interactions, and heavy-tailed noise while holding the main-benchmark draw fixed apart from the redrawn outcome.The interaction variant retains strong heredity while raising interaction variance to match main-effect variance.Source: §3.3 · p. 7
  • On seven published tables with actual outcomes, LINEARPFN is compared with an MCMC sampler targeting the same posterior through PIP and posterior-mean coefficient correlations and RMSE.These tables do not provide a known active set for scoring selection against truth.Source: §3.3 · p. 8

4 RESULTS

On prior-generated outcomes over real social-science predictor matrices, LinearPFN leads the five baselines in selection across benchmark slices and matches the best baseline in coefficient error. Its selection advantage persists across three outcome-model variants, while comparisons on published outcomes show agreement with a sampler targeting the same posterior.

Fidelity to the exact posterior

  • LinearPFN meets all six exact-posterior validation criteria; its PIP AUC gap is −0.0026 and its posterior-mean coefficients correlate 0.9931 with exact values.On 2,299 datasets within enumeration reach, its F1 differs from the exact posterior by −0.030 and coefficient RMSE by +0.015.Source: §4 · p. 8

Real predictor matrices, outcomes from the prior

  • Against HIERNET, the best baseline, the RMSE difference is −0.001 [−0.002, 0.000] for the dense posterior mean and 0.000 [−0.001, +0.001] when hard-zeroed.For p ≤5, LinearPFN's coefficient RMSE exceeds the exact posterior by +0.015.Source: §4 · p. 8

Real predictor matrices, outcomes outside the prior

  • LinearPFN retains the highest AUC and F1 in every variant; under many large interactions, its F1 margin over SUSIE is +0.115 [+0.112, +0.119], but its RMSE of 0.064 trails three baselines.The variants test fixed-magnitude coefficients, many large interactions and heavy-tailed noise.Source: §4 · p. 8

Real predictor matrices, outcomes from the prior

  • Selection leads every baseline in every slice; against closest baseline SUSIE, its per-dataset F1 difference is +0.118 [+0.114, +0.122].Figure 4 slices the 10,000 datasets by predictor count, row count, density and R2.Source: §4 · p. 8, Figure 4 · p. 9

Real predictor matrices, outcomes outside the prior

  • Many large interactions raise every method's coefficient error and lower its AUC; heavy-tailed noise changes LinearPFN's coefficient error by −0.001 and slightly raises F1 for every method.Under heavy-tailed noise, stability selection alone loses coefficient accuracy, with mean error rising +0.123.Source: §4 · p. 9

Real outcomes

  • On seven published tables with actual outcomes, LinearPFN's PIPs correlate 0.972 with those of the same-posterior MCMC sampler, and their MPMs match on five tables.These tables have no known active set, so they do not measure selection accuracy.Source: §4 · p. 9

5 DISCUSSION

LinearPFN amortizes Bayesian variable selection for main effects and pairwise interactions, but its performance and scope depend on a fixed prior. The discussion highlights selection gains, practical limitations, and the method’s data and reproducibility boundaries.

5 DISCUSSION

  • Four limitations are identified: near-collinearity can leave both PIPs below one half, CPU use is slow, coefficient error only matches the best baseline, and training assumes one fixed prior.That prior assumes strong heredity, Gaussian noise, numeric predictors, and independent rows; selection gains held when coefficients, interactions, or noise departed from it.Source: §5 · p. 9
  • On real predictor matrices with prior-drawn outcomes, LINEARPFN selects better than every baseline and returns variables and coefficients in 0.03 seconds at the median on one GPU.The MCMC sampler takes 57.43 seconds at the median on one CPU thread and meets its convergence criteria on 93.0% of runs.Source: §5 · p. 9
  • Generative AI assisted with implementation, dataset cleaning, related-work discovery, and wording, while authors reviewed the work and take responsibility for the final content.Synthetic datasets were drawn by code from the stated prior, not generated by a generative model.Source: AI USE STATEMENT · p. 10
  • The work uses public predictor tables and published regression outcomes, and its probabilities under a stated prior do not by themselves support causal claims.The code and scripts producing experiments, numbers, tables, and figures are publicly available, as is the trained model.Source: ETHICS STATEMENT · p. 10, REPRODUCIBILITY STATEMENT · p. 10

B THE PRIOR

  • The prior’s outcome scale is fixed, so outcomes from outside the prior are standardized before the forward pass.The specified prior also expands standardized predictors with pairwise products and draws active effects from type-dependent normal slabs.Source: B THE PRIOR · p. 15

C.2 WHY THE PREDICTIVE DISTRIBUTION NEEDS ITS OWN HEAD

  • The predictive variance needs its own head because posterior-mean coefficients determine the predictive mean but not the variance, which includes coefficient covariances.The variance also includes expected noise, which does not shrink as data grow, and uncertainty about the regression line, which shrinks along informed directions.Source: §C.2 · p. 18

C.3 TRAINING

Training uses batches whose datasets share their row and predictor counts, while the exact reference integrates over the slab scale by quadrature. A 20-node calculation closely matches a denser grid in the tested cases.

C.3 TRAINING

  • Training batches share dataset size (n, p), while the query-row share is drawn per batch and the remaining rows are context.Muon trains backbone weight matrices and AdamW trains other parameters; all three loss terms send gradients to the shared backbone.Source: §C.3 · p. 18

D.1 ENUMERATION AND QUADRATURE

  • For each model and quadrature node, the exact predictive distribution is a Student-t component, and the overall predictive distribution is their mixture.Its quantiles are found by bisection on the mixture distribution function.Source: §D.1 · p. 19
  • The reference enumerates every model with nonzero prior mass and integrates over c using 20 Gauss–Legendre nodes on log c; model probabilities and PIPs use every model, while coefficient means and prediction use pairs carrying all but 1.0 × 10^-9 of posterior mass.At p = 5, the largest p the reference enumerates, there are 1,450 such models.Source: §D.1 · p. 19
  • At 20 quadrature nodes, the largest tested PIP deviation from a 30-node grid was 1.4 × 10^-7, and the largest log-evidence deviation was 9.8 × 10^-6.The largest deviations occurred at n = 1,024 and p = 5; the largest PIP deviation across 12-, 15-, and 20-node grids fell from 4.1 × 10^-4 to 1.4 × 10^-7.Source: §D.1 · p. 19

D.2 THE MCMC SAMPLER

Beyond exact enumeration, the reference method uses Metropolis–Hastings MCMC over active sets to target the same posterior. Strong heredity gives zero target mass to models with ineligible interactions.

D.2 THE MCMC SAMPLER

  • The MCMC sampler proposes adding, dropping, or swapping one effect, starts several chains from prior draws, discards burn-in, and estimates PIPs from visit frequencies.It averages each visited model’s posterior-mean coefficients and declares convergence when the largest split R-hat is below 1.1 and the smallest effective sample size exceeds 200.Source: §D.2 · p. 20

E VALIDATION AGAINST THE EXACT POSTERIOR

Validation compares LinearPFN with an exact posterior on small-predictor prior draws, while broader coverage checks assess predictive and selection behavior. The real predictor matrices are cleaned from published social-science tables using fixed rules.

E VALIDATION AGAINST THE EXACT POSTERIOR

  • The validation panel crosses five context sizes from 32 to 2,000 rows with seven predictor counts from 3 to 30, using 100 datasets per cell and 64 query rows per dataset.Exact comparisons are made only in cells with p ≤ 5; coverage pools every cell, while the other five criteria use the exact cells.Source: E VALIDATION AGAINST THE EXACT POSTERIOR · p. 20
  • Six validation criteria assess predictive headroom, coefficient recovery, interval coverage, selection ranking, calibration, and resolution against the exact posterior.Selection ranking allows the PIP AUC to be at most 0.01 below the exact posterior in the worst pooled, main-effect, or interaction subset; calibration limits excess ECE to 0.01.Source: E VALIDATION AGAINST THE EXACT POSTERIOR · p. 20

F.1 THE REAL PREDICTOR MATRICES

  • Fixed cleaning rules leave 383 eligible predictor tables from the RDATASETS collection, after excluding tables that fail row, column, missingness, or other criteria.The 383 matrices come from 20 packages across seven fields, with median 758 rows and 9 predictors; 87 have at most 5 predictors.Source: §F.1 · p. 21

F.2 CONSTRUCTION OF THE MAIN BENCHMARK

The benchmark uses rejection sampling to construct datasets from real predictor matrices, then balances draws across selected density and R2 cells. This grid covers only part of the prior and its cells differ in the predictor matrices they contain.

F.2 CONSTRUCTION OF THE MAIN BENCHMARK

  • The 25 benchmark cells cover 19.1% of prior mass, with equal weighting within each cell and acceptance rates from 3.6% in the most frequent cell to 0.013% in the rarest.Draws are discarded when density or R2 falls between bins or below the lowest bin, excluding datasets without active main effects and those with R2 below the lowest bin.Source: §F.2 · p. 21
  • Comparisons across cells mix density and R2 differences with predictor-matrix differences, because cells contain distinct sets of predictor matrices.Density takes only multiples of 1/pℓ, so matrices with few predictors can reach only some bins; none of the 2,299 datasets with pℓ≤5 falls in the lowest density bin.Source: §F.2 · p. 21

F.3 METHODS

The section defines how classical baselines select and rank effects, and how their coefficient outputs are handled for comparison with LINEARPFN.

F.3 METHODS

  • The lasso selects effects at the penalty with the smallest cross-validated error and ranks them by their entry point on the regularization path.GLINTERNET and HIERNET use the same selection and ranking rules under their heredity constraints.Source: §F.3 · pp. 21–22
  • Stability selection fits lasso paths on 50 complementary pairs of half-samples, selects effects with frequency at least 0.75, and ranks them by frequency.Its coefficients come from a least-squares fit on selected effects, and fit time includes that refit.Source: §F.3 · pp. 21–22
  • SUSIE selects effects with PIP above one half and ranks them by PIP, while LINEARPFN’s hard-zeroed read-out keeps the intercept and involves no refit.The lasso, GLINTERNET and HIERNET return penalized coefficients; the exact posterior and LINEARPFN return posterior means over models, with no coefficient exactly zero.Source: §F.3 · p. 22

F.4 METRICS

The metrics evaluate selection, ranking, coefficient estimation, and paired differences between methods, with explicit conventions for ties, missing scores, and fit time.

F.4 METRICS

  • AUC ranks each dataset’s candidate effects by each method’s own score, and the average covers 9,974 datasets because it is undefined when all or no effects are active.Tied scores receive their mean rank.Source: §F.4 · p. 22
  • Table 6 reports paired differences as LINEARPFN minus each baseline on the same datasets, with bootstrap 95% intervals.All Wilcoxon signed-rank p-values are below 0.001 except RMSE against HIERNET (p=0.155) and hard-zeroed RMSE against GLINTERNET (p=0.775).Source: Table 6 · p. 23
  • Fit time is wall time for one fit on one CPU thread, including each method’s own tuning.Coefficient RMSE excludes the intercept; paired-difference intervals use 2,000 seeded bootstrap resamples.Source: §F.4 · p. 23

G.1 MAIN BENCHMARK: FULL RESULTS

LINEARPFN’s benchmark results include selection performance, inference speed, comparisons with exact and MCMC references, tests under prior violations, and agreement with MCMC on published outcomes.

G.1 MAIN BENCHMARK: FULL RESULTS

  • LINEARPFN’s F1 was 0.78 on panel predictor matrices and 0.77 on cross-sections, exceeding HIERNET by +0.241 and +0.244 and SUSIE by +0.110 and +0.120, respectively.The 61 panel matrices contribute 1,645 datasets, with outcomes independent across rows.Source: §G.1 · p. 23
  • The GPU forward pass took 0.03 seconds at the median and 0.42 seconds at most, 222 times faster than the single-thread median.The exact posterior took 0.04 seconds at the median on 2,299 datasets; MCMC took 57.43 seconds at the median and 866.31 at the 90th percentile on 300 datasets.Source: §G.1 · p. 23

G.2 DISTANCE TO THE EXACT POSTERIOR AND TO MCMC

  • Against exact inference on 2,299 datasets with pℓ≤5, LINEARPFN was lower by 0.009 AUC and 0.030 F1, and its coefficient RMSE was higher by 0.015.The exact posterior’s own RMSE was 0.050; against MCMC on 300 datasets, LINEARPFN was lower by 0.008 AUC and 0.050 F1 and higher by 0.009 RMSE.Source: §G.2 · p. 23

G.3 OUTCOMES OUTSIDE THE PRIOR

  • The three variants preserve the main-effect density; under many large interactions, the median dataset has 3 active interactions and 15.7% have none.The exact posterior on these variants is evaluated under a wrong prior.Source: §G.3 · p. 24
  • LINEARPFN’s ECE was 0.009 on the control and 0.008, 0.026, and 0.007 under fixed-magnitude coefficients, many large interactions, and heavy-tailed noise, respectively.On enumeration-reach datasets, F1 was 0.79 versus 0.82 for the exact posterior on the control, and 0.78 versus 0.82 under many large interactions; the grid covers 19.1% of prior mass, so ECE here does not test calibration.Source: §G.3 · p. 24

G.4 REAL OUTCOMES

  • Across 278 candidate effects from seven published tables, LINEARPFN’s PIPs correlated 0.972 with MCMC and its posterior-mean coefficients correlated 0.991.The two median probability models agreed on 272 of 278 effects; held-out negative log predictive density for a Gaussian fit to fitting rows was 1.385.Source: §G.4 · p. 26

G.5 TRAINING WITHOUT THE PREDICTIVE LOSS

The predictive-loss model exceeds its two-head ablation on selection and coefficient metrics, while the ablation still beats every baseline on F1.

G.5 TRAINING WITHOUT THE PREDICTIVE LOSS

  • +0.005 AUC per dataset, +0.026 F1 per dataset, and −0.005 coefficient RMSE distinguish the paper’s model from its two-head ablation (all Wilcoxon signed-rank p < 0.001).The paper’s model has smaller coefficient error on 67.0% of datasets.Source: §G.5 · p. 26
  • The two-head model’s F1 is 0.75 versus 0.66 for SUSIE, and it still exceeds every baseline on F1.Source: §G.5 · p. 26
  • On Appendix E’s validation panel, the two-head model’s AUC gap is −0.0055 versus −0.0026, while its ECE gap is 0.0005 versus 0.0034.Source: §G.5 · p. 26
  • The two-head model meets the three selection criteria, with a smallest resolution ratio of 0.9050 on interactions versus 0.9640 for the paper’s model.Source: §G.5 · p. 26
  • Because its regression head is never trained, the two-head model cannot meet headroom or coverage criteria; validation does not score its coefficient head.Source: §G.5 · p. 26
  • The ablation changed only predictive negative log-likelihood, using the same prior, architecture, seed, and 37,500 steps, with each model trained once.The remaining loss terms were not reweighted, so part of the gap may be recoverable.Source: §G.5 · p. 26

Related papers

Loading 2610.06580v1…