Technische Samenvatting: Subspace-Constrained Federated Learning met Low-Rank Adaptation

1. Probleemstelling

Federated Learning (FL) maakt collaboratieve training mogelijk over gedecentraliseerde cliënten zonder ruwe data te delen, maar het staat voor aanzienlijke uitdagingen wanneer de datadistributies van cliënten heterogeen zijn. In standaard Federated Averaging (FedAvg) leidt statistische heterogeniteit tot "client drift", waarbij lokale modelupdates divergeren, wat instabiliteit en trage convergentie veroorzaakt.

Met de opkomst van grote fundamentele modellen is Low-Rank Adaptation (LoRA) een populaire techniek geworden voor parameter-efficiënte fine-tuning in FL. LoRA beperkt gewichtsupdates ($\Delta W$) tot een low-rank decompositie ($BA$), wat communicatie- en rekenkosten vermindert. Dit papier veronderstelt echter dat onder heterogene datadistributies verschillende cliënten low-rank updates leren die distincte subspaces beslaan. Wanneer deze geometrisch niet-uitgelijnde updates naïef worden geaggregeerd op de server, kunnen ze elkaar gedeeltelijk opheffen—een fenomeen dat wordt aangeduid als "silent cancellation". Deze structurele annulering vermindert de effectieve magnitude van de globale update, wat convergentie belemmert, zelfs wanneer individuele cliënten substantiële lokale vooruitgang boeken.

2. Methodologie

De auteurs stellen Subspace-Regularized LoRA (Subspace-Reg) voor, een framework dat is ontworpen om lokale cliëntupdates uit te lijnen met een gedeelde globale referentie-subspace.

2.1 Doelfunctie

Elke cliënt $k$ optimaliseert een subspace-regularized doelfunctie die een proximale alignment penalty bevat. De doelfunctie is:
$$L_k^{SR}(w_k) = \min_{A_k, B_k} \left\{ L_k(w_k) + \frac{\mu}{2} \|B_k A_k - \Delta W_g^{(t)}\|_F^2 + \frac{\lambda}{2} (\|A_k\|_F^2 + \|B_k\|_F^2) \right\}$$
Waar:

• $L_k(w_k)$ de lokale loss is.
• $\Delta W_g^{(t)}$ de huidige globale low-rank update is (server referentie).
• $\mu$ de sterkte van de subspace alignment controleert.
• $\lambda$ de low-rank factoren regulariseert.
• De Frobenius-norm term de lokale updates dwingt dicht bij de referentie-update van de server te blijven, wat constructieve aggregatie stimuleert.

2.2 Algoritme

De training volgt een standaard FL-structuur met aanpassingen:

1. Initialisatie: De server initialiseert een globale low-rank basis $B_g$ met orthonormale kolommen.
2. Server Stap: De server selecteert cliënten, berekent een globale referentie-update via SVD van de gewogen som van voorgaande cliëntupdates, en broadcast het model en de referentie-basis.
3. Cliënt Stap: Elke cliënt initialiseert hun lokale factoren op de globale referentie en voert $E$ lokale epochs uit. De lokale gradiëntafdaling minimaliseert de subspace-regularized loss, die bestaat uit:
   • De primaire taak loss.
   • Een term die afwijking van het product $B_k A_k$ van de globale referentie $B_g A_g$ bestraft.
   • Een optionele term die afwijking van de basis $B_k$ van $B_g$ bestraft.
4. Aggregatie: Cliënten sturen hun updates ($\Delta W_k = B_k A_k$) naar de server, die hen aggregeert en de globale model bijwerkt.

3. Theoretische Analyse

Het papier biedt een theoretische analyse om te begrijpen hoe subspace alignment de convergentie onder heterogeniteit beïnvloedt.

• Geometric Cancellation Bound: De auteurs leiden af dat de gekwadrateerde norm van de geaggregeerde update afhankelijk is van de cross-termen tussen cliëntupdates. Onder heterogeniteit, indien updates niet uitgelijnd zijn, worden deze cross-termen negatief, wat de effectieve update magnitude vermindert (silent cancellation).
• Convergentie Garanties: Onder standaard aannames (smoothness, bounded heterogeneity, bounded variance, en bounded low-rank updates), bewijst het papier een convergentiebound voor de gemiddelde gekwadrateerde gradiëntnorm.
  • De bound bevat een term $C_1 \sigma^2 / \mu^2$ die de misalignment penalty representeert.
  • In tegen tegenstelling tot standaard LoRA-FedAvg, dat lijdt onder een vaste cancellation term, staat Subspace-Reg toe dat deze penalty willekeurig klein gemaakt kan worden door de alignment parameter $\mu$ te verhogen.
  • De auteurs leiden condities af voor $\mu$ en het aantal communicatierondes $T$ om een doelstelling van gradiënt nauwkeurigheid $\epsilon$ te bereiken.

