Federated learning improves with new robust multi-source data method

Distributionally Robust Federated Learning with Multi-Source Data

Machine Learning

Summary

Federated learning allows many devices to train a shared AI model without sharing their private data. However, the data on these devices can be quite different, and it is hard to know exactly how they vary. The authors propose a new way to handle this uncertainty by considering many possible mixes of data differences both between and within devices. Their method ensures strong performance even when the amounts of data per device are small, and they develop an algorithm that works well in this setting. They tested their approach in simulations and found it effective.

What this means in practice

  • For mobile app developers: Build federated learning models that remain accurate despite varying and uncertain user data distributions on different devices.
  • For healthcare data teams: Train shared medical models across clinics where patient data differs and sample sizes may be limited without compromising model reliability.

Authors

Yingzhu Liu, Zhongkui Li, Pengcheng You, Ashish Cherukuri

Abstract

Federated learning trains a shared model from private client data. In practice, data-generating distributions may differ, and the true mixture across clients is often unknown, making the underlying group distribution difficult to specify. Existing approaches address cross-client mixture uncertainty by optimizing against the worst-case mixture, yet assume accurate client-wise distribution estimates. However, these estimates can be unreliable when based on finite samples. To handle both cross-client mixture uncertainty and within-client distributional ambiguity, we construct a global ambiguity set as the union of admissible mixtures of local ambiguity sets. The construction allows client-specific ambiguity radii and admits a client-wise separable reformulation. Leveraging this structure, we establish a high-probability out-of-sample performance guarantee. We further develop a federated algorithm for a penalty-based reformulation and prove its convergence under milder regularity conditions. Simulations validate the algorithm's effectiveness.