Een Mathematisch Raamwerk voor Transformer-circuits

Inleiding

Transformer-taalmodellen zijn een opkomende technologie die steeds vaker in de praktijk wordt gebruikt, denk aan systemen zoals GPT-3, LaMDA, Codex, Meena en Gopher. Naarmate deze modellen schalen, zorgt hun open-ended karakter en hoge capaciteit echter voor een toenemende kans op onverwacht en soms schadelijk gedrag. Zelfs jaren na de training ontdekken makers en gebruikersroutinematig modelcapaciteiten — inclusief problematisch gedrag — waar ze voorheen niet van op de hoogte waren.

Een manier om deze problemen aan te pakken is via mechanistische interpreteerbaarheid. Hierbij wordt geprobeerd de gedetailleerde berekeningen die door transformers worden uitgevoerd te reverse-engineeren, vergelijkbaar met hoe een programmeur complexe binaire bestanden probeert om te zetten in leesbare broncode. Als dit mogelijk zou zijn, zou het een systematischere aanpak kunnen bieden voor het verklaren van huidige veiligheidsproblemen, het identificeren van nieuwe problemen en wellicht zelfs het anticiperen op veiligheidsrisico's van krachtigere toekomstige modellen. Een eerder project, de Distill Circuits thread, heeft geprobeerd visiemodellen te reverse-engineeren, maar tot nu toe was er geen vergelijkbaar project voor transformers of taalmodellen.

In dit artikel zetten we de eerste, zeer voorlopige stappen naar het reverse-engineeren van transformers. Gezien de enorme complexiteit en omvang van moderne taalmodellen, is het het meest vruchtbaar om te beginnen met de eenvoudigst mogelijke modellen en van daaruit op te bouwen. Ons doel is om eenvoudige algoritmische patronen, motieven of raamwerken te ontdekken die vervolgens kunnen worden toegepast op grotere en complexere modellen. Specifiek bestuderen we in dit artikel transformers met twee lagen of minder die alleen attention-blokken bevatten — in tegenstelling tot een modern model zoals GPT-3, dat 96 lagen heeft en afwisselt tussen attention-blokken en MLP-blokken.

We stellen vast dat we, door de werking van transformers op een nieuwe maar mathematisch equivalente manier te conceptualiseren, grip krijgen op deze kleine modellen en een aanzienlijk inzicht krijgen in hun interne werking. In het bijzonder ontdekken we dat specifieke attention heads (aandachtskoppen), die we "induction heads" noemen, in-context learning in deze kleine modellen kunnen verklaren. Deze heads ontwikkelen zich alleen in modellen met ten minste twee attention-lagen.

We passen deze inzichten in dit eerste artikel nog niet toe op grotere modellen, maar in een volgend artikel zullen we aantonen dat zowel ons mathematische raamwerk als het concept van induction heads gedeeltelijk relevant blijft voor veel grotere en realistischere modellen — hoewel we nog ver verwijderd zijn van het volledig kunnen reverse-engineeren van dergelijke modellen.

Samenvatting van de Resultaten

Resultaten van de Reverse Engineering

Om de uitdaging van het reverse-engineeren van transformers te verkennen, hebben we verschillende kleine attention-only modellen geanalyseerd. Hierbij hebben we het volgende vastgesteld:

  • Zero-layer transformers modelleren bigram-statistieken. De bigram-tabel kan direct uit de gewichten worden afgeleid.
  • One-layer attention-only transformers zijn een ensemble van bigram- en "skip-trigram"-modellen (sequenties van de vorm "A… B C"). De bigram- en skip-trigram-tabellen kunnen direct uit de gewichten worden afgeleid zonder het model te draaien. Deze skip-trigrammen kunnen verrassend expressief zijn, inclusief het implementeren van een eenvoudige vorm van in-context learning.
  • Two-layer attention-only transformers kunnen veel complexere algoritmen implementeren door compositie van attention heads. Deze compositionele algoritmen kunnen ook direct uit de gewichten worden gedetecteerd. Opvallend is dat modellen met twee lagen attention head-compositie gebruiken om "induction heads" te creëren, een zeer algemeen algoritme voor in-context learning.
  • Verschil in in-context learning: Modellen met één laag en twee lagen gebruiken zeer verschillende algoritmen voor in-context learning. De attention heads van modellen met twee lagen gebruiken kwalitatief meer geavanceerde algoritmen tijdens de inferentie — in het bijzonder de induction head — wat een belangrijk overgangspunt vormt dat relevant zal zijn voor grotere modellen.