4. Experimentele Resultaten

De methode werd geëvalueerd op twee pretrained modellen (RoBERTa-large en SmolLM-360M) met de HellaSwag dataset in een non-IID setting (10 cliënten, Dirichlet $\beta=0.5$). Experimenten werden uitgevoerd over 3 willekeurige seeds (42, 43, 44), wat neerkomt op 24 runs.

4.1 Prestaties op RoBERTa-large

Subspace-Reg presteerde significant beter dan alle baselines (FedAvg, SVD redistribution, en FedSVD):

• Nauwkeurigheid: Behaalde de hoogste gemiddelde beste nauwkeurigheid (0.454 ± 0.023) en gemiddelde finale nauwkeurigheid (0.429 ± 0.011), waarmee het de beste baseline (FedSVD) versloeg met een marge van ~0.118 in beste nauwkeurigheid.
• Loss: Behaalde de laagste finale loss (1.363).
• Alignment: Behaalde een bijna perfecte Basis Overlap (≈0.9999), vergeleken met 0.958–0.991 voor baselines.

4.2 Prestaties op SmolLM-360M

De resultaten op SmolLM-360M lieten zien dat de winst in nauwkeurigheid model-afhankelijk is:

• Nauwkeurigheid: FedAvg behaalde de hoogste nauwkeurigheid (0.581), terwijl Subspace-Reg tweede werd (0.564).
• Alignment: Ondanks het gebrek aan verbetering in nauwkeurigheid, behield Subspace-Reg een bijna perfecte Basis Overlap (0.99989), consistent met de RoBERTa-resultaten.
• Conclusie: De geometrische alignment was succesvol, maar dit vertaalde zich niet naar downstream nauwkeurigheidswinst voor dit specifieke model.

4.3 Belangrijke Diagnostische Bevindingen

• Basis Overlap vs. Cosine Similarity: Hoewel SVD-gebaseerde methoden hogere cosine-gebaseerde metrieken behaalden (Client-Global Cosine, Pairwise Cosine), behaalden zij niet dezelfde nauwkeurigheidsverbeteringen als Subspace-Reg op RoBERTa. De auteurs identificeren Basis Overlap (directe subspace alignment) als de kritieke diagnostische factor, in plaats van cosine similarity.
• Robuustheid: De bijna perfecte basis overlap was consistent over alle seeds en beide modellen, wat de hypothese van geometrische alignment valideert, zelfs wanneer de downstream prestaties varieerden.

5. Betekenis en Claims

Het papier claimt de volgende bijdragen en betekenis:

1. Identificatie van Silent Cancellation: Het benadrukt dat subspace misalignment in LoRA-gebaseerde FL leidt tot destructieve aggregatie, een fenomeen dat voorheen onvoldoende onderzocht was.
2. Subspace-Regularized Objective: Het stelt een proximale alignment penalty voor die expliciet aanmoedigt dat lokale updates dicht bij een globale referentie-subspace blijven, waarmee het concept van FedProx uitbreidt van parameter magnitude naar geometrische structuur.
3. Theoretische Verbetering: Het biedt een convergentiebound die aantoont dat de misalignment penalty instelbaar is, wat een strikterere garantie biedt dan standaard LoRA-FedAvg onder heterogeniteit.
4. Empirische Validatie: Het demonstreert dat de voorgestelde methode consistent een bijna perfecte basis overlap (≈0.9999) bereikt. Op RoBERTa-large vertaalt dit zich naar substantiële winst in nauwkeurigheid. Op SmolLM-360M bevestigt het dat geometrische alignment bereikt kan worden onafhankelijk van downstream nauwkeurigheid, wat suggereert dat de voordelen van alignment model-afhankelijk zijn.
5. Diagnostisch Inzicht: Het werk stelt vast dat basis overlap een relevantere indicator is voor aggregatiekwaliteit dan cosine similarity in federated LoRA settings.

De auteurs concluderen dat hoewel subspace alignment een robuust geometrisch mechanisme is, de impact ervan op de uiteindelijke taakprestaties afhangt van de specifieke modelarchitectuur en capaciteit, wat verder onderzoek rechtvaardigt naar de reden waarom deze voordelen variëren tussen modellen.