Hoppa till innehållet
9/30Kapitel 9 av 30

Attention och transformer-blocket, härlett från ett genomsnitt

Börja med den billigaste sammanfattningen av ett context — genomsnittet — och låt attention-formeln falla ut ur reparationen.

På den här sidan

Du kommer hit med en tokenizer från kapitel 7, en embedding-tabell från kapitel 8 och målet som hör till dem: givet alla tokens hittills, sätt en sannolikhet på nästa.

Det som saknas är mitten. För att förutsäga token tt behöver modellen en vektor som sammanfattar allt före den, och ingenting du har byggt producerar en sådan. Embedding för token t1t-1 är inte den — det är en bigrammodell, och den kan inte veta att meningen började med en fråga. En hopslagning av alla tidigare embeddings är inte heller svaret: deras antal ändras vid varje steg, och en fast viktmatris kan inte ta en indata med variabel längd.

Alltså: en vektor med fast storlek, som sammanfattar ett variabelt antal vektorer. Det är hela problemet, och attention är vad du får genom att lösa det på latast möjliga sätt och sedan reparera de två saker som går sönder.

Svaret fältet hade, och varför vi inte bygger det

Länk till avsnittet: Svaret fältet hade, och varför vi inte bygger det

Från 1997 till ungefär 2017 var sammanfattningen ett rekurrent tillstånd: behåll en vektor h\mathbf{h} och uppdatera den vid varje token, ht=f(ht1,xt)\mathbf{h}_t = f(\mathbf{h}_{t-1}, \mathbf{x}_t). Fast storlek, variabel indata, exakt rätt form.

Det misslyckades på tre sätt, och det här kapitlets arkitektur svarar på alla tre. Backpropagation genom TT steg multiplicerar TT Jacobianer, så gradienten försvinner eller exploderar — sjukdomen som kapitel 5 mätte inuti en enda tanh\tanh-nod. LSTM1 utformades just för att motverka det och flyttade det användbara intervallet från tiotals steg till hundratals, utan att ändra faktumet att information från token 5 når token 500 bara genom att överleva 495 sekventiella uppdateringar. Hela källan behövde få plats i en vektor: i sekvens-till-sekvens-översättning2 komprimerar en encoder indatan till sitt sluttillstånd. Bahdanau, Cho och Bengio namngav den flaskhalsen och fixade den 2014, tre år före transformer, genom att låta decoder ta en viktad summa av alla encoder-tillstånd med vikter den själv beräknade.3 Allt nedan är den idén, applicerad av en sekvens på sig själv, med rekurrensen borttagen. Och uppdateringen är sekventiell per konstruktion: ht\mathbf{h}_t behöver ht1\mathbf{h}_{t-1}, och en GPU med tiotusen kärnor kan inte göra något med det. Arkitekturen som vann är inte uppenbart smartare; det är den vars dyra steg är en matrismultiplikation.

Den andra klassiska induktiva biasen, konvolution — dra ett litet filter över hela indatan, så att en feature som upptäcks var som helst upptäcks överallt — byggs inte heller här; den är nästan exakt rätt för bilder och delegeras till en kurs i datorseende. Varken rekurrens eller konvolution dyker upp igen efter den här sidan, vilket är varför ingen av dem får ett kapitel: kapitel 1 lovade att utelämnandena skulle deklareras i stället för att ske tyst.

Den mest uppenbara funktionen av ett variabelt antal vektorer som returnerar en vektor är genomsnittet:

ct=1ti=1txi\mathbf{c}_t = \frac{1}{t}\sum_{i=1}^{t} \mathbf{x}_i

Godtyckligt antal indata, fast utdatastorlek, differentierbar, gratis. Embedding-tabell plus detta genomsnitt plus ett linjärt lager till vokabulären är en komplett språkmodell på femton rader. Den är också usel, och hur den är usel är hela härledningen.

Korpusen nedan är en megabyte Shakespeare, 1 115 394 tecken, genom en byte-nivå BPE-tokenizer av den typ som byggdes i kapitel 7 med en vokabulär på 1024: 459 760 tokens med 2,43 tecken vardera, delad 90/10. Varje modell är 128 bred, ser 128 tokens och tränar i 3000 steg med AdamW vid 10310^{-3} med en batch på 64. Perplexitet är på den undanhållna splitten.4

modellparametrarvalideringsperplexitet
endast aktuell token, ingen context alls263 16859.71
plus det enhetliga genomsnittet av allt före den263 168248.07
plus inlärda positions-embeddings279 552245.93
enhetligt genomsnitt adderat till token i stället för att ersätta den263 16860.45

Läs andra raden två gånger. Att ta genomsnittet av context hjälper inte lite; det gör modellen fyra gånger sämre än att ignorera context helt. Två skäl, båda bevisbara snarare än empiriska.

Genomsnittet kan inte se ordning. Addition kommuterar, så om fönstret blandas om lämnas sammanfattningen oförändrad — inte ungefär:

order.pyPYTHON
A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)          # rows of the averaging matrix
y = x[torch.randperm(T)]                # the same tokens, shuffled
print((A[-1] @ x - A[-1] @ y).abs().max().item())
TEXT
2.9802322387695312e-08

