|
59 | 59 | "> ##### Learning objectives\n", |
60 | 60 | ">\n", |
61 | 61 | "> - Understand what a transformer is used for\n", |
62 | | - "> - Understand causal attention, and what a transformer's output representsalgebra operations on tensors\n", |
| 62 | + "> - Understand causal attention, and what a transformer's output represents (algebraic operations on tensors)\n", |
63 | 63 | "> - Learn what tokenization is, and how models do it\n", |
64 | 64 | "> - Understand what logits are, and how to use them to derive a probability distribution over the vocabulary\n", |
65 | 65 | "\n", |
|
195 | 195 | "cell_type": "markdown", |
196 | 196 | "metadata": {}, |
197 | 197 | "source": [ |
198 | | - "Our tranformer's input is natural language (i.e. a sequence of characters, strings, etc). But ML models generally take vectors as input, not langage. How do we convert language to vectors?\n", |
| 198 | + "Our transformer's input is natural language (i.e., a string of characters). But ML models generally take vectors as input, not language. How do we convert language to vectors?\n", |
199 | 199 | "\n", |
200 | 200 | "We can factor this into 2 questions:\n", |
201 | 201 | "\n", |
|
211 | 211 | "source": [ |
212 | 212 | "### Converting sub-units to vectors\n", |
213 | 213 | "\n", |
214 | | - "We basically make a massive lookup table, which is called an **embedding**. It has one vector for each possible sub-unit of language we might get (we call this set of all sub-units our **vocabulary**). We label every element in our vocabulary with an integer (this labelling never changes), and we use this integer to index into the embedding.\n", |
| 214 | + "Text is split into small pieces called tokens (e.g., words or word parts) using a tokenizer. Each token gets a number, which we turn into a vector (a list of numbers) using a table called an **embedding**.\n", |
215 | 215 | "\n", |
216 | 216 | "A key intuition is that one-hot encodings let you think about each integer independently. We don't bake in any relation between words when we perform our embedding, because every word has a completely separate embedding vector.\n", |
217 | 217 | "\n", |
|
349 | 349 | "\n", |
350 | 350 | "#### **Step 1:** Convert text to tokens\n", |
351 | 351 | "\n", |
352 | | - "The sequence gets tokenized, so it has shape `[batch, seq_len]`. Here, the batch dimension is just one (because we only have one sequence).\n" |
| 352 | + "The sequence gets tokenized, so it has shape `[batch, seq_len]`. Here, the batch dimension is just one because we only have one sequence. Think of `batch` as how many sentences we process, and `seq_len` as sentence length in tokens.\n" |
353 | 353 | ] |
354 | 354 | }, |
355 | 355 | { |
|
592 | 592 | "source": [ |
593 | 593 | "#### Attention\n", |
594 | 594 | "\n", |
595 | | - "First we have attention. This moves information from prior positions in the sequence to the current token.\n", |
| 595 | + "Attention heads decide which words matter most for predicting the next one, copying info from those spots.\n", |
596 | 596 | "\n", |
597 | 597 | "We do this for *every* token in parallel using the same parameters. The only difference is that we look backwards only (to avoid \"cheating\"). This means later tokens have more of the sequence that they can look at.\n", |
598 | 598 | "\n", |
599 | 599 | "Attention layers are the only bit of a transformer that moves information between positions (i.e. between vectors at different sequence positions in the residual stream).\n", |
600 | 600 | "\n", |
601 | | - "Attention layers are made up of `n_heads` heads - each with their own parameters, own attention pattern, and own information how to copy things from source to destination. The heads act independently and additively, we just add their outputs together, and back to the stream.\n", |
| 601 | + "Attention layers are made up of `n_heads` attention heads - each with their own parameters, own attention pattern, and own information how to copy things from source to destination. The heads act independently and additively, we just sum their outputs together back into the residual stream.\n", |
602 | 602 | "\n", |
603 | 603 | "Each head does the following:\n", |
604 | 604 | "* Produces an **attention pattern** for each destination token, a probability distribution of prior source tokens (including the current one) weighting how much information to copy.\n", |
|
629 | 629 | "source": [ |
630 | 630 | "### MLP\n", |
631 | 631 | "\n", |
632 | | - "The MLP layers are just a standard neural network, with a singular hidden layer and a nonlinear activation function. The exact activation isn't conceptually important ([GELU](https://paperswithcode.com/method/gelu) seems to perform best).\n", |
| 632 | + "The MLP layers are just a standard machine learning, with a singular hidden layer and a nonlinear activation function. The exact activation function used isn't conceptually important, though [GeLU](https://paperswithcode.com/method/gelu) is often chosen for its performance.\n", |
633 | 633 | "\n", |
634 | 634 | "Our hidden dimension is normally `d_mlp = 4 * d_model`. Exactly why the ratios are what they are isn't super important (people basically cargo-cult what GPT did back in the day!).\n", |
635 | 635 | "\n", |
|
667 | 667 | "\n", |
668 | 668 | "#### LayerNorm\n", |
669 | 669 | "\n", |
670 | | - "* Simple normalization function applied at the start of each layer (i.e. before each MLP, attention layer, and before the unembedding)\n", |
671 | | - "* Converts each input vector (independently in parallel for each batch x position residual stream vector) to have mean zero and variance 1.\n", |
672 | | - "* Then applies an elementwise scaling and translation\n", |
| 670 | + "* Forces vectors into the range `[0, 1]` which helps the model learn more efficiently.\n", |
| 671 | + "* Applied at the start of each layer (i.e. before each MLP, attention layer, and before the unembedding)\n", |
| 672 | + "* Converts each input vector (independently in parallel for each batch x position residual stream vector) to have mean `0.0` and variance `1.0`\n", |
| 673 | + "* Then applies an element-wise scaling and translation\n", |
673 | 674 | "* Cool maths tangent: The scale & translate is just a linear map. LayerNorm is only applied immediately before another linear map. Linear compose linear = linear, so we can just fold this into a single effective linear layer and ignore it.\n", |
674 | 675 | " * `fold_ln=True` flag in `from_pretrained` does this for you.\n", |
675 | 676 | "* LayerNorm is annoying for interpertability - the scale part is not linear, so you can't think about different bits of the input independently. But it's *almost* linear - if you're changing a small part of the input it's linear, but if you're changing enough to alter the norm substantially it's not linear.\n", |
|
679 | 680 | "#### Positional embeddings\n", |
680 | 681 | "\n", |
681 | 682 | "* **Problem:** Attention operates over all pairs of positions. This means it's symmetric with regards to position - the attention calculation from token 5 to token 1 and token 5 to token 2 are the same by default\n", |
682 | | - " * This is dumb because nearby tokens are more relevant.\n", |
683 | | - "* There's a lot of dumb hacks for this.\n", |
684 | | - "* We'll focus on **learned, absolute positional embeddings**. This means we learn a lookup table mapping the index of the position of each token to a residual stream vector, and add this to the embed.\n", |
| 683 | + " * This is inefficient because nearby tokens are often more relevant.\n", |
| 684 | + "* One solution is **learned, absolute positional embeddings**. This involves a learned lookup table mapping each token's position to a residual stream vector, and adding this to the embed.\n", |
685 | 685 | " * Note that we *add* rather than concatenate. This is because the residual stream is shared memory, and likely under significant superposition (the model compresses more features in there than the model has dimensions)\n", |
686 | 686 | " * We basically never concatenate inside a transformer, unless doing weird shit like generating text efficiently.\n", |
687 | 687 | "* This connects to **attention as generalized convolution**\n", |
688 | | - " * We argued that language does still have locality, and so it's helpful for transformers to have access to the positional information so they \"know\" two tokens are next to each other (and hence probably relevant to each other).\n", |
| 688 | + " * Since language does have locality it is helpful for transformers to have access to positional information so they \"know\" whether two tokens are next to each other, and hence likely relevant to each other.\n", |
689 | 689 | "</details>" |
690 | 690 | ] |
691 | 691 | }, |
|
927 | 927 | " nn.init.normal_(self.W_pos, std=self.cfg.init_range)\n", |
928 | 928 | "\n", |
929 | 929 | " def forward(self, tokens: Int[Tensor, \"batch token\"]) -> Float[Tensor, \"batch token d_model\"]:\n", |
930 | | - " # Hint: You should use the einops.repeat or torch.reapeat function\n", |
| 930 | + " # Hint: You should use the einops.repeat or torch.repeat function\n", |
931 | 931 | " # to repeat batch-wise the positional embedding.\n", |
932 | 932 | " # Hide: hard\n", |
933 | 933 | " # The value of tokens is not important here, only the size of the tensor!\n", |
|
1125 | 1125 | " einops.einsum(\n", |
1126 | 1126 | " normalized_resid_pre,\n", |
1127 | 1127 | " self.W_Q,\n", |
1128 | | - " \"batch posn d_model, nheads d_model d_head -> ???\",\n", |
| 1128 | + " \"batch token d_model, head d_model d_head -> ???\",\n", |
1129 | 1129 | " )\n", |
1130 | 1130 | " + self.b_Q\n", |
1131 | 1131 | ")\n", |
|
1138 | 1138 | "attn_scores = einops.einsum(\n", |
1139 | 1139 | " q,\n", |
1140 | 1140 | " k,\n", |
1141 | | - " \"???,??? -> batch nheads posn_Q posn_K\",\n", |
| 1141 | + " \"???,??? -> batch head token_Q token_K\",\n", |
1142 | 1142 | ")\n", |
1143 | 1143 | "\n", |
1144 | 1144 | "# then scale and apply mask and apply softmax on the correct dimension to get probabilities\n", |
|
1149 | 1149 | "z = einops.einsum(\n", |
1150 | 1150 | " v,\n", |
1151 | 1151 | " attn_pattern,\n", |
1152 | | - " \"???,??? -> batch posn_Q nheads d_head\",\n", |
| 1152 | + " \"???,??? -> batch token_Q head d_head\",\n", |
1153 | 1153 | ")\n", |
1154 | 1154 | "\n", |
1155 | 1155 | "# Calculate output (by applying matrix W_O and summing over heads, then adding bias b_O)\n", |
1156 | 1156 | "attn_out = (\n", |
1157 | 1157 | " einops.einsum(\n", |
1158 | 1158 | " z,\n", |
1159 | 1159 | " self.W_O,\n", |
1160 | | - " \"batch posn_Q nheads d_head, nheads d_head d_model -> batch posn_Q d_model\",\n", |
| 1160 | + " \"batch token_Q head d_head, head d_head d_model -> batch token_Q d_model\",\n", |
1161 | 1161 | " )\n", |
1162 | 1162 | " + self.b_O\n", |
1163 | 1163 | ")\n", |
|
1646 | 1646 | "source": [ |
1647 | 1647 | "If you've finished this, congrats! \n", |
1648 | 1648 | "You should ask the TA what to do next. One option is \n", |
1649 | | - "to look at training and sampling from tranformers in the rest of the arena notebook, which you can find it at https://colab.research.google.com/github/EffiSciencesResearch/ML4G-2.0/blob/master/workshops/transformer/transformer-arena.ipynb." |
| 1649 | + "to look at training and sampling from transformers in the rest of the arena notebook, which you can find it at https://colab.research.google.com/github/EffiSciencesResearch/ML4G-2.0/blob/master/workshops/transformer/transformer-arena.ipynb." |
1650 | 1650 | ] |
1651 | 1651 | } |
1652 | 1652 | ], |
|
0 commit comments