Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

1. Probleemstelling

Het artikel adresseert de "drievoudige uitdaging" waar Federated Learning (FL) mee geconfronteerd wordt in het tijdperk van grootschalige modellen (honderden miljoenen tot miljarden parameters):

1. Berekening: Lokale trainingsstappen zijn computationally zwaar.
2. Geheugen: Het opslaan van volledige modeltoestanden en hulpvariabelen op hulpbronnen-beperkte randapparaten is verboden.
3. Communicatie: Het verzenden van volledig dimensionale updates via beperkte bandbreedte is inefficiënt.

Deze uitdagingen worden verergerd door data-heterogeniteit (Non-IID). In Non-IID-instellingen divergeren lokale doelen van het globale doel, wat leidt tot "client drift".

• Bestaande Oplossingen & Beperkingen:
  • Heterogeniteitscorrectie (bijv. SCAFFOLD): Gebruikt controlevariaties om drift te corrigeren, maar vereist het opslaan en communiceren van volledig dimensionale hulpstaten, wat de efficiencywinsten voor grote modellen tenietdoet.
  • Compressie (bijv. Quantisatie): Vermindert de verzonden bits maar vereist vaak volledig dimensionale foutbuffers, waardoor het geheugengebruik niet wordt verkleind.
  • Subruimte/Laag-rang Methoden: Verminderen de dimensionaliteit voor efficiency maar missen doorgaans principiële mechanismen om client drift te corrigeren. Naïeve combinaties gooien vaak historische correctie-informatie weg wanneer de subruimte verandert, wat leidt tot instabiliteit.

De Kloof: Er bestaat geen bestaand algoritme dat native heterogeniteitscorrectie binnen een laag-dimensionale subruimte uitvoert, terwijl het de robuustheid van volledig dimensionale methoden zoals SCAFFOLD behoudt.

2. Methodologie: Subspace-SCAFFOLD (SSF)

De auteurs stellen SSF voor, een algoritme dat heterogeniteitscorrectie direct integreert in een laag-dimensionale optimalisatiesubruimte.

Kernmechanismen

1. Subruimte-projectie:

   • In elke ronde $t$ wordt een gedeelde willekeurige orthonormale projectiematrix $P_t \in \mathbb{R}^{r \times d}$ (waarbij $r \ll d$) gegenereerd.
   • Optimalisatie (lokale updates en aggregatie) vindt strikt plaats in de $r$-dimensionale subruimte.
   • Gradienten worden geprojecteerd: $g_{proj} = P_t g$.

2. Geprojecteerde Controlevariaties (Heterogeniteitscorrectie):

   • SSF nabootst het SCAFFOLD-mechanisme maar werkt op geprojecteerde grootheden.
   • Clients onderhouden lokale controlevariaties $c_i$ en de server onderhoudt een globale controlevariatie $c$.
   • De lokale update-richting wordt gecorrigeerd met het verschil tussen geprojecteerde lokale en globale controlevariaties: $P_t g - c_{i, proj} + c_{proj}$.

3. Residu-bewarend Opvullen (Key Innovation):

   • In tegenstelling tot standaard subruimtemethoden die informatie buiten de actieve subruimte wegwerpen, onderhoudt SSF een volledig dimensionale controlestaat ($c_i$ en $c$).
   • Update-regel: Wanneer de subruimte verandert, worden de controlevariaties als volgt bijgewerkt:
     $$c^{t+1} = (I - P_t^\top P_t)c^t + P_t^\top P_t (\text{nieuwe gradientinformatie})$$
   • Dit betekent dat het orthogonale complement (residu) van de controlevariatie onveranderd behouden blijft, terwijl alleen het actieve subruimte-component wordt ververst.
   • Opvullen: Na de subruimte-update wordt het volledige-ruimte model gereconstrueerd door de nieuwe geprojecteerde coördinaten te combineren met het behouden residu: $x^{t+1} = P_t^\top x^{t+1}_{proj} + x^t_{res}$.

Algoritme-Flow

1. Decompositie: Server splitst model en controlevariaties in geprojecteerde en residu-delen.
2. Lokale Stappen: Clients reconstrueren het volledige model met behulp van het uitgezonden subruimtemodel en het residu van de server. Ze berekenen gradienten in volledige ruimte, projecteren deze, en voeren heterogeniteitsgecorrigeerde updates uit in de subruimte.
3. Aggregatie: Server aggregeert geprojecteerde eindpunten in de subruimte.
4. Opvullen & Verversen: Server tilt het model terug naar volledige ruimte en werkt controlevariaties bij, waarbij de residu-componenten behouden blijven.

