FedRot-LoRA: Mitigating Rotational Misalignment in Federated LoRA

Het paper introduceert FedRot-LoRA, een federatief leerframework dat de prestaties van LoRA-finetuning verbetert door rotatiemismatch tussen clients op te lossen via orthogonale transformaties, wat leidt tot een stabielere training en hogere nauwkeurigheid zonder extra communicatiekosten.

De kern

FedRot-LoRA: De "Draaiende Kompas" Oplossing voor Gezamenlijk Leren

Stel je voor dat een groep vrienden (de "clients") samen een enorm, slim boek wil schrijven (een groot taalmodel). Ze willen dit doen zonder hun eigen privé-notities (hun data) met elkaar te delen, uit angst dat iemand ze leest. Dit noemen we Federated Learning.

Om het boek te verbeteren, gebruiken ze een slimme truc genaamd LoRA. In plaats van het hele boek opnieuw te schrijven, maken ze slechts kleine, compacte aantekeningen op losse blaadjes. Dit is snel en bespaart ruimte.

Het Probleem: De Verwarde Kompasnaalden

Hier zit echter een addertje onder het gras.

Elke vriend schrijft zijn aantekeningen op een eigen manier. Soms gebruiken ze een ander "referentiekader".

• Vriend A denkt: "Noord is naar links."
• Vriend B denkt: "Noord is naar rechts."

Beiden hebben het precies over hetzelfde punt op de kaart (ze bedoelen dezelfde zin in het boek), maar hun aantekeningen wijzen in verschillende richtingen. Dit noemen de auteurs rotational misalignment (rotatie-misalignering).

Wanneer de centrale coördinator (de server) alle aantekeningen verzamelt om het boek te updaten, doet hij dat op de simpele manier: hij telt alles bij elkaar op.

• Als je een pijl die naar links wijst optelt bij een pijl die naar rechts wijst, krijg je een flauwe, onzinnige lijn in het midden.
• Het resultaat is dat de nieuwe versie van het boek verward raakt, slechter wordt, en de training instabiel wordt. De vrienden "vernielen" elkaars werk zonder het te merken.

De Oplossing: FedRot-LoRA

De auteurs van dit paper, FedRot-LoRA, hebben een elegante oplossing bedacht.

In plaats van de aantekeningen direct bij elkaar te tellen, laten ze elke vriend eerst zijn eigen kompas draaien zodat ze allemaal naar hetzelfde "Noorden" wijzen.

1. De Draai (Rotatie): Voordat een vriend zijn aantekeningen opstuurt, kijkt hij naar de laatste versie van het gezamenlijke boek. Hij draait zijn eigen aantekeningen (via een wiskundige transformatie) zodat ze perfect aansluiten bij de richting van het gezamenlijke boek.
2. Behoud van Betekenis: Belangrijk: door te draaien verandert de inhoud van de aantekening niet. Als je een tekening op een vel papier draait, is het nog steeds dezelfde tekening. Alleen de oriëntatie verandert.
3. Het Samenvoegen: Nu, wanneer de server alle gedraaide aantekeningen optelt, wijzen ze allemaal in dezelfde richting. Ze versterken elkaar in plaats van elkaar op te heffen.

Waarom is dit zo slim?

• Geen extra kosten: Het draaien van de kompasnaald kost bijna geen tijd of energie. Het is een lichte berekening die elke vriend lokaal doet.
• Geen geheimen: Ze hoeven hun privé-data niet te delen. Ze sturen alleen de al gedraaide aantekeningen.
• Stabiel leren: Omdat de richtingen nu overeenkomen, wordt het boek veel sneller en beter. Het voorkomt die "flauwe lijnen" in het midden.

Een Dagelijkse Analogie: Het Orkest

Stel je een orkest voor waar elke muzikant (client) een solo speelt.

• Zonder FedRot-LoRA: Iedereen speelt in een andere toonsoort of tempo. Als ze samen spelen, klinkt het als een luidruchtig, onbegrijpelijk geraas.
• Met FedRot-LoRA: Voor ze beginnen, stemt elke muzikant zijn instrument af op de dirigent (de server). Ze draaien hun toonhoogte en tempo een klein beetje aan zodat ze perfect in harmonie zijn met de groep. Nu, als ze samen spelen, klinkt het als een prachtige symfonie.

