Can a pretrained transformer amortize spike-and-slab variable selection with interactions?

LinearPFN amortizes spike-and-slab variable selection with interactions, returning inclusion probabilities in one forward pass.

Direct answer

Yes, within a stated prior. LinearPFN is a transformer pretrained once on synthetic datasets drawn from an explicit conjugate spike-and-slab prior over main effects and pairwise interactions. A single forward pass over a new dataset returns posterior inclusion probabilities, posterior-mean coefficients and posterior predictive distributions with no per-dataset fitting, and on real predictor matrices with outcomes drawn from that prior it attains a higher per-dataset selection AUC and a higher F1 under the median probability model rule than five classical baselines.

1sources cited

This article was generated with WisPaper-powered search and paper analysis.

What amortization replaces in spike-and-slab selection

Spike-and-slab regression does not return one selected subset. Its prior sets each coefficient either exactly to zero (the spike) or draws it from a Gaussian (the slab), so the posterior is a distribution over active sets, and every candidate effect carries a posterior inclusion probability (PIP). With p predictors the candidate effects are an intercept, p main effects and p choose 2 pairwise interactions, d = 1 + p + p choose 2 in total. With the intercept always included there are 2^(d-1) models: at p = 10 that is about 3.6 x 10^16 models.

Because the exact posterior is a sum over all models, it can be enumerated only at small p. Beyond that reach, practice relies on Markov chain Monte Carlo over the model space, adaptive sampling or variational approximations, which need a fresh fitting run for every dataset and, within a fixed budget of steps, may fail to converge. LinearPFN takes a different route: it is pretrained once on synthetic datasets drawn from an explicitly specified prior, and inference on a new dataset is one forward pass rather than a per-dataset fit.

Three outputs from one forward pass

The network has three read-out heads. The selection head returns a PIP for every candidate effect, one MLP scoring each main effect and a second scoring each interaction. The coefficient head returns one number per effect, the posterior mean E[beta_e | T], which is the marginal mean averaged over all models. The regression head returns a posterior predictive distribution for each query row.

The two coefficient quantities are easy to confuse. The conditional mean E[beta_e | T, gamma_e = 1] is the size of the effect given that it is active; the marginal mean is linked to it by E[beta_e | T] = Pr(gamma_e = 1 | T) E[beta_e | T, gamma_e = 1], so an effect with a PIP of one half receives half its conditional mean. The coefficient output is dense: every candidate effect receives a value and none is set to zero by a selection step. Hard selection happens separately, by the median probability model rule, which selects every effect whose PIP exceeds one half.

A prior built for conjugacy and social-science predictors

Training data come from a prior over data-generating processes. For row i the outcome follows y_i = beta_0 + sum_j beta_j x_ij + sum_{j<k} beta_jk x_ij x_ik + epsilon_i, with Gaussian noise of variance sigma^2. The expanded predictor matrix Z = [1, X, (x_j elementwise-times x_k)_{j<k}] has one column per candidate effect. An inactive effect has coefficient zero; an active main effect has variance c sigma^2, an active interaction c rho sigma^2, and the intercept tau_0^2 sigma^2, so the slab scale c measures effect size in units of the noise variance.

The prior is conjugate by design, so the posterior for each fixed set of active effects has a closed form, and wherever the exact posterior is still computable by enumeration the network's outputs are verified against it. The predictor prior mimics individual-level social-science data: a factor block standing in for multi-item scales, a generic block of correlated columns, shuffled column order so position carries no information, and a Gaussian copula mapping each column to Gaussian, Likert, binary, count or skewed marginals. The active-effects prior enforces strong heredity, so an interaction is eligible only when both parent main effects are active.

How the network is trained and scored

LinearPFN has 31.5M parameters and was trained for 37,500 optimizer steps at an effective batch of 64 datasets, 2.4M datasets in total, on two H200 GPUs. Training minimizes the sum of three terms: the negative log-likelihood of the regression head on query rows, a binary cross-entropy between the selection logits and the simulated active indicators, and a mean squared error between the coefficient outputs and the simulated coefficients. The second term is a proper scoring rule whose expected minimizer is the posterior probability that an effect is active, so the selection head is trained toward the exact PIPs; the third is trained toward the posterior mean of the coefficient.

