Interactive note · AAAI 2025
Federated Unsupervised Domain Generalization using Global and Local Alignment of Gradients
[Paper, arXiv, Code, BibTeX]Each client holds unlabeled data from its own domain and cannot share it. The model still has to work on a domain no client has seen.
Domain shift shows up as disagreement between gradients, and gradients are exactly what federated learning gets to see.
Clients drop batch gradients that disagree with the global direction, and the server down-weights clients that disagree. Best results on all four benchmarks.
Many domains, no labels, no sharing
Picture a network of wearable activity monitors. Every device records under its own conditions, so each one sees a slightly different world, a different domain. Nobody labels their activity, and privacy rules keep the raw data on the device. Yet the model they train together should work for a new user it has never seen.
Earlier work handled these constraints one at a time: federated domain generalization assumes labels, and federated unsupervised learning ignores domain shift. This paper puts them together.
Federated unsupervised domain generalization is learning general representations from decentralized, unlabeled datasets, each from a different domain, when data cannot be shared.
The benchmarks make this concrete. In PACS every domain is a client except one, which is held out for testing. Pick which domain to hold out.
Different domains, different gradients
Under these rules the server never sees data, means or variances. What it does see is every client's model update, and each client sees the global update. The paper's first result says that this is enough to detect domain shift.
Model two domains as features that are correlated across domains with covariance \(\sigma\). Their mutual information, a measure of how similar the domains are, is
For clients trained with self-supervised learning, with features modeled as standardized Gaussians and gradients approximated to first order, as the shift between two clients' domains grows, the covariance of their gradients, \(\mathrm{Cov}(g_i, g_j)\), shrinks.
Below, two domains are sampled with the similarity you choose, and the gradients of a one-layer sigmoid encoder, as in the theorem but with a simpler loss, are computed in your browser. Move the slider and watch the points line up.
The proof runs through a first-order Taylor expansion (Lemma 2): the gradient covariance is a sum of the feature covariances \(\sigma\), weighted by products of derivatives that are always positive in this setting (Claim 1). A corollary follows: as domains drift apart, the variance of the difference of their gradients grows. So gradient alignment is a signal of domain shift that can be read without sharing raw data, and that is what FedGaLA uses. (Sharing gradients is not by itself a formal privacy guarantee.)
One communication round
FedGaLA aligns gradients in two places.
On each client, training is self-supervised (SimCLR). For every batch and every layer, the gradient is compared with a reference: the global model's change over the last round, \(\hat g_{est} = \theta^{(t)} - \theta^{(t-1)}\). If its cosine with the reference is below a threshold \(\tau\), the batch gradient is discarded. Scaling it down would not help, since cosine ignores scale.
On the server, each client's update gets a weight \(w_i = \tfrac12\big(\cos(\hat g_i, \hat g) + 1\big)\), normalized across clients, where \(\hat g\) is the current aggregate. The aggregate is recomputed with the new weights, three times. Dropping a whole client would throw away a domain, so the server weighs softly.
Here is one round. Two domains are typical; make the third one different and press Play round.
The global weights are gentle by design: \(w_i\) only falls to zero when a client points exactly against the aggregate. Most of the correction happens on the clients, where unaligned batches are simply skipped. In the paper's ablation the two parts are complementary: either one alone barely changes PACS accuracy, while together they add 1.4 points.
What the experiments show
Best on PACS, DomainNet, Office-Home and TerraInc
ResNet-18 trained from scratch, 100 rounds. Linear evaluation with 10% of the target's labels, averaged over held-out domains. Every baseline is a self-supervised method made federated with FedAvg; FedGaLA is FedSimCLR plus the two alignments.
| Method | PACS | DomainNet | Office-Home | TerraInc |
|---|---|---|---|---|
| FedEMA | 41.9 | 32.4 | 13.5 | 53.2 |
| FedBYOL | 44.2 | 31.8 | 13.8 | 54.3 |
| FedMoCo | 42.1 | 27.2 | 10.7 | 45.7 |
| FedSimSiam | 39.8 | 36.9 | 18.9 | 47.9 |
| FedSimCLR | 58.6 | 39.5 | 22.0 | 55.1 |
| FedGaLA | 60.0 | 41.1 | 23.0 | 56.7 |
Pairs of PACS domains
Gradient covariance against mutual information for each pair of domains. Apart from one outlier, more shift means less covariance, as Theorem 1 predicts.
Fewer gradients discarded over time
The share of discarded batch gradients falls from about 68% to 37% over 100 rounds: the clients learn features that agree with the global model.
Both alignments matter
| Variant | Accuracy |
|---|---|
| FedGaLA | 60.0 |
| without global alignment | 58.8 |
| without local alignment | 58.5 |
| without both (FedSimCLR) | 58.6 |
Neither alignment helps much alone; together they add 1.4 points.
Better than training on pooled data
PACS, 10% labels. The centralized methods see all domains in one place with no privacy constraint; FedGaLA still leads by a wide margin, in line with prior evidence that federation helps domain generalization.
BibTeX
@inproceedings{pourpanah2025fedgala,
title = {Federated Unsupervised Domain Generalization
using Global and Local Alignment of Gradients},
author = {Pourpanah, Farhad and Molahasani, Mahdiyar and
Soltany, Milad and Greenspan, Michael and
Etemad, Ali},
booktitle = {Proceedings of the AAAI Conference on
Artificial Intelligence},
pages = {19948--19958},
year = {2025}
}
* Equal contribution. The models in §2 and §3 are small toys built and run in your browser. The numbers and plots in §4 are from the paper.