Flyttalsbrus på en omordnad summa: de två sammanfattningarna är samma vektor. En modell vars enda vy av context är ett genomsnitt kan inte skilja hunden bet mannen från mannen bet hunden. Rad tre bevisar att detta inte kan fixas genom att lägga till positioner till indatan — en inlärd position embedding på varje token före genomsnittet köpte 2,14 poäng av 188. Positionerna går in i summan, och summan glömmer dem.

Och genomsnittet dränker nuet. Vid position 100 är aktuell token en hundradel av sammanfattningen. Det har en billig lösning som du redan äger: behåll token och addera sammanfattningen till den — en residualkoppling, från kapitel 6, och rad fyra visar vad den gör. När utspädningen reparerats bidrar det enhetliga genomsnittet med ingenting alls: 60.45 mot en baseline på 59.71. Varje token finns där, viktad lika, och lika viktning är samma sak som ingen information.

Problemet är inte genomsnittsbildningen. Det är vikterna.

Genomsnittet är en matrismultiplikation, och masken är en softmax

Länk till avsnittet: Genomsnittet är en matrismultiplikation, och masken är en softmax

Att ta genomsnitt över ett växande prefix ser ut som en loop. Det är en multiplikation med en undertriangulär matris vars rader summerar till ett — och också, exakt, en softmax:

mechanics.pyPYTHON
loop = torch.stack([x[:t + 1].mean(0) for t in range(T)])   # the obvious version

A = torch.tril(torch.ones(T, T))
A = A / A.sum(1, keepdim=True)
mat = A @ x                                                  # the same thing

S = torch.zeros(T, T).masked_fill(torch.tril(torch.ones(T, T)) == 0, float("-inf"))
soft = F.softmax(S, dim=-1) @ x                              # and the same thing again
TEXT
loop vs matmul   max |diff| = 5.960464477539063e-08
loop vs softmax  max |diff| = 5.960464477539063e-08

the averaging matrix A (rows sum to 1, upper triangle is zero):
  1.000 0.000 0.000 0.000 0.000 0.000
  0.500 0.500 0.000 0.000 0.000 0.000
  0.333 0.333 0.333 0.000 0.000 0.000
  0.250 0.250 0.250 0.250 0.000 0.000
  0.200 0.200 0.200 0.200 0.200 0.000
  0.167 0.167 0.167 0.167 0.167 0.167

Tre namngivna komponenter i en transformer är nu på skärmen. Triangeln är causal mask, framtvingad av målet: om position tt kunde se position t+1t{+}1 skulle svaret finnas i indatan — läckan som kapitel 6 sa åt dig att granska efter, fast inuti arkitekturen. Softmax är hur masken implementeras: att sätta förbjudna poster till -\infty skickar dem till exakt noll och normaliserar det som återstår, så maskning och normalisering är en operation. (Använd -\infty, inte -1e9: det är värdet som maskningen betyder, det överlever en cast till float16 som -\infty, och det besparar dig beslutet om huruvida konstanten du valde är tillräckligt stor för intervallet du råkar befinna dig i — vilket är flyttalsrutan i kapitel 2 som ställer en fråga du inte behöver svara på.) Och scores är den fria parametern. Det enhetliga genomsnittet är vad du får när varje tillåten score är samma tal; sätt vilka tal som helst där och softmax gör om dem till giltiga vikter.

Resten av det här kapitlet är en fråga: varifrån kommer de talen?

De kan inte vara vanliga parametrar. En inlärd T×TT \times T-matris skulle vara identisk för varje mening — den kunde koda ”titta fyra tokens bakåt” men aldrig ”titta på substantivet som detta pronomen syftar på”. Vikten som länkar position tt till position ii måste bero på vad som finns på båda positionerna, eftersom relevans är en relation, inte en egenskap: ordet det är inte intrinsikalt relevant, det är relevant för något.

Den billigaste funktionen av två vektorer som returnerar ett tal är skalärprodukten från kapitel 1. Sätt score för position ii för position tt som xtxi\mathbf{x}_t \cdot \mathbf{x}_i och mekanismen fungerar — dåligt, på två sätt som tvingar fram allt annat. En vektors skalärprodukt med sig själv är dess kvadrerade norm, så varje token skulle mest attend till sig själv. Och relationen skulle vara symmetrisk: om det attend starkt till djur, då attend djur starkt till det, vilket är falskt om språk, där ett adjektiv behöver sitt substantiv mycket mer än substantivet behöver adjektivet.

Ge därför varje token två roller, som två inlärda linjära avbildningar av den: vad denna position letar efter, qt=Wqxt\mathbf{q}_t = W_q\mathbf{x}_t, query; och vad den erbjuder sig att hittas av, ki=Wkxi\mathbf{k}_i = W_k\mathbf{x}_i, key. Sätt score qtki\mathbf{q}_t \cdot \mathbf{k}_i och symmetrin är borta, eftersom WqWkW_q \neq W_k: en token kan annonsera en sak och söka efter en annan.

