Geometric Knowledge-Guided Localized Global Distribution Alignment for Federated Learning

Conference: CVPR 2025
arXiv: 2503.06457
Code: https://github.com/WeiDai-David/2025CVPR_GGEUR
Area: AI Security
Keywords: Federated Learning, Data Heterogeneity, Geometric Distribution Alignment, CLIP Embeddings, Covariance Reconstruction

TL;DR

Obtains the geometry of the global embedding distribution by accurately reconstructing the global covariance matrix from local covariance matrices in federated learning. It generates augmented samples along global principal directions to localize global distribution information, improving performance by 17 percentage points on CIFAR-100 under extreme heterogeneous scenarios (\(\beta=0.01\)).

Background & Motivation

Background: Federated Learning (FL) faces severe data heterogeneity, where divergent data distributions across clients lead to post-aggregation model performance degradation. Pre-trained model embeddings, such as CLIP, provide a unified feature space for FL.

Limitations of Prior Work: (1) FedAvg achieves only 58.71% accuracy on CIFAR-100 under extreme non-IID conditions (\(\beta=0.01\)). (2) FedFA approximates the global distribution using a Gaussian assumption, whereas the actual feature distribution is often non-Gaussian. (3) Individual clients cannot directly access the global data distribution.

Key Challenge: Global data distribution information is required for localized training, yet sharing raw data is prohibited. Common approximations using only mean or Gaussian distributions fail to capture the geometric structure of the distribution.

Goal: Accurately reconstruct the "geometry" (the principal directions of the covariance matrix) of the global embedding distribution for local data augmentation without sharing raw data.

Key Insight: The eigendecomposition of a covariance matrix reveals the distribution's "geometry"—specifically, its principal directions and scales. Since the global covariance can be mathematically reconstructed from local statistics without raw data sharing, perturbing local features along these global principal directions can successfully simulate samples from the global distribution.

Core Idea: Accurately reconstruct the geometry (eigenvectors and eigenvalues) of the global covariance matrix from local covariances, and generate augmented samples along the global principal directions to achieve localized global distribution alignment.

Method

Overall Architecture

Each client: Extracts embeddings using CLIP \(\to\) Computes local class covariance matrices \(\to\) Uploads statistical metadata to the server \(\to\) Server accurately reconstructs the global covariance via a weighted formula \(\to\) Server distributes the resulting eigenvectors/eigenvalues down \(\to\) Clients generate augmented samples along the global principal directions \(\to\) Train local MLP classifiers.

Key Designs

1. Accurate Global Covariance Reconstruction:

   • Function: Obtains the geometric structure of the global embedding distribution without relying on data sharing.
   • Mechanism: \(\Sigma_i = \frac{1}{N_i}(\sum_k n_k^i \Sigma_k^i + \sum_k n_k^i (\mu_k^i - \mu_i)(\mu_k^i - \mu_i)^T)\) is computed exactly using local covariances and means uploaded by clients. Eigendecomposition is then performed to obtain principal directions \(\xi\) and scales \(\lambda\).
   • Design Motivation: Exact reconstruction vs. Gaussian approximation—GGEUR outperforms FedFA by over 17 percentage points under \(\beta=0.01\), validating that precise geometric information is significantly more effective than coarse approximations.

2. Geometric-Guided Feature Augmentation (GGEUR):

   • Function: Generates augmented samples along the principal directions of the global distribution.
   • Mechanism: \(X_{(k,h)}^{(i,j)} = X_k^{(i,j)} + \sum_m \epsilon_m \lambda_i^m \xi_i^m\) injects perturbations scaled by eigenvalues along global eigenvector directions into the original embeddings. Each class is augmented to 2000 samples.
   • Design Motivation: Augmenting along global directions allows local training to adapt to the diversity of the global distribution, conceptually exposing the model to the data characteristics of other clients.

3. Multi-Domain Extension:

   • Function: Standardizes cross-domain federated learning paradigms.
   • Mechanism: Cross-domain geometries of the same category in the CLIP embedding space are highly consistent (empirically validated), which justifies sharing global prototypes. Each prototype is augmented to 500 samples.
   • Design Motivation: The zero-shot generalization of CLIP embeddings enables effective cross-domain geometric alignment.

Loss & Training

Standard cross-entropy loss is employed. The CLIP encoder is frozen, and only the MLP classifier is trained. This workflow acts as a preprocessing step and operates orthogonally to any other FL aggregation framework.

Key Experimental Results

Main Results

Method	CIFAR-100 \(\beta=0.5\)	\(\beta=0.1\)	\(\beta=0.01\)
FedAvg	81.41	68.22	58.71
FedFA	81.98	74.68	-
+GGEUR	83.31	77.70	75.72

Ablation Study

Configuration	Performance
FedAvg \(\beta=0.01\)	58.71
+GGEUR	75.72 (+17.01!)
Tiny-ImageNet \(\beta=0.01\)	53.03 → 64.27 (+11.24)

Key Findings

• Highest gains under extreme heterogeneity: At \(\beta=0.01\), performance increases by 17 percentage points. Severe distribution shifts amplify the utility of geometric alignment.
• Exact geometry outclasses Gaussian approximation: GGEUR significantly beats FedFA (75.72 vs. ~60), proving that distribution geometry is far more informative than simple mean/variance constructs.
• Plug-and-play property: As a preprocessing data augmentation technique, it can be seamlessly combined with any FL aggregation standard.

Highlights & Insights

• The mathematical guarantee of exact covariance reconstruction is pivotal—this is an exact calculation rather than an estimation, ensuring no privacy leakage occurs.
• Cross-domain geometric consistency in CLIP embeddings is an insightful observation that subsequent FL methodologies can exploit.

Limitations & Future Work

• Uploading covariance matrices introduces transmission overhead (\(D \times D\) matrices).
• Assumes high-fidelity CLIP embeddings; performance may degrade on low-quality visual inputs.
• Augmenting up to a static sample size (2000/500) may not be globally optimal across diverse scenarios.

Related Work & Insights

• vs. FedFA: While FedFA approximates the global distribution via a Gaussian assumption, GGEUR's exact geometric reconstruction demonstrates an immense advantage in extreme non-IID scenarios.
• vs. FedProx / SCAFFOLD: Unlike gradient- or model-level regularization techniques, GGEUR is a data-level augmentation method, making it highly orthogonal to existing frameworks.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Highly original formulation of exact geometric covariance reconstruction and distribution alignment.
• Experimental Thoroughness: ⭐⭐⭐⭐ Evaluated across multiple datasets, non-IID levels, and single-to-multi-domain tasks.
• Writing Quality: ⭐⭐⭐⭐ Mathematically sound and clearly explained.
• Value: ⭐⭐⭐⭐ Offers a powerful, elegant mitigation for severe data heterogeneity in federated environments.

Related Papers

• [CVPR 2025] Detecting Backdoor Attacks in Federated Learning via Direction Alignment Inspection
• [CVPR 2025] A Simple Data Augmentation for Feature Distribution Skewed Federated Learning
• [CVPR 2026] ProxyFL: A Proxy-Guided Framework for Federated Semi-Supervised Learning
• [CVPR 2026] Batman: Benign Knowledge Alignment Through Malicious Null Space in Federated Backdoor Attack
• [CVPR 2025] Infighting in the Dark: Multi-Label Backdoor Attack in Federated Learning