Problem: Let \(|\mathcal V|,N_c\in\mathbf Z^+\) be positive integers (where \(|\mathcal V|\) is the cardinality of an arbitrary set \(\mathcal V\) called the vocabulary and \(N_c\) will come to be seen as the number of codebooks), and let \(\mathcal T\) be a (finite or infinite) set whose elements are called tokens. What does it mean for a function \(\mathbf i\) to be a \((|\mathcal V|,N_c)\)-tokenizer on the token space \(\mathcal T\)? Give some examples of tokenizers.
Solution: It means that \(\mathbf i:\mathcal T\to\{1,…,|\mathcal V|\}^{N_c}\) quantizes each abstract token \(\tau\in\mathcal T\) as some concrete \(N_c\)-tuple of integers \(\mathbf i(\tau):=(i_1(\tau),…,i_{N_c}(\tau))\) where each of these \(N_c\) token IDs \(i_1(\tau),…,i_{N_c}(\tau)\in\mathbf N\) associated to the token \(\tau\in\mathcal T\) ranges from \(1\) to the vocabulary size \(|\mathcal V|\).
- For \(\mathcal T:=\{a, b, …, z, @, \&, ., uh, th, …\}\) the set of natural language tokens (could be finite or infinite depending on how one defines \(\mathcal T\)), tokenization \(\mathbf i\) is typically done via a single (\(N_c=1\)) lookup inside a dictionary (“codebook/vocabulary”) of size \(|\mathcal V|\sim 10^5\) (indeed, one can simply take \(\mathcal V:=\mathcal T\)).
- For \(\mathcal T\) the (infinite) set of \(20\text{ ms}\) Nyquist-downsampled audio waveform tokens, tokenization \(\mathbf i\) might be implemented by first passing an audio token \(\tau\in\mathcal T\) into some CNN encoder \(\mathbf e_{\text{CNN}}:\mathcal T\to\mathbf R^{<|\mathcal T|}\), thereby obtaining some hidden latent vector \(\mathbf e_{\text{CNN}}(\tau)\in\mathbf R^{<|\mathcal T|}\), followed by residual vector quantization \(\text{RVQ}:\mathbf R^{<|\mathcal T|}\to\{1,…,|\mathcal V|\}^{N_c}\) of \(\mathbf e_{\text{CNN}}(\tau)\) through a sequence of \(N_c\) ordered codebooks \(\mathcal V_1,…,\mathcal V_{N_c}\) each containing the same number \(|\mathcal V|:=|\mathcal V_1|=…=|\mathcal V_{N_c}|\) of vectors:
\[\text{RVQ}(\mathbf e):=(\text{argmin}_{1\leq i\leq |\mathcal V|}|\mathbf e-\mathbf e_{i}^{(1)}|,\text{argmin}_{1\leq i\leq |\mathcal V|}|\mathbf e-\text{argmin}_{\mathbf e^{(1)}\in\mathcal V_1}|\mathbf e-\mathbf e^{(1)}|-\mathbf e_{i}^{(2)}|,…)\]
(one could also consider codebooks of different sizes, though conceptually that would simply change the range of the tokenizer to \(\mathbf i:\mathcal T\to\{1,…,|\mathcal V_1|\}\times…\times\{1,…,|\mathcal V_{N_c}|\}\)).
Problem: Describe in painstaking detail the architectural mathematics of the forward pass through a (decoder-only) transformer.
Solution: Note that tokenization is simply a data pre-processing step unaffiliated with the transformer architecture itself which operates on the token IDs themselves. A useful terminology heuristic in what follows:
\[\text{Encoding}\Leftrightarrow\text{Deterministic}\]
\[\text{Embedding}\Leftrightarrow\text{Learnable}\]
Finally, one more useful way to visualize the whole residual stream flow is to note that there are essentially \(5\) key vector spaces between which vectors are being “tossed around”: \(\mathbf R^{|\mathcal V|},\mathbf R^d,\mathbf R^{d_{qk}},\mathbf R^{d_{ov}}\) and \(\mathbf R^{d_{ff}}\). Also keep in mind that there are \(2\) distinct softmax operations, one which might be thought of as the “attention softmax”, and the other as the “output softmax”. During training, both softmaxes have temperature \(T:=1\), it is only during inference that the output softmax might be tuned to have temperature \(T\neq 1\).
- (One-hot encoding) Take the sequence of \(n\leq n^*\) input token IDs from tokenization (where \(n^*\) is the context window of the transformer), and one-hot encode them with respect to the vocabulary size \(|\mathcal V|\).
- (Latent space embedding) Take each unit one-hot basis vector \(\hat{\mathbf e}_i\in\{0,1\}^{|\mathcal V|}\) (it is important that these are unit vectors not just with respect to the \(\ell^2\) metric, but also the \(\ell^1\) metric as will be seen later), and multiply it by the transformer’s (learnable) embedding matrix \(W_E\in\mathbf R^{d\times |\mathcal V|}\), thereby obtaining some latent vector \(\hat{\mathbf e}_i\mapsto\mathbf x_i:=W_E\hat{\mathbf e}_i\in\mathbf R^{d}\) (thus, the column vectors of \(W_E\) are the latent embeddings of the one-hots) that, after pre-training the transformer, should ideally learn a semantically rich latent space representation of the vocabulary \(\mathcal V\).
- (Positional encoding) Perturb each latent in some way that uniquely encodes the position of its corresponding token in the original input (the original Attention is All You Need paper used an additive sinusoidal encoding applied directly to the latents, but one can also use a rotary positional encoding (RoPE) applied to the key/query vectors below).
- (LayerNorm) Send each latent \(\mathbf x_i=(x_i^{(1)},…,x_i^{(d)})\in\mathbf R^d\) with \(\mu_i:=\frac{x_i^{(1)}+…+x_i^{(d)}}{d}\) and \(\sigma_i^2=\frac{(x_i^{(1)}-\mu_i)^2+…+(x_i^{(d)}-\mu_i)^2}{d}\) and \(\varepsilon\sim 10^{-5}\) onto the unit \((d-1)\)-sphere \(S^{d-1}\), and then apply a (learnable) affine transformation of \(2d\) parameters \(\mathbf w,\mathbf b\in\mathbf R^d\):
\[\mathbf x_i\mapsto \mathbf w\odot\frac{\mathbf x_i-\mu_i\boldsymbol{1}}{\sqrt{\sigma_i^2+\varepsilon}}+\mathbf b\]
(note: in the original Attention is All You Need paper, a LayerNorm was applied post-attention rather than pre-attention as is the modern standard. Also, an alternative known as RMSNorm is increasingly common, defined to be LayerNorm but with \(\mu_i=0\) (hence \(\sigma_i^2=\frac{(x_i^{(1)})^2+…+(x_i^{(d)})^2}{d}\)) and \(\mathbf b=\mathbf 0\), so only \(d\) learnable parameters \(\mathbf w\in\mathbf R^d\)). - (Multi-headed self-attention) For each self-attention head \(\eta=1,…,h\) (typical choice is \(h:=d/d_{ov}\) though not necessarily) and for each LayerNormed latent \(\mathbf x_i\in\mathbf R^{d}\), compute its query vector \(\mathbf q_i^{(\eta)}:=W_{\mathbf q}^{(\eta)}\mathbf x_i\in\mathbf R^{d_{qk}}\), its key vector \(\mathbf k_i^{(\eta)}:=W_{\mathbf k}^{(\eta)}\mathbf x_i\in\mathbf R^{d_{qk}}\), and its value vector \(\mathbf v_i^{(\eta)}:=W_{\mathbf v}^{(\eta)}\mathbf x_i\in\mathbf R^{d_{ov}}\). Form the bidirectional, variance-preserving attention scores matrix \(K^{(\eta)T}Q^{(\eta)}/\sqrt{d_{qk}}\in\mathbf R^{n\times n}\), apply causal masking \(M\in(\mathbf R\cup\{-\infty\})^{n\times n}\) if relevant, and then softmax it into a probability distribution to be used in a convex linear combination of the \(n\) value vectors, map back from \(\mathbf R^{d_{ov}}\to\mathbf R^d\) via the so-called output matrix \(W_o^{(\eta)}\), and sum the proposed contextualizing perturbations from each of the \(h\) heads to get the net displacement \(\Delta\mathbf x_i\) of the skip connection to be added to the residual stream latent \(\mathbf x_i\):
\[\Delta\mathbf x_i=\sum_{\eta=1}^{h}W_o^{(\eta)}V^{(\eta)}\mathbf p_{T=1}\left(\frac{K^{(\eta)T}\mathbf q_i^{(\eta)}}{\sqrt{d_{qk}}}+\mathbf M_i\right)\]
6. (Multilayer Perceptron) First another LayerNorm/RMSNorm, then a linear layer from \(\mathbf R^d\to\mathbf R^{d_{ff}}\) (typically \(d_{ff}=4d\)), followed by a nonlinear activation (originally \(\text{ReLU}\), GPT-\(2\) used \(\text{GeLU}\), and these days also gated units like \(\text{SwiGLU},\text{GeGLU}\), etc.), followed by another linear layer back down from \(\mathbf R^{d_{ff}}\to\mathbf R^{d}\).
7. (Layers) Alternate between the multi-headed self-attention blocks and MLP feedforwards for \(L\) layers
8. (Sampling) After the final output MLP layer, emerge from the “deep dive” in latent space \(\mathbf R^d\) back to the vocabulary space \(\mathbf R^{|\mathcal V|}\) via an unembedding matrix \(W_U\) (typically \(W_U=W_E^T\)) and take a \(T\)-softmax over these logits to predict a probability distribution over next tokens at each position \(i=1,…,n\). In particular, to get a prediction of the actual next token, just sample from that \(i=n\) position, append it to the input sequence to obtain a new sequence now of length \(n+1\), and repeat the whole thing.
Problem: Just to close some loose threads from above, show exactly how RoPE positional encoding works with an example. Also, explain and motivate the cross-entropy loss function on which autoregressive transformers are trained, and how this relates to the earlier comment about the one-hot encoded vectors \(\hat{\mathbf e}_i\) being \(\ell^1\)-unit vectors.
Solution: (to be added!)