En sak är fortfarande fel. Den viktade summan var över xi\mathbf{x}_i själva, vilket tvingar det som kopieras att vara det som matchas. Matchning vill ha de features som identifierar en token; kopiering vill ha de features som är användbara längre nedströms. Lär därför en tredje avbildning, vi=Wvxi\mathbf{v}_i = W_v\mathbf{x}_i, value, och summera dem.

Formeln är nu bokföring:

Attention(Q,K,V)=softmax ⁣(QKdk+M)V\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V

med MM som causal mask, noll på och under diagonalen och -\infty ovanför. I kod är det trettio rader, varav tjugo är former:

attention.pyPYTHON
class Head(nn.Module):
    """One head of causal self-attention."""

    def __init__(self, d_model, d_head, block):
        super().__init__()
        self.q = nn.Linear(d_model, d_head, bias=False)      
        self.k = nn.Linear(d_model, d_head, bias=False)      
        self.v = nn.Linear(d_model, d_head, bias=False)      
        self.d_head = d_head
        self.register_buffer("mask", torch.tril(torch.ones(block, block)).bool())

    def forward(self, x):
        T = x.shape[1]
        q, k, v = self.q(x), self.k(x), self.v(x)
        s = q @ k.transpose(-2, -1) / math.sqrt(self.d_head)          
        s = s.masked_fill(~self.mask[:T, :T], float("-inf"))          
        w = F.softmax(s, dim=-1)                                      
        return w @ v                                                  

Score, mask, normalisera, blanda. Allt annat är en projektion.

Divisionen med kvadratroten, och vad den försvarar mot

Länk till avsnittet: Divisionen med kvadratroten, och vad den försvarar mot

Nästan varje förklaring av dk\sqrt{d_k} säger ”för att hindra softmax från att mättas”, vilket är sant och inte förklarar något. Argumentet är två rader av variansen från kapitel 2. Om posterna i q\mathbf{q} och k\mathbf{k} är oberoende med medelvärde noll och varians ett, har varje produkt qjkjq_j k_j varians ett, och varianser för oberoende saker adderas:

Var(qk)=j=1dkVar(qjkj)=dk\mathrm{Var}(\mathbf{q}\cdot\mathbf{k}) = \sum_{j=1}^{d_k}\mathrm{Var}(q_j k_j) = d_k

Alltså har scores standardavvikelse dk\sqrt{d_k}. Mätt över tjugotusen slumpmässiga par:

TEXT
     d     Var(q.k)         std   sqrt(d)
     4        3.975       1.994     2.000
    16       16.071       4.009     4.000
    64       64.249       8.016     8.000
   256      253.065      15.908    16.000
  1024     1015.562      31.868    32.000

Varför det spelar roll: softmax är skalkänslig på ett sätt som ett linjärt lager inte är. Att dubbla ett linjärt lagers indata dubblar dess utdata; att multiplicera scores med tio före en softmax gör en mjuk blandning till ett hårt val. En rad med 64 scores, med och utan divisionen:

dkd_kstörsta vikt, odeladentropieffektiva tokensstörsta vikt, deladentropieffektiva tokens
40.2052.94419.00.0813.75842.9
160.4381.6925.40.0753.84946.9
640.4890.8742.40.0853.67339.4
2560.99990.00071.00.1433.54734.7
10241.00000.00001.00.1323.64438.3

”Effektiva tokens” är exponentialen av entropin: hur många positioner raden egentligen tar genomsnitt över. Odelat, vid dk=256d_k = 256, attend ett nyinitialiserat head till exakt en token av 64, vald av ingenting annat än slumpdragningen.

Det är dåligt framåt och värre bakåt, i en form som kapitel 5 redan mätte på en tanh\tanh. En softmax som har låst sig vid en post har nästan ingen derivata: diagonalen i dess Jacobian är wi(1wi)w_i(1-w_i), noll i båda ändar. Över tvåtusen slumpmässiga rader:

dkd_kiwi(1wi)\sum_i w_i(1-w_i) odeladdeladrader mättade (största vikt över 0.99)
40.84270.95680.2 % → 0.0 %
640.29400.960917.9 % → 0.0 %
2560.14060.960949.1 % → 0.0 %
10240.06810.961170.4 % → 0.0 %

Vid dk=1024d_k = 1024 är sju rader av tio frusna innan träningen börjar, och ett head som börjar fruset kan inte lära sig vad det ska titta på. Delat är storheten platt vid 0.96 vid varje bredd och ingenting mättas.

Nu den del ingen publicerar: ändrar det den slutliga perplexiteten? Ta bort divisionen och träna, vid fyra head-bredder:

head-breddodeladdelad med dk\sqrt{d_k}delad med dkd_k
fyra heads, dk=32d_k = 3237.2938.0737.89
ett head, dk=128d_k = 12848.5146.1045.99
ett head, dk=256d_k = 25665.3747.53
ett head, dk=512d_k = 51267.0649.15
ett head, dk=1024d_k = 102476.6959.17

De första två raderna kommer från budgeten på 3000 steg ovan; de sista tre är en kortare körning — 1500 steg, batch på 32, ett head, ingen normalisering före projektionerna — med båda varianterna under identiska inställningar.

