High-dimensional Linear Regression Foundation Model

Abstract

Penalized linear regression is a standard tool for high-dimensional data, but in practice it requires solving many optimization problems and tuning the penalty by repeated cross-validation. In this talk, I present foundation models, pretrained on synthetic regression problems, that amortize both steps. One network returns the penalized coefficient estimate for a given penalty and tuning parameter in a single forward pass. A second network predicts the optimal tuning parameter and penalty function directly from the data.

Our goal is to make these models fast enough to run on an ordinary CPU, even when the number of features far exceeds the sample size. Attention-based tabular foundation models cannot meet this goal, because their cost grows quadratically with the number of tokens. We therefore develop new attention-free architectures whose cost grows linearly in the dimension. They combine a fast classical initialization with a few learned refinement steps, and they use compact summary statistics instead of the raw data. I will discuss the main difficulties, including instability in high dimensions and dependence on the synthetic prior. I will also present the current results and the remaining limitations.

Please click here for participating the talk online.