Fast and robust invariant generalized linear models

Abstract

Statistical integration of diverse data sources is an essential step in the building of generalizable prediction tools, especially in precision health. The invariant features model is a new paradigm for multi-source data integration which posits that a small number of covariates affect the outcome identically across all possible environments. Existing methods for estimating invariant effects suffer from immense computational costs or only offer good statistical performance under strict assumptions. In this work, we provide a general framework for estimation under the invariant features model that is computationally efficient and statistically flexible. We also provide a robust extension of our proposed method to protect against possibly corrupted or misspecified data sources. We demonstrate the robust properties of our method via simulations, and use it to build a transferable prediction model for end stage renal disease using electronic health records from the All of Us research program.

1. Introduction

Data integration has become a central challenge of contemporary biomedical research. As research networks spread across institutions and even geographical boundaries, practitioners need to think carefully about how to handle heterogeneity and similarity between diverse data sources. This problem is especially pertinent in precision medicine (Martínez-García & Hernández-Lemus, 2022). Distributional shifts between patient subgroups, manifesting themselves in the training and testing data, may severely impact downstream model performance, which in turn can adversely affect decision making at the patient level (Green et al., 2024). If statistical machine learning systems are to play a role in the clinic, statistical researchers must develop models and methods that enable domain experts to better understand and, in some cases, protect against discrepancies between sources of medical data.

One such model that is growing in popularity is the invariant prediction or features model proposed by Peters et al. (2016). Here, the authors presume that the analyst is given data as outcome-feature pairs from a discrete set of sources, referred to as environments, that may correspond to different sub-populations of observations, experimental conditions, or any other source of heterogeneity. Letting $ℰ$ denote this set of environments, the invariant features model assumes that there exists a subset $S$ of features such that for any two environments $e$, $f\in ℰ$, we have

$$\mathbb{E}[Y^e\mid X_S^e=x]=\mathbb{E}[Y^f\mid X_S^f=x]$$ (1)

where $(Y^e,X^e)\in \mathbb{R}\otimes {\mathbb{R}}^p$ represents the outcome and feature vector drawn from environment $e$. The features in $S$ are invariant with respect to the environments $ℰ$, and knowledge of $S$ can be used to build prediction tools that generalize well to previously unseen environments satisfying (1). This model is particularly amenable to large-scale electronic health records data. Due to differences across healthcare systems such as in clinical workflows, patient demographics, as well as EHR systems implementations, many features derived from EHR may admit effects that are idiosyncratic to their source institution, and hence disease prediction models built using these features will be useless if applied to data from a new clinical setting (Sarwar et al., 2022; Sauer et al., 2022; Singh et al., 2022). Knowledge of a set of features satisfying (1) would greatly improve EHR-based phenotyping and disease screening tools across institutions, and lead researchers to a better understanding of disease symptoms.

The interpretability and potential applicability of the invariant features model has inspired a litany of works exploring methods for estimating the conditional mean $\mathbb{E}[Y^e\mid X_S^e]$ and the set $S$ (Fan et al., 2024; Heinze-Deml et al., 2018; Pfister et al., 2019; Rojas-Carulla et al., 2018; Wang et al., 2024; Yin et al., 2024). While the community has made fundamental strides in building tools for estimation under (1), existing methods are routinely hampered by computational infeasibility. For instance, the original Invariant Causal Prediction algorithm of Peters et al. (2016) requires running a hypothesis test for each $S\subset [p]$, and the Environment Invariant Linear Least Squares estimator of Fan et al. (2024) is defined as the solution to novel subset selection problem which even recent advancements in integer optimization (Bertsimas et al., 2016) cannot solve efficiently. This computational cost totally prohibits these methods from being useful for reasonable data analysis tasks. Recent works such as Yin et al. (2024) and Wang et al. (2024) attempt to design computationally faster methods, but do so by imposing stricter assumptions on the data generating process or by requiring a-priori knowledge of the set $S$. When these assumptions are not satisfied or such information is unavailable, the analyst has no option but to use an exponentially slow method.

We make two contributions in the present work. First, we provide a computationally efficient procedure for estimating $\mathbb{E}[Y^e\mid X_S^e]$ that is valid for any generalized linear model satisfying (1). Our approach uses a linear relaxation of the estimator proposed in Fan et al. (2024), and we propose an alternating minimization scheme to compute our estimator quickly and accurately. We then propose a robust version of our method that is able to maintain good statistical performance even when some environments fail to satisfy the condition in (1). Robustness is invaluable in real-world settings, as data sources may suffer from measurement error, corruption, or outright misspecification. We demonstrate through an analysis of synthetic data that our proposed method and its robust extension perform well under a variety of settings. Then, we use a real EHR dataset from the All of Us program (of Us Research Program Investigators, 2019) to build a risk prediction model for end stage renal disease. We show that our proposed methods give the best predictive performance on previously unseen environments even when some of the training data is corrupted.

1.1. Related literature

Our work directly builds on the invariant features model first proposed by the seminal work of Peters et al. (2016). The authors of Peters et al. (2016) propose a multiple testing based procedure for recovering the invariant features $S$ under a linear assumption on (1). In Heinze-Deml et al. (2018) and Pfister et al. (2019), these authors extend this idea to nonlinear and time series models respectively. The first theoretical guarantees on estimating the conditional mean $\mathbb{E}[Y^e\mid X_S^e]$ under a linear model are given in Fan et al. (2024), and these authors extend their results to nonlinear models in Gu et al. (2024). The work of Yin et al. (2024) also aims to estimate $\mathbb{E}[Y^e\mid X_S^e]$ under a linear model and does so using a relaxation similar to the one proposed in the present work. They provide theoretical results showing that their CoCo procedure can identify the invariant effects if the analyst has prior knowledge on the support $S$. The recent work of Wang et al. (2024) uses a non-convex optimization program to estimate $\mathbb{E}[Y^e\mid X_S^e]$, requiring that distinct environments are generated from additive interventions. Finally, Li and Zhang (2024) studies an invariant features-type model under fairness constraints, and provides algorithms for estimation and feature screening under this model. In the sense that invariant features can be understood as causal features (Peters et al., 2016) for the outcome $Y^e$, the cited works address causal discovery and causal effect estimation under particular parameterizations of the conditional mean $\mathbb{E}[Y^e\mid X_S^e]$. While we do not focus on causality in the present work, the reader may keep this interpretation of our proposed algorithm in mind, as it constitutes a fast algorithm for estimating causal effects under the invariant generalized linear model.

