Pretraining and the Lasso

TL;DR

This paper explores how pretraining techniques can enhance the performance of the lasso regression method by developing a framework for transfer learning from large to small datasets, improving support recovery.

Contribution

It introduces a novel framework for applying pretraining to the lasso, enabling effective transfer learning and support recovery across related datasets and models.

Findings

Pretraining improves support recovery of common coefficients.

The framework applies to stratified, multinomial, and multi-response models.

Pretraining enhances scientific interpretability of model coefficients.

Abstract

Pretraining is a popular and powerful paradigm in machine learning to pass information from one model to another. As an example, suppose one has a modest-sized dataset of images of cats and dogs, and plans to fit a deep neural network to classify them from the pixel features. With pretraining, we start with a neural network trained on a large corpus of images, consisting of not just cats and dogs but hundreds of other image types. Then we fix all of the network weights except for the top layer(s) (which makes the final classification) and train (or "fine tune") those weights on our dataset. This often results in dramatically better performance than the network trained solely on our smaller dataset. In this paper, we ask the question "Can pretraining help the lasso?". We develop a framework for the lasso in which a model is fit to a large dataset, and then fine-tuned using a smaller dataset. This latter dataset can be a subset of the original dataset, or it can be a dataset with a different but related outcome. This framework has a wide variety of applications, including stratified models, multinomial responses, multi-response models, conditional average treatment estimation and even gradient boosting. In the stratified model setting, the pretrained lasso pipeline estimates the coefficients common to all groups at the first stage, and then estimates the group-specific coefficients at the second "fine-tuning" stage. We show that under appropriate assumptions, the support recovery rate of the common coefficients is superior to that of the usual lasso trained only on individual groups. This separate identification of common and individual coefficients can also be useful for scientific understanding.