3. Belangrijkste Bijdragen

1. Algoritme-ontwerp (SSF):

   • Het eerste FL-algoritme dat native heterogeniteitscorrectie uitvoert in een laag-dimensionale subruimte.
   • Introduceert een opvullende-update die residu-componenten van controlevariaties behoudt, zodat historische drift-correctie-informatie niet verloren gaat wanneer de subruimte roteert of verandert.

2. Theoretische Analyse:

   • Bewijst dat onder standaard gladheids- en begrensd-variatie-aannames, SSF een niet-asymptotische convergentiesnelheid bereikt van $\tilde{O}(1/T + 1/\sqrt{NKT})$.
   • Deze snelheid komt overeen met de lineaire versnelling van volledig dimensionale SCAFFOLD, wat aantoont dat subruimte-beperkingen het asymptotische convergentiegedrag met betrekking tot client drift niet verslechteren.
   • De analyse geldt zonder de onafhankelijkheid aan te nemen tussen de willekeurige projector en stochastische gradienten.

3. Empirische Validatie:

   • Gedemonstreerd op zowel een gecontroleerd matrix-regressie toy-probleem als een deep learning-benchmark (CIFAR-100 met ResNet-110).
   • Toonde aan dat SSF de kosten voor geheugen, communicatie en berekening aanzienlijk verlaagt terwijl het een nauwkeurigheid behoudt die vergelijkbaar is met volledig dimensionale SCAFFOLD.

4. Experimentele Resultaten

Toy-probleem (Matrix-regressie)

• Robuustheid: Onder variërende heterogeniteitsniveaus (laag, medium, hoog) presteerde SSF consequent beter dan FedSub (een baseline subruimtemethode zonder residu-bewaring) en FedAvg.
• Stabiliteit: FedSub vertoonde numerieke instabiliteit en divergentie (NaN-fouten) bij hogere subruimte-dimensies ($r=50$) door het verlies van orthogonale informatie. SSF bleef stabiel over alle dimensies.
• Nauwkeurigheid: Met een subruimte-ratio van $r/d = 0.2$ bereikte SSF relatieve fouten die zeer dicht bij Full-SCAFFOLD lagen, aanzienlijk beter dan Full-FedAvg.

Deep Learning (CIFAR-100 / ResNet-110)

• Rangschikking: Full-SCAFFOLD > SSF > FedAvg > FedSub.
• Prestaties: SSF behaalde een finale testnauwkeurigheid van 45,43%, wat FedAvg (34,35%) en FedSub (23,17%) met een grote marge overtrof.
• Efficiency: SSF slaagde erin de communicatie- en geheugenoverhead te verminderen terwijl de robuustheid van SCAFFOLD tegen data-heterogeniteit behouden bleef.

5. Betekenis en Impact

• De Kloof Dichtmaken: SSF lost de fundamentele spanning op tussen robuustheid (het hanteren van Non-IID-data) en efficiency (het hanteren van grote modellen). Het bewijst dat men geen volledig dimensionale hulpstaten nodig heeft om client drift effectief te corrigeren.
• Schaalbaarheid: Door de communicatie- en geheugenfootprint te verkleinen van $O(d)$ naar $O(r)$ (waarbij $r \ll d$) terwijl theoretische convergentiegaranties behouden blijven, maakt SSF grootschalig Federated Learning haalbaar op hulpbronnen-beperkte randapparaten.
• Theoretisch Inzicht: Het "residu-bewarende" mechanisme biedt een nieuw paradigma voor subruimte-optimalisatie, waarbij aantoont dat het behouden van volledig dimensionaal geheugen voor controlevariaties (zelfs als dit slechts gedeeltelijk wordt benaderd) cruciaal is voor stabiliteit in dynamische subruimte-omgevingen.

Concluderend biedt SSF een praktische en theoretisch onderbouwde oplossing voor het implementeren van grootschalige, heterogene Federated Learning, en bereikt het het "beste van twee werelden": de efficiency van subruimtemethoden en de robuustheid van controlevariatie-correctie.