Table of Contents
Fetching ...

Handling Covariate Mismatch in Federated Linear Prediction

Alexis Ayme, Rémi Khellaf

TL;DR

This work addresses covariate mismatch in federated linear prediction under blockwise missing data modeled as MCAR. It introduces two modular approaches: a low-dimensional plug-in method that aggregates cross-site second moments to build site-specific predictors, and an impute-then-regress (ItR) strategy for high-dimensional regimes that imputes missing covariates via exchangeable methods and then applies ridge regression. Theoretical results provide both asymptotic and finite-sample rates, showing that the plug-in approach is consistent and generalizes to unseen clients, while ItR with optimal linear imputation can achieve the global optimum with rates that scale with the total sample size across sites. Practically, the paper offers guidance on when to prefer plug-in versus ItR and demonstrates that federated covariance pooling often outperforms isolated local training in fragmented or high-dimensional settings, thereby enabling privacy-preserving, cross-institutional predictive modeling.

Abstract

Federated learning enables institutions to train predictive models collaboratively without sharing raw data, addressing privacy and regulatory constraints. In the standard horizontal setting, clients hold disjoint cohorts of individuals and collaborate to learn a shared predictor. Most existing methods, however, assume that all clients measure the same features. We study the more realistic setting of covariate mismatch, where each client observes a different subset of features, which typically arises in multicenter collaborations with no prior agreement on data collection. We formalize learning a linear prediction under client-wise MCAR patterns and develop two modular approaches tailored to the dimensional regime and communication budget. In the low-dimensional setting, we propose a plug-in estimator that approximates the oracle linear predictor by aggregating sufficient statistics to estimate the covariance and cross-moment terms. In higher dimensions, we study an impute-then-regress strategy: (i) impute missing covariates using any exchangeability-preserving imputation procedure, and (ii) fit a ridge-regularized linear model on the completed data. We provide asymptotic and finite-sample learning rates for our predictors, explicitly characterizing their behaviour with the global dimension, the client-specific feature partition, and the distribution of samples across sites.

Handling Covariate Mismatch in Federated Linear Prediction

TL;DR

This work addresses covariate mismatch in federated linear prediction under blockwise missing data modeled as MCAR. It introduces two modular approaches: a low-dimensional plug-in method that aggregates cross-site second moments to build site-specific predictors, and an impute-then-regress (ItR) strategy for high-dimensional regimes that imputes missing covariates via exchangeable methods and then applies ridge regression. Theoretical results provide both asymptotic and finite-sample rates, showing that the plug-in approach is consistent and generalizes to unseen clients, while ItR with optimal linear imputation can achieve the global optimum with rates that scale with the total sample size across sites. Practically, the paper offers guidance on when to prefer plug-in versus ItR and demonstrates that federated covariance pooling often outperforms isolated local training in fragmented or high-dimensional settings, thereby enabling privacy-preserving, cross-institutional predictive modeling.

Abstract

Federated learning enables institutions to train predictive models collaboratively without sharing raw data, addressing privacy and regulatory constraints. In the standard horizontal setting, clients hold disjoint cohorts of individuals and collaborate to learn a shared predictor. Most existing methods, however, assume that all clients measure the same features. We study the more realistic setting of covariate mismatch, where each client observes a different subset of features, which typically arises in multicenter collaborations with no prior agreement on data collection. We formalize learning a linear prediction under client-wise MCAR patterns and develop two modular approaches tailored to the dimensional regime and communication budget. In the low-dimensional setting, we propose a plug-in estimator that approximates the oracle linear predictor by aggregating sufficient statistics to estimate the covariance and cross-moment terms. In higher dimensions, we study an impute-then-regress strategy: (i) impute missing covariates using any exchangeability-preserving imputation procedure, and (ii) fit a ridge-regularized linear model on the completed data. We provide asymptotic and finite-sample learning rates for our predictors, explicitly characterizing their behaviour with the global dimension, the client-specific feature partition, and the distribution of samples across sites.
Paper Structure (39 sections, 11 theorems, 88 equations)

This paper contains 39 sections, 11 theorems, 88 equations.

Key Result

Proposition 2.1

Under as:mcar_iid, the best predictor in $\mathcal{F}^{(k)}_{\mathrm{lin}}$ for client $k$, denoted by $\theta_\star^{(k)}$, is the solution to the local least-squares problem. It admits the following two equivalent representations: where $\Sigma_{\mathrm{obs}(k)}$ denotes the principal submatrix $\Sigma_{\mathrm{obs}(k),\mathrm{obs}(k)}$. The associated optimal global risk is the weighted sum of

Theorems & Definitions (26)

  • Proposition 2.1
  • Remark 2.2: Singular covariance
  • Remark 2.3: Allowing site-dependent covariate distributions
  • Theorem 3.1: Consistency
  • Remark 3.2: Finite-Sample Upper Bound
  • Proposition 4.1
  • Theorem 4.2
  • Remark 4.3
  • Remark 4.4
  • Proposition 5.1: Risk of local learning
  • ...and 16 more