MambaPFN learns to select relevant variables before seeing a new dataset

A dataset can contain dozens of measurements, yet only a few may provide information needed to predict an outcome. Identifying those variables becomes harder when their relationships are curved, interactive or otherwise nonlinear. In a preprint, Jin Zhu and Junxian Zhu propose MambaPFN, a Transformer-based system that learns variable selection across synthetic datasets, then applies that learning to new datasets without updating its parameters.

The approach moves computation rather than eliminating it. Pretraining is expensive, the authors acknowledge, but its cost can be shared across subsequent selection tasks. Their experiments suggest that this trade-off can deliver fast selection and useful predictive subsets within the settings studied.

Learning what makes a variable relevant

MambaPFN targets the response’s minimal Markov blanket: the smallest group of variables that makes the remaining variables uninformative about the response once that group is known. This is a conditional-information target, not simply a list of columns with strong individual correlations.

The paper assumes that the joint distribution has a strictly positive density, excluding exact deterministic relationships between variables. Under this assumption, its target set is unique. Although synthetic causal graphs help generate training examples, selecting this set should not be confused with establishing which variables cause an outcome in real data.

Training requires examples with known answers. The researchers generate directed acyclic graphs, extract the response’s graph-based blanket, and simulate observations using nonlinear functions and independent noise. Because a graph’s blanket can differ from the distribution’s minimal blanket when dependencies cancel, they impose restrictions on local graph structures, functional dependence and noise distributions to ensure that the training labels match their statistical target.

The network learns a score for each variable using a binary log-loss objective. The authors show that, under their assumptions and with a sufficiently expressive function class, the population-level optimum represents posterior inclusion probabilities. Those probabilities depend on the synthetic data-generating family; the result does not guarantee that a trained finite network produces exact probabilities for arbitrary real datasets. Thresholding the scores produces a selected subset.

A Transformer built for tables

The architecture processes information along both the observation and variable axes. A special embedding distinguishes the response from predictors, while the design respects the fact that reordering observations should not change the selection and reordering predictors should reorder the outputs accordingly.

Despite its name, MambaPFN uses a Transformer. The paper mentions selective state-space models only as a possible direction for future efficiency improvements.

On held-out synthetic datasets generated through the same procedure, MambaPFN typically achieves the highest F1-score and ranking AUC among the tested methods. Comparisons include linear selectors, nonlinear approaches and CauScale, an amortized causal-graph learner whose estimated graph is used to extract a blanket. These results support performance within the tested synthetic family, rather than demonstrating unrestricted generalization.

The authors also find decreasing validation loss as they increase pretraining datasets up to 320,000 and model size to roughly 11 million parameters. This is empirical scaling behavior over the examined ranges, not proof that improvements continue indefinitely.

Three variables in a Facebook example

The real-data experiment uses 40,949 Facebook posts with 53 predictors. Across ten repetitions, the researchers separate observations for variable selection, downstream neural-network training and final testing. Because the true relevant set is unknown, predictive performance serves as a surrogate measure of selection quality.

The authors report that MambaPFN selects the fewest variables and achieves the highest test R², with an improvement approaching 20% over the best baseline. Its checkpoint was pretrained with at most 35 predictors, so this experiment also tests extrapolation to a wider table.

A separate full-data descriptive analysis selects elapsed time since publication, comments in the latest 24 hours, and the change in comment counts between consecutive 24-hour periods. The latter shows a roughly V-shaped association with the log-transformed response, illustrating why a straight-line relationship may miss useful structure. This descriptive analysis is distinct from the held-out predictive evaluation.

Where the evidence stops

Fast selection is not the same as cheap end-to-end learning: pretraining remains computationally costly, and prediction requires a separately trained model. Runtime comparisons also involve different hardware: MambaPFN and CauScale use an NVIDIA A100 GPU, while other baselines run on CPUs.

The Facebook example provides encouraging evidence beyond synthetic data, but its unknown true blanket means better prediction cannot establish exact variable recovery. As a preprint, the work presents an early case for amortizing variable selection—not a universal replacement for existing methods.