Vid dk=32d_k = 32 är divisionen värd ingenting och körningen utan den ligger mycket lite före. Det är inte en licens att slopa den, eftersom den vid 256 är värd 18 perplexitetspoäng och vid 1024 är värd 17. Mekanismen syns i själva scores:

dkd_kscore-std vid initefter 1500 steg, odeladefter 1500 steg, deladrader mättade, odeladdelad
25610.49121.672.1391.9 %0.8 %
51215.13836.852.6698.7 %1.3 %
102421.155147.463.4499.9 %16.5 %

Det odelade head återhämtar sig inte. Det skenar: standardavvikelsen för dess scores går från 21 vid initialisering till 5147, attention-entropin faller till noll, och 99,9 % av raderna lägger mer än 0.99 av sin vikt på en enda token. När ett head väl är en hård selektor är dess gradient nästan noll och ingenting drar det tillbaka, så kollapsen är stabil. Det delade head ligger på en score-standardavvikelse på 3.44 efter samma träning, vilket är en mjuk blandning som fortfarande kan ändras.

Vaswani et al. säger exakt detta och inget mer — de misstänker att produkterna ”grow large in magnitude for large values of dkd_k” och delar.5 Ordet large bär lasten, och tabellerna säger var stort börjar: ingenting vid 32, allt vid 256.

Mer än en åsikt, och de två tredjedelar ingen pratar om

Länk till avsnittet: Mer än en åsikt, och de två tredjedelar ingen pratar om

Ett head är en softmax-rad per position, så det rymmer ett svar på ”vad är relevant här”. Att förutsäga ordet efter den i det djur som korsade den våta gatan behöver den syntaktiska luckan, subjektet och föregående token samtidigt, och en sannolikhetsfördelning kan inte vara koncentrerad på tre platser. Kör därför flera heads parallellt, vart och ett med bredd dmodel/hd_{\text{model}}/h, konkatenera och blanda med ytterligare en matris WoW_o: du har partitionerat bredden, inte lagt till den.

Attention gör också exakt en sak — den flyttar information mellan positioner. Varje operation i koden ovan är linjär längs feature-axeln, och kapitel 5 bevisade vad en stack av linjära avbildningar är. Därför har varje block också en liten MLP applicerad på varje position oberoende, som expanderar bredden med fyra och kommer tillbaka, med en GELU i mitten. Arbetsdelningen är värd att memorera: attention blandar över positioner, feed-forward-nätverket beräknar inom en position.

Hela stegen, där varje rad lägger till en del till raden ovanför:

modellparametrarvalideringsperplexitet
enhetligt genomsnitt, adderat279 55260.45
ett attention head, ersätter token328 70455.47
ett attention head, adderat328 70446.10
fyra heads i stället för ett345 21643.21
plus feed-forward-nätverket476 92839.87
plus LayerNorm — det kompletta blocket477 69638.07

Inlärda vikter slår enhetliga med 14 perplexitetspoäng, vilket är hela det här kapitlets argument i en rad. Fyra heads köper ytterligare 3 för 16 512 extra parametrar. Och samma head är värt 9 poäng mer adderat än ersättande: attention för in information, det bestämmer inte vad en position är.

Nu var parametrarna faktiskt sitter, vilket överraskar folk som bara har sett diagrammet:

breddheadsattentionfeed-forwardtotalt per block
128465 664 (33.2 %)131 712 (66.6 %)197 888
768122 360 064 (33.3 %)4 722 432 (66.6 %)7 085 568
40963267 112 960 (33.3 %)134 238 208 (66.7 %)201 367 552

Två tredjedelar av varje transformer-block är feed-forward-nätverket, i varje skala, eftersom attention har fyra d×dd \times d-matriser och MLP:n har motsvarigheten till åtta. Vad en modell än vet, sitter de flesta parametrar som håller det i den per-positionella MLP:n.

Residualer och LayerNorm, ärvda från kapitel 6

Länk till avsnittet: Residualer och LayerNorm, ärvda från kapitel 6

LayerNorm byggdes och mättes i kapitel 6, och det här kapitlet använder det som det lämnades där; residualkopplingar namngavs och ablaterades där, och byggs här. Raderna ”adderat, inte ersättande” ovan är residualkopplingar, värda 188 perplexitetspoäng för genomsnittet och 9 för ett head. LayerNorm7 normaliserar varje exempel över dess features, och kapitel 6 gav skälen till att det och inte BatchNorm överlevde här — inget beroende av batchen, inga löpande statistiker, identiskt i träning och inferens, likgiltigt för sekvenslängd — varje punkt blir ett krav när du genererar en token åt gången för en användare, vilket är där kapitel 13 hamnar. Det kostar 768 parametrar och köper 1,8 perplexitetspoäng.

block.pyPYTHON
class Block(nn.Module):
    def forward(self, x):
        x = x + self.att(self.ln1(x))     
        x = x + self.ff(self.ln2(x))      
        return x

Titta på var normaliseringen sitter: på indatan till varje sub-layer, med residualvägen från indata till utdata aldrig normaliserad. Det är pre-norm. 2017-artikeln gör motsatsen, x = LayerNorm(x + Att(x))post-norm, vilket sätter en LayerNorm på själva residualvägen.

