Je had transformers uit je geheugen kunnen uitvinden

Grote taalmodellen (LLM's) kunnen tegenwoordig obscure feiten uit alle hoeken van de wereld ophalen.onlangs was ik te laat voor een spelletjesavond met vrienden, dus vroeg ik Astra me de regels van Dune: Imperium te herinneren. Het model vertelde ze me ter plekke terwijl ik onderweg was, waardoor ik me genoeg herinnerde om weer mee te kunnen spelen.

We zijn zo gewend om modellen alles te vragen dat we hun enorme kennis voor lief nemen. Maar hoe onthouden modellen zoveel informatie? En hoeveel kunnen ze precies opslaan?

In dit artikel bespreek ik een intuïtieve manier om lineaire afbeeldingen, MLP's (Multi-Layer Perceptrons) en self-attention te beschouwen als verschillende soorten geheugenopslag. Vanuit deze intuïtie kunnen we heuristieken afleiden voor de hoeveelheid geheugen die elke laag kan opslaan.

De belangrijkste eigenschap van een "goede" geheugenopslag is het vermogen om opgeslagen informatie correct terug te halen. Zelfs als we oneindig veel feiten zouden kunnen opslaan, is dat niet nuttig als we ze niet kunnen ophalen wanneer dat nodig is.

Hoe meer herinneringen we opslaan, hoe meer ze elkaar storen (interferentie) wanneer we proberen er één op te halen. Daarom focussen we in dit artikel op één grootheid: de interferentie-signaalverhouding, die meet in hoeverre andere herinneringen in de weg zitten tijdens het ophalen. Het aantal herinneringen dat een opslag kan bevatten voordat de interferentie even groot wordt als het signaal, noemen we de capaciteit.

We zullen zien hoe je van een lineaire laag naar een MLP gaat, wat zowel de capaciteit als het ophalen verbetert, en uiteindelijk naar self-attention, die een zeer nauwkeurige recall heeft en een potentieel onbeperkt aantal herinneringen kan opslaan. Bij 64-dimensionale sleutels is dat een sprong van ongeveer 64 herinneringen, naar ongeveer 2.700, tot over een miljoen.

Definities

• Geheugenopslag: Een functie $f$ die $M$ herinneringen bevat, voor $i = 1, \ldots, M$.
• Herinnering: Een sleutel-waarde-paar: een sleutel $ki \in \mathbb{R}^{dk}$ (gebruikt voor opzoeking op basis van gelijkenis) en een waarde $vi \in \mathbb{R}^{dv}$ (hetgeen we willen ophalen).
• Query: Een input $x$ waarmee we de functie $f$ aanroepen om een output $f(x)$ te verkrijgen.
• Associatief geheugen: Een geheugenopslag die informatie ophaalt op basis van gelijkenis.

---

Een geheugenopslag op basis van buitenproducten

De meest basale vorm van geheugenopslag is lineair: het accumuleert herinneringen simpelweg via buitenproducten:

$$W = \sum{i=1}^{M} vi k_i^\top = V^\top K$$

waarbij de rijen van $K \in \mathbb{R}^{M\times dk}$ en $V \in \mathbb{R}^{M\times dv}$ de sleutels $ki^\top$ en de waarden $vi^\top$ zijn. Om een nieuwe herinnering $(k, v)$ op te slaan, voegen we simpelweg $vk^\top$ toe aan de bestaande gewichten.

Terzijde: een gradient descent-stap schrijft een herinnering

Neem een lineaire laag $y = Wx$ in een netwerk met een scalaire loss $L$, en laat $\bar y = \nabla L(y)$ de gradiënt van de loss zijn ten opzichte van de output van de laag. De gradiënt ten opzichte van $W$ is $\bar y x^\top$, dus een stap met leersnelheid $\eta$ is:

$$W \leftarrow W - \eta\, \bar y x^\top.$$

Deze recursie is een geheugenupdate: de gewichten worden bijgewerkt met een nieuwe herinnering met sleutel $x$ en waarde $-\eta\bar y$.

Om de geheugenopslag te bevragen met een input $x$, vermenigvuldigen we de input met de matrix:

$$f(x) = V^\top K x = \sum{i=1}^{M} vi\,(k_i^\top x).$$

Het ophalen is dus een opzoeking op basis van gelijkenis met de opgeslagen sleutels: we vergelijken de query met elke sleutel en retourneren de som van de waarden, gewogen naar die gelijkenissen.

Interferentie bij het ophalen

Het opslaan van herinneringen is eenvoudig: blijf buitenproducten toevoegen aan de gewichten. Maar wat gebeurt er als we er één proberen op te halen?

Stel voor de eenvoud dat elke sleutel is genormaliseerd naar lengte 1 ($kj^\top kj = 1$). Dan geldt:

$$f(kj) = \underbrace{vj}{\text{signaal}} + \underbrace{\sum{i\neq j} vi\,(ki^\top kj)}{\text{interferentie}},$$

Dit betekent dat we de waarde van de doelherinnering terugkrijgen (het signaal), plus interferentie van elke andere herinnering waarvan de sleutel gelijkaardig is aan de query.

Als de sleutels allemaal orthonormaal waren, zou de interferentie nul zijn en zou het ophalen exact zijn. Maar $\mathbb{R}^{dk}$ heeft slechts $dk$ orthogonale richtingen, dus we kunnen geen perfecte recall krijgen zodra we meer herinneringen opslaan dan de dimensie van de sleutel: $M > d_k$.

Wat is de capaciteit van de buitenproduct-opslag? Er zijn $M-1$ interferentietermen. Als de sleutels willekeurige eenheidsvectoren zijn, is elke gelijkenis $ki^\top kj$ ongeveer $\pm 1/\sqrt{dk}$. Als ook de waarden willekeurig zijn met vergelijkbare lengtes, tellen de $M-1$ interferentietermen op als een random walk. De gekwadrateerde lengte van de interferentie is dan ongeveer $(M-1)/dk$ keer die van het signaal:

$$\frac{\|\text{interferentie}\|}{\|\text{signaal}\|} \sim \sqrt{\frac{M}{d_k}}.$$

Wanneer we deze verhouding op 1 zetten, vinden we een capaciteit van: $$M \approx d_k.$$ Zodra we evenveel herinneringen opslaan als de dimensie van de sleutel, is de interferentie even groot als het signaal en wordt het signaal overstemd.

Hoe verminderen we interferentie?

Zelfs onder de capaciteit van $d_k$ produceert de buitenproduct-opslag aanzienlijke interferentie. Om de capaciteit te verhogen of cleaner op te halen, moeten we de interferentie verminderen.

De brute-force methode is het vergroten van de sleuteldimensie $d_k$, maar dit wordt beperkt door hardware. Alternatief kunnen we algoritmische verbeteringen aanbrengen:

1. Zorgvuldiger opslaan: De buitenproduct-opslag voegt elke herinnering toe zonder naar de anderen te kijken.
2. Selectiever ophalen: Elke herinnering in een lineaire opslag draagt bij aan interferentie, zelfs bij lage gelijkenis. We kunnen dit verbeteren door minder gelijkenis te onderdrukken.

De optimale lineaire geheugenopslag

In plaats van simpelweg buitenproducten op te tellen, kunnen we de gewichten $W$ zo kiezen dat de recall alle opgeslagen herinneringen tegelijkertijd past. Dit is een lineair regressieprobleem:

$$\text{minimaliseer}W \sum{i=1}^M \|W ki - vi\|_2^2.$$

Wanneer $M > d_k$, is de least-squares oplossing:

$$W = \left(\sumi vi ki^\top\right)\left(\sumi ki ki^\top\right)^{-1} = V^\top K\,(K^\top K)^{-1}.$$

Dit is de buitenproduct-opslag $V^\top K$ met één aanpassing: voordat we de query met de sleutels vergelijken, vermenigvuldigen we deze met $(K^\top K)^{-1}$. Dit heft de covariantie van de sleutels op (whitening), waardoor richtingen die door veel sleutels worden gedeeld, niet langer overgewogen worden.

• Bij $M \leq d_k$: De regressie kan elke herinnering exact passen (mits de sleutels lineair onafhankelijk zijn). Er is geen interferentie.
• Bij $M > d_k$: Er ontstaat interferentie. Voor willekeurige sleutels en waarden wordt de verhouding:

$$\frac{\|\text{interferentie}\|}{\|\text{signaal}\|} \approx \sqrt{\frac{M}{d_k} - 1}.$$

Deze verhouding is 1 bij $M = 2d_k$. De capaciteit van de optimale lineaire opslag is dus twee keer zo groot als die van de buitenproduct-opslag.

MLP als nonlinear geheugenopslag

Een lineaire geheugenopslag haalt informatie op via: $$f(x) = \sum{i=1}^M vi (k_i^\top x).$$

Als we een manier hadden om het gewicht te onderdrukken wanneer $k_i^\top x$ klein is, zouden we de interferentie van sleutels met lage gelijkenis kunnen verminderen. We laten elke gelijkenis door een nonlinear functie $\sigma$ gaan:

$$f(x) = \sum{i=1}^{M} vi\,\sigma(k_i^\top x) = V^\top \sigma(Kx).$$

Hiermee hebben we in feite een twee-laags MLP afgeleid. De rijen van de eerste gewichtenmatrix zijn de sleutels, de kolommen van de tweede matrix zijn de waarden.

Ideaal gezien onderdrukt $\sigma$ termen waarbij de gelijkenis laag of negatief is. Een ReLU-activatie zet sleutels met negatieve gelijkenis op nul. Een gekwadrateerde ReLU ($\text{ReLU}^2$) onderdrukt sleutels met een gelijkenis dicht bij nul nog sterker (een gelijkenis van 0,1 wordt 0,01, terwijl 1 gelijk blijft aan 1).

In moderne taalmodellen worden vaak SwiGLU MLP's gebruikt, waarbij de poort zijn eigen gewichten heeft en een Swish-activatie gebruikt in plaats van ReLU.

Interferentie verminderen met hogere machten

Het kwadrateren van gelijkenissen werkt omdat het lage gewichten veel sterker krimpt dan hoge. We kunnen dit verder doortrekken met $\sigma(s) = \text{ReLU}(s)^p$. De interferentie-signaalverhouding wordt dan:

$$\frac{\|\text{interferentie}\|}{\|\text{signaal}\|} \approx \sqrt{\frac{M}{d_k^p}}.$$

Dit betekent dat de capaciteit ongeveer $M \approx d_k^p$ is, wat een geometrische groei is ten opzichte van de lineaire opslag.

Bij $d_k = 64$ geeft dit de volgende geschatte capaciteiten:

• Lineair: $M \approx 64$.
• ReLU: $M \approx 128$.
• Squared ReLU (GLU MLP): $M \approx 2.700$.
• Cubed ReLU: $M \approx 35.000$.

Self-attention als non-parametrische geheugenopslag

Wat als we deze trend tot het uiterste drijven? In plaats van één macht te kiezen, kunnen we er oneindig veel gebruiken. De exponentiële functie is namelijk een gewogen som van elke macht:

$$e^{\beta s} = 1 + \beta s + \frac{(\beta s)^2}{2} + \frac{(\beta s)^3}{6} + \cdots = \sum_{p=0}^{\infty} \frac{(\beta s)^p}{p!},$$

waarbij $\beta$ (de sharpness) bepaalt hoeveel gewicht er op de hogere machten ligt.

De capaciteit van de exponentiële functie is: $$M \approx e^{2\beta(1 - \beta/d_k)}.$$

Dit groeit exponentieel met $\beta$, in plaats van polynomiaal in $dk$. Met $\beta = \sqrt{dk} = 8$ is de capaciteit ongeveer $e^{14} \approx 1,2$ miljoen herinneringen.

Om numerieke stabiliteit te garanderen (voorkomen van overflow/underflow), delen we elk gewicht door de som van alle gewichten. Dit verandert de interferentie-signaalverhouding niet:

$$f(x) = \sum{i=1}^{M} vi\,\frac{\exp(\beta\,ki^\top x)}{\sumj \exp(\beta\,k_j^\top x)} = V^\top \text{softmax}(\beta K x).$$

Dit is self-attention. In tegenstelling tot een MLP, waarbij de sleutels en waarden getrainde parameters zijn, berekent self-attention deze uit de tokens in de context. Hierdoor is het non-parametrisch: de opslag groeit mee met de contextlengte.

Omgaan met veel sleutel-waarde-paren in lange contexten

Omdat $\beta$ bepaalt hoe sterk we minder gelijkenis onderdrukken, kunnen we $\beta$ verhogen naarmate het aantal paren ($M$) toeneemt. We kunnen $\beta$ echter niet oneindig groot maken, omdat de softmax dan verzadigt en de gradiënten verdwijnen. Daarom wordt $\beta$ vaak geschaald met $\log M$.

Verschillende methoden voor lange contexten gebruiken variaties hiervan:

• Su (2021) / Scalable-Softmax: Vermenigvuldigt de logits met $\log M / \log L$ (waarbij $L$ de trainingslengte is).
• Qwen LogN-Scaling: Past deze factor alleen toe zodra de context langer is dan $L$.
• YaRN (Qwen2.5-1M): Gebruikt een schaling die groeit als $(\log M)^2$.

Door $\beta$ te schalen met $\log M$ kan in principe een onbeperkt aantal herinneringen met verschillende sleutels worden opgeslagen.

Slotbeschouwingen

Hieronder volgt een overzicht van de besproken geheugenopslagen en hun capaciteiten (bij willekeurige sleutels en waarden):

Geheugenopslag $f(x)$	Capaciteit	Bij $d_k = 64$
Buitenproduct $V^\top K x$	$d_k$	64
Optimale lineair $V^\top K (K^\top K)^{-1} x$	$2d_k$	128
ReLU MLP $V^\top \text{ReLU}(Kx)$	$2d_k$	128
Squared ReLU (GLU MLP) $V^\top \text{ReLU}(Kx)^2$	$2d_k^2/3$	$\approx 2.700$
SwiGLU MLP $V^\top (\text{Swish}(Kx) \odot Kx)$	$\approx 2d_k^2/3$	$\approx 2.900$
Cubed ReLU $V^\top \text{ReLU}(Kx)^3$	$2d_k^3/15$	$\approx 35.000$
Self-attention ($\beta = \sqrt{d_k}$) $V^\top \text{softmax}(\beta K x)$	$e^{2\sqrt{d_k} - 2}$	$\approx 1,2$ miljoen
Self-attention ($\beta \propto \log M$) $V^\top \text{softmax}(\beta K x)$	(niet bereikt in exp.)	$> 2$ miljoen

We zijn begonnen bij een eenvoudige heuristiek en zijn zo opgebouwd naar de standaardlagen in moderne transformers. Deze benadering laat ons schatten hoe ze verschillen in hun capaciteit om als vector-waardige sleutel-waarde-opslagen te fungeren.

Natuurlijk is dit slechts één manier om naar architecturen te kijken. Deze schatting telt alleen hoeveel willekeurige vectoren een laag kan ophalen voordat interferentie domineert; het houdt geen rekening met de werkelijke trainingsdata, wat invloed heeft op wat er daadwerkelijk in de sleutels en waarden wordt gerepresenteerd.

Gerelateerde richtingen

Voor wie verder wil lezen, zijn er diverse fascinerende onderwerpen:

• Informatiecapaciteit: Onderzoeken naar hoeveel bits per parameter worden opgeslagen.
• Hopfield-netwerken: Associatief geheugen met exponentiële capaciteit.
• Kortetermijngeheugen: Fast weights, Mamba, TTT-lagen en test-time regressie.
• Sparse parametrische opslag: Product-key memories, Mixture of a Million Experts.
• Non-parametrische opslag: kNN-LM, RETRO en Memorizing Transformers.
• Geheugenbewerking: Het vinden en bewerken van herinneringen in getrainde modellen (bijv. ROME en MEMIT).
• Compressie: KVzip, gist tokens en attention matching.