Optimal Transport-Based Heterogeneous Federated Learning for Chest X-Rays

Abstract

In some federated learning (FL) scenarios, discrepancies in local client devices result in inconsistent image resolutions, which motivates clients to adopt models with different depths and widths. Existing heterogeneous federated learning methods struggle to maintain model accuracy while preserving computational efficiency. To tackle this issue, this paper proposes a heterogeneous federated learning framework based on optimal transport (OT) and cross-layer alignment. The framework addresses the inconsistency of model depth via cross-layer alignment, fuses parameters of layers with different widths using optimal transport, and develops an aggregation strategy for multiple heterogeneous models. Experiments demonstrate that our method can improve model accuracy by up to 1.65% while maintaining satisfactory efficiency.