Technische Samenvatting: RADLADS

Probleemstelling

Hoewel varianten van lineaire aandacht (zoals RWKV en Mamba) aanzienlijke voordelen bieden ten opzichte van traditionele softmax-aandachtstransformers — specifiek $O(1)$ inferentietijd per token en constant geheugengebruik door het vermijden van Key-Value caches — blijft het vanaf nul trainen van grootschalige lineaire modellen prohibitief duur. State-of-the-art (SoTA) transformer-modellen vereisen vaak training op meer dan tien biljoen tokens, een kostenpost die voor de meeste organisaties ontoegankelijk is. Bestaande methoden voor het converteren van vooraf getrainde softmax-transformers naar lineaire of recurrente modellen (bijv. T2R, SUPRA, LoLCats, MOHAWK) hebben historisch gezien enorme aantallen tokens vereist (variërend van 20 miljard tot 100 miljard+ tokens) of resulteerden in een matige prestatie op downstream-benchmarks, met name op MMLU. Bovendien vertrouwen veel conversiepogingen op hybride architecturen (waarbij enige softmax-aandacht behouden blijft) of slagen ze er niet in om de kwaliteit van de originele docentmodellen te evenaren.

Methodologie

De auteurs presenteren RADLADS (Rapid Attention Distillation to Linear Attention Decoders at Scale), een protocol ontworpen om softmax-aandachtstransformer-modellen te converteren naar lineaire aandacht-decodermodellen met minimale computationele kosten. Het proces omvat drie primaire stadia en specifieke architecturale aanpassingen:

1. Het RADLADS-protocol

De conversie is een driestaps-proces dat slechts 350–700 miljoen tokens vereist (minder dan 0,005% van de pre-training data van de docent):

• Setup (Overdracht van aandachtgewichten): Aandachtsgerelateerde gewichten ($W_q, W_k, W_v, W_o$) van het docentmodel worden direct overgedragen naar de sequence mixing-lagen van het studentmodel. Andere gewichten worden geïnitialiseerd via standaard pre-training methoden of worden ingesteld om het gedrag van de docent na te bootsen zonder onmiddellijk effect.
• Stap 1: Uitlijning van de Attention Hidden State: De sequence mixing-lagen van het studentmodel worden getraind om de hidden state-outputs van de overeenkomstige aandachtlagen van de docent te benaderen. Dit wordt uitgevoerd met een L2-afstand (of MSE) objectief. De auteurs ontdekten dat het gebruik van een Gated Linear Attention kernel (het verwijderen van off-by-one decay en bonustermen) zorgt voor een nauwere aansluiting bij de hidden states van de docent dan standaard RWKV-6.
• Stap 2: Kennisdistillatie: Het volledige studentmodel wordt getraind om de output-logits van de docent te benaderen met behulp van een Kullback-Leibler (KL) divergentie-verlies. De auteurs hypothetiseren dat feitelijke kennis zich voornamelijk in de MLP's en embeddings van de docent bevindt; daarom worden de leersnelheden voor deze componenten laag gehouden of vastgezet om catastrofale vergetelheid te voorkomen, terwijl de sequence mixer agressiever wordt getraind.
• Stap 3: Contextlengte-extensie (Optioneel): Het model wordt gefinetuned op langere sequenties (tot 16k tokens) met behandeling van standaard cross-entropy verlies zonder docentmodel om lang-context capaciteiten te verbeteren. Alternatief stellen de auteurs Stap 2a voor, die de sequentielengte tijdens de distillatiefase zelf uitbreidt naar 4096, wat potentieel de noodzaak voor een aparte Stap 3 wegneemt.

2. Nieuwe Architecturen

De auteurs stelden vast dat standaard RWKV-architecturen niet altijd optimaal waren voor conversie. Ze introduceerden twee nieuwe varianten:

• RAD-RWKV6 ("RADFinch"): Een modificatie van RWKV-6 die gebruikmaakt van een Gated Linear Attention kernel en state balancing technieken om stabiliteit en fit te verbeteren.
• RAD-RWKV7 ("RADGoose"): Een modificatie van RWKV-7 die het "tokenshift"-mechanisme verwijdert (dat geen voordeel bood in deze context) en Rotary Positional Embeddings (RoPE) direct toepast. Deze architectuur vertoonde snellere convergentie en een lager distillatieverlies vergeleken met ongewijzigde RWKV-7.

3. Hyperparameters en Data

