Overview
This is my on-going project of writing a transformer from scratch with minimal help from AI. I only use AI for extremely specific conceptual questions that I can't find a direct explanation to, and for sanity checks on my code (basically I ask "did I implement this concept correctly? yes/no only!").
In the end, I'm only getting out what put in to this project, and I'm only cheating myself.
I was inspired by this post to take on the challenge. Here I will list the resources I found helpful as well as some personally written summaries that help me remember/understanding the concepts.
Understanding Backpropagation
Take-aways
- We represent the process of backpropagation with a computation directed acyclic graph (DAG), where one side contains all training data samples and model weights, and the other side resolves to a scalar loss function that takes compares the model predictions to the training labels. Each node in the graph represents a mathematical operation.
- In the forward pass, the training data and weights are passed through the DAG topologically (meaning every node of one layer must finish computing before the next layer begins) where the mathematical operations manipulate the values until they reach the final loss calculation node, where the computed values are compared to the training data sample's assigned label values. The labels are technically inputs to the DAG, however they are fed directly into the loss node and gradients are not computed for them.
- The loss value is a representation of the difference between: a) the DAG's ability to model the relationship between samples and labels via parameterized transforms and b) the actual relationship.
- Backpropagation uses the loss value to determine the "effect" that a particular input value has on the loss calculation. This "effect", or gradient, is used in the optimization step to update the model's weights in a direction that will ultimately reduce the loss value the next time the inputs are passed through the DAG. The gradients for the training data and labels are not computed, only model parameters.
- Backpropagation is a recursive process that moves topologically from loss function node to input. For a given node, a gradient is computed for each input parameter by multiplying the input parameter's partial derivative w.r.t. output (which is the "local gradient") by the gradient passed to the output side of the node, typically represented as the partial derivative of the output w.r.t the loss function. To kick off the process, the loss node's output contains a gradient value of 1, produced by the partial derivative of the loss node (L) w.r.t L, which is an identity function.
- For nodes that have multiple outputs, the gradient value is accumulated via summation.
- Typically, inputs to the computational graph are tensors, meaning that backpropagation on a computational node would produce a Jacobian matrix; a 2D matrix that shows the influence of every element from the input on every element of the output. The gradient calculation for the input tensor would then be the Jacobian matrix (input tensor vs output tensor) multiplied by the partial derivative of node output w.r.t loss node.
- It would be impractical to actually compute the full Jacobian due to its size, so in practice each gate's backward implementation untilizes the sparse structure of a Jaconbian and hand-codes its effect directly, without ever materializing the matrix.
Take-aways
- Sigmoid / Tanh activation functions can produce vanishing gradients (where the value is 0) if initialized in a fully saturated state (close to the tails). Additionally, the local maximum of the sigmoid derivative will diminish the gradient to a quarter of its original value. This means that layers towards the input of the DAG may hardly change in value when using SGD, due to the compounded degradation.
- ReLU activation function can produce "dead relu", where it has been initialized to the left of 0, or it's value gets pushed so far to the left where the neuron never fires. (thats why leaky ReLu is preferred is some cases)
Architecture
- Transformer Architecture in General
- Self Attention
My current understanding of Q, K, V
-
The Query (Q), Key (K) and Value (V) tensors are learned tensors that project the input sequence into a space that enables the model to determine which tokens are most relevant to each other.
-
This architecture is modeled after traditional retrieval systems. Think of Youtube for example: you type in your query to the search bar, the algorithm best matches it with the video titles (keys), and selecting a title (key) retrieves the video contents (value). In this case, the attention mechanism does the "selecting" of the most relevant information to the query.
-
The query vector represents the information the current token is seeking. The key vectors represent the kinds of information each existing token contains. The attention scores measure how well the query aligns with each key, and the value vectors provide the actual information that is blended together to produce the token's updated representation. It's important to remember that the queries are always looking at existing tokens, not considering future tokens!
-
All tokens in the sequence are updated with a new representation after going through the attention mechanism. This is analogous to how words change meaning when new words are introduced to the end of a sentence.
Process:
-
Every input token's tensor is multiplied by learned weight tensors to create the Q, K, and V vectors.
-
The model measures how well a token's Query matches every other word's Key using a dot product (directional similarity measure). This is produces a "compatibility score".
-
These scores are converted into attention weights via softmax.
-
The attention weights are multiplied by the value vectors, producing the new representations of the token embeddings.
- Position Wise Feed Forward Networks
- Positional Encoding I decided to go with Rotary Position Encodings due to their superiority
- RoPE Paper
- Rotary Positional Encodings vs Absolute vs Relative
- Implementations of Pos. Enc. functions
- Decoder
- Source 1
- One piece of information I couldn't find basically anywhere was how the we go from a sequence of tokens passing through the decoder stack, to a constant sized vector to input into the final linear layer and softmax. The answer? Its an architectural choice that is up to you. The standard way is to use the last token embedding in the sequence of tokens1.