# Distribution Alignment for One-Shot Federated Learning via Optimal Transport

arXiv: https://arxiv.org/abs/2606.16655 · HTML: https://mlsanity.com/paper/2606.16655/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 51 references that are in the map.
- First posted on arXiv 15 Jun 2026.

## Builds on (most cited first, 30 of 35)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Moment Matching for Multi-Source Domain Adaptation](https://mlsanity.com/paper/1812.01754/) (2018)
- [On the Convergence of FedAvg on Non-IID Data](https://mlsanity.com/paper/1907.02189/) (2019)
- [Deep Hashing Network for Unsupervised Domain Adaptation](https://mlsanity.com/paper/1706.07522/) (2017)
- [Ensemble Distillation for Robust Model Fusion in Federated Learning](https://mlsanity.com/paper/2006.07242/) (2020)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Optimal Transport for Domain Adaptation](https://mlsanity.com/paper/1507.00504/) (2015)
- [Federated Learning on Non-IID Data: A Survey](https://mlsanity.com/paper/2106.06843/) (2021)
- [Joint Distribution Optimal Transportation for Domain Adaptation](https://mlsanity.com/paper/1705.08848/) (2017)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Generalized Federated Learning via Sharpness Aware Minimization](https://mlsanity.com/paper/2206.02618/) (2022)
- [One-Shot Federated Learning](https://mlsanity.com/paper/1902.11175/) (2019)
- [Distilled One-Shot Federated Learning](https://mlsanity.com/paper/2009.07999/) (2020)
- [Federated Learning for Generalization, Robustness, Fairness: A Survey and Benchmark](https://mlsanity.com/paper/2311.06750/) (2023)
- [Practical One-Shot Federated Learning for Cross-Silo Setting](https://mlsanity.com/paper/2010.01017/) (2020)
- [Revisiting Weighted Aggregation in Federated Learning with Neural Networks](https://mlsanity.com/paper/2302.10911/) (2023)
- [Federated Learning via Decentralized Dataset Distillation in Resource-Constrained Edge Environments](https://mlsanity.com/paper/2208.11311/) (2022)
- [Effects of Degradations on Deep Neural Network Architectures](https://mlsanity.com/paper/1807.10108/) (2018)
- [Preserving Privacy in Federated Learning with Ensemble Cross-Domain Knowledge Distillation](https://mlsanity.com/paper/2209.04599/) (2022)
- [Multi-Level Branched Regularization for Federated Learning](https://mlsanity.com/paper/2207.06936/) (2022)
- [Exploring One-shot Semi-supervised Federated Learning with A Pre-trained Diffusion Model](https://mlsanity.com/paper/2305.04063/) (2023)
- [Enhancing One-Shot Federated Learning Through Data and Ensemble Co-Boosting](https://mlsanity.com/paper/2402.15070/) (2024)
- [FRAug: Tackling Federated Learning with Non-IID Features via Representation Augmentation](https://mlsanity.com/paper/2205.14900/) (2022)
- [FedDEO: Description-Enhanced One-Shot Federated Learning with Diffusion Models](https://mlsanity.com/paper/2407.19953/) (2024)
- [FedBiP: Heterogeneous One-Shot Federated Learning with Personalized Latent Diffusion Models](https://mlsanity.com/paper/2410.04810/) (2024)
- [Revisiting Ensembling in One-Shot Federated Learning](https://mlsanity.com/paper/2411.07182/) (2024)
- [An Optimal Transport Approach to Personalized Federated Learning](https://mlsanity.com/paper/2206.02468/) (2022)
- [Advances in Robust Federated Learning: A Survey with Heterogeneity Considerations](https://mlsanity.com/paper/2405.09839/) (2024)
- [One-shot Federated Learning via Synthetic Distiller-Distillate Communication](https://mlsanity.com/paper/2412.05186/) (2024)

## Built on by (most cited first, 1 of 1)

- [Understanding Head Geometry and Dynamics in Federated Regression through a Natural Solution Selection Rule: An Unconstrained Feature Model Analysis](https://mlsanity.com/paper/2609.39464/) (2026)

## Abstract (arXiv)

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.