• Dataset: De auteurs kozen voor DCLM (DataComp-LM) voor alle conversiestappen, en vonden deze superieur aan FineWeb of FineWeb-Edu voor het converteren van Qwen-modellen.
• Leersnelheden: Er wordt een cosine annealing schema gebruikt in Stap 1, beginnend hoog ($10^{-3}$) om hidden states uit te lijnen en eindigend nabij de finale pre-training leersnelheid van de docent ($10^{-5}$). Stappen 2 en 3 gebruiken een vlakke leersnelheid.

Belangrijkste Bijdragen

1. RADLADS Distillatie Recept: Een gedetailleerd, stapsgewijs protocol inclusief specifieke hyperparameters, tokenaantallen en datasetkeuzes dat een hoogwaardige conversie met minimale data mogelijk maakt.
2. Nieuwe Architecturen: De introductie van RAD-RWKV6 en RAD-RWKV7, die zijn geoptimaliseerd voor het conversieproces, wat snellere inferentie en een betere uitlijning met docentmodellen oplevert dan hun ongewijzigde voorgangers.
3. Grootschalige Modellen: De succesvolle conversie van populaire open-source Qwen2.5 modellen naar lineaire aandacht-varianten op schalen van 7B, 32B en 72B parameters.
4. Open Source Release: De release van code en geconverteerde modellen (QRWKV6/7-7B/32B/72B) onder de Apache 2.0 licentie (met Qwen Licentie beperkingen voor het 72B model), waardoor anderen het proces kunnen repliceren.

Resultaten

De geconverteerde modellen bereiken state-of-the-art prestaties onder pure recurrente modellen van hun omvang:

• Efficiëntie: Het converteren van een 72B model kost minder dan $2.000 USD en vereist slechts ~700M tokens.
• Prestaties: Op standaard benchmarks (Lambada, MMLU, ARC, etc.) presteren de RADLADS-modellen consistent beter dan andere conversiemethoden (bijv. SUPRA, LoLCats, MOHAWK, ARWKV).
  • De QRWKV7-7B-Instruct bereikt een relatieve MMLU-score van 92,4% ten opzichte van zijn docent, wat aanzienlijk beter is dan andere 7B conversies.
  • De QRWKV6-72B-Instruct bereikt een relatieve MMLU-score van 89,9%, wat een nieuwe SoTA vestigt voor pure RNN taalmodellen op die schaal.
• Inferentiesnelheid: Vanwege het lineaire aandachtmechanisme vertonen de geconverteerde modellen significante versnellingen ten opzichte van de docentmodellen naarmate de contextlengte toeneemt. Bijvoorbeeld, bij 8k input tokens en 256 output tokens is de 32B QRWKV7 model 1,61x sneller dan de Qwen3-32B docent; bij 6k input en 2k output bereikt de versnelling 3,41x.

Beperkingen en Claims

Het artikel erkent bescheiden enkele beperkingen:

• Redeneren en Lange Context: Hoewel de prestaties sterk zijn op standaard benchmarks, vertonen de modellen beperkingen in complexe redeneertaken (bijv. Minerva Math) en zeer lange-context taken (bijv. RULER), waarbij ze niet even effectief verbeteren met extra output tokens als de docentmodellen.
• Architecturale Gevoeligheid: Elk nieuw architecturaal ontwerp vereist nauwgezet testen om compatibiliteit met het RADLADS-protocol te garanderen. Bijvoorbeeld, GroupNorm/LayerNorm, gebruikelijk in sommige lineaire aandachtvarianten, veroorzaakte trainingsinstabiliteit op 14B+ schaal en moest worden vervangen door pre-scaling technieken.
• Dataset Uitlijning: De auteurs merken op dat redeneringsdatasets nauwer moeten worden afgestemd op de distributie van het docentmodel om repetitief loopgedrag te voorkomen, een uitdaging die zij gedeeltelijk hebben aangepakt door DCLM te mengen met OpenThoughts.

Betekenis

Het artikel claimt dat RADLADS een kosteneffectief pad biedt om toegang tot grootschalige lineaire aandachtmodellen te democratiseren. Door de vereiste trainingstokens te verminderen van biljoenen naar honderden miljoenen, stelt het onderzoekers en kleinere organisaties in staat om nieuwe varianten van RNN-architecturen op grote schaal te testen, te trainen en te implementeren zonder de extreme kosten die gepaard gaan met pre-training. De auteurs positioneren dit als een hulpmiddel om de ontwikkeling van de volgende generatie compressieve state attention varianten te versnellen.