Sheaf-FRL Enables Robust Federated Learning with Data Heterogeneity

Federated learning systems often face challenges when agents have heterogeneous data distributions, sensing modalities, or model architectures. Existing approaches typically assume a shared global latent space, which can be restrictive. Sheaf-based Federated Representation Learning (SFRL) is a novel framework designed to overcome these limitations, enabling agents in heterogeneous federated systems to learn and exchange informative representations effectively. SFRL jointly optimizes local learning objectives with a geometric alignment regularizer that is constrained by manifolds and based on learnable sheaf restriction maps. Unlike many current methods, SFRL does not require a shared global latent space. Instead, global consistency emerges from the alignment of neighboring latent representations through orthogonal transformations and isometric embeddings. This alignment is enforced by a quadratic gluing regularizer, derived from the sheaf Laplacian, whose learnable restriction maps dynamically adapt the geometry to the observed data. The penalty for this alignment is efficiently evaluated on a small set of shared pilot samples, ensuring both scalability and communication efficiency. The researchers developed a decentralized algorithm, Sheaf-FRL, which alternates between gradient updates for local models and closed-form Procrustes updates for the edge-wise restriction maps. Experiments in cooperative classification tasks under model and data heterogeneity demonstrate that Sheaf-FRL outperforms baseline approaches in terms of local and post-communication classification accuracy, showing greater robustness to latent-space dimensionality compression.