This line of work is situated within the broader statistical data integration literature. The invariant features model is closely related to the multi-task learning problem (Zhang & Yang, 2018), as distinct environments can be thought of as different tasks that are modeled simultaneously. The papers Maurer et al. (2013) and Maurer et al. (2016) initiate the contemporary theoretical study of multi-task learning using classical tools from empirical process theory. The more recent work of Tripuraneni et al. (2021) gives sharper results in the multi-task linear model under a low rank feature representation. This model is fully explored in Duchi et al. (2022), Niu et al. (2024), and Tian et al. (2023). In Duan and Wang (2023), the authors provide a general framework for robust multi-task learning when some tasks may be misspecified. Knight and Duan (2024) gives theoretical results for sparse and low rank multi-task learning under data sharing constraints. The invariant features model can also be conceived of as a relaxation of the covariate shift assumption as explored in Liu et al. (2023) and Ma et al. (2023). Here, the end goal is typically to adjust for differences between a source dataset and target data in a transfer learning task. Transfer learning has grown in popularity in its own right, as recent works provide tools for transfer learning under high-dimensional linear models (Li et al., 2022), generalized linear models (Li et al., 2024), as well as in federated settings (Li et al., 2023).

In the machine learning literature, researchers study the data integration problem under the umbrella of domain adaptation (Ben-David et al., 2006). A recent seminal work in this area is the Invariant Risk Minimization framework of Arjovsky et al. (2020), which aims to find a feature representation that mitigates domain-specific spurious correlations. This framework is conceptually similar to the invariant features model, although the community has recently pointed out its shortcomings (Kamath et al., 2021; Rosenfeld et al., 2021). We refer readers to the survey Farahani et al. (2021) for a general overview of developments in this field.

1.2. Structure

The rest of the paper is structured as follows. In Section 2.1 we outline the version of the invariant features model that we consider in detail, as well as describe the primary statistical challenges arising under this model. In Section 2.2 we describe our proposed Fast Invariant Generalized Linear Model (FILM) estimator and provide an alternating optimization procedure to compute it. Next, we describe the Robust-FILM method in Section 2.3. Section 3 gives simulation results demonstrating our method in comparison other methods in the literature, and in Section 4 we use FILM to build a risk prediction model for chronic kidney disease using electronic health records for the All of Us program.

2. Method

2.1. Model and problem setup

We let $ℰ$ denote the set of environments, and for each $e\in ℰ$, we observe iid pairs ${(Y_i^e,X_i^e)}_{i=1}^{n_e}$ for $Y_i^e\in \mathbb{R}$ and $X_i^e\in {\mathbb{R}}^p$. We will use ${\mathbf{Y}}^e\in {\mathbb{R}}^{n_e}$ to refer to the vector of outcomes and ${\mathbf{X}}^e\in {\mathbb{R}}^{n_e\times p}$ to indicate the matrix of covariates from environment $e$. Following the invariant features model proposed by Peters et al. (2016), we assume that there exists a set of features $S\subset [p]$ such that the following generalized linear model holds for all $e\in ℰ$:

$$g(\mathbb{E}[Y_i^e\mid X_{i,S}^e])={\theta }^e+{\beta }_S^{\mathsf{T}}{X}_{i,S}^e$$ (2)

where $g$ is a known canonical link function and ${\theta }^e\in \mathbb{R}$ denotes an environment specific intercept term. We collect $\theta ={({\theta }^1,\dots ,{\theta }^{\mid ℰ\mid })}^{\mathsf{T}}\in {\mathbb{R}}^{\mid ℰ\mid }$. Our aim is to estimate the vector $\beta \in {\mathbb{R}}^p$ which has support $\text{supp}(\beta )=S$. We emphasize that $\beta$ is shared across all environments; for this reason, we refer to $\beta$ as the vector of invariant effects and the features $j\in S$ as invariant features.

Let

$${\epsilon }_i^e=Y_i^e-\mathbb{E}[Y_i^e\mid X_{i,S}^e]=Y_i^e-g^{-1}({\theta }^e+{\beta }_S^{\mathsf{T}}{X}_{i,S}^e).$$

Under the moment condition assumed in Model (2), we have that $\mathbb{E}[X_{i,j}^e{\epsilon }_i^e]=0$ for each $j\in S$, but we have no guarantee that $\mathbb{E}[X_{i,j}^e{\epsilon }_i^e]=0$ for any $j\notin S$. We call any $j$ with $\mathbb{E}[X_{i,j}^e{\epsilon }_i^e]\ne 0$ a spurious feature in environment $e$ with spurious correlation $\mathbb{E}[X_{i,j}^e{\epsilon }_i^e]$. Including a spurious feature in a prediction model may enhance accuracy within a specific environment but will likely reduce the model’s transferability to other environments. Importantly, we will not presuppose knowledge of the set $S$; rather we observe the pairs ${(Y_i^e,X_i^e)}_{i=1}^{n_e}$ blindly, without knowing which features are invariant and which are spurious. This distinguishes the invariant features model in (2) from, for instance, the instrumental variables model (Angrist et al., 1996), which requires knowledge of which features are uncorrelated with the residual term. This adds further motivation to our aim of estimating $\beta$, as a good estimate of these invariant effects may lead us to the features that affect the outcome identically across environments.

The primary difficulty in estimating $\beta$ under Model (2) arises from the potential presence of spurious features across the environments. Naively using a maximum likelihood or empirical risk minimization estimator in the presence of spurious features will lead to highly biased estimates of the invariant effects. To demonstrate this, suppose that $g(x)=x$ which admits the linear model for $Y_i^e$ in $X_{i,S}^e$, and suppose that $Y^e$ is appropriately centered so that ${\theta }^e=0$. If we use the single-environment ordinary least squares estimator for $\beta$, an elementary calculation reveals

$${\widehat{\beta }}_{\text{OLS}}=\beta +{\left(\frac{1}{n_e}\sum _{i=1}^{n_e}{X}_i^e{(X_i^e)}^{\mathsf{T}}\right)}^{-1}\frac{1}{n_e}{({\mathbf{X}}^e)}^{\mathsf{T}}{\epsilon }^e\stackrel{\mathbb{P}}{\to }\beta +{({\Sigma }^e)}^{-1}\mathbb{E}[X^e{\epsilon }^e]$$ (3)

where the limit in probability uses the WLLN and Slutksy’s theorem, and assumes that the covariates have a finite second moment ${\Sigma }^e$. When $\mathbb{E}[X^e{\epsilon }^e]\ne 0$ (i.e., if our data includes any spurious features), the OLS estimator is not even consistent. A natural idea to improve this estimator is to run OLS on the pooled data from all the environments $ℰ$, resulting in the estimator

