Attention Is All You Need: Deconstructing the Transformer Architecture

The Transformer is a neural network architecture that revolutionized Natural Language Processing by abandoning recurrence in favor of a powerful mechanism called self-attention. This fundamental shift allows for massive parallelization and a superior ability to capture long-range dependencies within text. Its core innovation, self-attention, enables each word in a sentence to dynamically assess the importance of all other words, generating a deeply contextualized representation. This is achieved by projecting each word's embedding into Query, Key, and Value vectors and calculating a weighted sum of Values based on Query-Key similarity scores. The full architecture typically consists of an encoder and a decoder. The encoder processes the entire input sequence to build a rich set of numerical representations. The decoder then uses these representations, along with the output it has generated so far, to produce the target sequence one token at a time. Key components include Multi-Head Attention for focusing on diverse contextual patterns, Positional Encodings to provide a sense of word order, and masked attention in the decoder to prevent it from 'cheating' during training. This flexible and powerful design has become the foundation for most modern large language models, driving progress in everything from machine translation and text summarization to chatbots and content generation.

1. Introduction: A Paradigm Shift in Sequence Modeling

For years, Recurrent Neural Networks (RNNs), particularly Long Short-Term Memory (LSTM) and Gated Recurrent Unit (GRU) networks, were the state-of-the-art for sequence modeling tasks in Natural Language Processing (NLP). Their sequential nature, processing one token at a time while maintaining a hidden state, seemed intuitive for language. However, this sequential processing had two major drawbacks: it was computationally slow and parallelization was difficult, and it struggled to capture long-range dependencies in text (e.g., connecting a pronoun at the end of a long paragraph to its subject at the beginning).

In 2017, a seminal paper from Google titled "Attention Is All You Need" introduced the Transformer architecture. It completely discarded recurrence and relied solely on a mechanism called self-attention. This new model not only achieved superior performance on tasks like machine translation but also allowed for massive parallelization, enabling the training of much larger models on unprecedented amounts of data. This breakthrough laid the foundation for modern large language models (LLMs) like GPT, BERT, and T5.

2. The High-Level Architecture: Encoder-Decoder Stacks

The original Transformer model is an Encoder-Decoder architecture, primarily designed for sequence-to-sequence (seq2seq) tasks like machine translation.

• The Encoder: Its job is to read and understand the input sentence, creating a rich numerical representation (a set of context-aware embeddings) for each word. It processes the entire input sequence at once.
• The Decoder: Its job is to take the encoder's representation and generate the output sequence one token at a time, using the previously generated tokens as additional context.

Imagine translating the sentence "The cat sat on the mat." The encoder would process this entire sentence to build a deep contextual understanding. The decoder would then generate "Le" (French for "The"), then use the encoder's output and "Le" to generate "chat" (cat), and so on, until the full translated sentence is produced.

Both the encoder and decoder are composed of a stack of identical layers (N=6 in the original paper).

3. The Building Blocks: A Deep Dive

3.1 Input Processing: Embeddings and Positional Encoding

Before entering the encoder, two things happen to the input tokens:

1. Input Embedding: Each word is converted into a high-dimensional vector using a learned embedding matrix. This vector represents the word's meaning.
2. Positional Encoding: Since the model contains no recurrence, it has no inherent sense of word order. To solve this, we inject information about the position of each token in the sequence. The original paper used a clever combination of sine and cosine functions of different frequencies. This allows the model to learn relative positions, as the positional encoding for any position can be represented as a linear function of the encodings for other positions. This is crucial for understanding grammar and syntax.

3.2 The Core Mechanism: Scaled Dot-Product Self-Attention

This is the heart of the Transformer. For each word in a sequence, self-attention allows it to look at all other words in the same sequence and calculate a weighted "attention score" for each. This score determines how much focus to place on other words when encoding the current word.

To achieve this, three vectors are created for each input embedding:

