LinearPFN: Amortized Variable Selection for Linear Models with Interactions

2026-10-05

AI Summary

This study addresses the high computational cost, per-dataset fitting requirements, and convergence issues of Markov Chain Monte Carlo (MCMC) methods in Bayesian variable selection by proposing an amortized inference approach based on Prior-Data Fitted Networks (PFNs). Leveraging a Transformer architecture and conjugate prior design, the model is pretrained on synthetic data to perform spike-and-slab posterior inference for linear models with interaction terms. This represents the first application of PFNs to such settings, enabling posterior estimation via a single forward pass without dataset-specific training. Experimental results demonstrate that the proposed method significantly outperforms five classical baselines in AUC and F1 scores on real-world social science datasets, while maintaining strong robustness under deviations from prior assumptions.

Abstract

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.

Problem

• Variable Selection
• Spike-and-slab Regression
• Linear Models with Interactions
• Bayesian Inference
• Amortized Inference

Innovation

• Amortized Inference
• Prior-data Fitted Networks
• Spike-and-slab Regression
• Variable Selection
• Transformer