Xiong et al. förklarade skillnaden genom gradienten vid initialisering, som i ett post-norm-nätverk är dåligt skalad med djup — skälet till att den ursprungliga transformer behövde en learning-rate warmup för att överhuvudtaget träna.8 Tolv block, 1000 steg, learning rate 3×1033 \times 10^{-3}:

TEXT
gradient norm per block at initialisation, before any step
  pre-norm    block 1 0.0498 ... block 12 0.0657   ratio last/first  1.32
  post-norm   block 1 0.0977 ... block 12 0.1613   ratio last/first  1.65

  pre-norm,  no warmup          perplexity   37.82
  pre-norm,  200-step warmup    perplexity   37.62
  post-norm, no warmup          perplexity  308.05
  post-norm, 200-step warmup    perplexity   37.88

Post-norm utan warmup är åtta gånger sämre, och post-norm med warmup matchar pre-norm exakt. Warmup är inte en generell god praxis här; det är en patch för ett specifikt arrangemang av normaliseringen, och att flytta LayerNorm tar bort behovet av den. Det är därför i princip varje modell sedan 2019 är pre-norm, och därför 2017-diagrammet bör läsas som historia snarare än som specifikation.

Ta bort position embeddings och modellen tränar fortfarande; den kan helt enkelt inte se var något är, och det är en symmetri snarare än ett träningsfel. Ingenting i attention-score nämner tt eller ii själva, så att permutera indatan permuterar utdatan: self-attention är permutationsekvivariant. Det är genomsnittets ordningsblindhet i en bättre förklädnad — causal mask återställer viss ordning, eftersom varje position ser ett annat prefix, men inom ett prefix är alla ordningar lika.

Fyra sätt att injicera position, tränade på 64-token-fönster och utvärderade vid 64, 128 och 256 — bortom alla längder de såg:

positionerperplexitet vid 64vid 128vid 256
inga alls48.7952.6357.52
inlärda absoluta embeddings38.63108.47181.94
fasta sinusoider42.9695.26152.25
RoPE44.1250.5284.84
ALiBi44.9543.5142.49

Inlärda absoluta embeddings — en vektor per position, adderad till token — vinner vid den tränade längden och faller sedan från en klippa, eftersom position 100 aldrig var i en batch och dess embedding fortfarande är den slumpvektor den började som. Sinusoider, originalvalet, beräknas snarare än lärs, från sinus och cosinus vid geometriskt fördelade frekvenser; 2017-artikeln hoppades att det skulle extrapolera, och tabellen säger att det inte gör det — funktionen är definierad vid position 200, men modellen lärde sig aldrig läsa den där. RoPE9 lägger inte till något och roterar i stället query och key med en vinkel proportionell mot position, i tvådimensionella skivor; eftersom lika rotation av båda sidor av en skalärprodukt lämnar den oförändrad, kommer score till slut bara bero på tit - i, så position blir relativ gratis och det finns ingen tabell som kan ta slut. Den degraderar, men den degraderar. ALiBi10 är det enklaste och märkligaste resultatet här: ett linjärt straff på score proportionellt mot avstånd, med olika lutning per head. Dess perplexitet förbättras när fönstret växer förbi träningslängden, från 44.95 till 42.49, eftersom straffet är definierat på vilket avstånd som helst och varje head fortsätter göra vad det tränades att göra.

Lärdomen överlever tabellen: en arkitektur som inte kan representera något är ett annat problem än en som aldrig lärde sig det intervallet, och det andra är det som biter. Det är också maskineriet bakom varje ”vi utökade context till 128K”-annonsering — de är nästan alltid omskalningar av en rotary encoding, och de är varför kapitel 16 säger att context-gränsen flyttar sig snarare än försvinner.

Dropout ärvs på samma sätt: den dyker upp på attention-vikterna efter softmax, på varje sub-layer:s utdata före residualaddition och på embedding-summan, och gör exakt vad kapitel 6 beskrev. I stora pretraining-körningar sätts den ofta till noll, eftersom en modell som ser varje token en gång inte är i en position att overfit.

Två tensorer i lagret har formen n×nn \times n, där nn är antalet tokens: scores och vikterna efter softmax. Allt annat — varje projektion, hela MLP:n — är linjärt i nn.

Ett attention-lager, 512 brett, 8 heads, batch på ett, float32, på en laptop-GPU. Läs de två millisekundkolumnerna endast för deras kvoter: de är väggklocka på ett 8 GB laptop-kort som stryps från 1 785 MHz till under 300 MHz när det blir varmt, så en kall körning av samma kod kommer tillbaka sju till tio gånger snabbare och en upptagen ännu långsammare. Megabytekolumnerna är allocator-byteantal och rör sig inte.

TEXT
  tokens   ms total    ms x4   ms projections   attn matrix MB    peak MB    MB x4
     128      2.246        -            1.324              0.5       14.6        -
     256      2.855     1.27            2.113              2.0       19.2     1.31
     512      5.761     2.02            3.105              8.0       34.4     1.79
    1024     16.414     2.85            4.008             32.0       89.1     2.59
    2048     51.573     3.14            9.989            128.0      296.1     3.32
    4096    225.432     4.37           20.176            512.0     1100.1     3.72
    8192    832.838     3.69           40.106           2048.0     4300.1     3.91
   16384   OUT OF MEMORY                                 8192.0

