Distribution Alignment for One-Shot Federated Learning via Optimal Transport

Technische Samenvatting: Distributie-alignement voor One-Shot Federated Learning via Optimal Transport

1. Probleemstelling

Het artikel behandelt de uitdagingen die inherent zijn aan One-Shot Federated Learning (OSFL), een communicatieregime waarbij cliënten precies één keer met een centrale server interageren. Hoewel standaard Federated Learning (FL) datatheterogeniteit mitigeert door middel van iteratieve client-server optimalisatie, mist OSFL deze verfijningscapaciteit.

Het kernprobleem dat wordt geïdentificeerd, is het cumulatieve effect van twee specifieke soorten statistische heterogeniteit:

1. Domain Shift: Cliënten observeren verschillende inputdistributies $P_k(x)$ als gevolg van variërende acquisitiecondities of omgevingen.
2. Label Shift: Cliënten vertonen verschillende labelmarginals $P_k(y)$, wat leidt tot variërende klassenproporties.

Bij gebrek aan iteratieve correctie leiden deze verschuivingen tot misalignede feature-representaties en inconsistente klasse-posterior beslissingsfuncties bij cliënten. Bestaande OSFL-methoden (die vertrouwen op distillatie, datasynthese of ensembling) gaan vaak uit van gealigneerde representaties of pakken deze verschuivingen afzonderlijk aan. Bijgevolg blijven discrepanties in de geometrie van features en klasse-conditionele biases ongecorrigeerd, wat de prestaties van het geaggregeerde globale model verslechtert.

2. Methodologie: SLOT-Align

De auteurs stellen SLOT-Align (Single-round, Learning-free Optimal Transport Alignment) voor, een geometrie-bewust framework ontworpen om de client feature-distributies te harmoniseren voorafgaand aan downstream OSFL-training. De methode werkt als een preprocessing-stap en vertrouwt op een gedeelde bevroren encoder $f_\theta$ (voorgetraind op een externe dataset) om alle lokale inputs naar een gemeenschappelijke latente ruimte te mappen.

De workflow bestaat uit vier sleutelfases:

A. Extractie van Lokale Feature-statistieken

Elke client $k$ berekent compacte eerste- en tweede-orde statistieken van de features $z = f_\theta(x)$ die zijn geëxtraheerd uit hun lokale data:

• Gemiddelde: $\mu_k = \mathbb{E}[f_\theta(x)]$
• Covariantie: $\Sigma_k$ (geregulariseerd via Ledoit–Wolf shrinkage om positieve definitie te garanderen).
  Deze statistieken dienen als een Gaussische proxy $N(\mu_k, \Sigma_k)$ voor de geïnduceerde feature-distributie, waarbij de dominante discrepanties in locatie en spreiding veroorzaakt door domain- en label-shifts worden gevangen.

B. Constructie van een Globale Referentie

De server aggregeert deze lokale statistieken om een globale referentiedistributie te construeren:

• Globaal Gemiddelde: Een gewogen gemiddelde van de client-gemiddelden.
• Globale Covariantie: Berekend als de Bures–Wasserstein barycenter van de client-covarianties. Dit wordt opgelost via een fixed-point iteratie om de unieke minimizer te vinden van de som van gekwadrateerde Bures-afstanden, wat de variabiliteit van de individuele client effectief balanceert met globale consistentie.

C. Geodetische Optimal Transport Mapping

Elke client construeert een transportmap om de lokale Gaussische benadering uit te lijnen met de globale referentie.

• Optimale Map: Onder een kwadratische kostenfunctie is de optimale transportmap tussen twee Gaussische verdelingen affien en heeft deze een gesloten vorm die de wortels van de covariantie-matrices bevat.
• Interpolatie: Om overcorrectie te voorkomen (aangezien de ware feature-distributie niet strikt Gaussisch is), introduceert de methode een interpolatieparameter $\tau \in [0, 1]$. De uiteindelijke transformatie $T^{(\tau)}_k$ is een convexe combinatie van de identiteitsafbeelding en de volledige optimal transport map. Dit creëert een Wasserstein displacement interpolatie, waarbij de distributie van de client langs een constante-snelheid geodeet naar de globale referentie beweegt.

