Technische Samenvatting: MobileMoE: Schalen van Mixture of Experts op het apparaat

Probleemstelling

Hoewel Mixture-of-Experts (MoE)-architecturen de standaard zijn geworden voor het schalen van grote taalmodellen (LLM's) naar honderden miljarden parameters op cloudservers, blijft hun potentieel op schalen onder de miljard voor implementatie op het apparaat grotendeels onontgonnen. Bestaande LLM's voor op het apparaat (bijv. MobileLLM, MobileLLM-Pro) zijn overweldigend dicht, waarbij een verhoogd aantal actieve parameters en inferentieberekening nodig is om de kwaliteit te verbeteren. Dit creëert een knelpunt voor rand-ai (edge AI), waarbij apparaten te maken hebben met strikte beperkingen op het gebied van geheugen (DRAM) en rekenkracht (batterijduur, latentie). Hoewel de capaciteiten van smartphone-DRAM recent zijn toegenomen (bijv. 12–16 GB), is er nog geen gestructureerde methodologie vastgesteld voor het schalen van MoE-architecturen onder deze specifieke mobiele beperkingen—waarbij een evenwicht wordt gezocht tussen de totale modelgrootte, actieve parameters en inferentiekosten.

Methodologie

1. Schalingswet voor MoE op het apparaat

De auteurs formuleren een gegeneraliseerde schalingswet voor MoE op het apparaat om de architectuur gezamenlijk te optimaliseren onder geheugen- en rekenbeperkingen. De verliesfunctie $L$ wordt gemodelleerd als een functie van actieve parameters ($N_{act}$), trainingsdatagrootte ($D$), een monotoone transformatie van het aantal experts ($\hat{E}$), en architecturale keuzes ($x$, inclusief expert-granulariteit $g$ en gedeelde experts $s$).

De optimalisatiedoelstelling minimaliseert het verlies onder de volgende voorwaarden:

• Rekenbeperkingen: Trainings-FLOPs ($F_{train} = 6N_{act}D$) en per-token inferentie-FLOPs ($F_{inf} = 2N_{act}$).
• Geheugenbeperkingen: Totale geheugenvoetafdruk $M(N_{total}, T)$, inclusief gekwantiseerde gewichten ($b_w$-bit) en KV-cache ($b_{kv}$-bit) bij contextlengte $T$.

2. Architecturaal ontwerp via Divide-and-Conquer

Met behulp van de schalingswet identificeren de auteurs een optimaal "sweet spot" voor MoE op het apparaat via drie gecontroleerde ablatiestudies:

• Aantal experts ($E$): Matige sparsiteit is optimaal. Het verhogen van $E$ verlaagt het verlies, maar met afnemende meerwaarde boven $E=8$. $E=8$ wordt geselecteerd als het praktische sweet spot voor geheugengelimiteerde apparaten.
• Expert-granulariteit ($g$): Fijnkorrelige experts (het splitsen van elke expert in $g$ sub-experts) zorgen voor meer diverse routeringspaden. De studie vindt $g=8$ als de rekenoptimaal granulariteit, wat aanzienlijke verliesreductie biedt ten opzichte van grofkorrelige experts zonder excessieve overhead.
• Gedeelde expert ($s$): Het opnemen van een gedeelde, altijd-actieve expert (een dichte pathway) vult de gerouteerde specialisten aan. De configuratie met een gedeelde expert ($s=\checkmark$) bereikt een lager verlies dan de uitsluitend gerouteerde tegenhanger onder vaste rekenkracht.

Resultaat Architectuur: De afgeleide MobileMoE-familie bestaat uit drie schalen (S, M, L) met actieve parameters van 0,3B, 0,5B en 0,9B, en totale parameters van respectievelijk 1,3B, 2,8B en 5,3B. Het ontwerp gebruikt $E=8$, $g=8$ (64 fijnkorrelige gerouteerde experts) en een gedeelde expert met top-4 routering.

3. Vier-staps trainingsrecept

Om de architecturale voordelen te realiseren, wordt een uitgebreide trainingspijplijn toegepast op open-source datasets:

1. Pre-training: ~6T tokens met een web-zware mix om brede linguïstische dekking en expert-specialisatie te waarborgen. Technieken omvatten auxiliair-verliesvrije balansering, router z-loss en gegroepeerde MLP (GMM)-kernels voor efficiëntie.
2. Mid-training: ~500B tokens met focus op hoogwaardige domeinspecifieke data (wiskunde, code, kennis) om de context uit te breiden naar 8k en de expert-specialisatie te scherpen.
3. Instructie-finetuning (SFT): ~80M samples over diverse domeinen met behulp van dropless token-dispatching om gestructureerde instructiesignalen te behouden.
4. Kwantiseringsbewuste training (QAT): INT4-gewichtskwantisering (groepsgrootte 32) met INT8-dynamische activeringen. Cruciaal zijn de routeringsgewichten in FP32 bewaard om routeringsstabiliteit te behouden, wat verwaarloosbare geheugenoverhead toevoegt.

4. Efficiënte inferentie op het apparaat

Om de kloof naar implementatie te overbruggen, implementeren de auteurs een aangepaste gefuseerde MoE-kernel in ExecuTorch. Deze kernel:

• Zet sparse expert-dispatch om in dichte gegroepeerde GEMM's door tokens te herschikken.
• Fuseert alle sub-operaties (routering, dispatch, projecties, activatie, scatter) tot één enkele operator om de overhead van kernelstart en kwantisatie te amortiseren.
• Maakt efficiënte inferentie mogelijk op gangbare smartphone-CPU's (Snapdragon 8 Elite, Apple A18 Pro) en GPU's.

Belangrijkste resultaten

Benchmarkprestaties

Over 14 fundamentele benchmarks (gezond verstand, kennis, wetenschap, lezen, redeneren) vestigt MobileMoE een nieuwe Pareto-grens:

• Efficiëntie vs. Dichte modellen: MobileMoE-S/M komen overeen met of overtreffen toonaangevende dichte LLM's voor op het apparaat (bijv. MobileLLM-Pro, Llama 3.2 1B) met 2–4× minder inferentie-FLOPs bij vergelijkbare geheugenvoetafdrukken.
• Efficiëntie vs. MoE-modellen: MobileMoE-M komt overeen met de nauwkeurigheid van de state-of-the-art OLMoE-1B-7B met ~60% minder actieve en totale parameters. MobileMoE-L overtreft OLMoE-1B-7B met 30% minder actieve parameters en een 23% kleinere geheugenvoetafdruk.
• Token-efficiëntie: MobileMoE-L bereikt superieure prestaties met slechts ~6T pre-training tokens, en presteert beter dan dichte baselines die zijn getraind op 1,5–2× meer tokens (bijv. Llama 3.2 1B op 9T).

Runtime-profileren op het apparaat

Geïmplementeerd op Samsung Galaxy S25 en iPhone 16 Pro met INT4-gewichten:

• Latentie: MobileMoE-S bereikt 1,8–3,8× snellere prefill en 2,2–3,4× snellere decode in vergelijking met de dichte MobileLLM-Pro-baseline bij vergelijkbare INT4-gewichtsgeheugen.
• Geheugen: Bij 8k context verbruikt MobileMoE-S tot 22% minder piek-RAM dan de dichte baseline. MobileMoE-L past binnen 5 GB DRAM bij 8k context, wat de haalbaarheid voor moderne smartphones bevestigt.
• Nauwkeurigheid: MobileMoE-L levert aanzienlijk hogere nauwkeurigheid (Gemiddeld 57,8 vs. 45,5 voor MobileLLM-Pro) met slechts een gematigde runtime-kost, die afneemt naarmate de contextlengte toeneemt.

Betekenis en claims

Het artikel claimt de eerste familie van MoE's met minder dan een miljard actieve parameters te presenteren die specifiek is geoptimaliseerd voor implementatie aan de rand, afgeleid van een nieuwe schalingswet voor op het apparaat. De betekenis ligt in drie hoofdcontributies:

1. Theoretische basis: Het biedt een gestructureerd kader voor MoE-ontwerp onder gezamenlijke geheugen- en rekenbeperkingen, waarbij matige sparsiteit ($E=8$), fijnkorrelige experts ($g=8$) en gedeelde experts worden geïdentificeerd als de optimale configuratie voor mobiele apparaten.
2. Trainings-efficiëntie: Het demonstreert dat MoE-architecturen state-of-the-art prestaties op het apparaat kunnen bereiken met aanzienlijk minder trainings-tokens (~6T) in vergelijking met dichte tegenhangers, wat de leerefficiëntie van MoE op schalen onder de miljard valideert.
3. Praktische implementatie: Het overbrugt de "laatste mijl" naar mobiele implementatie door de eerste efficiënte MoE-inferentie op gangbare smartphones te bieden. De resultaten bevestigen dat MoE niet alleen een cloud-schaal architectuur is, maar een praktische, efficiënte route voor next-generation rand-ai, waardoor low-latency, privacy-bewarende en kosteneffectieve toepassingen op smartphones en wearables mogelijk worden.

De auteurs concluderen dat hoewel bestaande modellen voor op het apparaat overwegend dicht blijven, MobileMoE aantoont dat MoE een superieur alternatief biedt voor de rand, en een nieuwe Pareto-grens vestigt die nauwkeurigheid, geheugen en latentie in evenwicht brengt.