Technische Samenvatting: Voortdurende LLM Upcycling via Predictor-Gated Bank-Wise Sparsity

Probleemstelling

De computationele kosten van Large Language Models (LLMs), met name decoder-only Transformers, worden steeds meer gedomineerd door de parametergrootte en de per-token berekening vereist door dichte Feed-Forward Networks (FFNs). In moderne SwiGLU-architecturen breidt de FFN de verborgen toestand (hidden state) uit naar een grote intermediaire dimensie, verwerkt deze, en projecteert deze vervolgens terug. Tijdens training en inferentie met lange contexten wordt deze dichte berekening toegepast op alle tokens, ook al is slechts een subset van de intermediaire kanalen cruciaal voor een specifieke token.

Bestaande oplossingen kampen met beperkingen:

• Mixture-of-Experts (MoE): Vermindert de berekening door tokens te routeren naar specifieke expert-pools, maar introduceert aparte parameterpools, complexe eisen voor load-balancing en all-to-all communicatie-overhead.
• Post-hoc Sparsity: Methoden die sparse masks afleiden uit natuurlijke FFN-activaties komen vaak te laat om de dominante gate en up-projectie berekeningen over te slaan, waardoor ze geen hardware-niveau acceleratie bieden.
• Contextuele Sparsity: Hoewel eerder werk aantoont dat LLM-computatie input-afhankelijk is, kunnen training-vrije methoden fragiel zijn bij complexe generatie-taken.

Het artikel adresseert de uitdaging om een vooraf getraind dicht LLM te converteren naar een kanaal-sparse model door middel van voortdurende training (upcycling) zonder significante prestatievermindering, met als doel een hardware-georiënteerd sparse pad te creëren waarbij het routing-signaal beschikbaar is voordat de hoofd-FFN berekening plaatsvindt.

Methodologie

De voorgestelde aanpak, Predictor-Gated Bank-Wise Sparsity, transformeert een dichte Qwen2.5-8B model naar een sparse variant via een specifiek recept voor voortdurende training.

1. Architectuur: Predictor-Gated Sparse SwiGLU

In plaats van de natuurlijke gate-outputs te gebruiken om sparsity te bepalen, introduceren de auteurs een low-rank predictor ($P_l$) die opereert op de token hidden state ($x_{l,t}$) voordat de FFN-berekening plaatsvindt.

• Predictor: Een lichtgewicht twee-laags module ($R^d \to R^{448} \to R^{d_i}$) die de hidden state mapt naar logits over de FFN intermediaire kanalen.
• Routing Mechanisme: De predictor genereert logits die worden gebruikt om een routing-signaal te genereren. Dit vervangt de standaard dichte SwiGLU gate $\sigma(u_{l,t})$.
• Bank-wise Top-k Regel: De intermediaire dimensie ($d_i = 14336$) wordt gepartitioneerd in vaste banks van grootte 64. Voor elke bank wordt een harde top-k regel ($k=16$) toegepast. Dit behoudt 16 actieve kanalen per 64-kanaals bank, wat een 4x sparsity oplevert in de FFN intermediaire activatiebreedte ($d_s = 3584$).
• Executie:
  • Forward Pass: Een hard binair mask wordt toegepast. Inactieve kanalen zijn exact nul.
  • Backward Pass: De niet-differentieerbare top-k selectie wordt afgehandeld via een soft surrogate gradient (sigmoid-achtige lokale gradiënt) rond de selectiedrempel, waardoor de taalmodellering-loss de predictor-parameters kan updaten.

2. Trainingsrecept (Dense-to-Sparse Continual Training)

Het proces converteert een Qwen2.5-8B model getraind op 8K context naar een 32K sparse model via vier onderscheidende stadia:

1. Dense 8K Stage: Train vanaf scratch met de standaard Qwen2.5-8B architectuur.
2. Dense 32K Stage: Resumeer vanaf de 8K checkpoint, breid de context uit naar 32K, en ga door met dichte training voor ongeveer 3T tokens.
3. Predictor Warm-up: Introduceer de low-rank predictor als een auxiliary branch om het natuurlijke gate-signaal te imiteren. Deze branch heeft geen invloed op de hoofd-loss.
4. Predictor-Gated Dense FFN: Routeer de hoofd FFN gate via de predictor (met gebruik van sigmoid), zodat de predictor de loss beïnvloedt, maar zonder nog een harde top-k selectie toe te passen.
5. Predictor-Gated Top-k Sparse FFN: Activeer de bank-wise top-k selectie om de 4x sparsity op te leggen.

Noot: De sparse branch wordt laat in de 32K stage geïntroduceerd (na ~3T tokens) om een praktisch upcycling-scenario te reflecteren waarbij een dicht model naar sparse wordt geconverteerd.