D. Integratie

Cliënten passen de transformatie $z' = T^{(\tau)}_k(z)$ toe op hun lokale features. Deze gealigneerde features worden vervolgens gevoed aan elk bestaand OSFL-algoritme (bijv. FedAvg, FedCGS, FedPFT) zonder de downstream optimalisatieprocedures te wijzigen.

3. Belangrijkste Bijdragen

1. Formalisering: Het artikel formaliseert OSFL als een distributie-alignement probleem onder gezamenlijke domain- en label-shift, waarbij wordt benadrukt hoe deze verschuivingen de effectiviteit van one-shot benaderingen beperken.
2. Algoritme Voorstel: Het introduceert SLOT-Align, een learning-free, optimal transport-gebaseerd framework dat de lokale feature-distributies harmoniseert via een enkele uitwisseling van compacte statistieken. Het vereist geen hulpdata, trainbare generatoren of iteratieve optimalisatie.
3. Empirische Validatie: Uitgebreide experimenten tonen aan dat SLOT-Align de nauwkeurigheid en robuustheid consistent verbetert over meerdere benchmarks (Office-Home, Digits, DomainNet) en diverse voorgetrainde backbones (ViT, ResNet) en state-of-the-art OSFL-baselines.

4. Experimentele Resultaten

De auteurs evalueerden SLOT-Align onder variërende graden van label-shift (gecontroleerd door de Dirichlet concentratieparameter $\alpha$) en domain diversiteit.

• Prestatiewinst: Onder gecombineerde domain- en label-shift ($\alpha=0.1$), verbeterde SLOT-Align de macro-gemiddelde nauwkeurigheid consequent voor alle geteste baselines (O-FedAvg, FedCGS, FedPFT). Op de Office-Home dataset verbeterde het bijvoorbeeld O-FedAvg met +7,86% en FedPFT met +6,55%.
• Robuustheid: De methode bleef effectief onder sterkere label-skew ($\alpha=0,05$ en $\alpha=0,01$), waarbij winst werd geboekt zelfs in regimes met extreme heterogeniteit.
• Ablatie op Alignement Sterkte: De interpolatieparameter $\tau$ bleek cruciaal. Intermediaire waarden (bijv. $\tau=0,4$) leverden de beste balans op, terwijl $\tau=1,0$ (volledige transport) de prestaties soms verslechterde, wat suggereert dat gedeeltelijke alignment te prefereren is boven volledige correctie.
• Alleen Domain Shift: In settings met gebalanceerde labels (alleen domain shift), leverde SLOT-Align marginale winst of neutrale prestaties, wat bevestigt dat de methode specifiek mikt op de complexe interactie tussen domain- en label-shifts.
• Efficiëntie: De communicatiekosten zijn minimaal ($O(m + m^2)$ per client voor gemiddelde en covariantie), en de computationele overhead is verwaarloosbaar vergeleken met de feature extractie en downstream training.

5. Betekenis en Claims

Het artikel claimt dat SLOT-Align een modulair, computationeel efficiënt oplossing biedt voor een kritieke bottleneck in OSFL: het onvermogen om distributie-misalignment te corrigeren in één enkele communicatieronde.

• Modulariteit: Het vervangt bestaande OSFL-pipelines niet, maar fungeert als een compatibele preprocessing-laag, waardoor het toepasbaar is op een breed scala aan huidige en toekomstige methoden die vertrouwen op bevroren encoders.
• Geometrie-bewustzijn: Door gebruik te maken van de geometrie van de 2-Wasserstein ruimte en Bures–Wasserstein barycenters, corrigeert het expliciet de discrepanties in zowel de locatie (gemiddelde) als de vorm (covariantie) van de feature-distributies.
• Praktische bruikbaarheid: De methode vermijdt de privacyrisico's en computationele overhead die gepaard gaan met datasynthese of iteratieve distillatie, door enkel te vertrouwen op compacte statistische samenvattingen.

De auteurs concluderen dat hoewel OSFL wordt beperkt door strikte communicatiebudgetten, SLOT-Align de prestatieverslechtering door heterogene datadistributies effectief mitigeert, met name wanneer domain- en label-shifts interageren om de feature-statistieken te verstoren.