Amortized Nonlinear Variable Selection with Prior-Data Fitted Networks
Abstract
Variable selection is an important task for interpreting potentially complex systems in fields such as biology, business, and sociology. The importance of variable selection has attracted considerable attention in statistics and machine learning. In this paper, we are mainly interested in model-agnostic variable selection for identifying variables that have general nonlinear dependence on the response. We propose MambaPFN, short for Model-Agnostic Markov BlAnket selection with Prior-data Fitted Networks, which is built upon the Transformer, a building block of many popular large language models. Specifically, we consider a functional-mapping perspective on variable-selection learning, in which the Transformer is treated as a learnable mapping from a dataset to its relevant variables. To learn this mapping, MambaPFN introduces three new components: (i) a faithful, rich, and scalable scheme for generating pretraining data, (ii) an objective that targets each variable's posterior inclusion probability, and (iii) a modified Transformer architecture tailored to variable selection. Once pretrained, MambaPFN is applied to new datasets without parameter updates. Empirically, MambaPFN achieves desirable variable-selection performance in simulation studies and exhibits a scaling law---more pretraining datasets and larger models can continuously improve its performance. Applied to a Facebook comments dataset, MambaPFN selects the fewest variables and achieves the highest computational efficiency and predictive performance, improving upon the best baselines by about 20%.
Related
- SelectPFN: Amortized Bayesian Variable Selection for Sparse Nonlinear Regression with Prior-Data Fitted Networks
- Prior-data fitted networks for variable selection
- Debiased Prior-Data Fitted Networks for Amortized Bayesian Semi-Parametric Inference
- LinearPFN: Amortized Variable Selection for Linear Models with Interactions
- Amortized Structured Stochastic Variational Inference for Gaussian Process Latent Variable Models
BibTeX
@misc{zhu2026amortized,
title = {{Amortized Nonlinear Variable Selection with Prior-Data Fitted Networks}},
author = {Jin Zhu and Junxian Zhu},
howpublished = {arXiv},
year = {2026},
month = nov,
doi = {10.48550/arxiv.2611.04848},
eprint = {2611.04848},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2611.04848},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Amortized Nonlinear Variable Selection with Prior-Data Fitted Networks
AU - Jin Zhu
AU - Junxian Zhu
PY - 2026
DA - 2026/11/03
PB - arXiv
DO - 10.48550/arxiv.2611.04848
UR - https://arxiv.org/abs/2611.04848
AB - Variable selection is an important task for interpreting potentially complex systems in fields such as biology, business, and sociology. The importance of variable selection has attracted considerable attention in statistics and machine learning. In this paper, we are mainly interested in model-agnostic variable selection for identifying variables that have general nonlinear dependence on the response. We propose MambaPFN, short for Model-Agnostic Markov BlAnket selection with Prior-data Fitted Networks, which is built upon the Transformer, a building block of many popular large language models. Specifically, we consider a functional-mapping perspective on variable-selection learning, in which the Transformer is treated as a learnable mapping from a dataset to its relevant variables. To learn this mapping, MambaPFN introduces three new components: (i) a faithful, rich, and scalable scheme for generating pretraining data, (ii) an objective that targets each variable's posterior inclusion probability, and (iii) a modified Transformer architecture tailored to variable selection. Once pretrained, MambaPFN is applied to new datasets without parameter updates. Empirically, MambaPFN achieves desirable variable-selection performance in simulation studies and exhibits a scaling law---more pretraining datasets and larger models can continuously improve its performance. Applied to a Facebook comments dataset, MambaPFN selects the fewest variables and achieves the highest computational efficiency and predictive performance, improving upon the best baselines by about 20%.
ER -