Understanding attention and transformer architectures

Three mental models

• RNNs (Storyteller): Reads a story word by word, remembering only what came before to predict what comes next.
• CNNs (Photographer): Focuses on details in patches of an image and combines them to form a complete picture.
• Transformers (Collaborator): Reads the entire story at once and learns relationships between all parts simultaneously, identifying which details to focus on globally.

Let's break this into two key parts: the attention mechanism and the transformer architecture. The architecture builds on attention, so mastering attention first is critical to understanding the full transformer model.

1. Attention mechanism

The attention mechanism is the foundation of Transformers. It allows the model to focus on relevant parts of the input sequence when generating outputs, addressing limitations of RNNs (like long-range dependencies).

Core idea

• Instead of processing sequences step by step (like RNNs), attention compares all tokens in the input to each other to determine relevance.
• It assigns weights (importance scores) to each token based on how much they contribute to understanding a given token or query.

Key components of attention

Attention is built around the Query (Q), Key (K), and Value (V) paradigm.

1. Query (Q):

   • Represents what you're looking for.
   • Example: In the sentence "The cat sat on the mat," if we're focusing on "cat," the query is related to "cat."

2. Key (K):

   • Represents the characteristics of each token.
   • Example: Every word in the sequence has a key, encoding its meaning.

3. Value (V):

   • Represents the actual content of the tokens.
   • Example: "The", "cat", "sat", etc., are passed as values.

Self-attention formula

Attention computes relevance scores for each token pair using dot products:

$$\text{Attention(Q, K, V)} = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

• QKᵀ: Measures similarity between query and key.
• dₖ: Scaling factor to avoid large dot products in high dimensions.
• softmax: Converts scores into probabilities (attention weights).
• V: Final weighted sum of values, representing the focused context.

Example of self-attention

Input sentence: "The cat sat on the mat."

If the focus is on "cat":

• Query (Q) = "cat".
• Each word in the sentence has a Key (K) and Value (V).
• The attention mechanism calculates how much each word contributes to "cat" using QKᵀ:
  • "The" → Low weight (irrelevant).
  • "sat" → High weight (verb related to "cat").
  • "mat" → Low weight (less connected).

2. Multi-head attention

A single attention mechanism may miss nuanced relationships (e.g., syntax vs. semantics). Multi-Head Attention solves this by:

• Running multiple attention mechanisms (heads) in parallel.
• Each head captures different relationships (e.g., word order, context, synonyms).
• Outputs are concatenated and projected into a final representation.

$$\text{MultiHead(Q, K, V)} = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h)W^O$$

3. Positional encoding

Transformers process sequences in parallel, so they lack inherent order information (unlike RNNs). Positional Encoding adds order back:

• Each token embedding is enriched with a positional vector.
• Common approach: Use sine and cosine functions for each position.

$$\text{PosEnc(pos, 2i)} = \sin(\frac{\text{pos}}{10000^{2i/d}}), \quad \text{PosEnc(pos, 2i+1)} = \cos(\frac{\text{pos}}{10000^{2i/d}})$$

4. Transformer architecture

The Transformer builds on attention by combining Self-Attention, Feed-Forward Networks (FFN), and Residual Connections in a modular, stackable architecture.

Transformer encoder

Encodes the input into a rich contextual representation.

One encoder block:

1. Multi-head self-attention:

   • Computes relationships between all tokens in the input.
   • Outputs context-aware embeddings for each token.

2. Feed-forward network (FFN):

   • Applies two dense layers with non-linear activation.
   • Refines the token representations.

3. Residual connections and layer normalization:

   • Residual connections stabilize training by adding the input back to the output of attention/FFN layers.
   • Layer normalization improves gradient flow.

Structure:

$$\text{Input Embeddings} \xrightarrow{\text{Self-Attention}} \xrightarrow{\text{FFN}} \xrightarrow{\text{Normalization}} \text{Output Embeddings}$$

Transformer decoder

Generates the output sequence step-by-step.

One decoder block:

1. Masked multi-head attention:

   • Ensures the decoder only attends to already-generated tokens (causal attention).

2. Cross-attention:

   • Allows the decoder to attend to the encoder's outputs for context (important for tasks like translation).

3. Feed-forward network and residual connections:

   • Similar to the encoder.

Structure:

$$\text{Generated Embeddings} \xrightarrow{\text{Masked Self-Attention}} \xrightarrow{\text{Cross-Attention}} \xrightarrow{\text{FFN}} \text{Next Token Embedding}$$

End-to-end transformer workflow

1. Input processing:
   • Tokenize input (e.g., "The cat sat").
   • Add positional encodings.
2. Encoder:
   • Processes input embeddings into contextual representations.
3. Decoder:
   • Predicts the next token based on encoder outputs and already-generated tokens.
4. Output:
   • Decodes token predictions into readable text.

Key strengths of transformers

1. Handles long-range dependencies: Self-attention captures relationships between distant tokens.
2. Parallelizable: Processes sequences in parallel, unlike RNNs.
3. Scalable: Stacking multiple encoder/decoder blocks improves performance on large datasets.
4. Adaptable: Works across NLP, vision, and multimodal tasks.

Next steps

1. Experiment with attention matrices (visualizing weights).
2. Implement a small Transformer from scratch to solidify understanding.
3. Explore pre-trained models (e.g., BERT, GPT) to see Transformers in action.