LinearPFN: Amortized Variable Selection for Linear Models with Interactions

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

cs.LGOct 5, 2026

LinearPFN: Amortized Variable Selection for Linear Models with Interactions

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

Organizations: Institute of Psychology, Humboldt-Universität zu Berlin, Berlin, Germany · Department of Education and Psychology, Freie Universität Berlin, Berlin, Germany · Department for AI in Society, Science, and Technology, Zuse Institute Berlin, Germany · Institute of Mathematics, Technische Universität Berlin, Germany

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.

Figures & tables

Figure 1: Pretraining and prediction of LinearPFN . Top: each prior draw fixes an active set

\gamma

and spike-or-slab coefficients

\beta

for one synthetic dataset, and the loss compares the network’s output with them. Bottom: in every layer, each cell (green) first attends to the other cells in its row and then to the other cells in its column (striped), and an MLP then updates each cell on its own. One forward pass over a new dataset returns inclusion probabilities and posterior-mean coefficients for every main effect and interaction, and a posterior predictive distribution for each query row.

Figure 2: Predictor prior. Left: the generic block is generated jointly from a random correlation matrix

R

(top), and the factor block loads on latent factors (bottom). Middle: the two blocks are merged and the column order is shuffled. Right: a Gaussian copula maps each column to its marginal, either unchanged (Gaussian) or transformed to a Likert, binary, count or skewed family.

Figure 3: The three read-out heads. Left: the final representation, one vector

\mathbf{r}_{ij}

per cell in row

i

and column

j

. Middle: the pooled vectors of Equation 3 . Path A averages each predictor column over the context rows into

\mathbf{h}_{j}

. Paths B and C are the combine-then-mean summaries , which multiply the cells of two predictor columns within each row, in C also with the outcome cell, and then average into

\mathbf{h}_{jk}

and

\mathbf{h}_{jky}

. Path D is the outcome cell

\mathbf{r}_{iy}

of a query row, not pooled. Right: the selection and coefficient heads each use a main MLP on

\mathbf{h}_{j}

and a pair MLP on

\mathbf{g}_{jk}

(Equation 4 ). The regression head maps

\mathbf{r}_{iy}

of each query row to a bar distribution.

predictive distribution and coefficients			inclusion probabilities
criterion	threshold	value	criterion	threshold	value
headroom	$\geq$ 0.95	0.9605	AUC gap	$\geq-0.01$	$-0.0026$
coverage error	$\leq$ 0.03	0.0005	ECE gap	$\leq$ 0.01	0.0034
coefficient correlation	$\geq$ 0.99	0.9931	resolution ratio	$\geq$ 0.90 / 0.85	0.9640

Table 1: Validation of LinearPFN against the exact posterior on held-out prior draws with

p\leq 5{}

(coverage on every

p

). The thresholds were fixed before the paper’s model was trained.

method	AUC/ds, mean	$F_{1}$ /ds, mean	coef RMSE, mean	fit s, median
LinearPFN	0.96 [0.95, 1.00]	0.77 [0.67, 1.00]	0.041 [0.011, 0.051]	7.40 [1.07, 24.59]
hierNet	0.92 [0.88, 1.00]	0.53 [0.38, 0.67]	0.042 [0.011, 0.051]	25.45 [9.17, 79.28]
glinternet	0.79 [0.68, 0.97]	0.57 [0.43, 0.71]	0.044 [0.011, 0.052]	4.39 [1.18, 15.05]
SuSiE	0.85 [0.76, 1.00]	0.66 [0.46, 0.89]	0.051 [0.009, 0.066]	0.16 [0.08, 0.40]
lasso	0.82 [0.72, 0.96]	0.48 [0.33, 0.63]	0.049 [0.012, 0.059]	0.14 [0.08, 0.71]
stability	0.82 [0.72, 0.95]	0.55 [0.40, 0.75]	0.068 [0.014, 0.074]	0.36 [0.33, 0.49]

