Multi-Environment GLAMP: Approximate Message Passing for Transfer Learning with Applications to Lasso-based Estimators

TL;DR

This paper introduces Multi-Environment GLAMP, an AMP-based framework for transfer learning across multiple data sources, providing precise risk characterizations for Lasso-based estimators with strong finite-sample validation.

Contribution

It develops a novel AMP framework for transfer learning with multiple environments and rigorously establishes its state evolution, enabling accurate risk analysis of new estimators.

Findings

Precise risk characterization of three Lasso-based transfer estimators.

Rigorous proof of state evolution for multi-environment GLAMP.

Strong finite-sample accuracy demonstrated through simulations.

Abstract

Approximate Message Passing (AMP) algorithms enable precise characterization of certain classes of random objects in the high-dimensional limit, and have found widespread applications in fields such as signal processing, statistics, and communications. In this work, we introduce Multi-Environment Generalized Long AMP, a novel AMP framework that applies to transfer learning problems with multiple data sources and distribution shifts. We rigorously establish state evolution for multi-environment GLAMP. We demonstrate the utility of this framework by precisely characterizing the risk of three Lasso-based transfer learning estimators for the first time: the Stacked Lasso, the Model Averaging Estimator, and the Second Step Estimator. We also demonstrate the remarkable finite sample accuracy of our theory via extensive simulations.