Technische Samenvatting: Efficiënte Zeroth-Order Federated Finetuning van Taalmodellen op Middelenbeperkte Apparaten

1. Probleemstelling

Federated Learning (FL) biedt een privacy-beschermend kader voor het finetunen van Large Language Models (LLMs) over gedistribueerde databronnen heen. Het implementeren van deze modellen op middelenbeperkte edge-apparaten stuit echter op aanzienlijke hindernissen:

• Geheugenbeperkingen: First-order methoden (bijv. standaard backpropagation) en parameter-efficiënte technieken zoals LoRA vereisen het opslaan van tussenliggende activaties voor gradiëntberekening. Dit resulteert in een geheugenvoetafdruk (bijv. ~4 GB voor RoBERTa-large) die de capaciteiten van de meeste edge-apparaten overstijgt.
• Communicatie-overhead: Het verzenden van volledige modelparameters of zelfs low-rank adaptatiematrices (LoRA) brengt substantiële uploadkosten met zich mee, wat met name problematisch is in omgevingen met een lage bandbreedte.
• Beperkingen van Zeroth-Order (ZO): Hoewel Zeroth-Order Optimalisatie (ZO) methoden zoals MeZO geen backpropagation nodig hebben door gebruik te maken van forward passes en pseudo-willekeurige perturbaties, lijden ze aan trage convergentie en een hoge computationele vraag. Om stabiele gradiënt-inschattingen te verkrijgen, vereisen bestaande ZO-FL methoden een groot aantal perturbaties per trainingsstap, wat leidt tot excessieve Forward-Only berekeningen (FLOPs).

Er bestaat een kritieke afweging: huidige ZO-methoden zijn geheugenefficiënt maar computationeel duur, terwijl first-order methoden computationeel efficiënt maar geheugenintensief zijn.

2. Methodologie: FedSPZO

De auteurs stellen Federated Split-Perturbation Zeroth-order Optimization (FedSPZO) voor, een nieuw framework ontworpen om de computationele last van ZO-FL te verminderen met behoud van de geheugen- en communicatievoordelen.

Kernmechanisme: Split-Perturbation Strategie

FedSPZO maakt gebruik van de sequentiële compositionele structuur van neurale netwerken door het model te splitsen in twee opeenvolgende blokken, $f_1$ (het eerste, grotere blok) en $f_2$ (het tweede, kleinere blok), zodanig dat $y = f_2(\theta_2; f_1(\theta_1; B))$.

• Asymmetrische Perturbatie-budgetten: In plaats van hetzelfde aantal perturbaties toe te passen op het gehele netwerk, wijst FedSPZO een hoger perturbatiebudget toe aan het tweede, kleinere blok ($f_2$).
  • Blok 1 ($\theta_1$): Wordt geperturbeerd met $P_1$ richtingen.
  • Blok 2 ($\theta_2$): Wordt geperturbeerd met $P_2$ richtingen, waarbij $P_2 = 2 \times P_1 \times P_s$.
• Activatie-hergebruik: De belangrijkste efficiëntiewinst komt voort uit het hergebruiken van tussenliggende activaties.
  1. De client perturbeert $\theta_1$ in één richting ($+z_1$) en voert een forward pass uit door $f_1$ om een tussenliggende output $+y_l$ te genereren.
  2. Deze enkele $+y_l$ wordt vervolgens gebruikt als input voor $P_s$ verschillende perturbaties van $\theta_2$ (zowel positieve als negatieve richtingen).
  3. Het proces wordt herhaald voor de negatieve richting van $\theta_1$ ($-z_1$), waarbij $-y_l$ wordt gegenereerd, die wordt hergebruikt voor een andere set van $P_s$ perturbaties van $\theta_2$.
• Gradiënt-inschatting:
  • De gradiënt voor $\theta_2$ wordt ingeschat met behulp van de verschillen in losses van de $P_2$ perturbaties, die allemaal dezelfde input van $f_1$ delen.
  • De gradiënt voor $\theta_1$ wordt ingeschat door de verschillen in losses over de meerdere perturbaties van $f_2$ te middelen, wat effectief de ruis vermindert die wordt geïntroduceerd door de stochasticiteit van de perturbaties van $f_2$.

Communicatieprotocol

FedSPZO handhaaft een communicatie-efficiënt ontwerp:

• Client-zijde: Clients voeren $K$ lokale trainingsstappen uit. Ze uploaden geen modelparameters. In plaats daarvan uploaden ze alleen scalaire gradiëntwaarden ($G_1, G_2$) en random seeds ($S_1, S_2$) die worden gebruikt om de perturbatievectoren te regenereren.
• Server-zijde: De server reconstrueert de exacte client-modellen door de perturbatievectoren te regenereren van de ontvangen seeds en de scalaire updates toe te passen. Dit elimineert de noodzaak voor clients om modelgewichten te verzenden, wat de upload-overhead met ordes van grootte vermindert vergeleken met parameter-gebaseerde FL.

