Prevention over Correction: Learning Aligned Representations in One-shot Federated Learning
Abstract
One-shot Federated Learning (OFL) reduces communication costs by requiring only a single round of model exchange. However, when clients train independently on heterogeneous data, their learned representations develop different means and variances, creating distributional mismatch that prevents effective server-side aggregation. Existing methods attempt post-calibration at the server, but struggle to reconcile representations that have already diverged into incompatible statistical spaces. We propose FACE (Federated Aligned-Consistent Ensemble) to prevent distributional divergence during client training rather than correcting it afterward. FACE uses a shared fixed reference matrix to guide representation means toward consistent directions, applies a variance stabilizer to maintain consistent feature variance, and performs featurefusion ensemble inference. Extensive experiments on CIFAR-10/100 and Tiny-ImageNet demonstrate that FACE achieves state-of-the-art performance, with particularly strong improvements under extreme data heterogeneity. Comprehensive ablation studies validate the necessity of both components for achieving distributional alignment.