SpecFed: Accelerating Federated LLM Inference with Speculative Decoding and Compressed Transmission

1. Probleemstelling

Het paper adresseert twee kritieke knelpunten in Federated Large Language Model (LLM) Inference:

• Sequentiële Latentie: Standaard autoregressieve inferentie vereist dat alle gedistribueerde workers een volledige forward pass uitvoeren voor elke gegenereerde token. Dit leidt tot frequente synchronisatie en beperkt de decoding-throughput ernstig.
• Communicatie-Overhead: Hoewel Speculative Decoding (het gebruik van een klein "draft"-model om kandidaat-tokens te genereren voor parallelle verificatie) het aantal forward passes reduceert, introduceert het een massaal communicatieknelpunt. In een federale setting moet elke worker de volledige waarschijnlijkheidsverdeling (over een vocabulaire van 32k+ tokens) voor elke draft-token naar de server transmitten voor aggregatie en verificatie. Dit uplink-verkeer domineert de end-to-end latentie, waardoor de throughput-winsten door parallelisatie worden tenietgedaan.

Bestaande oplossingen zoals kwantisatie of split-verificatie adresseren de unieke vereisten van federale inferentie niet adequaat, waarbij heterogene modellen gezamenlijk tokens moeten genereren via gewogen aggregatie die een nauwkeurige verdelingsalignatie vereist.

2. Methodologie

De auteurs stellen SpecFed voor, een framework dat speculative decoding combineert met een Top-K Gecomprimeerd Transmissieschema.

A. Systeemmodel

• Architectuur: Een centrale server coördineert $M$ gedistribueerde workers, elk met een lokale LLM. Een small language model (SLM) bevindt zich op de server.
• Werkwijze:
  1. Drafting: De SLM van de server genereert autoregressief een sequentie van $\gamma$ draft-tokens.
  2. Broadcast: De draft-sequentie wordt naar alle workers verzonden.
  3. Parallel Decoding: Workers berekenen conditionele waarschijnlijkheidsverdelingen voor de draft-tokens parallel.
  4. Transmissie & Aggregatie: In plaats van volledige verdelingen te sturen, transmitten workers alleen de Top-K tokens en hun waarschijnlijkheden. De server reconstrueert deze verdelingen en aggregeert ze met behulp van gewogen gemiddelden.
  5. Verificatie: De server verifieert de draft-tokens tegen de geaggregeerde verdeling, accepteert of herbemonstert op basis van acceptatiecriteria.

B. Compressie- en Reconstructiestrategieën

Om de communicatiekosten te mitigeren, introduceert het paper Top-K Truncatie, waarbij worker $i$ alleen de top-$k_i$ tokens met de hoogste waarschijnlijkheden transmittet. Om de ontbrekende waarschijnlijkheidsmassa te behandelen, worden twee reconstructiestrategieën aan de kant van de server voorgesteld:

1. Gerenormaliseerde Top-K Reconstructie: De server gerenormaliseert de ontvangen Top-K waarschijnlijkheden zodat ze optellen tot 1, en wijst een waarschijnlijkheid van nul toe aan alle andere tokens. Dit behoudt de relatieve waarschijnlijkheden tussen Top-K tokens, maar verwijdert de rest.
2. Residuele-Massa Herverdeling: De server behoudt de oorspronkelijke waarschijnlijkheden voor de Top-K tokens en herverdeelt de resterende waarschijnlijkheidsmassa ($\epsilon$) uniform over de resterende vocabulaire-tokens. Dit behoudt een geldige verdeling over het volledige vocabulaire.

3. Belangrijkste Bijdragen

• Nieuw Transmissieschema: De eerste toepassing van Top-K-compressie specifiek afgestemd op federale speculative decoding, die het communicatieknelpunt adresseert zonder het mechanisme voor gewogen aggregatie op te offeren.
• Theoretische Analyse: De auteurs bieden rigoureuze theoretische grenzen voor de robuustheid van hun methode:
  • Lokale Reconstructiefout: Bewezen dat de $\ell_1$-fout begrensd is door $2\epsilon$ (waarbij $\epsilon$ de residuele massa is).
  • Aggregatie-Bias: Afgeleid een bovengrens voor de bias die wordt geïntroduceerd bij het aggregeren van gereconstrueerde verdelingen van meerdere workers.
  • Variatie in Acceptatiepercentage: Bewezen dat de variatie in het acceptatiepercentage voor speculative decoding begrensd is door de aggregatie-bias, waardoor decoding-efficiëntie stabiel blijft zelfs onder compressie.
• Twee Reconstructiestrategieën: Twee onderscheiden methoden voorgesteld en geanalyseerd voor het behandelen van afgeknipte waarschijnlijkheidsmassa, waarbij een afweging wordt geboden tussen eenvoud (Gerenormalisatie) en verdelingsgetrouwheid (Residuele-Massa Herverdeling).

4. Experimentele Resultaten

De auteurs evalueerden het schema met LLaMA-7B en LLaMA-13B als workers en LLaMA-68M als server-SLM op het WMT14 En-De dataset.

• Aggregatie-Bias: Experimenten bevestigden dat de gemiddelde aggregatie-bias ($\bar{\Delta}$) strikt begrensd is door $2\bar{\epsilon}$. De strategie Residuele-Massa Herverdeling leverde consequent een lagere bias op vergeleken met de Gerenormaliseerde aanpak.
• Acceptatiepercentage: De variatie in acceptatiepercentage ($\Delta\bar{\alpha}$) bleek begrensd te zijn door de residuele massa $\bar{\epsilon}$.
• Compressie-efficiëntie: Zelfs onder agressieve compressie (bijvoorbeeld het transmitten van slechts 1% van het vocabulaire of $K=320$ van de 32.000 tokens), behield het systeem een lage bias en hoge generatiegetrouwheid over verschillende temperatuurinstellingen ($T=0.8, 1.0, 1.2$).
• Afweging: De resultaten tonen een gunstige afweging, waarbij het bandbreedtegebruik drastisch wordt gereduceerd (van honderden kilobits per token naar een fractie) terwijl de kwaliteit van de federale inferentie behouden blijft.

5. Betekenis

Dit werk biedt een praktische oplossing voor de schaalbare edge-implementatie van LLM's. Door het communicatieknelpunt inherent aan federale speculative decoding op te lossen, maakt SpecFed mogelijk:

• Hoge Throughput: Significante reductie in end-to-end latentie door uplink-verkeer te minimaliseren.
• Resource-efficiëntie: Maakt federale inferentie haalbaar op apparaten met beperkte bandbreedte of verbindingen met hoge latentie.
• Theoretische Garanties: Biedt wiskundige zekerheid dat compressie het speculative decoding-proces niet destabiliseert of de uiteindelijke generatiekwaliteit significant verslechtert.

Kortom, SpecFed overbrugt de kloof tussen de computationele efficiëntie van speculative decoding en de communicatiebeperkingen van federatief leren, waardoor gedistribueerde LLM-inferentie haalbaarder wordt voor real-world edge computing-scenario's.