Conceptuele Inzichten

We hebben ontdekt dat veel subtiele details van de transformer-architectuur vereisen dat we reverse engineering op een andere manier benaderen dan bij InceptionV1 Circuits. De belangrijkste punten zijn:

  • Onafhankelijkheid van Attention Heads: Attention heads kunnen worden begrepen als onafhankelijke operaties; elk resultaat wordt opgeteld in de residual stream.
  • End-to-end functies: Attention-only modellen kunnen worden geschreven als een som van interpreteerbare functies die tokens mappen naar veranderingen in de logits. Deze functies corresponderen met "paden" door het model.
  • Lineaire structuur: Transformers hebben een enorme hoeveelheid lineaire structuur. Veel kan worden geleerd door sommen op te splitsen en ketens van matrices met elkaar te vermenigvuldigen.
  • QK- en OV-circuits: Attention heads voeren twee grotendeels onafhankelijke berekeningen uit: een QK (Query-Key) circuit, dat het aandachtspatroon berekent, en een OV (Output-Value) circuit, dat berekent hoe elk token de output beïnvloedt als er aandacht aan wordt besteed.
  • Virtual Weights: Key-, query- en value-vectoren kunnen worden gezien als tussenresultaten in de berekening van de matrices $WQ^T WK$ en $WO WV$.
  • Compositie: De compositie van attention heads vergroot de expressiviteit van transformers aanzienlijk. Er zijn drie manieren waarop heads kunnen componeren (via keys, queries of values), waarbij key- en query-compositie sterk verschillen van value-compositie.
  • Residual Stream als Communicatiekanaal: Alle componenten van een transformer (token embedding, attention heads, MLP-lagen en unembedding) communiceren met elkaar door te lezen uit en te schrijven naar verschillende subruimtes van de residual stream.

Transformer Overzicht

Model Simplificaties

Om de ideeën in dit artikel in hun zuiverste vorm te presenteren, richten we ons op "toy transformers" met enkele simplificaties:

  1. Attention-only: We focussen op transformers zonder MLP-lagen. Dit is een drastische simplificatie, maar stelt ons in staat de uitdagingen van attention heads in isolatie te bestuderen.
  2. Geen biases: We houden geen rekening met biases; een model met biases kan worden gesimuleerd door deze in de gewichten te vouwen.
  3. Geen layer normalization: We negeren layer normalization omdat dit veel complexiteit toevoegt, terwijl het (op een schalingsfactor na) vaak kan worden samengevoegd met aangrenzende gewichten.

Hoogwaardige Architectuur

We richten ons op autoregressieve, decoder-only transformer-taalmodellen (zoals GPT-3). Een transformer begint met een token embedding, gevolgd door een reeks "residual blocks", en eindigt met een token unembedding. Elk residu-blok bestaat uit een attention-laag, gevolgd door een MLP-laag. Beide lezen hun input uit de residual stream (via een lineaire projectie) en schrijven hun resultaat terug door dit op te tellen.

De Residual Stream als Communicatiekanaal

De residual stream is de som van de output van alle voorgaande lagen en de oorspronkelijke embedding. We beschouwen dit als een communicatiekanaal omdat het zelf geen verwerking uitvoert, maar alle lagen via dit kanaal communiceren.

Virtual Weights

Een gevolg van de lineaire structuur van de residual stream is dat men kan denken in impliciete "virtuele gewichten" die elk paar lagen direct verbinden. Deze virtuele gewichten zijn het product van de output-gewichten van de ene laag met de input-gewichten van de andere (bijv. $W{I}^2 W{O}^1$).

Subruimtes en Bandbreedte

