ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation

Each language version is independently generated for its own context, not a direct translation.

ZorBA: De Slimme Manier om Grootte Taalmodellen te Leren zonder je Computer te Verbranden

Stel je voor dat je een gigantische, superintelligente robot wilt leren om Nederlands te spreken. Deze robot is zo groot dat hij niet in één hoofd past; hij is verdeeld over duizenden kleine hersencellen (we noemen ze "blokken"). Om hem te leren, moeten we hem oefeningen geven.

In de wereld van kunstmatige intelligentie noemen we dit Fine-tuning. Maar er is een groot probleem: deze robot is zo zwaar dat hij niet op de gewone laptops van mensen past. Als je probeert hem lokaal te trainen, springt je computer-geheugen (VRAM) eruit, alsof je probeert een olifant in een Mini Cooper te proppen.

Daarnaast willen we de data van mensen niet delen (privacy), dus we trainen de robot op vele verschillende computers tegelijk. Dit heet Federated Learning. Maar het uitwisselen van de enorme robot tussen al die computers kost zoveel tijd en internetbandbreedte, dat het bijna onmogelijk wordt.

De auteurs van dit paper, Chuiyang Meng, Ming Tang en Vincent Wong, hebben een oplossing bedacht genaamd ZorBA. Laten we uitleggen hoe het werkt met een paar simpele vergelijkingen.

1. Het Probleem: De "Grote Robot" en de "Kleine Werkbanken"

Stel je voor dat je een bouwteam hebt met 50 mensen (de klanten), maar ze werken allemaal in kleine schuurtjes met weinig ruimte. Ze moeten samen een enorme kathedraal bouwen (het taalmodel).

Het oude probleem: Om de kathedraal te bouwen, moesten ze eerst de hele blauwdruk in hun hoofd houden en elke steen berekenen. Dat paste niet in hun kleine schuurtjes.
Het communicatie-probleem: Ze moesten ook constant de hele blauwdruk naar elkaar mailen. Dat duurde eeuwen.

2. De Oplossing: ZorBA (De Slimme Bouwmeester)

ZorBA lost dit op met drie slimme trucjes:

Truc 1: "Kijken zonder te onthouden" (Zeroth-Order Optimalisatie)

Normaal gesproken moet een bouwer elke steen leggen, kijken of het goed zit, en dan onthouden waar hij fout zat om het de volgende keer beter te doen. Dat onthouden kost veel ruimte in het hoofd (VRAM).
ZorBA gebruikt een andere methode: Zeroth-Order.

De analogie: In plaats van te onthouden waar je fout zat, doet de bouwer een proefje. Hij legt een steen, kijkt of het goed is, en zegt: "Oké, als ik deze steen een beetje naar links schuif, wordt het beter." Hij onthoudt de richting van de verbetering, maar niet de hele complexe berekening erachter.
Het resultaat: Je hebt veel minder ruimte nodig in je hoofd (VRAM) omdat je niets hoeft op te slaan. Je doet alleen maar "proefjes" (voorwaartse passes).

Truc 2: De "Kieswijzer" (Heterogene Blok Activering)

Niet iedereen heeft even veel ruimte in zijn schuurtje. Sommigen hebben een grote garage, anderen een kleine schuur.

Het oude probleem: Iedereen probeerde dezelfde grote taak te doen, en de mensen met kleine schuurtjes kwamen in de knel.
De ZorBA-oplossing: De hoofdbouwer (de server) kijkt naar elke werknemer en zegt: "Jij, met de grote garage, mag de muren bouwen. Jij, met de kleine schuur, mag alleen de ramen plaatsen."
Dit heet Heterogene Blok Activering. Iedereen doet een ander stukje van de robot. Zo past het allemaal netjes in de beschikbare ruimte.

Truc 3: De "Geheime Code" (Gedeelde Willekeurige Zaden)

Normaal moeten mensen hun berekeningen naar elkaar sturen. Dat is veel post.

De ZorBA-oplossing: De hoofdbouwer geeft iedereen dezelfde geheime code (een willekeurige zaden-lijst).
Omdat iedereen dezelfde code heeft, weten ze allemaal precies welke "proefjes" ze moeten doen, zonder dat ze elkaar hoeven te bellen. Ze sturen alleen het resultaat van hun proefje (een klein getalletje) terug, in plaats van de hele blauwdruk.
Dit bespaart enorm veel internetverkeer.

3. Het Grote Raadsel: Wie doet wat?