$${\widehat{\beta }}_{\text{Pooled-OLS}}=\underset{\beta }{\text{argmin}}\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}\frac{1}{n_e}{\Vert {\mathbf{Y}}^2-{\mathbf{X}}^e\beta \Vert }_2^2$$ (4)

However, this approach also incurs bias from spurious features, as Proposition 4.1 in Fan et al. (2024) shows

$${\Vert {\widehat{\beta }}_{\text{Pooled-OLS}}-\beta \Vert }_2\asymp {\Vert \frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}\mathbb{E}[X^e{\epsilon }^e]\Vert }_2$$ (5)

Thus, simply pooling data from distinct environments is insufficient to address the bias from spurious features. To overcome this, a strategy commonly adopted in the literature (Arjovsky et al., 2020; Fan et al., 2024; Yin et al., 2024) is to augment a pooled loss function with a gradient-based penalty term which exploits our assumption that the data from each environment satisfies the moment condition in Model (2). This is the path that we will take to derive our estimator of $\beta$ using the general framework outlined in Fan et al. (2024).

Under Model (2), our data from environment $e$ admit a negative log-likelihood of the form

$$L^e(\beta ,{\theta }^e)=\frac{1}{n_e}\sum _{i=1}^{n_e}\{\psi ({\beta }^{\mathsf{T}}{X}_i^e+{\theta }^e)-Y_i^e\cdot ({\beta }^{\mathsf{T}}{X}_i^e+{\theta }^e)\}$$ (6)

where the function $\psi$ is uniquely determined by $g$ via $\psi =\int g^{-1}$. Note that we consider $L^e$ as a function on ${\mathbb{R}}^p\times \mathbb{R}$; for $j=\{1,\dots ,p\}$, we use the notation ${\nabla }_j{L}^e(\beta ,{\theta }^e)=\frac{d}{d{\beta }_j}L(\beta ,{\theta }^e)$. The key observation (Fan et al., 2024; Yin et al., 2024) is that for each $j\in S$, we have $\mathbb{E}[{\nabla }_j{L}^e(\beta ,{\theta }^e)]=0$ for all $e\in ℰ$. This follows directly from (2). Since ${\nabla }_j{L}^e$ is an average of iid observations, it follows that ${\nabla }_j{L}^e(\beta ,{\theta }^e)\approx 0$ for all $j\in S$, $e\in ℰ$, indicating that any reasonable estimators of $\beta$ and $\theta$ should satisfy

$${\nabla }_j{L}^e(\widehat{\beta },{\widehat{\theta }}^e)\approx 0$$ (7)

for all $j\in \text{supp}(\widehat{\beta })$ and $e\in ℰ$. This is the intuition behind the Environment Invariant Linear Least Squares (EILLS) method proposed by Fan et al. (2024). Their EILLS estimator is defined as

$${\widehat{\beta }}_{\text{EILLS}}=\underset{\beta }{\text{argmin}}\left\{\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}{L}^e(\beta ,0)+{\lambda }_1\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}\sum _{j=1}^{p}1({\beta }_j\ne 0){\left({\nabla }_j{L}^e(\beta ,0)\right)}^2+{\lambda }_2{\Vert \beta \Vert }_0\right\}$$ (8)

where they assume that ${\theta }^e=0$ for all $e\in ℰ$. The first penalty term phases out spurious features by enforcing the moment condition (7) across all of the environments and the second penalty removes features that do not admit any spurious correlation but have no effect on the outcome. The pooled loss function ensures that the resulting estimator maintains good predictive performance which excludes trivial solutions to (7), such as the zero vector in ${\mathbb{R}}^p$. The authors of Fan et al. (2024) provide a complete set of theoretical results for ${\widehat{\beta }}_{\text{EILLS}}$ in the linear model case (i.e. when Model (2) is satisfied with $g(x)=x$). They show that, under mild conditions, the true vector of invariant effects $\beta$ is identified by the solution to the population-level analog of the EILLS optimization problem in (8), and give statistical results demonstrating that ${\widehat{\beta }}_{\text{EILLS}}$ is an efficient estimator of $\beta$. We refer the reader to Theorems 4.2, 4.3, and 4.4 in Fan et al., 2024 for details.

Unfortunately, computing ${\widehat{\beta }}_{\text{EILLS}}$ is infeasible for even moderate $p$ due to the dependence on the support of $\beta$ in the invariance-inducing penalty. The authors of Fan et al. (2024) propose a brute force approach, in which the user of the EILLS method iterates through every possible support set $S\subseteq [p]$ and solves Equation (8) over vectors with support $S$, then taking the vector that gives the best loss value as the resulting estimator. In other words, computing ${\widehat{\beta }}_{\text{EILLS}}$ requires solving a quadratic program for every set $S\subseteq [p]$, incurring a runtime of at least $\Omega (2^p)$. This prevents the EILLS estimator from being a practical tool for building transferable models in any realistic data analysis setting, as $p$ may be in the hundreds, thousands, or even millions in contemporary datasets.

2.2. Proposed method

Our first contribution is a new, computationally efficient estimator of $\beta$ that builds on the EILLS framework. We derive our method via a linear relaxation of the minimization problem in (8) as follows. We parametrize $\beta$ as an element-wise product $\beta =a\odot b$ where $a\in {[0,1]}^p$ is a relaxation of the support of $\beta$ and $b\in {\mathbb{R}}^p$ contributes the magnitude of each entry. Although the true coefficient vector $\beta$ does not admit a unique parametrization of this form, relaxing the support indicator to lie in [0, 1] circumvents the combinatorial nature of the EILLS program in Equation (8). Using this re-parametrization, we define a relaxed loss function as

$$Q_{ℰ}(a,b,\theta )=\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}{L}^e(a\odot b,\theta )+{\lambda }_1\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}\sum _{j=1}^p{a}_j^2{\left({\nabla }_j{L}^e(a\odot b,\theta )\right)}^2+{\lambda }_2\sum _{j=1}^p{a}_j(1-a_j)$$ (9)

Our Fast Invariant Generalized Linear Model (FILM) estimator is defined as

$${\widehat{\beta }}_{\text{FILM}}=\widehat{a}\odot \widehat{b};(\widehat{a},\widehat{b},\widehat{\theta })=\underset{a\in {[0,1]}^p,b\in {\mathbb{R}}^p,\theta \in {\mathbb{R}}^{\mid ℰ\mid }}{\text{argmin}}{Q}_{ℰ}(a,b,\theta )$$ (10)