fitted exponent (log-log slope, last four rows):  time ~ n^1.91   memory ~ n^1.87

Kolumnerna x4 är kvoten till raden ovanför, och en fördubbling av nn konvergerar mot exakt 4 för både tid och minne — 3.91 vid sista steget mot teoretiska 4. Projektionskolumnen är kontrollen: 4.0 ms vid 1024 tokens till 40.1 ms vid 8192, en faktor tio för en faktor åtta. Linjär, som utlovat.

Sedan sista raden. Ett attention-lager, en sekvens, ingen modell runt det, får slut på minne på en 8 GB GPU vid 16 384 tokens — bara score-matrisen skulle vara 8 GB, eftersom den är 8 heads gånger 16 384 gånger 16 384 gånger 4 byte. Inte modellen; en mellanliggande tensor i ett lager.

Det är det fysiska faktumet under tre senare kapitel. Det är varför en context window överhuvudtaget har en gräns, vilket kapitel 16 gör till ett pris. Det är varför FlashAttention finns, som beräknar samma resultat i tiles utan att någonsin lagra matrisen — en minnesoptimering innan det är en hastighetsoptimering.11 Och det är aritmetiken bakom priset för en lång prompt, som kapitel 24 betalar i en agent-loop — en separat sak från det kapitlets andra fynd, att en modell också använder en lång context sämre, vilket det mäter och avstår från att skylla på den här formeln.

Visa detaljer

De två cache-krympande varianterna, namngivna här och betalda för i kapitel 13.

Generering cachar keys och values för de tokens som redan har processats — en key och en value per token, per head per lager. Multi-query attention12 behåller hh query-projektioner men en enda key- och value-projektion som delas av alla heads, vilket delar den cachen med hh. Grouped-query attention13 interpolerar: heads grupperas, varje grupp delar en key och en value, så g=hg = h är vanlig attention och g=1g = 1 är multi-query. Nästan varje öppen modell sedan 2023 använder det med 4 eller 8 grupper. Ingen av dem finns för kvalitet; båda finns för storleken på den cachen, och kapitel 13 gör aritmetiken som blir ”vilken modell får plats i din GPU”.

2017-artikeln beskriver en encoder-decoder: en stack som läser källan med omaskad attention, en andra som genererar målet kausalt, och en tredje sorts attention i mitten där decoderns queries möter encoderns keys. Det är rätt för översättning, där indata och utdata är två sekvenser.

Det som vann var decoder-only-halvan — en stack, kausal överallt, indata och utdata i samma sekvens — och skälet är inte elegans. ”Förutsäg nästa token” körs på vilken text som helst, så träningsmängden är internet snarare än en parallell korpus, och allt blir den enda uppgiften: en översättning är ett dokument som innehåller källa och sedan mål, en fråga och dess svar är ett dokument, en konversation med ett tool call i mitten är ett dokument. Kapitel 11 handlar om hur den sista tillverkas. Encoders försvann inte — en ser hela indatan på en gång, vilket är vad du vill när jobbet är att representera en text snarare än fortsätta den, och det är varför kapitel 19:s retrieval embeddings kommer från encoders och inte från modellen som chattar.

Med blocket definierat är modellstorlek aritmetik. Per block, med bredd dd och en fyrfaldig expansion: 4d2+4d4d^2 + 4d för Wq,Wk,Wv,WoW_q, W_k, W_v, W_o med biaser på alla fyra, som GPT-2 har dem — tabellen ovan lämnar bias borta från tre av dem, därav 2 304 färre per block vid d=768d = 768; 8d2+5d8d^2 + 5d för MLP:n; 4d4d för två LayerNorms — 12d2+13d12d^2 + 13d, plus en token-tabell på V×dV \times d och, för absoluta positioner, nctx×dn_{\text{ctx}} \times d. För formen hos GPT-2 small — d=768d = 768, 12 block, en vokabulär på 50 257, en context på 1024, utdatalagret delar embedding-vikterna:

TEXT
  token embeddings     50,257 x 768 = 38,597,376
  position embeddings   1,024 x 768 =    786,432
  one block                             7,087,872
  12 blocks                            85,054,464
  final LayerNorm         2 x 768 =        1,536
  total (weights tied)                124,439,808

Vilket är den publicerade storleken på den modellen. Formeln är inte en approximation; den är modellen. Notera också att nästan en tredjedel av en liten modell är embedding-tabellen, vilket är varför vokabulärstorlek är ett arkitekturbeslut och inte ett preprocessing-beslut — avvägningen som kapitel 7 satte upp.

Perplexitet är ett tal om en korpus. Vad ett head gör är en annan fråga, och en modell tränad på en megabyte Shakespeare är fel instrument för den: det ärligt korrekta att säga om attention-kartan för en modell med 500 000 parametrar är att den mestadels inte är tolkbar. Alltså: ett språk där frågan har ett rätt svar.

