DeepFusion: Accelerating MoE Training via Federated Knowledge Distillation from Heterogeneous Edge Devices

Probleemstelling

De opkomst van grote taalmodellen (LLMs) op basis van Mixture-of-Experts (MoE)-architecturen (zoals Qwen-MoE en DeepSeek-MoE) heeft de prestaties in natuurlijke taalverwerking (NLP) aanzienlijk verbeterd. Deze modellen vereisen echter enorme hoeveelheden diverse trainingsdata om alle "expert"-subnetwerken effectief te laten leren. Tegelijkertijd dreigt er een tekort aan hoogwaardige publieke tekstdata.

Federated Learning (FL) biedt een oplossing door data van randapparaten (edge devices) te benutten zonder de privacy te schenden. Echter, bestaande FL-methoden voor MoE-training hebben twee fundamentele beperkingen:

1. Resource-beperkingen: Traditionele FL vereist dat elk randapparaat een lokaal kopie van het globale MoE-model (of een gesneden expert-versie daarvan) draait. Voor resource-beperkte apparaten (zoals embedded systemen of oudere smartphones) is dit onmogelijk vanwege de grote modelgrootte en complexiteit.
2. Heterogeniteit en View-Mismatch: Bestaande methoden gaan er vaak van uit dat alle apparaten dezelfde modelarchitectuur hebben. In de praktijk zijn randapparaten echter zeer heterogeen (verschillende hardware, verschillende LLM-architecturen). Wanneer een lokaal model (de "leraar") en het globale MoE-model (de "student") verschillende architecturen hebben, ontstaan er discrepanties in hun latente ruimtes en voorspellingsperspectieven. Dit maakt traditionele kennisdistillatie (die direct op logits of features in een gedeelde ruimte afgestemd is) inefficiënt of onwerkbaar.

Methodologie: DEEPFUSION

DEEPFUSION is een schaalbaar framework dat federated kennisdistillatie gebruikt om heterogene randapparaten toch te laten deelnemen aan MoE-training. Het systeem werkt in drie hoofdfasen:

1. One-Shot Federated Learning & Lokale Kennisclustering:

   • In plaats van meerdere communicatierondes, trainen randapparaten hun eigen, op maat gemaakte LLM's lokaal op privé-data.
   • Zodra de training voltooid is, worden de modellen één keer geüpload naar de centrale server (One-Shot FL), wat communicatiekosten drastisch verlaagt.
   • De server cluster de geüploade modellen in $K$ lokale kennis-domeinen (bijv. medisch, financieel) op basis van semantische embeddings van de trainingsdata. Voor elk domein wordt een proxy-model gegenereerd door de gewichten van de geklusterde modellen te middelen.

2. Cross-Architectuur Kennisdistillatie met VAA:

   • Dit is de kerninnovatie. De proxy-modellen (leraren) hebben vaak een compacte architectuur, terwijl het globale MoE-model (student) complexer is.
   • Om het probleem van view-mismatch (verschil in voorspellingsperspectief) op te lossen, introduceert het paper de View-Aligned Attention (VAA) module.
   • De VAA-module gebruikt een multi-head attention-mechanisme om features van het student-model uit verschillende representatiefasen te integreren. Hierdoor leert het student-model een voorspellingsperspectief aan dat overeenkomt met dat van de leraar, ondanks de architecturale verschillen.
   • De distillatie gebeurt via een combinatie van cross-entropy loss, feature-matching loss (via VAA) en KL-divergentie loss op de logits.

3. Globale MoE Model Tuning:

   • De $K$ gedistilleerde basismodellen worden samengevoegd tot het globale MoE-model.
   • De FFN-lagen (Feed-Forward Networks) van de basismodellen worden direct toegewezen aan de specifieke experts in het MoE-model.
   • De andere lagen (embedding, self-attention, gating) worden gefuseerd door gemiddelde berekening en vervolgens fijn afgestemd (fine-tuning) op publieke benchmark-data, terwijl de experts bevroren blijven.

Belangrijkste Bijdragen

• Eerste Schaalbare Framework: DEEPFUSION is het eerste framework dat heterogene, resource-beperkte randapparaten toelaat om bij te dragen aan MoE-training via kennisdistillatie, zonder dat deze apparaten zware MoE-modellen hoeven te hosten.
• View-Aligned Attention (VAA): Een nieuwe module die het probleem van architecturale heterogeniteit oplost door de voorspellingsperspectieven van leraar en student uit te lijnen, waardoor effectieve cross-architectuur distillatie mogelijk wordt.
• Efficiëntie: Door gebruik te maken van een "one-shot" FL-benadering en het benutten van bestaande lichte LLM's op apparaten, worden communicatiekosten en lokale rekenlast geminimaliseerd.

Resultaten

Het framework is geëvalueerd met industriële MoE-modellen (Qwen-MoE en DeepSeek-MoE) op real-world datasets (medische meerkeuzevragen en financiële open-ended vragen).

• Prestaties: DEEPFUSION bereikt prestaties die zeer dicht in de buurt komen van centraal getrainde MoE-modellen (zoals met DeepSpeed).
• Verbetering t.o.v. Baselines: In vergelijking met state-of-the-art methoden (FedJETS, FedKMT, OFA-KD) toont DEEPFUSION superieure resultaten, met name bij uitdagende open-ended vragen.
  • Perplexiteit: Tot 5,28% verbetering in token-perplexiteit.
  • Communicatiekosten: Tot 71% reductie in communicatiekosten door het vermijden van frequente updates en het gebruik van lichte modellen.
• Hardware Compatibiliteit: In tegenstelling tot FedJETS (dat zware lokale expert-modellen vereist), kan DEEPFUSION draaien op apparaten met beperkt geheugen (bijv. NVIDIA Jetson Nano, iPhone 16), wat de deelname van een veel bredere reeks apparaten mogelijk maakt.

Significantie

DEEPFUSION lost een kritieke knelpunt op in de ontwikkeling van toekomstige AI-modellen: het tekort aan publieke trainingsdata en de onmogelijkheid om deze data te verzamelen via privacy-bevorderende methoden op beperkte hardware. Door heterogene randapparaten effectief te benutten als bron van gespecialiseerde kennis, maakt het framework schaalbare, privacy-bewuste training van enorme MoE-modellen mogelijk. Dit opent de deur voor robuuste, domeinspecifieke AI-modellen die trainen op data die anders nooit beschikbaar zou zijn geweest voor centrale servers.