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

Die Arbeit stellt ZorBA vor, ein Framework für das federierte Fine-Tuning von LLMs, das durch den Einsatz von Zeroth-Order-Optimierung und einer heterogenen Blockaktivierung den VRAM-Verbrauch erheblich senkt und gleichzeitig die Konvergenz beschleunigt.

Das Wesentliche

Stellen Sie sich vor, Sie und eine Gruppe von Freunden wollen gemeinsam ein riesiges, komplexes Puzzle lösen. Dieses Puzzle ist ein Künstlicher Intelligenz-Modell (ein sogenanntes "Large Language Model" oder LLM), das so groß ist, dass es den gesamten Speicher (VRAM) eines einzelnen Computers sprengen würde.

Das Ziel ist es, das Modell zu verbessern, ohne dass jemand seine privaten Notizen (die Daten) mit anderen teilt. Das nennt man Federated Learning (verteiltes Lernen).

Das Problem bisher: Um das Puzzle zu verbessern, mussten alle Freunde ihre gesamten Puzzleteile (die Berechnungen) aufwendig hin- und herschicken. Das brauchte viel Speicherplatz auf ihren Geräten und viel Zeit für den Datentransfer.

Hier kommt ZorBA ins Spiel – eine neue, clevere Methode, die wie ein genialer Puzzle-Leiter funktioniert.

1. Das Problem: Der "Speicher-Überlauf"

Stellen Sie sich vor, jeder Freund hat nur einen kleinen Arbeitstisch (begrenzter VRAM). Wenn sie versuchen, das ganze riesige Puzzle auf einmal zu bearbeiten, stürzt ihr Tisch um. Außerdem ist es mühsam, das ganze Puzzle immer wieder zu fotografieren und per E-Mail zu versenden (hoher Kommunikationsaufwand).

2. Die Lösung: ZorBA (Der clevere Leiter)

ZorBA löst dieses Problem mit drei genialen Tricks:

Trick A: "Raten statt Rechnen" (Zeroth-Order Optimization)

Normalerweise müssen die Freunde genau berechnen, warum ein Puzzleteil falsch liegt (das nennt man "Gradienten berechnen"). Das ist wie eine komplizierte mathematische Analyse, die viel Speicher braucht.
ZorBA sagt: "Vergessen wir die komplizierte Analyse! Wir machen einfach einen kleinen Test."

• Die Analogie: Statt die genaue Ursache eines Fehlers zu berechnen, probieren die Freunde einfach eine winzige Verschiebung eines Puzzleteils aus. Sie schauen nur, ob das Bild dadurch besser oder schlechter aussieht (Vorwärtsdurchlauf).
• Der Vorteil: Sie müssen keine aufwendigen Berechnungen speichern. Das spart enorm viel Platz auf dem Arbeitstisch (VRAM).

Trick B: "Jeder bearbeitet nur einen Teil" (Heterogene Block-Aktivierung)

Das Puzzle besteht aus vielen Abschnitten (Blöcken). Früher musste jeder Freund versuchen, alle Abschnitte gleichzeitig zu bearbeiten. Das war zu viel für die kleinen Tische.
ZorBA sagt: "Wir teilen die Arbeit auf! Du bearbeitest nur die Landschaftsteile, du nur die Gesichter, und er nur die Gebäude."

• Die Analogie: Der Leiter (Server) weist jedem Freund genau zu, welche Puzzleteile er gerade anfassen darf, basierend darauf, wie groß sein Tisch ist.
• Der Vorteil: Jeder braucht weniger Platz. Aber das Puzzle wird trotzdem schneller fertig, weil alle gleichzeitig an verschiedenen Teilen arbeiten, ohne sich im Weg zu stehen.

Trick C: "Der geheime Code" (Shared Random Seeds)

Normalerweise müssten die Freunde ihre Verschiebungen (die Testbewegungen) genau beschreiben und verschicken. Das sind riesige Datenmengen.
ZorBA sagt: "Wir nutzen alle denselben geheimen Zufalls-Code (Seed)."

• Die Analogie: Alle Freunde haben das gleiche Buch mit Zufallszahlen. Wenn der Leiter sagt "Nimm Seite 5", wissen alle sofort, welche Zahl sie nehmen müssen, ohne dass er sie ihm schicken muss.
• Der Vorteil: Statt riesige Datenpakete zu senden, reicht es, nur zu sagen: "Wir haben Seite 5 benutzt." Das spart enorm viel Zeit und Bandbreite.

3. Das Ergebnis: Schneller und schlanker

In Tests hat sich gezeigt, dass ZorBA:

• Bis zu 62 % weniger Speicherplatz auf den Geräten der Freunde braucht.
• Viel schneller lernt als andere Methoden, weil die Arbeit intelligent aufgeteilt wird.
• Kaum Datenverkehr verursacht, weil nur kleine "Hinweise" statt ganzer Datenpakete ausgetauscht werden.

Zusammenfassung