3. Hardware Implicaties

De training wordt uitgevoerd op GPU's als een functionele simulatie. Hoewel de huidige GPU-implementatie dichte gate en up-projecties materialiseert voordat er gemaskeerd wordt (wat directe wall-clock snelheidswinst beperkt), staat het algoritmische principe toe dat toegewijde sparse hardware de berekeningen in alle drie de FFN-stadia (gate, up, en down projecties) kan overslaan voor inactieve kanalen, wat theoretisch de effectieve intermediaire breedte met 4x vermindert.

Belangrijkste Bijdragen

1. Predictor-Gated Channel Sparsity Pattern: Een nieuwe architectuur waarbij een low-rank predictor kanaal-routing logits genereert vóór de FFN-berekening, gecombineerd met een harde bank-wise top-k regel. Dit maakt het overslaan van inactieve gate, up, en down kanalen mogelijk, mits gespecialiseerde kernels of hardware worden gebruikt.
2. Dense-to-Sparse Continual Training Recept: Een praktisch workflow om een Qwen2.5-8B dicht model (getraind op 8K, voortgezet op 32K) te converteren naar een predictor-gated 4x sparse FFN model met beperkte prestatievermindering.
3. Trainingslessen en Foutmodi:
   • Kanaalbalans: Het artikel demonstreert dat MoE-stijl expert-balancing biases niet zuiver overdraagbaar zijn naar bank-wise kanaal-routing. Het expliciet afdwingen van kanaalbalans (via online updates) werd gevonden als nadelig of neutraal vergeleken met geen balans.
   • Layer-Local Long-Context Failure: Evaluatie op RULER-CWE onthulde een scherpe prestatie-cliff (12K–16K) in het sparse model die onzichtbaar was in short-context benchmarks. De auteurs identificeerden dit als een layer-local failure (specifiek in de 7e laag) en stelden een single-layer repair algorithm voor (dense fallback interventie) die de prestaties in het getroffen lengtebereik substantieel herstelt.

Resultaten

Het sparse model werd geëvalueerd tegen twee baselines: een Dense Companion (volledig dicht, gelijke budget) en een Dense-4x baseline (naïve 4x sparsity toegepast post-hoc op basis van activatie-amplitude).

• Prestatienicidade: Over een brede set benchmarks (Factual QA, Code, Reasoning, Math, Composite) bleef het predictor-gated sparse model consistent dichter bij de dense companion dan de naïeve Dense-4x baseline.
• Specifieke Winsten:
  • Het sparse model matchen of overtrof de dense baseline op MBPP, MultiPL-E, en MMLU.
  • Het verminderde substantieel de degradatie gezien in de Dense-4x baseline op Chinese SimpleQA, MBPP+, MBXP, en KOR-Bench.
• Long-Context Prestaties:
  • Op RULER benchmarks (4K, 8K, 16K), presteerde het sparse model vergelijkbaar met het dense model.
  • Bij 32K, presteerde het sparse model (80.4) beter dan de Dense-4x baseline (73.5) en overtrof het licht de dense companion (79.2).
  • RULER-CWE Analyse: Het sparse model vertoonde een "cliff" in recall tussen 12K en 16K. Het toepassen van de voorgestelde single-layer (Layer 7) reparatie verbeterde de gemiddelde recall in het 8.7K–12.2K bereik van 0.856 naar 0.947 en in het 12.8K–16.3K bereik van 0.583 naar 0.802.

Betekenis en Claims

Het artikel claimt dat sparse upcycling niet verliesvrij is, maar een levensvatbaar pad is voor het construeren van hardware-georiënteerde sparse modellen vanuit dense checkpoints.

• Training-Centrische Aanpak: In tegenstelling tot post-hoc compressie behandelt deze methode kanaal-sparsity als een integraal onderdeel van het taalmodellering pad, geoptimaliseerd tijdens voortdurende training.
• Hardware Alignment: Door de routing predictor vóór de FFN-berekening te plaatsen en een harde mask te gebruiken, sluit de methode aan bij de vereisten voor toegewijde sparse hardware om echte acceleratie te bereiken, in plaats van alleen theoretische sparsity.
• Diagnostische Waarde: Het werk benadrukt dat long-context evaluatie cruciaal is voor het identificeren van foutmodi (zoals de layer-local cliff) die brede short-context benchmarks missen. Het biedt ook een negatief resultaat met betrekking tot de overdraagbaarheid van MoE-balancing strategieën naar kanaal-niveau routing in continual training settings.

De auteurs concluderen dat hoewel het sparse model geen perfecte drop-in vervanging is voor het dense model in alle scenario's, het een praktisch recept biedt voor het upcyclen van dense modellen naar efficiënte, sparse varianten met beheersbare prestatie-trade-offs.