9 A Line-by-Line Implementation of Attention and Transformer

Transformers are presented as a breakthrough architecture for sequence-to-sequence tasks, surpassing RNNs and CNNs by modeling long-range dependencies efficiently and enabling parallel training. The chapter motivates this shift by showing how attention allows models to weigh relationships among tokens, unlocking the performance behind modern large language models. It also positions encoder-decoder Transformers within a broader family that includes encoder-only and decoder-only variants, and frames the chapter’s goal: to build a translation-ready Transformer from the ground up.

The core mechanism—scaled dot-product attention—maps input embeddings to queries, keys, and values, computes similarity scores, normalizes them into attention weights, and aggregates information via the values. This is extended with multi-head attention so different heads capture complementary relationships and senses. Inputs are tokenized (often into subwords), embedded, and augmented with positional encodings to inject order information. The encoder stacks identical layers of multi-head self-attention and position-wise feed-forward networks, each wrapped with layer normalization and residual connections, to convert a source sentence into rich vector representations. The decoder mirrors this stack but adds masked self-attention to preserve autoregressive generation and cross-attention to condition on encoder outputs, producing the target sequence token by token from BOS to EOS.

Implementationally, the chapter builds each piece line by line in PyTorch: an attention function, a MultiHeadedAttention module, a PositionwiseFeedForward layer, LayerNorm, and a SublayerConnection that combines normalization with residuals. These form EncoderLayer and DecoderLayer blocks, which are stacked into Encoder and Decoder modules. A Transformer wrapper orchestrates encode/decode steps, while a Generator head projects decoder outputs to target-vocabulary probabilities for next-token prediction. Finally, a create_model function assembles encoder, decoder, embeddings plus positional encoding for source and target, and the generator, yielding a complete encoder-decoder Transformer ready for training on bilingual data, with the next chapter dedicated to training it for English–French translation.

A diagram of the self-attention mechanism. To calculate attention, the input embedding X is first passed through three neural layers with weights, WQ, WK, and WV, respectively. The outputs are query Q, key K, and value V. The scaled attention score is the product of Q and K divided by the square root of the dimension of K, dk. We apply the softmax function on the scaled attention score to obtain the attention weight. The attention is the product of the attention weight and value V.

Steps to calculate attention weights. The input embedding is passed through two neural networks to obtain query Q and key K. The scaled attention scores are calculated as the dot product of Q and K divided by the square root of the dimension of K. Finally, we apply the softmax function on the scaled attention scores to obtain attention weights, which demonstrate how each element is related to all other elements in the sequence.

Use attention weights and the value vector to calculate the attention vector. The input embedding is passed through a neural network to obtain value V. The final attention is the dot product of the attention weights that we calculated earlier and the value vector V.

The Transformer architecture. The encoder in the Transformer (left side of the diagram), which consists of N identical encoder layers, learns the meaning of the input sequence and converts it into vectors that represent its meaning. It then passes these vectors to the decoder (right side of the diagram), which consists of N identical decoder layers. The decoder constructs the output (e.g., the French translation of an English phrase) by predicting one token at a time, based on previous tokens in the sequence and vector representations from the encoder. The generator on the top right is the head attached to the output from the decoder so that the output is the probability distribution over all tokens in the target language (e.g., the French vocabulary).

The structure of the encoder in the Transformer. The encoder consists of N=6 identical encoder layers. Each encoder layer contains two sublayers. The first sublayer is a multi-head self-attention layer and the second is a feed-forward network. Each sublayer uses layer normalization and residual connection.

The structure of the decoder in the Transformer. The decoder consists of N=6 identical decoder layers. Each decoder layer contains three sublayers. The first sublayer is a masked multi-head self-attention layer. The second is a multi-head cross attention layer to calculate the cross attention between the output from the first sublayer and the output from the encoder. The third sublayer is a feed-forward network. Each sublayer uses layer normalization and residual connection.

An example of multi-head attention. This diagram uses the calculation of the multi-head self attention for the phrase “How are you?” as an example. We first pass the embedding through three neural networks to obtain query Q, key K, and value V, each with a size of (1, 6, 256). We split them into 8 heads, each with a set of Q, k, and V, with a size of (1, 6, 32). We calculate the attention in each head. The attention vectors from the 8 heads are then joined back into one single attention vector, with a size of (1, 6, 256).

An example of how cross-attention weights are calculated between the input to the decoder and the output from the encoder. The input to the decoder is passed through a neural network to obtain query Q. The output from the encoder is passed through a different neural network to obtain key K. The scaled cross attention scores are calculated as the dot product of Q and K divided by the square root of the dimension of K. Finally, we apply the softmax function on the scaled cross-attention scores to obtain the cross-attention weights, which demonstrate how each element in Q is related to all elements in K.

The structure of the generator in the Transformer. The generator converts the output from the decoder stack to a probability distribution over the targe language’s vocabulary, so that the Transformer can use the distribution to predict the next token in the French translation of an English phrase. The generator contains a linear layer so that the number of outputs is the same as the number of tokens in the French vocabulary. The generator also applies a softmax activation to the output so that the output is a probability distribution.

• Transformers are advanced deep-learning models that excel in handling sequence-to-sequence prediction challenges. Their strength lies in effectively understanding the relationships between elements in input and output sequences over long distances.
• The revolutionary aspect of the Transformer architecture is its attention mechanism. This mechanism assesses the relationship between words in a sequence by assigning weights, determining how closely words are related based on the training data. This enables Transformer models like ChatGPT to comprehend relationships between words, thus understanding human language more effectively.
• To calculate SDPA, the input embedding X is processed through three distinct neural network layers, query (Q), key (K), and value (V). The corresponding weights for these layers are W^Q, W^K, and W^V. WE can calculate Q, K, and V as Q=X*W^Q, K=X*Q^K, and V=X*W^V. SDPA first calculates the attention score as follows:

• where d_k represents the dimension of the key vector K. The next step is to apply the softmax function to these attention scores, converting them into attention weights. This ensures that the total attention a word gives to all words in the sentence sums to 100%. The final attention is then calculated as the dot product of these attention weights and the value vector V:

• Furthermore, instead of using one set of query, key, and value vectors, Transformer models use a concept called multi-head attention. The query, key, and value vectors are split into multiple heads. Each head pays attention to different parts or aspects of the input, enabling the model to capture a broader range of information and form a more detailed and contextual understanding of the input data. Multi-head attention is especially useful when a word has multiple meanings in a sentence.

FAQ