De residual stream is een hoog-dimensionale vectorruimte. Lagen kunnen verschillende informatie naar verschillende lagen sturen door deze in verschillende subruimtes op te slaan. Omdat de bandbreedte van de residual stream beperkt is in vergelijking met het aantal neuronen in MLP-lagen, ontstaat er een "bottleneck". We vermoeden dat sommige MLP-neuronen en attention heads een rol spelen in "geheugenbeheer" door informatie uit de residual stream te verwijderen.

Attention Heads als Informatiebeweging

Onafhankelijkheid en Additiviteit

We beschouwen attention layers als een verzameling onafhankelijke attention heads $h \in H$ die parallel werken en hun output onafhankelijk toevoegen aan de residual stream. Hoewel implementaties vaak gebruikmaken van concatenatie voor computationele efficiëntie, is dit mathematisch equivalent aan onafhankelijke optelling.

Informatiebeweging

De fundamentele actie van attention heads is het verplaatsen van informatie. Ze lezen informatie uit de residual stream van het ene token en schrijven dit naar de residual stream van een ander token.

Met behulp van tensorproducten kunnen we de werking als volgt beschrijven: $$h(x) = (A \otimes WO WV) \cdot x$$

Waarbij:

  • $A$ (het aandachtspatroon) bepaalt van welk token informatie wordt verplaatst naar welk token.
  • $WO WV$ (de OV-matrix) bepaalt welke informatie wordt gelezen uit het bron-token en hoe dit wordt geschreven naar het doel-token.

Belangrijke Observaties over Attention Heads

  1. Scheiding van operaties: De beweging van informatie (A) en de inhoud van de informatie ($W_{OV}$) zijn onafhankelijk.
  2. Half-lineair: Als het aandachtspatroon $A$ wordt vastgezet, voert de attention head een lineaire operatie uit.
  3. Low-rank matrices: Hoewel ze als aparte matrices ($WQ, WK, WV, WO$) worden geparametriseerd, werken ze altijd samen als low-rank matrices $W{OV}$ en $W{QK}$.
  4. Virtuele Attention Heads: Producten van attention heads gedragen zich als attention heads zelf, wat leidt tot "virtuele" heads.

Zero-Layer Transformers

Een "zero-layer" transformer neemt een token, embedt het en unembedt het direct om logits te produceren: $$T = WU WE$$

Omdat het model geen informatie tussen tokens kan verplaatsen, voorspelt het het volgende token puur op basis van het huidige token. Dit betekent dat de optimale werking van $WU WE$ een benadering is van de bigram-statistieken.

One-Layer Attention-Only Transformers

Een transformer met één laag kan worden begrepen als een ensemble van een bigram-model en verschillende "skip-trigram"-modellen (sequenties van de vorm "A… B C").

De Path Expansion Trick

We kunnen de transformer representeren als een product van termen (embedding → attention → unembedding). Door dit product uit te breiden, transformeren we het in een som van end-to-end paden.

  • Het directe pad: Draagt bij aan de bigram-statistieken.
  • De attention-head termen: Zorgen voor de skip-trigrammen.

QK- en OV-circuits

Elke attention head $h$ kan worden gesplitst in twee circuits:

  1. Het Query-Key (QK) circuit ($WE^T W{QK}^h W_E$): Bepaalt de aandachtsscore voor elk paar van query- en key-tokens. Het bepaalt naar welk "bron-token" het "doel-token" kijkt.
  2. Het Output-Value (OV) circuit ($WU W{OV}^h W_E$): Bepaalt hoe een token de output-logits beïnvloedt als er aandacht aan wordt besteed.

Samen vormen deze drie tokens een skip-trigram: [bron]... [doel][uitkomst].

Kopieergedrag en In-Context Learning

Veel attention heads in modellen met één laag zijn gericht op kopiëren. Het OV-circuit zorgt ervoor dat tokens hun eigen waarschijnlijkheid verhogen, terwijl het QK-circuit zoekt naar tokens die plausibel zijn als volgende token.

Dit is een primitieve vorm van in-context learning. Voorbeelden van complexer gedrag via skip-trigrammen:

  • Python: Voorspellen van else, elif of except na een vermindering van inspringing (indentatie).
  • HTML: Voorspellen van <td> na <tbody>.
  • Algemeen: Matchen van openende en sluitende haken of quotes.
  • LaTeX: Voorspellen van \right na een \left.

