Technische Samenvatting: Conditioneel Geheugen via Schaalbare Opzoekfuncties (Engram)

1. Probleemstelling

Huidige Large Language Models (LLMs), met name die gebruikmaken van Mixture-of-Experts (MoE) architecturen, vertrouwen op conditionele berekening om capaciteit te schalen. Echter, standaard Transformers missen een natuurlijke primitief voor kennisopzoeking (knowledge lookup). Bijgevolg moeten modellen het ophalen inefficiënt simuleren via dynamische berekening.

De auteurs stellen dat linguïstische signalen intrinsieke heterogeniteit vertonen:

• Compositionele Redenering: Vereist diepe, dynamische berekening.
• Kennisretrieven: Een aanzienlijk deel van de tekst (bijv. benoemde entiteiten, formuleuze patronen) is lokaal, statisch en sterk gestypeerd.

Standaard Transformers dwingen het model om meerdere vroege lagen van aandacht en feed-forward netwerken (FFN's) te consumeren om statische opzoekfuncties te reconstrueren (bijv. het oplossen van "Diana, Princess of Wales") die theoretisch via een eenvoudige tabelopzoekfunctie opgehaald zouden kunnen worden. Dit verspilt waardevolle sequentiële diepte aan triviale operaties, wat voorkomt dat het model middelen kan toewijzen aan hogere niveau redeneringen.

2. Methodologie: De Engram-module

Om dit aan te pakken, introduceert het artikel Conditioneel Geheugen als een complementaire schaarste-as (sparsity axis) voor conditionele berekening. Dit wordt geïnstantieerd via Engram, een module die klassieke $N$-gram embeddings moderniseert om $O(1)$ constant-tijd opzoekfuncties mogelijk te maken.

2.1. Architectuur

Engram breidt de Transformer-backbone uit door de opslag van statische patronen structureel te scheiden van dynamische berekening. Het werkt in twee fasen:

1. Sparse Retrieval via Gehashte $N$-grams:
   • Tokenizer Compressie: Een surjectieve functie brengt ruwe token-ID's naar canonieke identificatoren (normalisatie van hoofdletters, interpunctie, etc.) om de semantische dichtheid te maximaliseren en de vocabulaire te verkleinen.
   • Multi-Head Hashing: Om de combinatorische ruimte van $N$-grams aan te pakken, gebruikt de module $K$ verschillende hash-heads per $N$-gram orde. Deze mappen gecomprimeerde contexten naar indices in embedding-tabellen via deterministische functies, waarbij statische vectoren worden opgehaald.
   • Concatenatie: De opgehaalde embeddings van verschillende $N$-gram orden en hash-heads worden geconcateneerd om de geheugenvector te vormen.
2. Context-Bewuste Gating & Fusie:
   • Omdat statische geheugens geen context hebben, moduleert een gating-mechanisme de opgehaalde geheugenvector met behulp van de huidige verborgen toestand (geaggregeerde globale context) als Query.
   • Een scalaire gate $\alpha_t$ wordt berekend via een dot product tussen de genormaliseerde verborgen toestand en de geheugen Key. Als het geheugen in strijd is met de context, onderdrukt de gate het geheugen (ruisonderdrukking).
   • Een lichtgewicht depthwise causal convolution breidt het receptieve veld uit voordat het gated geheugen via een residuele verbinding in de backbone wordt gefuseerd.

2.2. Systeemefficiëntie & Infrastructuur Co-design

In tegenstelling tot MoE, dat dynamische routing en communicatie van verborgen toestanden vereist, gebruikt Engram deterministische adressering gebaseerd op enkel de input token-ID's.

• Prefetching: Indices zijn bekend vóór de forward pass, wat het asynchroon prefetchen van embeddings vanuit het hostgeheugen (CPU/DRAM) over PCIe mogelijk maakt.
• Overlap: Dit maakt het mogelijk om communicatie te overlappen met de berekening van voorafgaande Transformer-blokken, waardoor de overhead verwaarloosbaar blijft.
• Offloading: Massieve embedding-tabellen (bijv. 100B parameters) kunnen naar het hostgeheugen worden verplaatst (offloaded), waardoor de beperkingen van GPU-geheugen worden omzeild.

2.3. Sparsity Allocatie

De auteurs formuleren het Sparsity Allocation probleem: gegeven een vast parameterbudget, hoe moet de capaciteit worden verdeeld tussen MoE-experts (dynamische berekening) en Engram-geheugen (statische opzoekfunctie)?

• Ze definiëren een allocatieratio $\rho$ die het fractie van de inactieve parameters vertegenwoordigt die aan MoE wordt toegewezen.
• Experimenten onthullen een U-vormige schaalwet: Pure MoE ($\rho=100\%$) en pure Engram ($\rho=0\%$) zijn suboptimaal. De beste prestaties worden bereikt met een hybride allocatie (bijv. $\rho \approx 75\%–80\%$), waarbij ongeveer 20–25% van het schaarse budget wordt geheralloceerd naar Engram.

3. Belangrijkste Resultaten

3.1. Pre-training Prestaties

De auteurs trainden vier modellen op 262B tokens, waarbij ze strikt dezelfde geactiveerde parameters (3.8B) en trainings-FLOPs gebruikten:

• Dense-4B: Baseline.
• MoE-27B: Standaard MoE baseline.
• Engram-27B: Iso-parameter aan MoE-27B, met 5.7B parameters geheralloceerd naar Engram (reductie van experts van 72 naar 55).
• Engram-40B: Schaalt Engram-geheugen naar 18.5B parameters (Totaal 39.5B).

Bevindingen:

• Superieur aan MoE: Engram-27B presteerde beter dan de iso-parameter MoE-27B baseline over diverse domeinen.
• Winst in Redenering: In tegen wordt van de intuïtie dat geheugen alleen helpt bij kennisretrieven, toonde Engram significante winsten in algemene redenering (BBH +5.0, ARC-Challenge +3.7) en code/wiskunde (HumanEval +3.0, MATH +2.4).
• Winst in Kennis: Verbeterde prestaties op kennisintensieve taken (MMLU +3.0, CMMLU +4.0).
• Schaalbaarheid: Engram-40B reduceerde de loss verder, wat suggereert dat de geheugen-as effectief schaalt zonder extra rekenkracht.

3.2. Long-Context Capaciteiten

Door lokale afhankelijkheden te verplaatsen naar opzoekfuncties, bevrijdt Engram de aandachtscapaciteit voor de globale context.

• RULER Benchmark: Engram-27B behaalde 97.0% nauwkeurigheid op Multi-Query Needle-in-a-Haystack (NIAH) vergeleken met 84.2% voor de MoE-baseline.
• Efficiëntie: Engram-27B behaalde de LongPPL-prestaties van de baseline met slechts 82% van de pre-training FLOPs.

3.3. Mechanistische Analyse

• Effectieve Diepte: Gebruikmakend van LogitLens en CKA (Centered Kernel Alignment), vonden de auteurs dat de vroege lagen van Engram representaties produceren die semantisch equivalent zijn aan veel diepere lagen in de MoE-baseline. Dit suggereert dat Engram het netwerk "verdiept" door de noodzaak voor vroege lagen om statische kenmerken te reconstrueren te omzeilen.
• Functionele Dichotomie: Ablatie-studies toonden aan dat benchmarks voor feitelijke kennis instortten (slechts 29–44% prestatie behoudend) wanneer Engram werd uitgeschakeld, terwijl leesvaardigheid robuust bleef (81–93%), wat bevestigt dat Engram fungeert als de primaire opslagplaats voor parametrische kennis.

3.4. Systeemefficiëntie

• Doorvoersnelheid (Throughput): Het offloaden van een 100B-parameter Engram-tabel naar het hostgeheugen resulteerde in een verwaarloosbaar straf op de doorvoersnelheid van <3% (2.8% op een 8B backbone), wat de prefetch-en-overlap strategie valideert.

4. Betekenis en Claims

Het artikel claimt dat Conditioneel Geheugen een onmisbare modelleringsprimitief is voor de volgende generatie schaarse modellen. De betekenis ligt in drie hoofdzaken:

1. Structurele Complementariteit: Het stelt vast dat statische opzoekfuncties en dynamische berekening niet wederzijds uitsluitend zijn, maar structureel complementair. Het optimaliseren van de allocatie tussen beide levert een U-vormige schaalwet op die de pure MoE overtreft.
2. Efficiëntie via Ontkoppeling: Door parameteropslag te ontkoppelen van berekening via deterministische adressering, stelt Engram modellen in staat om naar enorme parameteraantallen te schalen (bijv. 40B+) zonder beperkt te worden door GPU-geheugen of rekenkosten, waardoor hardware-bottlenecks effectief worden omzeild.
3. Verdiepen van het Netwerk: De mechanistische analyse suggereert dat het bieden van expliciete kennisopzoekfuncties de vroege lagen van de backbone ontlast van statische reconstructietaken. Dit vergroot effectief de beschikbare diepte van het netwerk voor complexe redeneringen, wat de onverwachte winsten in redenerings- en code-taken verklaart.

De auteurs positioneren Engram niet louter als een vorm van retrieval augmentation, maar als een fundamentele architecturale verschuiving die $N$-gram concepten moderniseert voor het tijdperk van grootschalige schaarse modellen, en daarmee een nieuwe as van schaarste biedt die de heersende MoE-paradigma aanvult.