We include the second penalty in Equation (9) to encourage the entries of $a$ to live near either 0 or 1, which improves interpretability of the resulting FILM estimator as well as our method’s agreement with the original EILLS approach. Additionally, we use $a_j^2$ in the first penalty instead of $a_j$ to further down-weight features that have a non-zero gradient while maintaining a functional form that is easy to minimize.

While our proposed loss function $Q_{ℰ}(a,b)$ removes the combinatorial aspect of the optimization problem, it is still non-trivial to compute ${\widehat{\beta }}_{\text{FILM}}$. This is because $Q_{ℰ}$ is jointly non-convex in $a$ and $b$ due to the element-wise products in the pooled loss function and the first penalty term. Solving the optimization problem in Equation (10) directly, i.e. via gradient descent, is not guaranteed to be a global minimizer of $Q_{ℰ}$. To overcome this, we propose an iterative alternating minimization scheme that optimizes over $a$, $b$, and $\theta$ separately, and leverages a quadratic surrogate objective function that improves computational efficiency.

To describe the ${(k+1)}_{\text{th}}$ step of our algorithm, let ($a^{(k)}$, $b^{(k)}$, ${\theta }^{(k)}$) denote the values of $a$, $b$, and $\theta$ obtained from the $k_{\text{th}}$ iteration. Our first aim is to compute $a^{(k+1)}$. Letting ${\beta }^{(k)}=a^{(k)}\odot b^{(k)}$, we introduce the function