Stellen Sie sich ZorBA wie einen effizienten Koch vor, der eine riesige Suppe für viele kocht:

• Statt dass jeder Koch den ganzen Topf umrührt (was den Herd sprengen würde), bekommt jeder nur einen kleinen Löffel und rührt einen bestimmten Teil der Suppe um.
• Statt komplizierte Rezepte zu schreiben, nutzen alle denselben Timer und dieselben Zutatenlisten.
• Das Ergebnis: Die Suppe ist schneller fertig, niemand hat den Herd überhitzt, und es wurde kaum Papier für Rezepte verschwendet.

ZorBA macht es also möglich, dass auch kleine, schwache Geräte gemeinsam an riesigen KI-Modellen lernen können, ohne dass der Speicherplatzt oder die Internetverbindung kollabiert.

Technische Zusammenfassung

1. Problemstellung

Das Fine-Tuning von Large Language Models (LLMs) in einem federierten Lernumfeld (Federated Learning, FL) stößt auf zwei wesentliche Herausforderungen, insbesondere bei ressourcenbeschränkten Clients (z. B. Geräten mit begrenztem VRAM):

1. Hoher VRAM-Verbrauch: Herkömmliche FL-Ansätze nutzen Backpropagation (BP), um Gradienten erster Ordnung zu berechnen. Dies erfordert das Speichern von Aktivierungen und Gradienten für alle Transformer-Blöcke des Modells. Bei großen Modellen (Milliarden von Parametern) übersteigt dies oft die VRAM-Kapazität lokaler GPUs.
2. Kommunikations-Overhead: Der Austausch vollständiger Modell-Updates oder Gradienten zwischen Clients und dem Server führt zu erheblichem Kommunikationsaufwand, besonders bei hochdimensionalen Modellen.
3. Limitationen existierender Nullter-Ordnung-Methoden: Zwar eliminiert die Optimierung nullter Ordnung (Zeroth-Order Optimization, ZOO) den Bedarf an Gradientenspeicherung, indem sie nur Vorwärtsdurchläufe nutzt, jedoch führen bestehende ZOO-FL-Ansätze zu langsameren Konvergenzraten (aufgrund von Varianz in den geschätzten Gradienten) und einem immer noch hohen VRAM-Verbrauch durch die Speicherung aller Aktivierungen. Zudem wird oft der gesamte Modellraum aktualisiert, was ineffizient ist.

Ziel: Entwicklung eines Frameworks, das den VRAM-Verbrauch minimiert, den Kommunikations-Overhead reduziert und gleichzeitig eine schnelle Konvergenzrate bei heterogenen Client-Ressourcen gewährleistet.

2. Methodik: Das ZorBA-Framework

ZorBA (Zeroth-order Federated Fine-tuning with Heterogeneous Block Activation) kombiniert drei Kernkomponenten:

A. Zeroth-Order Optimierung (ZOO)

Statt Backpropagation verwendet ZorBA eine „Forward-Pass-only"-Methode.

• Gradientenschätzung: Die Gradienten werden durch endliche Differenzen von Verlustfunktionswerten geschätzt, die durch zufällige Perturbationsvektoren erzeugt werden.
• Vorteil: Es müssen keine Gradienten gespeichert werden, was den VRAM-Verbrauch für Gradienten eliminiert.

B. Heterogene Block-Aktivierung (Heterogeneous Block Activation)

Dies ist der zentrale Innovationsschritt. Anstatt dass alle Clients das gesamte Modell aktualisieren:

• Der zentrale Server weist jedem Client eine Teilmenge der Transformer-Blöcke zu, die aktualisiert werden sollen.
• Die Zuweisung basiert auf den VRAM-Kapazitäten der einzelnen Clients und der Notwendigkeit, die Konvergenzrate zu optimieren.
• Mechanismus: Ein Client $n$ aktiviert nur die Blöcke $m$, für die $a_{m,n} = 1$ gilt. Die nicht aktivierten Blöcke bleiben eingefroren (frozen). Dies reduziert den VRAM-Verbrauch für Aktivierungen drastisch, da nur für die aktiven Blöcke Zwischenergebnisse gespeichert werden müssen.

C. Gemeinsame Zufallssamen (Shared Random Seeds) zur Reduktion des Kommunikations-Overheads

• Anstatt Perturbationsvektoren oder Gradienten zu übertragen, initialisiert der Server eine Menge zufälliger Samen ($S$) und teilt diese einmalig mit allen Clients.
• Da alle Parteien dieselben Samen verwenden, können Clients und Server unabhängig aber identisch dieselben Perturbationsvektoren generieren.
• Kommunikation: Clients senden nur die endlichen Differenzen der geschätzten Gradienten (Skalarwerte) zurück, nicht die hochdimensionalen Vektoren selbst. Der Server rekonstruiert die Gradienten basierend auf den empfangenen Differenzen und den lokalen Samen. Dies reduziert den Kommunikationsaufwand signifikant.

3. Theoretische Analyse und Optimierung

