SelectPFN: Amortized Bayesian Variable Selection for Sparse Nonlinear Regression with Prior-Data Fitted Networks
Organizations: School of Data Science and Analytics, Kennesaw State University
Abstract
Tabular foundation models based on prior-data fitted networks (PFNs) predict on a new data set in a single forward pass, but they cannot report which inputs their predictions depend on. We propose SelectPFN, a PFN that performs Bayesian variable selection and prediction jointly for sparse nonlinear regression. In one forward pass it returns posterior inclusion probabilities for every variable, and its predictions are computed from sampled subsets only, so excluded variables provably cannot affect them; we show this is exactly what the Bayesian model prescribes. The network is pretrained once on synthetic tasks from a meta-prior built around a group spike-and-slab sparse-input neural network. Column-permutation equivariance combined with random interleaving of original and knockoff columns yields knockoff statistics with exact finite-sample false discovery rate control, even when the meta-prior is wrong. We develop theory for amortized selection: a Bayes-regret bound for the amortized median probability model, a characterization showing that the information needed to detect a variable grows like in the prior weight of structures that can use it, where is an efficiency, and a negative result showing that average-case training does not transfer to frequentist guarantees at a fixed truth. On out-of-prior nonlinear designs with p up to 500, SelectPFN improves the area under the precision-recall curve by 0.23-0.28 over the lasso, BART, DART, LassoNet and a group-penalized sparse neural network, at a false discovery proportion of about 0.07 and 0.04 seconds per data set. The advantage persists on real covariate matrices, and, as the theory predicts, detection of interaction-only variables rises monotonically with their prior weight.
Figures & tables
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
| Component | Distribution |
|---|---|
| Sample size | rows per task; a fraction are training (context) rows, the remaining are query rows |
| Dimension | with curriculum in thirds of training |
| Sparsity | Beta–Bernoulli, , , ; ; an empty is redrawn w.p. 0.9 |
| Dependence | independent; AR(1); block equicorrelation; latent factor (equal weights). Active variables at random positions |
| Marginals | Gaussian copula with per-column transform: identity, uniform, log-normal skew, Tukey- heavy tails ( ), discretized; columns standardized |
| Function family | sparse-input MLP with group spike-and-slab first layer, eq. 3 (0.30); additive (0.12); additive + pairwise, with no main effect for a variable w.p. 0.3 (0.12); trees/step functions (0.10); GP with ARD kernel, length-scales (0.10); linear (0.14); interaction-only pairs/triples through functions odd in each argument (a single active variable gets a main effect), plus weak main effects w.p. 0.3 (0.12) |
Feature Selection
Feature selection identifies compact, informative variables or input dimensions to improve prediction, efficiency, interpretability, and statistical inference.
See all 194 papers