$${\overline{Q}}_{ℰ}(a;b^{(k)},{\theta }^{(k)},{\beta }^{(k)})=\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}{L}^e(a\odot b^{(k)},{\theta }^{(k)})+\frac{{\lambda }_1}{\mid ℰ\mid }\sum _{e\in ℰ}\sum _{j=1}^p{a}_j^2({\nabla }_j{L}^e{({\beta }^{(k)},{\theta }^{(k)})}^2+{\lambda }_2\sum _{j=1}^p{a}_j(1-a_j)$$ (11)

We then define $a^{(k+1)}$ as

$$a^{(k+1)}=\underset{a\in {[0,1]}^p}{\text{argmin}}\overline{Q}(a;b^{(k)},{\theta }^{(k)},{\beta }^{(k)})$$ (12)

Our rationale for optimizing over ${\overline{Q}}_{ℰ}$ instead of $Q_{ℰ}$ with $b^{(k)}$ and ${\theta }^{(k)}$ held constant is purely computational. By fixing ${\beta }^{(k)}$ in the gradient term of the invariance-inducing penalty, we simplify both penalties to be quadratic in $a$, granting us a function that is very efficient to minimize. We also avoid potential spurious minima that may be introduced by optimizing over the products $a\mapsto a_j^2{({\nabla }_j{L}^e(a\odot b^{(k)},{\theta }^{(k)}))}^2$ that arise in the form of $Q_{ℰ}$. To compute $b^{(k+1)}$ and ${\theta }^{(k+1)}$, we optimize $Q_{ℰ}$ directly, granting

$$b^{(k+1)}=\underset{b\in {\mathbb{R}}^p}{\text{argmin}}{Q}_{ℰ}(a^{(k+1)},b,{\theta }^{(k)})$$ (13)

$${\theta }^{(k+1)}=\underset{\theta \in {\mathbb{R}}^{\mid ℰ\mid }}{\text{argmin}}{Q}_{ℰ}(a^{(k+1)},b^{(k+1)},\theta )$$ (14)

We point out that the runtime of each constituent optimization problem is quadratic, meaning that optimizing over $b$ and $\theta$ separately incurs a runtime of $O(p^2+{\mid ℰ\mid }^2)$, where optimizing over them jointly runs in $O({(p+\mid ℰ\mid )}^2)$. When $p$ is large, optimizing over $b$ and $\theta$ separately may therefore provide a nontrivial improvement in runtime, which motivates us to adopt this approach rather than solving for $b$ and $\theta$ jointly. We also allow the user to leverage domain specific knowledge of the data generating mechanism to improve the performance of our algorithm. In many application areas, for instance healthcare or biomedical science, prior studies can provide evidence that particular features indicated by the set $J$ have invariant or causal effects on the outcome of interest. In this case, we want to enforce $a_j=1$ for all $j\in J$ in the optimization process. This is straightforward to achieve in our computation of $a^{(k+1)}$ by adding an additional constraint to the optimization problem in Equation (12). We demonstrate via simulation studies in Section (3) that such prior knowledge of the invariant features can grant our algorithm better estimation performance, although it is not strictly necessary. We summarize our iterative procedure in Algorithm (2.1). In the supplement, we provide simulation results demonstrating the advantage of Algorithm (2.1) over direct minimization of the objective function $Q_{ℰ}(a,b,\theta )$.

Remark 2.1. The form of the relaxation used in the definition of our FILM estimator is similar conceptually to that of the constrained causal optimization (CoCo) method of Yin et al. (2024). Given data from a set of environments $ℰ$, the CoCo algorithm solves

$${\widehat{\beta }}_{CoCo}=\underset{b}{\text{argmin}}\frac{1}{\mid ℰ\mid }\sum _{e\in ℰ}{\Vert \tilde{b}\odot \nabla L^e(b)\Vert }_{2}for{\tilde{b}}_J=1,{\tilde{b}}_{J^c}=b_{J^c}$$ (15)

where $J$ is a set of known exogenous variables. The key difference between our FILM approach and the CoCo method is our decoupling of the support and magnitude of $\beta$ via the parameters $a$ and $b$. This severely reduces our dependence on knowledge of a set $J$ of exogenous variables. If the set $J$ is misspecified or missing, the statistical performance of the CoCo algorithm may suffer. Furthermore, we include the pooled loss function by default in Equation (9) to phase out uninformative solutions. While the authors of Yin et al. (2024) provide a risk-regularized variant of CoCo, they do not evaluate this approach for generalized linear models systematically via simulations in their work. Finally, the CoCo method is not robust to misspecified or corrupted environments, unlike our Robust-FILM estimator described in Section 2.3. We validate these points empirically in Section 3.

Algorithm 2.1 Alternating minimization for computing ${\widehat{\beta }}_{\text{FILM}}$

$$\begin{array}{cc}\hfill & \text{Given data}{(Y_i^e,X_i^e)}_{i=1}^{n_e}\text{from a set of environment}ℰ\hfill \\ \hfill & \text{Given a set of known exogenous variables}J\subset [p]\hfill \\ \hfill & \text{Initialize}{a}^{(0)}\in {[0,1]}^p,b_0\in {\mathbb{R}}^p,{\theta }^{(0)}\in {\mathbb{R}}^{\mid ℰ\mid },\text{and}{\beta }^{(0)}=a^{(0)}\odot b^{(0)}\hfill \\ \hfill & \mathbf{for}k\leftarrow 0\text{to}\mathsf{max}.\mathsf{iter}\mathbf{do}\hfill \\ \hfill & \mathbf{if}J\ne \varnothing \mathbf{then}\hfill \\ \hfill & a^{(k+1)}\leftarrow {\text{argmin}}_{a\in {[0,1]}^p}{\overline{Q}}_{ℰ}(a;b^{(k)},{\theta }^{(k)},{\beta }^{(k)})\text{s.t.}{a}_j=1\forall j\in J\hfill \\ \hfill & \mathbf{else}\hfill \\ \hfill & a^{(k+1)}\leftarrow {\text{argmin}}_{a\in {[0,1]}^p}{\overline{Q}}_{ℰ}(a;b^{(k)},{\theta }^{(k)},{\beta }^{(k)})\hfill \\ \hfill & \mathbf{end}\mathbf{if}\hfill \\ \hfill & b^{(k+1)}\leftarrow {\text{argmin}}_b{Q}_{ℰ}(a^{(k+1)},b,{\theta }^{(k)})\hfill \\ \hfill & {\theta }^{(k+1)}\leftarrow {\text{argmin}}_{\theta }{Q}_{ℰ}(a^{(k+1)},b^{(k+1)},\theta )\hfill \\ \hfill & {\beta }^{(k+1)}\leftarrow a^{(k+1)}\odot b^{(k+1)}\hfill \\ \hfill & \mathbf{end}\mathbf{for}\hfill \\ \hfill & \mathbf{return}(a^{(\max .\mathsf{iter})},b^{(\mathsf{max}.\mathsf{iter})},{\theta }^{(\mathsf{max}.\mathsf{iter})}\hfill \end{array}$$

We also propose a two-stage cross validation procedure to choose the tuning parameters ${\lambda }_1$ and ${\lambda }_2$, as well as the initial iterates ($a^{(0)}$, $b^{(0)}$, ${\theta }^{(0)}$). We evaluate CV error on data from a validation environment that is distinct from environments used in the training process to ensure that the chosen tuning parameters do not over-fit to the training environments, as ${\lambda }_1$ needs to be calibrated to balance in-sample predictive performance and invariance. Furthermore, comparing the performance of multiple initializations helps guarantee that our iterative algorithm does not get stuck in a spurious local minima, which would admit a poor estimate ${\widehat{\beta }}_{\text{FILM}}$. This cross validation procedure is detailed in Algorithm (2.2).

Remark 2.2. Our relaxation does not retain the identification and estimation results given in Theorems 4.2 and 4.3 of Fan et al. (2024), because the resulting non-convex objective function Q admits many local solutions that are not necessarily globally optimal. We over-come this challenge by running cross-validation over a grid of initializations as described in Algorithm 2.2. We emphasize that our objective is not to identify and estimate the invariant effects, but rather to build a transferable prediction model across the environments under the invariant generalized linear model. Treating estimation of $\beta$ as a surrogate for invariant prediction, our proposed FILM estimator achieves excellent performance as demonstrated in simulations and real data analysis in Sections 3 and 4 respectively, indicating our non-standard relaxation is able to recover a transferable prediction rule across a variety of settings.

Algorithm 2.2 Alternating minimization with cross-validation

$$\begin{array}{cc}\hfill & \text{Given data from a set of training environment}{ℰ}_{\text{tr}}\text{and validation environment}{ℰ}_{\text{val}}\hfill \\ \hfill & \text{Given a set of initial point}\{(a_l^{(0)},b_l^{(0)},{\theta }_l^{(0)}):l=1,\dots ,L)\}\hfill \\ \hfill & \text{Given a set of tuning parameters}\{({\lambda }_{1,m},{\lambda }_{2,m}):m=1,\dots ,M\}\hfill \\ \hfill & \mathbf{for}l\leftarrow 1\text{to}L\mathbf{do}\hfill \\ \hfill & \mathbf{for}m\leftarrow 1\text{to}\mathbf{do}\hfill \\ \hfill & •\text{Compute}({\widehat{a}}_{l,m},{\widehat{b}}_{l,m},{\widehat{\theta }}_{l,m})\text{by running Algorithm (2.1) over the data from}{ℰ}_{\text{tr}}\text{with}\hfill \\ \hfill & \text{initial points}(a_l^{(0)},b_l^{(0)},{\theta }_l^{(0)})\text{and tuning parameter}({\lambda }_{1,m},{\lambda }_{2,m}).\hfill \\ \hfill & \mathbf{end}\mathbf{for}\hfill \\ \hfill & \mathbf{end}\mathbf{for}\hfill \\ \hfill & \text{Compute}(m^{\ast },l^{\ast })={\text{argmin}}_{m,l}{Q}_{{ℰ}_{\text{val}}}({\widehat{a}}_{m,l},{\widehat{b}}_{m,l},{\widehat{\theta }}_{m,l})\text{with tuning parameters}({\lambda }_{1,m},{\lambda }_{2,m})\hfill \\ \hfill & \mathbf{return}({\widehat{a}}_{m^{\ast },l^{\ast }},{\widehat{b}}_{m^{\ast },l^{\ast }},{\widehat{\theta }}_{m^{\ast },l^{\ast }})\hfill \end{array}$$

2.3. Robust extension

Our assumption that Model (2) holds for all $e\in ℰ$ may be too strong in practice. In real data settings, some environments may fail to satisfy Model (2) for a litany of reasons, including model shifts, feature misalignment, or data corruption. To protect against possibly misspecified environments, we introduce the robust loss function

$$Q_{ℰ}^{\text{Robust}}(a,b)=\Psi \left(L^e(a\odot b):e\in ℰ\right)+{\lambda }_1\Psi (\sum _{j=1}^p{a}_j^2{({\nabla }_j{L}^e(a\odot b))}^2:e\in ℰ)+{\lambda }_2\sum _{j=1}^p{a}_j(1-a_j)$$ (16)

where $\Psi :{\mathbb{R}}^{\mid ℰ\mid }\mapsto \mathbb{R}$ is a robust measure of centrality, such as the median or trimmed mean.

Our Robust-FILM estimator is then defined as

$${\widehat{\beta }}_{\text{Robust-FILM}}=\widehat{a}\odot \widehat{b};(\widehat{a},\widehat{b})=\underset{a\in {[0,1]}^p,b\in {\mathbb{R}}^p}{\text{argmin}}{Q}_{ℰ}^{\text{Robust}}(a,b)$$ (17)

which we also compute via alternating minimization using a slightly modified version of Algorithm 2.1. We provide details in the appendix. This approach is inspired by the median-of-means literature (Devroye et al., 2016; Lecué & Lerasle, 2020), as we replace a ‘global’ sample average (in this case, over all the training environments) with a robust accumulation of ‘local’ sample averages. In Section 3, we demonstrate empirically that an appropriate choice of $\Psi$ grants excellent statistical performance in the presence of outlier environments with only a minor cost in runtime.

3. Simulations

We investigate the properties of the FILM and Robust-FILM estimators using data simulated from the invariant features model in (2) with the identity and logit link functions, corresponding to the linear and logistic regression models respectively. We again let $ℰ$ denote the set of environments and form the partition $ℰ={ℰ}_I\cup {ℰ}_O$ where ${ℰ}_I$ denotes the set of inlier environments and ${ℰ}_O$ the outliers. For each $e\in {ℰ}_I$, we generate $n_e$ i.i.d. copies from the following data generating process:

$$\begin{array}{cc}\hfill & \mathbf{Data}\mathbf{generating}\mathbf{process}\mathbf{for}e\in {ℰ}_I\hfill \\ \hfill & \text{Given}\beta ,{\gamma }^e,{\theta }^e,{\sigma }^e:\hfill \\ \hfill & •X_S^e\sim N(0,I)\in {\mathbb{R}}^s\hfill \\ \hfill & •\text{For linear models:}\hfill \\ \hfill & -{\epsilon }^e\sim N(0,{({\sigma }^e)}^2\in \mathbb{R}\hfill \\ \hfill & -Y^e={\theta }^e+{\beta }_S^{\mathsf{T}}{X}_S^e+{\epsilon }^e\in \mathbb{R}\hfill \\ \hfill & •\text{For logistic models:}\hfill \\ \hfill & -Y^e\sim \text{Bern}(\text{expit}({\theta }^e+{\beta }_S^{\mathsf{T}}{X}_S^e))\hfill \\ \hfill & •X_{S^c}^e={\gamma }^e{Y}^e\cdot 1_q+N(0,I)\in {\mathbb{R}}^q\hfill \\ \hfill & •Z^e\sim N(0,I)\in {\mathbb{R}}^{p-(s+q)}\hfill \end{array}$$

where $\beta$, ${\gamma }^e$, ${\Sigma }^e$, ${\theta }^e$, and ${\sigma }^e$ are varied in each simulation setting. We collect the covariates into the vector $(X_S^e,X_{S^c}^e,Z^e)\in {\mathbb{R}}^p$. $X_{S^c}^e$ denotes the spurious features with corresponding spurious correlation ${\gamma }^e$. The $Z^e$ vector represents null features that are not related to the outcome, and are included to increase the complexity of the problem. In the results that follow, we set $s=3$ and $q=1$ with ${ℰ}_I=\{1,2,3\}$ and $\beta ={(2,3,4,0,0,\dots ,0)}^{\mathsf{T}}$. Note that this defines $S=\{1,2,3\}$ as the set of invariant features.

We will evaluate four versions of our FILM method. The first runs Algorithm (2.2) directly with $J=\varnothing$, representing a naive application of FILM that incorporates no prior information nor protects against misspecified or corrupted environments. Our second implementation modifies this naive approach by setting $J=\{1\}$ to study the effect of partial oracle knowledge of the support of $\beta$ on the performance of FILM. Additionally, we fit robust versions of both of these methods by solving Equation (17) with Algorithm (2.2) and $\Psi$ as the median function. In all cases, we use ${ℰ}_{\text{tr}}=\{1,2\}\cup {ℰ}_O$ and ${ℰ}_{\text{val}}=\{3\}$. We fix ${\lambda }_2=0.1$ and choose ${\lambda }_1$ with cross validation over a grid of evenly spaced values from 50 to 125. We initialize our method by running cross validation over 30 randomly generated candidate initial points.

To understand the performance of FILM relative to alternative methods, we also fit an oracle generalized linear model using only the invariant features $X_S^e$ with samples only from the inlier environments ${ℰ}_I$ with an environment-level intercept. This model acts as our gold standard, as it uses knowledge of the data-generating process that is typically unavailable to the statistician. We additionally fit the CoCo method by solving Equation (15) with $J=\{1\}$, and in the linear model setting we fit EILLS with ${\lambda }_1=1$ and ${\lambda }_2=0.1$. Finally, we fit a basic pooled generalized linear model with all available features and data from all environments.

3.1. Results without outliers

In the first set of results, we set ${ℰ}_O=\varnothing$, and we investigate the performance of each method discussed above while varying the sample size per environment $n$, number of features $p$, number of environments $\mid ℰ\mid =E$, and the minimum strength of spurious correlations within each environment $\gamma$. Upon fixing the parameter $\gamma$, we choose each ${\gamma }^e$ for $e\in ℰ$ along an evenly spaced grid from $\gamma$ to $\gamma +4$. By default, we set $n=500$, $p=10$, $E=3$, and $\gamma =2$. We measure the performance of each method via its estimation error ${\Vert \widehat{\beta }-\beta \Vert }_2$ and accuracy of variable selection. Here, accuracy is defined as the sum of the number of true positives (${\widehat{\beta }}_j\ne 0$ for $j\in S$) and true negatives (${\widehat{\beta }}_j=0$ for $j\notin S$), divided by $p$. The results are given in Figures 1 and 2 respectively.

Figure 1:

Median estimation error over 50 replications. Results are given for the linear and logistic models along the rows. The columns of the figure delineate the parameter varied along the x-axis (sample size, number of features, number of environments, and $\gamma$). The y-axis is on the logarithmic scale.

Figure 2:

Median accuracy over 50 replications. Results are given for the linear and logistic models along the rows. The columns of the figure delineate the parameter varied along the x-axis (sample size, number of features, number of environments, and $\gamma$).

In these figures, we can see that all four versions of FILM are able to mimic the performance of the gold standard oracle estimator and EILLS over a variety of settings. The Robust-FILM estimators perform nearly identically to their non-robust analogs, which indicates that Robust-FILM is safe to use in all cases, even when the analyst does not suspect that any outlying environments are present in their data. This concordance between FILM and Robust-FILM begins to fail as the number of environments increases, but this is reasonable, as minimizing over the median of a set of values becomes more challenging as the number of inputs increases. We also see that including knowledge of an invariant feature by setting $J=\{1\}$ can improve both estimation error and accuracy, but this improvement is not substantial. The CoCo method admits relatively poor performance in these simulations since we only give it knowledge of one entry of the set $S$. Finally, the naive Pooled GLM is badly biased due to its ignoring of cross-environment heterogeneity and the spurious correlations. This bias becomes especially severe as $\gamma$ increases.

3.2. Results with outliers

Now we add one environment to ${ℰ}_O$ to investigate the effectiveness of Robust-FILM when non-trivial outliers are present in the training data. We draw the outlier environment from each of the following four schemes:

1. Pure noise: We draw $X^e\stackrel{\text{iid}}{\sim }{t}_2$ and $Y^e\stackrel{\text{iid}}{\sim }\text{Unif}(-5,5)$ (for linear models) or $Y^e\stackrel{\text{iid}}{\sim }\text{Bern}(0.1)$ (for logistic models).

2. Permuted features: Draw ($X^e$, $Y^e$) as if $e\in {ℰ}_I$, but return $(\pi (X^e),Y^e)$ for a permutation $\pi$.

3. Rotated features: Generate ($X^e$, $Y^e$) as if $e\in {ℰ}_I$, but return ($O{X}^e$, $Y^e$) for a random rotation matrix $O$.

4. Flipped effects: Generate ($X^e$, $Y^e$) as if $e\in {ℰ}_I$ with causal effects $-1\cdot \beta$.

We keep $n=500$, $E=3$, $p=10$, and $\gamma =2$, and evaluate each method by estimation error. The oracle method is computed using only the data from ${ℰ}_I$ so that it acts as our gold standard once more. Results are given in Table 1.

Table 1:

Median estimation error over 50 replications under each of the four outlier schemes. For the sake of space, we abbreviate ‘Robust-FILM’ as ‘R-FILM’.

Model	Outlier scheme	Oracle GLM	Pooled GLM	FILM	$\text{FILM}(J=\{1\})$	R-FILM	$\mathrm{R}\text{-FILM}(J=\{1\})$	EILLS	CoCo
Linear	Pure noise	0.037	5.059	5.385	5.244	0.100	0.082	5.385	5.368
-	Permuted	0.040	3.871	5.385	5.049	0.081	0.082	2.413	3.784
-	Rotated	0.038	4.840	5.385	5.240	0.088	0.077	5.392	5.157
-	Flipped effects	0.035	6.152	5.385	5.183	0.091	0.084	5.388	4.025
Logistic	Pure noise	0.277	5.035	5.385	5.294	0.572	0.470	5.385	5.275
-	Permuted	0.229	4.289	5.385	5.255	0.550	0.422	5.388	5.235
-	Rotated	0.232	4.663	5.385	5.334	0.669	0.506	5.357	5.213
-	Flipped effects	0.265	5.396	5.385	5.317	0.501	0.414	5.388	5.291

We can see that under each outlier setting, our Robust-FILM procedure achieves the best performance outside of the oracle estimator. Similar to the outlier-free results, we see that including $J=\{1\}$ is helpful, but Robust-FILM performs very well even without this additional information.

3.3. Numerical runtime

We also run a set of benchmarking simulations to evaluate the runtime of EILLS, FILM, and CoCo over varying values of $p$. For $p\in \{5,10,15,20\}$, we generate data from 2 environments with $n_1=n_2=100$ from Model (2) with $g(x)=x$. We compute ${\widehat{\beta }}_{\text{EILLS}}$ by solving Equation (8) via brute force as suggested by Fan et al., 2024. We compute the FILM estimator via Algorithm (2.1), and we compute the CoCo estimator by solving Equation (15). We also include results for Robust-FILM, which are obtained by solving (17) with $\Psi =\text{median}(\cdot )$. We fit each method 5 times for each value of $p$ and provide the mean runtime in seconds in Figure 3.

Figure 3:

Average runtime over 5 replications in seconds by number of features for EILLS, FILM, Robust-FILM, and CoCo. The y-axis is on the logarithmic scale.

We see in Figure 3 that the runtime of computing the EILLS estimator increases exponentially in the dimension $p$, as expected. We also see that CoCo is consistently the computationally fastest method. This is not surprising, as CoCo only requires solving one optimization problem, where FILM and Robust-FILM involve a nested series of alternating minimization problems. However, FILM and Robust-FILM still run comparably fast to CoCo and offer superior statistical performance as discussed in Section 3. Furthermore, Robust-FILM performs similarly to FILM in terms of runtime, suggesting that the flexibility and robustness granted by the Robust-FILM method can be achieved with very little computational cost.

4. Real data analysis

4.1. Background

We demonstrate the real world effectiveness of our FILM method on kidney disease data from the All of Us research program (Mayo et al., 2023). The All of Us program is a nationwide initiative by the National Institutes of Health that aims to integrate multimodal biomedical data, including genetic profiles, electronic health records (EHR), and wearable device measurements, from diverse patient cohorts at partner institutions across the United States. Such a wealth of information offers great opportunities for researchers to learn from heterogeneous populations, and in particular build disease risk prediction models with transferable performance across patient subgroups.

Chronic kidney disease (CKD) is a critical and rapidly growing global health challenge, currently affecting approximately 10% of the global population—equivalent to over 850 million individuals worldwide (International Society of Nephrology, 2023). CKD accounted for over 1.5 million deaths annually (American Society of Hematology, 2024) and is projected to become the fifth leading cause of death globally by 2040 (Moura et al., 2024). Advanced CKD may lead to end stage renal disease (ESRD), characterized by complete kidney failure, which requires treatment by dialysis or a kidney transplant (Valderrábano et al., 2001). Early screening of ESRD among patients with CKD can greatly improve quality of life and life expectancy (Lin et al., 2018). Estimated glomular filtration rate (eGFR) is the most common measure of kidney function and is often used for screening or diagnosis of kidney disease (Kalantar-Zadeh et al., 2021). However, screening based on eGFR alone may be sub-optimal in practice, as formulas for computing eGFR are not always reliable estimators of the underlying glomular filtration rate (Glassock & Winearls, 2008). Supplementing eGFR-based screening with predictions using EHR derived features offers a promising route to improve renal disease detection (Haris et al., 2024). To this end, we aim to leverage our proposed FILM method to build a transferable prediction model for ESRD among patients with severe eGFR values. Our outcome is defined as the occurrence of an ESRD diagnosis (SNOMED code 46177005). After careful data processing, we are left with 51 demographic and EHR-derived features with which to train our prediction model. We provide complete details on our data processing pipeline in the appendix.

As healthcare system quality and patient demographics vary greatly across geography in the United States (Bethell et al., 2011; Weaver et al., 2021), we use the state of residence for each patient as the environment variable. We restrict our attention to states in the northeast (New York, Pennsylvania, and Massachusetts) and one representative state from the midwest (Illinois) and the south (Texas), each chosen based on population size. We use the northeast states as our training environments, and aim to build a prediction model whose performance transfers well to Illinois and Texas. Table 2 gives the number of cases and controls that were used in our data analysis from each state.

Table 2:

Case and control totals by state.

Role	State	n cases	n controls
Training	Pennsylvania	14	266
-	Massachusetts	10	207
-	New York	19	118
Testing	Illinois	41	378
-	Texas	16	256

4.2. Analysis and results

We use the data from the northeast states to train four variants of the logistic regression model: a ridge regularized logistic regression using the glmnet package that pools data across all three states, the CoCo method using a logit link function, and FILM and Robust-FILM using the median as our $\Psi (\cdot )$ function, both with the logit link. Tuning parameters for each method are chosen using cross-validation (for the FILM methods, we use Massachusetts as the validation environment). For each method, we use the estimated effect vector to compute predicted case probabilities on the patients from Illinois and Texas, with which we compute the AUC. These results are given in the first two rows of Table 3.

Table 3:

AUC for each method computed on the Illinois and Texas testing data. The methods are trained on data from Pennsylvania, Massachusetts, and New York. Results are given for the models trained with and without the outlier environment separately.

Outlier?	State	Pooled Logistic	FILM	Robust-FILM	CoCo
No	Illinois	0.779	0.837	0.831	0.661
-	Texas	0.691	0.782	0.771	0.479
Yes	Illinois	0.603	0.531	0.759	0.340
-	Texas	0.513	0.607	0.748	0.465

To demonstrate the effectiveness of our Robust-FILM approach, we use the patient data from California ($n=338$) to construct an outlier environment. Leaving the covariate data unchanged, we flip the case-control indicator from 0 to 1 or vice-versa for each patient from California with probability 0.5 to mimic data corruption that results from outcome mislabeling. We then include this corrupted environment in the training set and retrain each of the four methods described above. Performance is computed again as AUC on the Illinois and Texas data. These results are given in the second two rows of Table 3.

Our FILM and Robust-FILM methods achieve the best performance in both the outlier-free and outlier-present settings respectively. The CoCo method suffers in this analysis due to the lack of sufficient oracle knowledge of the true invariant features, which prohibits its use in practical data analysis settings, and the Pooled Logistic regression incurs bias from ignoring the cross-environment heterogeneity. The EILLS approach of Fan et al. (2024) is not computationally feasible even in this moderate dimensional ($p=51$) setting.

In the outlier-free setting, our FILM method selects the following five features as having a nonzero regression coefficient: serum creatinine, acquired hemolytic anemia, bacteremia, renal disorder due to type 1 diabetes mellitus, and systemic sclerosis. In the appendix, we provide a table of each feature’s estimated regression coefficient and corresponding SNOMED code. These features comprise our estimate of the set $S$ of invariant features in Model (2). Each of these features has a clinical relevance to end stage renal disease. Creatinine levels are a well-established measure of kidney function (Fink et al., 1999; Levey et al., 1988) and are used to calculate eGFR (Grams et al., 2023). Hemolytic anemia is known to induce kidney damage by damaging red blood cells, releasing large amounts of haem into circulation which place undue stress on the kidneys (Van Avondt et al., 2019). Recent studies provide evidence that bacteremia is prevalent among kidney transplant recipients, and is associated with subsequent kidney failure and increased rates of mortality (Ito et al., 2016; Jamil et al., 2016; Tsikala-Vafea et al., 2020). Patients with Type 1 diabetes are at high risk for developing end stage renal disease (Skupien et al., 2012; Stadler et al., 2006). Finally, renal disorder is common in cases of systemic sclerosis (Woodworth et al., 2016), and renal failure is a major cause of death among patients affected by systemic sclerosis (Cole et al., 2023). Knowledge of these features as provided by our FILM estimator can allow researchers to build transferable and parsimonious prediction models for ESRD.

5. Discussion

In this work, we present the FILM framework for computationally efficient estimation of generalized linear models from multiple environments under the invariant features model. We demonstrate via simulations that our approach performs as well as state of the art methods while greatly improving runtime. Furthermore, we propose a robust extension of our method called Robust-FILM that protects against environments that may be corrupted or misspecified, and we show that Robust-FILM successfully preserves statistical performances under a variety of corruption settings. We then use our new methods to build a transferable risk prediction model for end stage renal disease using electronic health records from the All of Us database.

The limitations of our FILM method suggest promising directions for future work. First of all, FILM is only designed for generalized linear models with a canonical link function. A natural next step is to extend our approach to a more general class of parametric models and then to nonparametric regression models, perhaps leveraging ideas from Gu et al. (2024). Another possible direction for future work is to extend our FILM approach to handle more general, possibly nonlinear, invariant representations, as is considered in the works (Jiang et al., 2023; Nguyen et al., 2021; Shi et al., 2021; Zhao et al., 2022). This may allow us to integrate data sources in which which the different environments correspond to different modalities of data collection, which is an increasingly common setting in real-world applications. Invariant representation learning also provides an alternate idea for enforcing robustness beyond the median-of-means approach that we take in this work, as data from a potentially outlying environment may satisfy Model (2) only after a suitable transformation (Maurer et al., 2016; Tian et al., 2023). Another potential direction is to extend the current model to incorporate random effects, either to relax the definition of invariant features or to account for correlated or dependent samples by modeling between- and within-environment variability within the FILM framework. We see this as a valuable piece of future work that will bring about its own computational challenges, as is typical when handling generalized linear mixed effects models. Finally, more research is needed to better understand the trade-offs between transferability, robustness, and computational efficiency. The recent work of Gu et al. (2025) explores the fundamental computational difficulty of estimating invariant effects under the linear invariant features model, but further work is needed to characterize the additional cost of robust estimation under this model.

Supplementary Material

R script for FILM method: An R script implementing our FILM and Robust-FILM methods, along with the cross validation procedure (.R file).

R script for data simulation: An R script for simulating data from the invariant features generalized model as used in our simulation studies (.R file).

Additional simulation and details on computation and data analysis: A PDF file describing additional simulation results, the algorithm used to compute the Robust-FILM estimator, and additional information pertaining to the real data analysis including data availability (.pdf file).