Technical Summary: Distribution Alignment for One-Shot Federated Learning via Optimal Transport

1. Problem Statement

The paper addresses the challenges inherent in One-Shot Federated Learning (OSFL), a communication regime where clients interact with a central server exactly once. While standard Federated Learning (FL) mitigates data heterogeneity through iterative client-server optimization, OSFL lacks this refinement capability.

The core problem identified is the compounding effect of two specific types of statistical heterogeneity:

1. Domain Shift: Clients observe different input distributions $P_k(x)$ due to varying acquisition conditions or environments.
2. Label Shift: Clients exhibit different label marginals $P_k(y)$, leading to varying class proportions.

In the absence of iterative correction, these shifts induce misaligned feature representations and inconsistent class-posterior decision functions across clients. Existing OSFL methods (relying on distillation, data synthesis, or ensembling) often assume aligned representations or address these shifts separately. Consequently, discrepancies in feature geometry and class-conditional biases remain uncorrected, degrading the performance of the aggregated global model.

2. Methodology: SLOT-Align

The authors propose SLOT-Align (Single-round, Learning-free Optimal Transport Alignment), a geometry-aware framework designed to harmonize client feature distributions prior to downstream OSFL training. The method operates as a preprocessing step and relies on a shared frozen encoder $f_\theta$ (pretrained on an external dataset) to map all local inputs into a common latent space.

The workflow consists of four key stages:

A. Local Feature Statistics Extraction

Each client $k$ computes compact first- and second-order statistics of the features $z = f_\theta(x)$ extracted from their local data:

• Mean: $\mu_k = \mathbb{E}[f_\theta(x)]$
• Covariance: $\Sigma_k$ (regularized via Ledoit–Wolf shrinkage to ensure positive definiteness).
  These statistics serve as a Gaussian proxy $N(\mu_k, \Sigma_k)$ for the induced feature distribution, capturing the dominant discrepancies in location and spread caused by domain and label shifts.

B. Global Reference Construction

The server aggregates these local statistics to construct a global reference distribution:

• Global Mean: A weighted average of client means.
• Global Covariance: Computed as the Bures–Wasserstein barycenter of the client covariances. This is solved via a fixed-point iteration to find the unique minimizer of the sum of squared Bures distances, effectively balancing client-specific variability with global consistency.

C. Geodesic Optimal Transport Mapping

Each client constructs a transport map to align its local Gaussian approximation to the global reference.

• Optimal Map: Under quadratic cost, the optimal transport map between two Gaussians is affine and has a closed-form solution involving the square roots of the covariance matrices.
• Interpolation: To prevent over-correction (since the true feature distribution is not strictly Gaussian), the method introduces an interpolation parameter $\tau \in [0, 1]$. The final transformation $T^{(\tau)}_k$ is a convex combination of the identity map and the full optimal transport map. This creates a Wasserstein displacement interpolation, moving the client's distribution along a constant-speed geodesic toward the global reference.

D. Integration

Clients apply the transformation $z' = T^{(\tau)}_k(z)$ to their local features. These aligned features are then fed into any existing OSFL algorithm (e.g., FedAvg, FedCGS, FedPFT) without modifying the downstream optimization procedures.

3. Key Contributions

1. Formalization: The paper formalizes OSFL as a distribution alignment problem under joint domain and label shift, highlighting how these shifts limit the efficacy of one-shot approaches.
2. Algorithm Proposal: It introduces SLOT-Align, a learning-free, optimal transport-based framework that harmonizes client feature distributions via a single exchange of compact statistics. It requires no auxiliary data, trainable generators, or iterative optimization.
3. Empirical Validation: Extensive experiments demonstrate that SLOT-Align consistently improves accuracy and robustness across multiple benchmarks (Office-Home, Digits, DomainNet), various pretrained backbones (ViT, ResNet), and state-of-the-art OSFL baselines.

4. Experimental Results

The authors evaluated SLOT-Align under varying degrees of label shift (controlled by the Dirichlet concentration parameter $\alpha$) and domain diversity.

• Performance Gains: Under joint domain and label shift ($\alpha=0.1$), SLOT-Align consistently improved macro-averaged accuracy for all tested baselines (O-FedAvg, FedCGS, FedPFT). For instance, on the Office-Home dataset, it improved O-FedAvg by +7.86% and FedPFT by +6.55%.
• Robustness: The method remained effective under stronger label skew ($\alpha=0.05$ and $\alpha=0.01$), showing gains even in severe heterogeneity regimes.
• Ablation on Alignment Strength: The interpolation parameter $\tau$ was found to be critical. Intermediate values (e.g., $\tau=0.4$) yielded the best trade-off, while $\tau=1.0$ (full transport) sometimes degraded performance, suggesting that partial alignment is preferable to full correction.
• Domain-Only Shift: In settings with balanced labels (domain shift only), SLOT-Align provided marginal gains or neutral performance, confirming that the method specifically targets the complex interaction of domain and label shifts.
• Efficiency: The communication cost is minimal ($O(m + m^2)$ per client for mean and covariance), and the computational overhead is negligible compared to the feature extraction and downstream training.

5. Significance and Claims

The paper claims that SLOT-Align offers a modular, computationally efficient solution to a critical bottleneck in OSFL: the inability to correct distributional misalignment in a single communication round.

• Modularity: It does not replace existing OSFL pipelines but acts as a compatible preprocessing layer, making it applicable to a wide range of current and future methods relying on frozen encoders.
• Geometry-Awareness: By leveraging the geometry of the 2-Wasserstein space and Bures–Wasserstein barycenters, it explicitly corrects discrepancies in both the location (mean) and shape (covariance) of feature distributions.
• Practicality: The method avoids the privacy risks and computational overhead associated with data synthesis or iterative distillation, relying solely on compact statistical summaries.

The authors conclude that while OSFL is constrained by strict communication budgets, SLOT-Align effectively mitigates the performance degradation caused by heterogeneous data distributions, particularly when domain and label shifts interact to distort feature statistics.