Distribution Alignment for One-Shot Federated Learning via Optimal Transport

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.

Federated learning enables multiple participants to train a shared machine-learning model through a coordinating server, without sharing their original data. In one-shot federated learning, a setting with strict communication limits, each participant can communicate with the server only once. Under this constraint, learning an effective shared model is particularly difficult when participants use local datasets that differ in visual style or acquisition conditions, or that contain substantially different proportions of categories, because the feature representations extracted from their data can be misaligned. We introduce SLOT-Align, a lightweight preprocessing method that aligns these representations before global model training. Specifically, it operates in three steps: (i) each participant communicates only compact summaries of its local features; (ii) these summaries are combined to construct a shared global reference representation; and (iii) each participant’s local representations are moved toward this reference. Experiments on three image benchmarks and with several existing one-shot federated learning methods show consistent improvements in prediction accuracy, particularly when participants’ data distributions differ substantially. These results show that aligning feature representations can make federated learning more reliable when communication is extremely limited.