Den klassiska illustrationen är djuret korsade inte gatan eftersom det var för trött, där det är djuret, mot …eftersom den var för våt, där ett ord flyttar referenten till gatan. Dessa är Winograd-scheman14 — meningspar som är identiska utom ett ord, där det ordet avgör vad ett pronomen syftar på.

De är också lösbara genom fusk, vilket är den del handledningarna hoppar över. Om de två kandidaterna är ett djur och en plats identifierar trött och våt referenten via kategori, och en modell som bara vet vilka ord som finns närvarande får rätt utan att veta något om ordning. Mätt på den versionen av uppgiften, med undanhållna djur/plats-par:

TEXT
uniform causal average           held-out referent accuracy 100.0 %
one transformer block            held-out referent accuracy  91.7 %

Bag of words slår transformer. Varje demonstration byggd på den meningen bevisar ingenting om attention.

Stäng därför hålet: dra båda kandidater från en pool av sexton substantiv, där vilket som helst kan förekomma i vilken lucka som helst, och dela adjektiven efter roll i stället för kategori — fyra som gör det till den som korsar (trött, rädd, långsam, svag), fyra som gör det till den som korsas (våt, bred, trafikerad, brant).

TEXT
the {x} did not cross the {y} because it was too {adj} , so the {ref} waited .

Träna som en vanlig next-token-predictor, sätt score på en position — ordet efter så det — och bygg den undanhållna mängden från substantivpar vars omvända ordning fanns i träning, så allt som vet vilka två substantiv som finns men inte vilket som kom först måste svara baklänges.

modellparametrarundanhållennamnger det andra substantivet
endast aktuell token5 7965.2 %5.2 %
enhetligt kausalt genomsnitt5 79627.9 %50.0 %
ett head av inlärd attention18 08435.4 %64.6 %
fyra heads22 24475.0 %15.6 %
ett transformer-block55 71692.7 %4.2 %
två transformer-block105 508100.0 %0.0 %

Slumpen bland de två substantiv som finns närvarande är 50 %. Det enhetliga genomsnittet landar på 27.9 % och svarar med fel substantiv i paret exakt halva tiden — signaturen för något som vet vilka ord som finns där och ingenting om deras ordning, som shuffle-testet förutsade tre avsnitt sedan.

Nu kartan: attention vid positionen som måste namnge referenten, medelvärdesbildad över de fyra heads i varje block, för de två meningar som skiljer sig med ett ord. Ett enhetligt genomsnitt skulle lägga 0.067 på var och en av de femton synliga tokens.

TEXT
the animal did not cross the street because it was too tired , so the animal waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00 tired:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.00
         because:0.00 it:0.00 was:0.00 too:0.00 tired:1.00 ,:0.00 so:0.00 the:0.00

the animal did not cross the street because it was too wet , so the street waited .
  blk 1  the:0.00 animal:0.70 did:0.00 not:0.00 cross:0.00 the:0.00 street:0.06
         because:0.00 it:0.00 was:0.00 too:0.00   wet:0.00 ,:0.05 so:0.00 the:0.19
  blk 2  the:0.00 animal:0.00 did:0.00 not:0.00 cross:0.03 the:0.00 street:0.49
         because:0.00 it:0.00 was:0.00 too:0.20   wet:0.03 ,:0.00 so:0.00 the:0.25

Block 1 är identiskt i båda meningarna — 0.70 på det första substantivet, oavsett adjektivet. Det är inte ett misslyckande utan ett bevis: i första lagret är query vid en position en funktion av positionens egen token och index, och det vid position 14 är samma token i båda meningarna. Ett first-layer head kan inte condition på ett ord det ännu inte har hämtat. Så block 1 gör det enda användbara som är tillgängligt och drar det första substantivet framåt.

Block 2 är där meningarna delar sig, och samma rad över alla åtta adjektiv visar regeln modellen hittade:

adjektivblock 2 på djurgatapå adjektivetsvar
trött, rädd, långsam, svag0.0000.0001.000djur
våt, bred, trafikerad, brant0.0000.4910.00–0.03gata

För ett korsar-adjektiv lägger det andra blocket hela sin vikt på adjektivet, eftersom svaret redan finns i residualströmmen — block 1 lade det där — och allt det behöver är bekräftelse. För ett korsad-adjektiv går det och hämtar det andra substantivet i stället. Det är en tvåhoppkrets: ett head flyttar en kandidat framåt, ett head i ett senare lager läser en token som avgör om den ska behållas. Komposition över lager är mekanismen, och det är varför ett block nådde 92.7 % och två nådde 100 %.

Det är också formen på den bäst dokumenterade kretsen i riktiga modeller. Induction heads — ett previous-token head som matar ett head i nästa lager som kompletterar mönstret [A][B] … [A] → [B] — är vad Anthropics tolkningsbarhetsarbete identifierar bakom en stor del av in-context learning, och de bildas vid ett identifierbart ögonblick under pretraining. Det här kapitlet försöker inte göra den analysen: den delegeras, med båda artiklarna i referenserna, eftersom att läsa ut kretsar ur en riktig modell är ett forskningsfält och inte ett avsnitt.