Each transformer block has 16 attention heads. Within a layer, each cell first attends to the other cells in its row and then to the other cells in its column, and an MLP then updates each cell on its own. The per-effect heads read the context rows only, and their outputs do not depend on the order of those rows. The regression head uses a bar distribution with 512 bins of equal width on [-8, 8] and a half-normal tail on each side.

Selection lead on real predictor matrices

The main benchmark uses real predictor matrices from published social-science datasets, with outcomes drawn from the prior so that the true active set is known. The 383 predictor matrices come from 20 packages in 7 fields: economics (237), political science (58), general social science (32), education (24), criminology (17), psychology (14) and demography (1). Their median numbers of rows and predictors are 758 and 9, and 87 of them have at most 5 predictors.

On these datasets LinearPFN attains a higher per-dataset selection AUC and a higher F1 under the median probability model rule than five classical baselines, and the lead holds when the coefficients, the interactions or the noise depart from the prior. The two metrics answer different questions: AUC measures per-dataset ranking of effects, while F1 measures the harmonic mean of precision and recall of a hard selection against the true active set. The paper reports that in coefficient error LinearPFN only matches the best baseline, so its advantage is in selection rather than in coefficient estimation.

Speed against an MCMC sampler over the same posterior

Amortization reduces variable selection to one forward pass, which returns the selected variables and their coefficients in 0.03 seconds at the median on one GPU. The MCMC sampler over the same posterior needs 57.43 seconds at the median on one CPU thread and meets its convergence criteria on 93.0% of runs. The comparison is between one GPU forward pass and one CPU thread of sampling, so it is a comparison of hardware as well as of method.

The speed claim does not carry over to every setting. Without a GPU, LinearPFN is slow: on one CPU thread it is the second slowest of the six methods, ahead of HIERNET only. The paper notes that neither the CPU cost nor the fixed prior is intrinsic to the approach, and that pruning, quantization and knowledge distillation are established routes to a cheaper forward pass that were not pursued here.

Where the selection lead stops

The network is trained on a single prior that assumes strong heredity, Gaussian noise, numeric predictors and independent rows. Its selection lead held when the coefficients, the interactions or the noise departed from this prior, but the prior itself is fixed, so a user cannot state their own prior without retraining. The paper suggests follow-up work could condition the network on the prior's hyperparameters.

Collinearity is the sharpest boundary. When two predictors are nearly collinear, the posterior mass splits across them, both PIPs can fall below one half, and the median probability model rule then selects neither. Unlike SUSIE, LinearPFN reports no credible sets that group such predictors. The benchmark also has a structural feature worth noting: the density k_main / p takes only multiples of 1/p, so none of the 2,299 datasets with p <= 5 lies in the lowest density bin, and a comparison across grid cells mixes the effect of density and R^2 with the effect of the predictor matrices those cells hold.

When one forward pass pays off

The decision this paper supports is about repetition. Bayesian model averaging with interactions becomes feasible where it has to be repeated many times, because the per-dataset cost collapses from a sampling run to a single forward pass. The trade is a fixed prior: the network returns probabilities under a stated prior, and its outputs do not by themselves support causal claims.

The threshold that matters most in practice is the one-half PIP rule. It is what converts the dense coefficient output into a selected subset, and it is exactly where near-collinear predictors fail, because split posterior mass can push both PIPs below one half and leave neither selected. A reader weighing LinearPFN against an MCMC sampler should therefore compare them on the regime they actually face: many repeated datasets with a prior close to the one the network was trained on favor the amortized forward pass, while near-collinear predictors or a prior the user needs to set themselves favor a method that reports credible sets or accepts a user-specified prior.

Sources used in this answer

1

LinearPFN: Amortized Variable Selection for Linear Models with Interactions

LinearPFN is a prior-data fitted transformer that amortizes spike-and-slab inference for linear models with main effects and pairwise interactions; one forward pass returns posterior inclusion probabilities, posterior-mean coefficients and posterior predictive distributions, and on real predictor matrices with outcomes drawn from its prior it attains a higher per-dataset selection AUC and a higher F1 under the median probability model rule than five classical baselines.