Reinforced Fast Weights with Next-Sequence Prediction

Titel: Reinforced Fast Weights with Next-Sequence Prediction (REFINE)

Auteurs: Hee Seung Hwang, Xindi Wu, Sanghyuk Chun, Olga Russakovsky (Princeton University)

---

1. Het Probleem

Lang-context modellering is essentieel voor moderne Large Language Models (LLM's), maar traditionele attention-based transformers hebben een fundamentele beperking: hun reken- en geheugenkosten schalen kwadratisch met de contextlengte. Fast Weight-architecturen (zoals DeltaNet en LaCT) bieden een oplossing door globale attention te vervangen door een vaste, dynamisch bijgewerkte geheugenmatrix. Dit zorgt voor constante geheugenoverhead, ongeacht de contextlengte.

Echter, deze modellen worden momenteel getraind met het standaard Next-Token Prediction (NTP) doelwit. Het paper identificeert twee kritieke tekortkomingen van NTP voor fast weight-modellen:

1. Gebrek aan semantische coherentie: NTP optimaliseert alleen de voorspelling van het volgende token en negeert de kwaliteit van de daaropvolgende reeks tokens.
2. Suboptimale representaties: Omdat fast weights hun parameters dynamisch aanpassen om context op te slaan, leidt token-level feedback tot updates die alleen korte-termijn waarschijnlijkheid maximaliseren. Dit beperkt het vermogen van het model om lange-afhang afhankelijkheden effectief te leren en op te slaan.

2. Methodologie: REFINE

De auteurs introduceren REFINE (Reinforced Fast weIghts with Next sEquence prediction), een Reinforcement Learning (RL) framework dat het trainingsdoel verschuift van NTP naar Next-Sequence Prediction (NSP). In plaats van één token te voorspellen, leert het model een coherent reeks van toekomstige tokens te voorspellen.

Het framework bestaat uit vier kernstappen:

1. Entropy-Based Token Selection:

   • Niet elke positie in de sequentie is even informatief. REFINE berekent de entropie (onzekerheid) van de token-voorspellingen.
   • De inputsequentie wordt opgedeeld in chunks. Uit elke chunk wordt een tokenpositie geselecteerd met een kans evenredig aan de entropie. Dit zorgt ervoor dat het model zich richt op moeilijke, onzekere gebieden in de context.

2. Rollout Generation:

   • Voor elke geselecteerde positie genereert het model een "rollout": een voorspelling van een reeks $k$ tokens (een continuatie) gebaseerd op de huidige prefix.
   • Het model extrahert ook de hidden states van de ground-truth continuatie.

3. Reward Assignment:

   • In plaats van een harde cross-entropy loss, wordt een sequence-level reward toegewezen.
   • Cosine Similarity Reward ($R_\phi$): Voor mid-training en post-training wordt de reward berekend op basis van de cosine-ähnelijkheid tussen de hidden states van de voorspelde reeks en de ground-truth reeks. Dit beloont semantisch vergelijkbare antwoorden, zelfs als ze niet exact letterlijk overeenkomen.
   • Binary Exact Match Reward ($R_{binary}$): Voor Test-Time Training (TTT) wordt een striktere exacte match gebruikt om contextuele memorisatie te maximaliseren.
   • Een hybride reward wordt gebruikt voor post-training om zowel generalisatie als memorisatie te balanceren.

4. Optimalisatie met RL (GRPO):

   • Het model wordt geoptimaliseerd met Group Relative Policy Optimization (GRPO).
   • De totale loss is een gewogen som van de NSP-loss (via RL) en de standaard NTP-loss (SFT) om te voorkomen dat het model vergeet wat het al wist (catastrophic forgetting).

Toepassingsfasen: REFINE is fase-onafhankelijk en kan worden toegepast tijdens:

• Mid-training: Aanpassing aan lange contexten op pre-training-achtige data.
• Post-training: Verfijning van fast weights voor specifieke taken (bijv. instructie-volgen).
• Test-Time Training (TTT): Directe aanpassing van de weights tijdens inferentie op de prompt zelf, zonder extra labels.

3. Belangrijkste Bijdragen

1. Introductie van NSP: Het introduceren van Next-Sequence Prediction als een doelwit specifiek voor fast weight-modellen om de beperkingen van token-level feedback te overwinnen.
2. RL Framework voor Fast Weights: Een nieuw framework dat entropy-based selectie en sequence-level rewards combineert om NSP te optimaliseren, wat eerder niet was onderzocht voor deze architectuur.
3. Levenscyclus-integratie: Het aantonen dat REFINE effectief is in alle fasen van het trainingsproces (mid-, post- en test-time), wat een flexibele aanpak biedt voor lange-context modellering.

4. Resultaten

De auteurs testen REFINE op twee modellen: LaCT-760M en DeltaNet-1.3B. De resultaten tonen een consistente verbetering ten opzichte van standaard Supervised Fine-Tuning (SFT) met NTP:

• Mid-training:
  • Op de RULER-benchmark (Needle-in-a-Haystack taken) verbeterde REFINE de prestaties aanzienlijk. Voor DeltaNet-1.3B was er een stijging van +23,5% in Multi-key NIAH ten opzichte van geen mid-training, en +8,8% ten opzichte van SFT mid-training.
  • Opvallend: REFINE verbeterde zelfs de NTP-accuraatheid op het Booksum-dataset, terwijl SFT daar stagneerde. Dit suggereert dat NSP unieke leersignalen biedt.
• Post-training:
  • Bij taken zoals SQuADQA en HotpotQA (multi-document QA) overtrof "Nested REFINE" (waarbij REFINE wordt toegepast op de prompt voor het genereren van het antwoord) zowel standaard SFT als Nested SFT.
  • Voor LaCT-760M was de verbetering op SQuADQA 17% hoger dan bij Nested SFT.
• Test-Time Training (TTT):
  • REFINE toonde superieure prestaties op de LongBench-benchmark (12 taken) vergeleken met SFT TTT.
  • Dit bevestigt dat NSP sterkere adaptatie-signalen biedt om contextuele informatie in fast weights te comprimeren en op te slaan.

5. Betekenis en Conclusie

Dit paper is significant omdat het een fundamenteel probleem in de training van efficiënte, lange-context modellen aanpakt. Het toont aan dat het standaard NTP-doelwit suboptimaal is voor architecturen die dynamisch geheugen gebruiken.

Door Reinforcement Learning toe te passen op Next-Sequence Prediction, slaagt REFINE erin om fast weight-modellen te trainen die beter in staat zijn om semantische coherentie over lange reeksen te behouden. De methode is niet alleen effectief voor het verbeteren van de initiële gewichten (mid-training), maar ook voor het verfijnen van het model voor specifieke taken en voor adaptatie tijdens het gebruik (test-time).

De auteurs concluderen dat REFINE een veelzijdig en effectief kader biedt om de lange-context modellering van fast weight-architecturen te verbeteren, zonder de noodzaak van kwadratische schaling in geheugen of rekenkracht. De code is openbaar beschikbaar gemaakt.