De moeilijkste vraag was: Welke werknemer moet welk stukje doen om het snelst klaar te zijn, zonder dat iemand zijn schuurtje overvol raakt?

Als je te weinig mensen op een taak zet, duurt het lang (langzame convergentie). Als je te veel mensen op een taak zet, hebben ze geen ruimte meer (te veel VRAM).

De auteurs hebben een wiskundig algoritme bedacht (een soort slimme planner) dat dit perfect in evenwicht brengt. Ze noemen dit een $\epsilon$ -constraint lexicographic algoritme.

In het kort: Het algoritme zoekt de perfecte verdeling. Het zorgt ervoor dat niemand overbelast raakt, maar dat de robot toch zo snel mogelijk leert. Het is alsof je een orkest dirigeert waarbij elke muzikant een ander instrument speelt, maar ze allemaal precies op het juiste moment in de maat spelen.

Wat is het resultaat?

De tests tonen aan dat ZorBA wonderen doet:

Ruimtebesparing: Het gebruikt tot 62% minder geheugen dan andere methoden. Je kunt dus een veel grotere robot trainen op dezelfde kleine computer.
Snelheid: Het is sneller dan andere methoden die ook proberen ruimte te besparen.
Internet: Het verkeer tussen de computers is verwaarloosbaar klein.

Samenvatting in één zin

ZorBA is een slimme manier om enorme AI-modellen te trainen op duizenden kleine computers door ze alleen te laten "proeven" in plaats van "onthouden", door iedereen een ander stukje werk te geven, en door een geheime code te gebruiken zodat ze niet hoeven te bellen.

Het is alsof je een gigantisch puzzelstuk oplost met een team dat in kleine kamers zit, zonder dat je de hele puzzel hoeft mee te nemen, en zonder dat je de hele tijd hoeft te schreeuwen om instructies.

Each language version is independently generated for its own context, not a direct translation.

Probleemstelling

Het fine-tunen van grote taalmodellen (LLMs) in een federated learning (FL) omgeving stuit op twee fundamentele uitdagingen:

VRAM-beperkingen: Traditionele FL-methoden gebruiken backpropagation (BP) om gradiënten te berekenen. Bij LLMs met miljarden parameters vereist het opslaan van deze gradiënten en de bijbehorende forward-pass activaties aanzienlijke videogeheugen (VRAM) op de client-apparaten. Dit overschrijdt vaak de capaciteit van resource-beperkte clients.
Communicatie- overhead: De frequente uitwisseling van volledige modelupdates of gradiënten tussen clients en de centrale server leidt tot enorme bandbreedte-eisen, vooral bij hoge dimensies.
Beperkingen van bestaande Zeroth-Order (ZO) methoden: Hoewel zeroth-order optimalisatie (die gradiënten schat via forward passes zonder backpropagation) de VRAM-gebruik voor het opslaan van gradiënten elimineert, hebben eerdere federale ZO-aanpakken nog steeds drie nadelen:
- Langzamere convergentie door de hoge variantie in geschatte gradiënten in hoge dimensies.
- Zeer hoge VRAM-gebruik omdat ze alle transformer-blokken activeren, wat de opslag van forward-pass activaties lineair doet toenemen.
- Hoge communicatiekosten door het uitwisselen van geschatte gradiënten.

Methodologie: Het ZorBA Framework

ZorBA (Zeroth-order Federated Fine-tuning with Heterogeneous Block Activation) is een nieuw framework dat de volgende kerncomponenten combineert:

1. Zeroth-Order Optimalisatie (Forward-pass only):
In plaats van backpropagation gebruikt ZorBA zeroth-order optimalisatie. Clients schatten de gradiënten door het verlies te evalueren bij het huidige punt en bij een punt dat verstoord is met een willekeurige vector. Dit elimineert de noodzaak om gradiënten op te slaan, wat de VRAM-gebruik drastisch verlaagt.

2. Heterogene Blok-activering:
In plaats van dat elke client alle transformer-blokken van het model update, deelt de centrale server een subset van blokken toe aan elke client, gebaseerd op hun individuele VRAM-capaciteit.

Mechanisme: Een client activeert alleen een subset van blokken ( $a_{m,n} = 1$ ) en "bevriest" de rest. Dit reduceert de VRAM-gebruik voor forward-pass activaties, aangezien alleen de geactiveerde blokken tussenstanden moeten opslaan.
Doel: Dit zorgt voor een evenwicht tussen VRAM-efficiëntie en convergentiesnelheid.

3. Gedeelde Willekeurige Zaden (Shared Random Seeds):
Om de communicatie- overhead te minimaliseren, worden geen volledige gradiënten of perturbatievectoren uitgewisseld.