Die Autoren leiten eine theoretische Konvergenzschranke für ZorBA her, die zeigt, dass die Konvergenzrate stark von der Block-Aktivierungsmatrix $A$ abhängt.

• Konvergenzmetrik: Die Konvergenz wird durch den Term $\Lambda(A)$ beeinflusst, der von der „Beliebtheit" (Popularity) der Blöcke abhängt. $\Lambda(A)$ ist definiert als die Summe der reziproken Quadrate der minimalen Beliebtheit jedes Clients.
• Erkenntnis: Um die Konvergenz zu maximieren (d.h. $\Lambda(A)$ zu minimieren), muss die Verteilung der aktivierten Blöcke so gestaltet sein, dass kein Client eine extrem geringe Anzahl an gemeinsamen Blöcken mit anderen hat. Es ist nicht nur die Anzahl der aktivierten Blöcke entscheidend, sondern deren Verteilung über die Clients.
• Optimierungsproblem: Es wird ein multi-kriterielles Optimierungsproblem formuliert, das $\Lambda(A)$ minimiert (für schnelle Konvergenz) und gleichzeitig den VRAM-Verbrauch jedes Clients unter einer Kapazitätsgrenze hält.
• Lösungsalgorithmus: Da das Problem NP-schwer ist, wird ein $\epsilon$-Constraint-Lexicographic-Algorithmus vorgeschlagen:
  1. Schritt 1 (Maximierung der minimalen Beliebtheit): Bestimmung des maximal möglichen Wertes für die geringste Beliebtheit eines Clients unter VRAM-Beschränkungen (unter Verwendung von Max-Flow-Algorithmen wie Dinic).
  2. Schritt 2 (Greedy-Anpassung): Ein gieriger Algorithmus aktiviert zusätzliche Blöcke, um die Anzahl der Clients zu minimieren, die noch die minimale Beliebtheit aufweisen, ohne die VRAM-Grenzen zu verletzen.
  3. Pareto-Front: Durch Variation der VRAM-Einsparungsziele ($\epsilon$) wird eine Pareto-Front zwischen Konvergenzgeschwindigkeit und VRAM-Nutzung erstellt, aus der die optimale Konfiguration gewählt wird.

4. Ergebnisse

Die Evaluation erfolgte auf den Datensätzen AG-News, SST-2 und SNLI mit den Modellen OPT-125M und OPT-1.3B. ZorBA wurde mit FedIT (First-Order), FedZO und DeComFL verglichen.

• VRAM-Effizienz: ZorBA reduziert den gesamten VRAM-Verbrauch aller Clients im Vergleich zu den Baselines um bis zu 62,41%. Dies ermöglicht das Fine-Tuning auf Geräten mit deutlich geringerer VRAM-Kapazität.
• Konvergenzgeschwindigkeit: ZorBA konvergiert schneller als FedZO und DeComFL (bis zu 23,76% weniger Runden für das Erreichen der Zielgenauigkeit), was die Wirksamkeit der optimierten Block-Aktivierung unterstreicht.
• Kommunikations-Overhead: Durch den Austausch von Samen und endlichen Differenzen statt ganzer Gradienten ist der Kommunikations-Overhead von ZorBA vernachlässigbar im Vergleich zu FedIT und FedZO und vergleichbar mit DeComFL.
• Trade-off: Die Experimente zeigen, dass eine suboptimale Block-Zuweisung (z. B. zu viele Blöcke bei einem Client, zu wenige bei anderen) die Konvergenz trotz gleicher Gesamtanzahl aktivierter Blöcke verschlechtern kann. ZorBA löst dieses Problem durch die gezielte Optimierung der Verteilung.

5. Bedeutung und Fazit

ZorBA stellt einen bedeutenden Fortschritt im Bereich des federierten Fine-Tunings von LLMs dar:

1. Demokratisierung des LLM-Fine-Tunings: Durch die drastische Reduktion des VRAM-Bedarfs können ressourcenbeschränkte Edge-Geräte an der Anpassung großer Modelle teilnehmen, was bisher aufgrund von Speicherbeschränkungen unmöglich war.
2. Skalierbarkeit: Die Kombination aus Zeroth-Order-Optimierung und intelligenter Block-Aktivierung macht federiertes Lernen für Modelle mit Milliarden von Parametern praktikabel.
3. Theoretische Fundierung: Die Arbeit liefert nicht nur ein praktisches Framework, sondern auch tiefgehende theoretische Einblicke in den Zusammenhang zwischen Block-Aktivierungsstrategien und Konvergenzraten in nicht-konvexen, verteilten Umgebungen.
4. Effizienz: Die Reduktion des Kommunikations-Overheads macht das System für Netzwerke mit begrenzter Bandbreite geeignet.

Zusammenfassend bietet ZorBA eine robuste Lösung, um die „Dreiecksbeziehung" zwischen Rechenressourcen (VRAM), Kommunikationskosten und Konvergenzgeschwindigkeit beim Fine-Tuning von LLMs in federierten Umgebungen zu optimieren.