• Query (Q): Represents the current word's request for information. It's like asking, "Who am I, and what context do I need?"
• Key (K): Represents the label or role of all words in the sequence. It's like a set of signposts saying, "I am a verb," or "I am the subject."
• Value (V): Represents the actual content or meaning of all words in the sequence. It's the information to be extracted.

The process works as follows:

1. Score: For a given word's Query vector, you calculate a dot product with the Key vectors of all other words in the sequence. This dot product measures similarity or compatibility.
2. Scale: The scores are scaled by dividing by the square root of the dimension of the key vectors (√dₖ). This prevents the dot products from becoming too large, which would result in extremely small gradients after the softmax step, hindering training.
3. Softmax: A softmax function is applied to the scaled scores, converting them into a probability distribution. These are the attention weights, summing to 1. A high weight means the word is highly relevant.
4. Weighted Sum: The Value vectors are multiplied by their corresponding attention weights and summed up. The result is a new vector for the current word, which is a blend of its own information and the information from other words it paid attention to.

Example: In the sentence "The animal didn't cross the street because it was too tired," the self-attention mechanism for the word "it" would learn to assign a very high attention weight to "animal" and a low weight to "street," correctly resolving the pronoun's antecedent.

3.3 Multi-Head Attention

Instead of performing a single self-attention operation, the Transformer performs it multiple times in parallel. This is called Multi-Head Attention. Each "head" has its own set of learned Q, K, and V weight matrices. This allows the model to jointly attend to information from different representation subspaces at different positions. For instance, one head might focus on subject-verb relationships, another on pronoun-antecedent relationships, and a third on syntactic structure. The outputs of all heads are then concatenated and passed through a final linear layer to produce the final output of the attention block.

3.4 The Encoder Block

Each encoder block has two main sub-layers:

1. Multi-Head Self-Attention: As described above.
2. Position-wise Feed-Forward Network (FFN): A simple, fully connected feed-forward network applied to each position's vector independently. It consists of two linear transformations with a ReLU activation in between. Its purpose is to process the attention output further and introduce additional non-linearity, increasing the model's representational power.

Each of these two sub-layers is wrapped with a residual connection (the input to the sub-layer is added to its output) followed by layer normalization. This "Add & Norm" step is crucial for training deep networks by preventing vanishing gradients and stabilizing the activations.

3.5 The Decoder Block

The decoder block is similar to the encoder block but with one key addition. It has three sub-layers:

1. Masked Multi-Head Self-Attention: This operates on the output sequence generated so far. The key difference is the "mask." During training, it prevents the decoder from attending to subsequent (future) positions in the output sequence. This is essential because during inference, the model only knows the tokens it has already generated. The mask ensures the training condition mimics the inference condition, preventing the model from "cheating."
2. Encoder-Decoder Attention: This is where the two stacks connect. In this layer, the Queries (Q) come from the previous decoder layer, but the Keys (K) and Values (V) come from the final output of the encoder. This allows every token in the decoder to attend to every token in the input sentence, enabling it to base its output on the relevant parts of the input.
3. Position-wise Feed-Forward Network: Identical to the one in the encoder.

Like the encoder, each sub-layer has an "Add & Norm" step.

4. Final Output and Real-World Applications

After the final decoder block, the resulting vector is passed through a linear layer and a softmax function. This produces a probability distribution over the entire vocabulary. The token with the highest probability is chosen as the next output, and this process repeats until an end-of-sequence token is generated.

The Transformer's architecture has proven to be incredibly versatile. Different configurations have led to major breakthroughs:

• Encoder-Only (e.g., BERT): Excellent for understanding tasks like text classification, sentiment analysis, and named entity recognition.
• Decoder-Only (e.g., GPT series): Excellent for generative tasks like text completion, creative writing, and dialogue systems.
• Encoder-Decoder (e.g., T5, BART): Excellent for sequence-to-sequence tasks like translation, summarization, and question answering.

Browse public study notes