Distribution Alignment for One-Shot Federated Learning via Optimal Transport

2026-06-15

AI Summary

This work addresses the challenge of feature representation misalignment in single-round federated learning caused by the coexistence of domain shift and label shift across clients. To tackle this issue, the authors propose SLOT-Align, a framework that leverages a frozen pretrained encoder to extract feature statistics, constructs a global reference via the Bures-Wasserstein barycenter, and achieves learning-free, geometry-aware feature alignment through a closed-form geodesic optimal transport mapping. As the first method to jointly handle both types of shifts within a single-round setting, SLOT-Align seamlessly integrates into existing one-shot federated learning (OSFL) pipelines without altering their training procedures. Extensive experiments across multiple benchmarks, backbone architectures, and base algorithms demonstrate consistent improvements in accuracy and robustness, confirming the method’s effectiveness and broad applicability.

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.

Problem

• One-Shot Federated Learning
• domain shift
• label shift
• feature misalignment
• data heterogeneity

Innovation

• Optimal Transport
• One-Shot Federated Learning
• Distribution Alignment
• Bures-Wasserstein Barycenter
• Feature Harmonization