Table 2: Benchmark on the real predictor matrices, 10,000 datasets with outcomes drawn from the prior. Over datasets, AUC and

F_{1}

computed per dataset (/ds), with the first and third quartiles in brackets. Fit seconds are one fit on one CPU thread. On one NVIDIA L40 GPU, the same forward pass of LinearPFN takes 0.03 seconds at the median [0.01, 0.10].

Figure 4: Benchmark on the 383 real predictor matrices, 10,000 datasets with outcomes drawn from the prior. Left: selection,

F_{1}

per dataset at each method’s own selection rule. Middle: ranking, AUC per dataset. Right: estimation, coefficient RMSE over every candidate effect. Each panel slices the same datasets by the number of predictors, the number of rows, the density

k_{\mathrm{main}}/p

and the

R^{2}

; every point is the mean over the datasets of that slice. Density and

R^{2}

are binned in the 5 bins of the benchmark grid, each 0.1 wide, and the axis marks the lower edge of each bin (Appendix F.2 ).

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Predictor prior
rows $n$ , predictors $p$	U{20, …, 2,000}, U{2, …, 30}
factor share $f$	$0$ with probability 0.15, else Beta(1.5, 3)
factors $m$	U{1, …, $\lceil q/3{}\rceil$ }, $q$ the size of the factor block
communality: mean, concentration	U(0.1, 0.92), U(4, 40), cap 0.98
high band: probability, mean	0.15, U(0.9, 0.995), cap 0.995
share of $h^{2}_{j}$ on other factors	0.15

Appendix

Table 3: The prior of the paper’s model: its fixed values, and the standardized effects

|\beta_{e}|\,\mathrm{sd}(\mathbf{z}_{e})/\mathrm{sd}(\mathbf{y})

it implies over the active effects of 10,000 prior draws, where

\mathbf{z}_{e}

is the column of effect

e

\mathbf{Z}

Figure 5: The active-effects and effect-size priors, one draw with

p=5

. Top: the active main effects and then the active interactions are drawn from the candidate set, under the count prior (the number first, then a uniform subset) or the rate prior (a rate first, then one inclusion per candidate). Bottom: the noise variance and the slab scale are drawn, active effects receive slab coefficients and inactive ones a spike at zero, and the outcome is

\mathbf{y}=\mathbf{Z}\beta+\epsilon

layers / heads / embedding	8 / 16 / 512	parameters	31.5M (31,541,766)
MLP hidden units	1,024	GPUs	2 $\times$ H200
optimizer steps	37,500	wall time	14.6 h
effective batch	64 (4 $\times$ 16)	GPU-hours	29.1
datasets seen	2,400,000	throughput	0.72 steps/s
learning rate	Muon 0.02, AdamW 0.001	peak GPU memory	57.3 GiB

Appendix

Table 4: Model card of LinearPFN . Wall time sums the segments of the requeued job, and GPU-hours are wall time times the number of GPUs.

reason	entries	reason	entries
package outside the frame	2,908	described as simulated	17
index: fewer than 30 rows	93	fewer than 3 columns after cleaning	15
index: fewer than 3 usable columns	81	more than 50.0% of rows lost	13
time-series or contingency-table object	69	same title as another table	8
pure time series	32	download failed	3
same columns as a larger table	23	fewer than 30 rows after cleaning	3

Appendix

Table 5: Rdatasets entries that did not become a predictor matrix, by reason.

Figure 6: Share of 60,000 datasets drawn from the prior in each cell of the plane of density and $R^{2}$ , on a log color scale, with the row and column totals in the margins. Rings mark the lower-left corners of the 25 cells of the main benchmark, which together hold 19.1 % of the draws. A dot marks a cell whose share rounds to zero.