3. Belangrijkste Bijdragen

1. FedSPZO Framework: Een nieuwe ZO-FL methode die het model verdeelt in twee blokken en een groter perturbatiebudget toewijst aan het kleinere, diepere blok. Dit maakt gebruik van tussenliggend activatie-hergebruik om het volledige netwerk bij te werken met minder forward evaluaties.
2. Computationele Efficiëntie: De methode bereikt een reductie in computation van wel 3× vergeleken met state-of-the-art ZO-FL baselines (specifiek Fang et al., 2022, en Li et al., 2025) terwijl de geheugen- en communicatievoordelen van ZO behouden blijven.
3. Uitgebreide Evaluatie: Uitgebreide experimenten op RoBERTa-large, OPT-1.3B, en LLaMA-3-3.2B modellen over meerdere datasets (SST-2, RTE, WiC, MultiRC, SQuAD, BoolQ) tonen aan dat FedSPZO de totale GFLOPs aanzienlijk vermindert terwijl de nauwkeurigheid vergelijkbaar blijft met first-order backpropagation methoden (met minimale degradatie).
4. Ablatie-analyse: De auteurs hebben studies uitgevoerd naar modelsplitting, perturbatie-budgetten en onafhankelijke blok-wijze gradiëntberekening, waarmee wordt gevalideerd dat de split-perturbatie strategie superieur is aan onafhankelijke blok-inschatting of uniforme perturbatie-schema's.

4. Experimentele Resultaten

De evaluatie vergelijkt FedSPZO met first-order methoden (FedAvg, FedAvg+LoRA) en andere ZO-FL methoden (FedZO, DecomFL).

• Geheugenefficiëntie: FedSPZO behoudt een inference-achtige geheugenvoetafdruk. Voor RoBERTa-large met een contextlengte van 256 vereist het ~1,75 GB, terwijl FedAvg ~4,9 GB vereist en LoRA ~3,9 GB. Dit maakt FedSPZO haalbaar voor edge-apparaten waar LoRA en standaard FL onmogelijk zijn.
• Communicatie-overhead: FedSPZO vermindert de uploadkosten met meer dan drie ordes van grootte vergeleken met LoRA. Terwijl LoRA low-rank matrices uploadt (~100s van MB's), uploadt FedSPZO alleen scalaire gradiënten en seeds (kilobytes tot megabytes).
• Computationele Kosten:
  • Vergeleken met FedAvg(LoRA): FedSPZO vereist hogere totale computation (ongeveer 17× meer) vanwege de tragere convergentiesnelheid die inherent is aan ZO-methoden.
  • Vergeleken met Andere ZO Methoden: FedSPZO is aanzienlijk efficiënter. Het vermindert de computation met 2,1× tot 3× vergeleken met DecomFL en FedZO over diverse modellen en datasets.
• Nauwkeurigheid: FedSPZO bereikt een nauwkeurigheid die vergelijkbaar is met FedAvg(LoRA), met slechts minimale degradaties (bijv. <1% op SST-2, ~4% op RTE voor RoBERTa). Het presteert iets slechter dan volledige FedAvg (backpropagation), wat als een upper bound dient, maar het verschil wordt als minimaal beschouwd in de context van de middelenbeperkingen.
• Trainingstijd: In een gesimuleerde edge-omgeving (Jetson Orin Nano) is FedSPZO (2,6 uur voor RoBERTa/SST-2) trager dan LoRA (0,45 uur) maar drastisch sneller dan DecomFL (6 uur) en FedZO (166 uur).

5. Betekenis en Claims

Het artikel positioneert FedSPZO als een pragmatische afweging voor omgevingen met beperkte middelen.

• Voor Edge-apparaten: Het wordt gepresenteerd als een voorkeurskeuze wanneer geheugen en communicatiebandbreedte de primaire knelpunten zijn, scenario's waarin LoRA en standaard backpropagation falen.
• Balansact: De auteurs erkennen dat FedSPZO de computationele kosten van ZO niet elimineert, maar deze wel aanzienlijk mitigeert. Het biedt een levensvatbaar pad voor het finetunen van LLM's op apparaten waar het opslaan van activaties voor backpropagation onmogelijk is en het verzenden van volledige modelupdates te kostbaar is.
• Toekomstperspectief: De auteurs merken op dat hoewel FedSPZO de computationele efficiëntie over bestaande ZO-methoden heeft verbeterd, het nog steeds achterloopt op first-order methoden qua pure snelheid. Toekomstig werk richt zich op het verkennen van ZO op inference-only versnellers met lagere precisie om dit gat verder te dichten.

Het artikel concludeert dat FedSPZO er succesvol in slaagt de hoge geheugen/communicatiekosten van first-order FL te ontkoppelen van de hoge computationele kosten van traditionele ZO, waardoor het finetunen van LLM's mogelijk wordt in scenario's die voorheen als onhaalbaar werden beschouwd.