Conclusie

FedRot-LoRA lost een verborgen probleem op in het gezamenlijk trainen van AI: het feit dat iedereen zijn werk op een iets andere manier "richt". Door die richtingen evenwijdig te maken voordat ze worden samengevoegd, krijgen we een veel slimmer, sneller en betrouwbaarder model, zonder dat we onze privacy in gevaar brengen of enorme rekenkracht nodig hebben.

Het is simpelweg het verschil tussen een groep mensen die in verschillende richtingen duwen (niets gebeurt) en een groep die allemaal in precies dezelfde richting duwt (het voertuig beweegt snel vooruit).

Technische samenvatting

1. Het Probleem: Rotational Misalignment in Federated LoRA

Federated Learning (FL) stelt organisaties in staat om gezamenlijk grote taalmodellen (LLMs) te finetunen zonder dat ruwe data centraal wordt opgeslagen, wat privacy en regelgeving respecteert. LoRA (Low-Rank Adaptation) is een populaire methode voor parameter-efficiënt finetunen, waarbij gewichtsupdates worden geparametriseerd als het product van twee laag-rang matrices: $\Delta W = BA$.

In een federale setting berekent elke client $i$ lokale updates $B_i A_i$. Het server-side aggregatieprobleem ontstaat wanneer deze lokale updates worden samengevoegd:

• Ideale aggregatie: Het gemiddelde nemen van de volledige updates $\frac{1}{N}\sum B_i A_i$. Dit is echter vaak van een hogere rang dan $r$, wat de communicatiekosten en het geheugengebruik onacceptabel maakt.
• Naïeve aggregatie (Factoren gemiddeld): Om de laag-rang structuur te behouden, worden de matrices $B$ en $A$ apart gemiddeld: $\bar{B} = \frac{1}{N}\sum B_i$ en $\bar{A} = \frac{1}{N}\sum A_i$, waarna het resultaat $\bar{B}\bar{A}$ wordt gebruikt.

De Kern van het Probleem:
LoRA-decompositie is rotatie-invariant. Dit betekent dat voor elke orthogonale matrix $R$ geldt dat $(B_i R)(R^\top A_i) = B_i A_i$. Verschillende clients kunnen dus semantisch identieke updates genereren, maar deze weergeven in verschillende, niet-overeenkomende latente subruimtes (verschillende oriëntaties).
Wanneer deze niet-uitgelijnde factoren direct worden gemiddeld (zoals bij naïeve aggregatie), interfereren ze destructief. De subruimtes van de clients "stoten" elkaar af in plaats van elkaar te versterken, wat leidt tot:

1. Een aanzienlijke aggregatiefout.
2. Instabiel trainingsgedrag.
3. Een verslechtering van de prestaties van het globale model.

2. Methodologie: FedRot-LoRA

Om dit probleem op te lossen, stellen de auteurs FedRot-LoRA voor. Dit is een framework dat lokale LoRA-updates uitlijnt via orthogonale transformaties voordat ze worden samengevoegd.

Kernstappen van het algoritme:

1. Lokaal Trainen: Clients trainen hun lokale factoren $A_i$ en $B_i$ op hun lokale data.
2. Rotational Alignment (Uitlijning): Voordat de factoren naar de server worden gestuurd, past elke client een rotatiematrix $R_i$ toe om hun lokale factoren uit te lijnen met een globaal referentiepunt (meestal de geaggregeerde factoren van de vorige ronde, $\bar{A}_{t-1}$ en $\bar{B}_{t-1}$).
   • Dit wordt geformuleerd als een Orthogonal Procrustes-probleem: minimaliseer $\|R_i^\top A_i - A_{ref}\|_F^2$ onder de voorwaarde dat $R_i$ orthogonaal is ($R_i^\top R_i = I$).
   • Dit probleem heeft een gesloten oplossing via Singular Value Decomposition (SVD).
3. Alternating Alignment: Om zowel $A$ als $B$ effectief uit te lijnen, wisselen de clients per communicatieronde af tussen het uitlijnen van $A$ of $B$.
4. Soft Rotation: In de vroege fasen van training kan het referentiepunt ruis bevatten. Om overcorrectie te voorkomen, wordt een "soft rotation" gebruikt. Dit interpolatieert tussen geen uitlijning en volledige Procrustes-uitlijning via een parameter $\lambda \in [0, 1]$.
   • De transformatie behoudt de semantische update: $\tilde{B}_i \tilde{A}_i = B_i A_i$.