Slutligen implementationen. De trettio raderna ovan, med sina vikter kopierade från PyTorchs egna:

TEXT
ours vs nn.MultiheadAttention           max |diff| = 1.7881393432617188e-07
ours vs F.scaled_dot_product_attention  max |diff| = 1.7881393432617188e-07

1.8×1071.8 \times 10^{-7} på utdata vars medelmagnitud är 0.159: samma aritmetik i en annan ordning, med float32-precision.

Du har arkitekturen som varje modell i resten av kursen byggs från, och den är mindre än sitt rykte: ett viktat genomsnitt vars vikter är inlärda, en per-position-MLP som håller två tredjedelar av parametrarna, två normaliseringar och två additioner, staplade.

Det du inte har är en modell som vet något, och stapling fixar inte det av sig självt. Två block på denna korpus når en träningsperplexitet på 14.49 och en valideringsperplexitet på 40.57, mot ett blocks 18.77 och 38.07 — mer kapacitet, bättre på vad det har sett, sämre på vad det inte har sett, vilket är kapitel 6:s tabell med en transformer i den. Avståndet mellan denna modell och dem kapitel 14 till 30 pratar med är inte arkitektoniskt. Det är samma block, upprepat fler gånger, över mycket mer text.

Vilket gör det till ett bokföringsproblem, och bokföringen är märkligare än den ser ut. Hur mycket text, och var får någon tag på den? Hur mycket aritmetik, och hur uppskattar du den innan pengarna är spenderade? Givet en fast budget, är det bättre att göra modellen större eller visa den mer data — och finns det ett korrekt svar, eller bara ett mode? Kapitel 10 svarar på alla tre genom mätning, och sätter ett pris på den billigaste användbara formen av frågan: vad kostar det, i dag, att träna en modell som GPT-2 från ingenting?


Tre förklaringar av detta material är bättre än denna på vad de är till för, och detta kapitel är skrivet för att läsas tillsammans med dem. Jay Alammars The Illustrated Transformer är den bästa bilden av dataflödet som någonsin ritats. Harvard NLP:s The Annotated Transformer är 2017-artikeln med körbar kod inflätad rad för rad. Andrej Karpathys Let's build GPT: from scratch, in code, spelled out bygger samma modell live på två timmar, och ablationsstegen ovan är samma ryggrad mätt på en annan korpus. För tolkningsbarhetsfrågan som detta kapitel bara nuddar vid är primärkällorna Elhage et al., A Mathematical Framework for Transformer Circuits (2021) och Olsson et al., In-context Learning and Induction Heads (2022), båda från Anthropics tolkningsbarhetsgrupp.

  1. Hochreiter, S. och Schmidhuber, J. Long Short-Term Memory. Neural Computation 9(8), s. 1735–1780 (1997).

  2. Sutskever, I., Vinyals, O. och Le, Q. V. Sequence to Sequence Learning with Neural Networks. arXiv:1409.3215 (2014). Encoder-decoder-modellen vars enda context-vektor är flaskhalsen.

  3. Bahdanau, D., Cho, K. och Bengio, Y. Neural Machine Translation by Jointly Learning to Align and Translate. arXiv:1409.0473 (2014). Attention, tre år före transformer.

  4. Perplexitet är exponentialen av den genomsnittliga korsentropin per token, från kapitel 8. Varje tal här använder samma tokenizer och samma valideringssplit, vilket är det enda villkoret under vilket två perplexiteter överhuvudtaget kan jämföras.

  5. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł. och Polosukhin, I. Attention Is All You Need. arXiv:1706.03762 (2017). Avsnitt 3.2.1 är den enda meningen om dk\sqrt{d_k} som detta kapitel ägnar ett avsnitt åt att mäta.

  6. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G. och Dean, J. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv:1701.06538 (2017).

  7. Ba, J. L., Kiros, J. R. och Hinton, G. E. Layer Normalization. arXiv:1607.06450 (2016). Introducerad och mätt i kapitel 6; används här oförändrad.

  8. Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L. och Liu, T.-Y. On Layer Normalization in the Transformer Architecture. arXiv:2002.04745 (2020). Gradientanalysen bakom pre-norm, och argumentet att warmup är ett symptom.

  9. Su, J., Lu, Y., Pan, S., Murtadha, A., Wen, B. och Liu, Y. RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864 (2021).

  10. Press, O., Smith, N. A. och Lewis, M. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation. arXiv:2108.12409 (2021). Extrapolationsresultatet som reproduceras ovan.

  11. Dao, T., Fu, D. Y., Ermon, S., Rudra, A. och Ré, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135 (2022).

  12. Shazeer, N. Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150 (2019).

  13. Ainslie, J., Lee-Thorp, J., de Jong, M., Zemlyanskiy, Y., Lebrón, F. och Sanghai, S. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245 (2023).

  14. Levesque, H. J., Davis, E. och Morgenstern, L. The Winograd Schema Challenge. KR (2012). Konstruktionen bakom meningen djur / gata som varje attention-handledning använder.

Redo att låta LIA välja åt dig?

Bygg med alla AI-modeller på ett ställe – kom igång gratis i dag.