Detectie van Kopieergedrag

Kopieergedrag kan worden gedetecteerd via de eigenwaarden van de OV-matrix. Matrices die kopiëren vertonen vaak positieve eigenwaarden, wat aangeeft dat bepaalde sets tokens hun eigen waarschijnlijkheid wederzijds verhogen.

Two-Layer Attention-Only Transformers

Het belangrijkste verschil tussen één en twee lagen is compositie. Diepte stelt het model in staat om algoritmen uit te voeren in plaats van simpelweg tabellen met skip-trigrammen op te zoeken.

Drie Soorten Compositie

  1. Q-Compositie: $W_Q$ leest een subruimte die is beïnvloed door een vorige head.
  2. K-Compositie: $W_K$ leest een subruimte die is beïnvloed door een vorige head.
  3. V-Compositie: $W_V$ leest een subruimte die is beïnvloed door een vorige head.

Q- en K-compositie beïnvloeden het aandachtspatroon, terwijl V-compositie effectief "virtuele attention heads" creëert.

Induction Heads

In kleine modellen met twee lagen wordt compositie voornamelijk gebruikt voor het creëren van induction heads. Dit is een krachtig mechanisme voor in-context learning.

Werking van Induction Heads: Een induction head zoekt in de context naar eerdere voorbeelden van het huidige token. Als het dit vindt, kijkt het naar het token dat volgde op dat eerdere voorbeeld en kopieert dat naar de huidige positie.

Vergelijking:

  • One-layer kopieer-head: [b] … [a] → [b]
  • Two-layer induction head: [a][b] … [a] → [b]

Dit stelt het model in staat om patronen te herhalen, zelfs als het gaat om volledig willekeurige tokens die niet in de trainingsdata voorkwamen.

Mechanistische Theorie van Induction Heads

Voor het functioneren van een induction head zijn twee dingen nodig:

  1. Een kopiërend OV-circuit.
  2. Een QK-circuit dat gebruikmaakt van K-compositie met een "previous token head" om de key-vector één positie te verschuiven.

Conclusie en Vooruitblik

Door kleine attention-only modellen te bestuderen, hebben we een raamwerk ontwikkeld om delen van transformers te reverse-engineeren. Hoewel we nog geen volledige modellen kunnen begrijpen (vooral vanwege de complexiteit van MLP-lagen), zien we dat concepten zoals induction heads en K-compositie ook voorkomen in zeer grote taalmodellen.

Dit suggereert dat mechanistische interpreteerbaarheid een levensvatbare route is naar het begrijpen van de interne werking van AI, mits we voortgang boeken in het analyseren van MLP-lagen en het omgaan met fenomenen zoals superpositie.

---

Technische Bijlagen

MLP Lagen en Superpositie

MLP-lagen zijn moeilijker te analyseren omdat ze niet-lineaire activatiefuncties (zoals GeLU) gebruiken. We vermoeden dat neuronen in MLP-lagen vaak polysemantisch zijn (meerdere functies hebben), wat wijst op het fenomeen van superpositie: het model probeert meer kenmerken te representeren dan er dimensies beschikbaar zijn in de vectorruimte.

Virtuele Gewichten en Convolutie

Wanneer we pad-expansie toepassen, zien we structuren die lijken op convoluties. Een standaard convolutie kan worden geschreven als een som van tensorproducten waarbij de attention heads fungeren als relatieve posities.

Notatieoverzicht

  • $T(t)$: Logits voor tokens $t$.
  • $x^n$: Residual stream op laag $n$.
  • $W_E$: Token embedding matrix.
  • $W_U$: Unembedding matrix.
  • $A^h$: Aandachtspatroon van head $h$.
  • $W{OV}^h$: Product van $WO^h W_V^h$.
  • $W{QK}^h$: Product van $WQ^{hT} W_K^h$.

Omgaan met Low-Rank Matrices

Vanwege de enorme omvang van de matrices (bijv. $50.000 \times 50.0