5. Aggregatie: De server ontvangt de uitgelijnde factoren $\tilde{A}_i$ en $\tilde{B}_i$ en berekent het gemiddelde.

Complexiteit:
De extra rekentijd is minimaal en hangt alleen af van de LoRA-rang $r$ (complexiteit $O(d \cdot r^2 + r^3)$), wat verwaarloosbaar is vergeleken met de lokale trainingstijd. Er is geen extra communicatiekosten omdat alleen de reeds bestaande factoren worden verzonden.

3. Belangrijkste Bijdragen

1. Identificatie van Rotational Noise: De auteurs identificeren rotatie-invariantie als een onderbelichte, maar kritieke bron van aggregatiefouten in federale LoRA, los van de algebraïsche niet-commutativiteit.
2. FedRot-LoRA Framework: Een nieuw protocol dat client-updates uitlijnt via orthogonale transformaties, waardoor subruimte-mismatch wordt verminderd zonder de expressiviteit van het model te beperken (in tegenstelling tot methoden die één factor bevriezen).
3. Theoretische Convergentie-analyse: De paper levert een wiskundige analyse die aantoont dat naïeve factor-gemiddelde aggregatie een foutterm introduceert. FedRot-LoRA vermindert deze fout en levert een strikt strakkere bovengrens op voor de aggregatiefout, wat leidt tot betere convergentiegaranties.
4. Uitgebreide Experimenten: Validatie op diverse taken (Natural Language Understanding en Generatie) en modellen (RoBERTa-Large, Llama 3-8B).

4. Resultaten

De experimenten tonen aan dat FedRot-LoRA consistent beter presteert dan bestaande baselines (zoals FedIT, FFA-LoRA en RoLoRA) onder verschillende omstandigheden:

• NLP-taken (GLUE Benchmark): FedRot-LoRA behaalt de hoogste gemiddelde nauwkeurigheid op taken zoals SST-2, MNLI en RTE. Het reduceert de standaardafwijking (instabiliteit) aanzienlijk, vooral bij een groot aantal clients ($N=10$) en heterogene data.
• Generatieve taken (GSM8K & HumanEval): Op Llama 3-8B voor wiskundig redeneren en codegeneratie overtreft FedRot-LoRA alle concurrenten, met name door een lagere variantie in resultaten tussen verschillende runs.
• Robuustheid:
  • Rank: Prestaties blijven stabiel bij toenemende LoRA-rang ($r=4, 8, 16$), terwijl andere methoden (zoals RoLoRA) sterk degraderen bij hogere rangen.
  • Heterogeniteit: FedRot-LoRA presteert goed zowel bij homogene (IID) als sterk heterogene (Non-IID) dataverdelingen.
• Ablatie-studies:
  • Rotatie vs. Rescaling: Eenvoudige schaling (rescaling) werkt niet in hogere dimensies; rotatie is noodzakelijk om subruimtes uit te lijnen.
  • Soft Rotation: Een optimale $\lambda$ (vaak rond 0.4 - 0.6) is cruciaal; te agressieve uitlijning ($\lambda=1$) kan in vroege rondes schadelijk zijn door ruis.
  • Alternating: Het afwisselend uitlijnen van $A$ en $B$ werkt beter dan het uitlijnen van slechts één factor.

5. Betekenis en Conclusie

FedRot-LoRA biedt een lichtgewicht, communicatie-efficiënt oplossing voor een fundamenteel probleem in federatief finetunen van LLMs. Door de rotatie-invariantie van LoRA actief te benutten in plaats van erdoor te worden beperkt, verbetert het de stabiliteit en nauwkeurigheid van het trainingsproces.

De methode maakt het mogelijk om federale parameter-efficiënte finetuning (PEFT) praktischer toe te passen in privacygevoelige en hulpbronnenbeperkte omgevingen, zonder de communicatiekosten te verhogen of de modelcapaciteit te beperken. Dit is een belangrijke stap naar de bredere adoptie van federatief leren voor grote taalmodellen.