Unveiling Invariant and Transferable Latent Factors Across Heterogeneous Environments via ATLAS

2026-07-20Machine Learning

Machine Learning
AI summary

The authors study a model where data come from different environments, each with many variables, and some environments include extra labels. They assume that underlying hidden factors include parts that are the same across environments and parts that vary between them. By using a principle that some factors remain consistent (invariant), they create a method called ATLAS to separate these common and environment-specific factors, improving prediction even in new environments. Their approach works well for regression tasks and comes with strong guarantees about how accurately they can find these factors.

multi-environment factor modellatent factorsinvariance principletransfer learninghigh-dimensional covariateslatent factor regressionauxiliary labelsrobust predictionnon-asymptotic error bounds
Authors
Yihong Gu, Katherine Liao, Tianxi Cai
Abstract
This paper considers a multi-environment factor model in which high-dimensional covariates are collected from heterogeneous environments, with auxiliary labels available in a subset of these environments. The joint distribution of the covariates may vary across environments, whereas the latent structure is decomposed into invariant factors with shared loadings and heterogeneous factors with environment-specific loadings. Such a model is motivated by transfer learning and latent factor regression, where one seeks stable low-dimensional representations for both interpretation and robust out-of-sample prediction of the response $Y$. Leveraging the invariance principle, we show that the invariant and heterogeneous factors are disentangled under a minimal structural condition. Based on this, we propose ATLAS, an Auxiliary-label and invariance-guided Transfer via Latent Alignment across heterogeneous environmentS. ATLAS is a unified procedure that leverages the invariance principle to separate aligned invariant and unaligned heterogeneous factors, and further exploits supervision from auxiliary labels to extract prediction-invariant and transferable factors from those unaligned heterogeneous factors. ATLAS yields near-oracle performance for downstream latent factor regression, enables transferable prediction in new environments through the full latent signal when auxiliary labels are available, and reduces to robust invariant-factor-only prediction otherwise. We establish sharp non-asymptotic error bounds for recovering invariant and heterogeneous factors, identifying all the response-invariant factors, and estimating the invariant signal in $Y$.