baseline	$F_{1}$ /ds	AUC/ds	coefficient RMSE	RMSE, hard-zeroed
hierNet	+0.244	+0.044	$-0.001$	0.000
	[+0.238, +0.249]	[+0.043, +0.046]	[ $-0.002$ , 0.000]	[ $-0.001$ , +0.001]
glinternet	+0.205	+0.173	$-0.004$	$-0.003$
	[+0.200, +0.210]	[+0.170, +0.176]	[ $-0.005$ , $-0.003$ ]	[ $-0.004$ , $-0.002$ ]
SuSiE	+0.118	+0.110	$-0.011$	$-0.010$
	[+0.114, +0.122]	[+0.107, +0.113]	[ $-0.012$ , $-0.010$ ]	[ $-0.011$ , $-0.009$ ]

Appendix

Table 6: Main benchmark: paired differences, LinearPFN minus each baseline on the same datasets, with bootstrap 95 % intervals. The last column pairs the hard-zeroed read-out of LinearPFN with each baseline. Every Wilcoxon signed-rank

p

is below

0.001

, except for the RMSE against hierNet (

p=0.155{}

) and the hard-zeroed RMSE against glinternet (

p=0.775{}

Figure 7: Outcomes outside the prior, one row per variant, each on 10,000 datasets. Left: selection,

F_{1}

per dataset at each method’s own selection rule. Middle: ranking, AUC per dataset. Right: estimation, coefficient RMSE over every candidate effect. Each panel slices the datasets by the number of predictors, the number of rows, the density

k_{\mathrm{main}}/p

and the

R^{2}

; every point is the mean over the datasets of that slice. Density and

R^{2}

are the bins of the main-benchmark dataset that the variant is built on, so the

R^{2}

of a variant dataset can lie outside its bin. A method with an off-scale value in a panel is left out of that panel, and the panel names it.

	control			fixed-magnitude coefficients
method	AUC/ds	$F_{1}$ /ds	RMSE	AUC/ds	$F_{1}$ /ds	RMSE
LinearPFN	0.96	0.77	0.041	0.98 (+0.011)	0.85 (+0.072)	0.044 (+0.003)
hierNet	0.92	0.53	0.042	0.94 (+0.018)	0.52 ( $-0.006$ )	0.045 (+0.003)
glinternet	0.79	0.57	0.044	0.85 (+0.062)	0.56 ( $-0.011$ )	0.048 (+0.004)
SuSiE	0.85	0.66	0.051	0.89 (+0.034)	0.70 (+0.044)	0.063 (+0.012)
lasso	0.82	0.48	0.049	0.86 (+0.040)	0.48 (+0.004)	0.056 (+0.007)

Appendix

Table 7: Outcomes outside the prior: every method on the control (the main benchmark) and on the three variants, means over datasets, with the paired change from the control in parentheses. AUC and

F_{1}

per dataset, coefficient RMSE over every candidate effect.

table	$n$	$p$	$r$ , mains	RMSE	max diff.	RMSE	agree	LinearPFN	MCMC	OLS
			PIP			coefficient	MPM	held-out NLL
attitude	30	6	0.998	0.016	0.043	0.006	21/21	1.095	1.122	1.177
swiss	47	5	0.854	0.074	0.157	0.019	15/15	1.277	1.312	1.765
prostate	97	8	0.990	0.028	0.107	0.009	36/36	0.927	0.894	0.876
tal_or	123	5	0.999	0.017	0.044	0.009	15/15	1.226	1.226	1.246
bodyfat	252	13	0.780	0.083	0.600	0.019	87/91	0.651	0.663	0.667

Appendix

Table 8: Real outcomes: LinearPFN against the MCMC sampler on 7 published tables with their actual outcome. PIP agreement over the candidate effects (correlation on the main effects, RMSE and largest difference over all effects), RMSE between the posterior-mean coefficients, candidate effects on which the two MPMs agree, and held-out negative log predictive density.

Explore similar work

Apr 29, 2026stat.ME

Linear Models, Variable Selection, Artificial Intelligence

Variable selection in linear regression models has been a problem since hypothesis testing began. Which variables to include or exclude from a model is not an easy task. Techniques such as Forward, Back ward, Stepwise Regression sequentially add or delete variables from a model. Penalized likelihood methods such as AIC, BIC, etc. seek to choose variables that have a significant contribution to the likelihood. Penalized sum of square methods such as LASSO and Elastic Net have been used to penalize small coefficients to only allow variables with large coefficients in the model. This work introduces an Artificial Intelligence approach to model selection where an ANN is trained to determine the significance of the variables based on OLS estimates. A simulation study shows the accuracy across various sample sizes and variances. Furthermore, a simulation study is conducted to compare the performance of the approach against Forward, Backward, AIC, BIC and LASSO. The approach is illustrated using a dataset from the World Health Organization regarding Life Expectancy. A github link is provided to the pretrained ANN that can handle up to 100 predictor variables, the original WHO dataset and the subset used in this work.

By Riyadh Alrawkan, Edward Boone, Ryad Ghanam +1

1Statistical Sciences and Operations Research, Virginia Commonwealth University, Richmond, Virginia, USA · 2Liberal Arts and Sciences, Virginia Commonwealth University School of the Arts in Qatar, Doha, Qatar · 3Research School of Finance, Actuarial Studies and Statistics, Australian National University, Canberra, AUS

Model Selection Lasso

Ask PDF

May 26, 2026cs.LG

Amortized Factor Inference Networks for Posterior Inference

Amortized inference promises fast test-time Bayesian inference, but existing methods are inherently tied to fixed models. Extending amortization to unseen models typically requires retraining or costly test-time finetuning. In this paper, we ask: is it possible to build a single inference network capable of generalizing across varying priors, likelihoods, and dimensionality? We introduce Amortized Factor Inference Networks (AFINs), a family of encode-merge-decode inference networks built on dimension-independent modules that map a model specification and its observations to the parameters of a variational posterior. Experimentally, a single trained AFIN achieves posterior accuracy comparable to NUTS and several variational inference methods, while requiring 2 to 4 orders of magnitude less test-time compute. Code is available at https://github.com/joohwanko/AFINs.

Joohwan Ko, Justin Domke

Manning College of Information and Computer Sciences University of Massachusetts Amherst

Bayesian Inference Posterior Neural Network Inference

Ask PDF

Jun 28, 2026cs.LG

A Mathematical Optimization Approach for Expert-Informed Bayesian Best Subset Selection

A central challenge in statistical modeling is identifying the subset of features that belong in the true regression model. The classical best subset selection problem, recently made tractable via mixed-integer optimization (MIO), finds the globally optimal sparse solution. It does not, however, make use of any information beyond the observed data. In many applied settings, domain experts can meaningfully rank or score the relevance of candidate predictors, yet no existing framework integrates such probabilistic expert assessments directly into the best-subsets objective. This paper presents Expert-Implied Bayesian Best Subsets (EBBS), a method that incorporates domain-expert probability estimates of feature relevance into the MIO best-subsets problem through a maximum a posteriori (MAP) framework. Expert views from multiple respondents are aggregated into a single prior probability per feature using the Poisson binomial distribution for marginal probability estimates, the pairwise win rate for pairwise comparisons, or the normalized mean rank for ordinal rankings. This probability enters the objective function as a log-odds penalty term that smoothly encourages or discourages the selection of each feature consistent with the expert consensus. This paper provides analytic derivations of the MAP formulation and characterizes its theoretical properties. The proposed model reduces to Best Subsets when experts all have no views. Empirical results on synthetic and real datasets are forthcoming.

Nolan Alexander, Henning Mortveit

Systems and Information Engineering University of Virginia

Model Selection Probability Constrained Optimization Mixed-Integer Programming Relevance

Ask PDF