Distribution Alignment for One-Shot Federated Learning via Optimal Transport
2026-06-15 • Machine Learning
Machine Learning
AI summaryⓘ
The authors address a problem in federated learning where many devices send data only once to a central server, making it hard to combine different data types properly. They propose SLOT-Align, a new method that uses math tools to better match and align data features from different devices without extra training. Their approach improves the way data is combined even when there are differences in both data types and labels across devices. Tests show that their method helps existing systems perform better and more reliably.
Federated LearningOne-Shot LearningDomain ShiftLabel ShiftFeature RepresentationOptimal TransportBures-Wasserstein BarycenterEncoderData HarmonizationGeodesic Map
Authors
Daniele Berardini, Vito Paolo Pastore, Vittorio Murino
Abstract
One-Shot Federated Learning (OSFL) addresses extreme communication regimes in which clients interact with the server only once, amplifying the impact of heterogeneous client data distributions. In particular, the interaction of domain shift and label shift across clients induces misaligned feature representations that cannot be corrected through iterative optimization. Existing OSFL methods rely on distillation, server-side generation or ensemble-based aggregation, but assume aligned representations or address domain and label shift separately. We introduce SLOT-Align (Single-round, Learning-free Optimal Transport Alignment), a geometry-aware feature harmonization framework for OSFL. SLOT-Align uses a shared frozen encoder to extract compact feature statistics, constructs a global reference via Bures-Wasserstein barycenters, and aligns local representations using closed-form geodesic optimal transport maps. The method is computationally efficient and can be combined with existing OSFL pipelines relying on frozen encoders without modifying their training procedures. Extensive experiments across multiple benchmarks, pretrained backbones, and OSFL methods show that SLOT-Align consistently improves accuracy and robustness under joint domain and label shift.