De server deelt een set willekeurige zaden ( $S$ ) met alle clients.
Zowel de server als de clients genereren onafhankelijk dezelfde perturbatievectoren op basis van deze zaden.
Clients sturen alleen de eindige verschillen (finite differences) van het verlies terug naar de server. De server kan hiermee de geschatte gradiënten reconstrueren zonder de hoge-dimensionele vectoren zelf te hoeven ontvangen.

4. Theoretische Analyse en Optimalisatie:
De auteurs analyseren de convergentie in een niet-convexe setting. Ze introduceren een metriek $\Lambda(A)$ die de convergentie- snelheid relateert aan de "populariteit" van de blokken (hoeveel clients een bepaald blok activeren).

Inzicht: Het minimaliseren van $\Lambda(A)$ (wat leidt tot snellere convergentie) is equivalent aan het maximaliseren van de minimale populariteit van blokken over de clients.
Optimalisatieprobleem: Ze formuleren een multi-objectief probleem om $\Lambda(A)$ te minimaliseren (voor convergentie) en het totale VRAM-gebruik te reduceren, onderworpen aan VRAM-beperkingen per client.
Oplossingsalgoritme: Een $\epsilon$ $ϵ$ -constraint lexicografisch algoritme wordt voorgesteld. Dit splitst het probleem op in twee subproblemen:
1. Het maximaliseren van de minimale populariteit van blokken (met behulp van een stroomalgoritme zoals Dinic's).
2. Een greedy-algoritme om extra blokken toe te voegen om het aantal clients met de laagste populariteit te minimaliseren, binnen de VRAM-beperkingen.

Belangrijkste Bijdragen

ZorBA Framework: Het eerste federale fine-tuning framework dat zeroth-order optimalisatie combineert met een heterogene blok-activeringsstrategie voor LLMs.
Theoretische Kader: Een rigoureuze convergentie-analyse die laat zien hoe blok-activeringsbeslissingen de convergentie- snelheid en VRAM-gebruik beïnvloeden. Ze tonen aan dat de convergentie- snelheid afhangt van de verdeling van de "minimale populariteit" van blokken, niet alleen van het totale aantal geactiveerde blokken.
Efficiënt Algoritme: Een $\epsilon$ -constraint lexicografisch algoritme dat een Pareto-optimaal compromis vindt tussen convergentie- snelheid en VRAM-gebruik, met een veel lagere computationele complexiteit dan het direct oplossen van het oorspronkelijke NP-hard probleem.
Privacy en Communicatie: Door gedeelde zaden en het uitwisselen van alleen eindige verschillen, wordt de communicatie- overhead drastisch verlaagd en wordt het risico op gradiënt-lekkage geminimaliseerd.

Resultaten

De auteurs hebben ZorBA getest op de AG-News, SST-2 en SNLI datasets met de OPT-125M en OPT-1.3B modellen. Ze vergeleken het met FedIT, FedZO en DeComFL.

VRAM-gebruik: ZorBA reduceerde het totale VRAM-gebruik met tot 62,41% vergeleken met de baselines (zoals FedIT en FedZO die alle blokken activeren).
Convergentie: ZorBA convergeerde sneller dan FedZO en DeComFL (tot 23,76% minder rondes nodig om de doelaccuracy te bereiken), dankzij de geoptimaliseerde blok-activering.
Communicatie: De communicatie- overhead was verwaarloosbaar in vergelijking met FedIT en FedZO, en vergelijkbaar met DeComFL, maar dan met aanzienlijk minder VRAM-gebruik.
Trade-off: De experimenten bevestigden de theoretische inzichten: een betere verdeling van de blok-activering (hoger $\Lambda$ ) leidt tot snellere convergentie, zelfs als het totale aantal geactiveerde blokken gelijk blijft.

Significantie

ZorBA is een doorbraak voor het toepassen van federated learning op grote taalmodellen in real-world scenario's waar clients beperkte hardware hebben (zoals edge devices). Het lost het fundamentele dilemma op tussen de noodzaak van grote rekenkracht voor LLM-fine-tuning en de beperkte middelen van gedistribueerde clients. Door het slim combineren van zeroth-order optimalisatie en een dynamische, heterogene blok-activering, maakt het fine-tuning van LLMs haalbaar zonder de privacy te schaden of de communicatie- kosten onbeheersbaar te maken. Dit opent de deur voor meer inclusieve en schaalbare AI-training in gedistribueerde omgevingen.