Self-attention and multi-head attention in NLP - Deep Dive

Overview - Self-attention and multi-head attention

What is it?

Self-attention is a way for a model to look at all parts of a sentence or sequence at once and decide which parts are important to understand each word. Multi-head attention takes this idea further by having several self-attention processes run in parallel, each focusing on different parts or aspects of the sequence. Together, they help models like transformers understand language better by capturing different relationships between words. This method is key to many modern language models.

Why it matters

Without self-attention and multi-head attention, models would struggle to understand context and relationships in sentences, especially long ones. Traditional methods looked at words one by one or only nearby words, missing important connections. These attention methods let models see the whole sentence at once and learn complex patterns, making language understanding much more accurate and flexible. This has led to breakthroughs in translation, summarization, and many AI tasks.

Where it fits

Before learning self-attention, you should understand basic neural networks and sequence models like RNNs or LSTMs. After mastering self-attention and multi-head attention, you can explore transformer architectures, pre-trained language models like BERT or GPT, and advanced NLP tasks such as question answering and text generation.

Mental Model

Core Idea

Self-attention lets each word in a sentence look at every other word to decide what matters most for understanding itself, and multi-head attention does this many times in parallel to capture different kinds of relationships.

Think of it like...

Imagine you are reading a group chat where each person listens to everyone else but focuses on different topics at the same time—one listens for jokes, another for plans, and another for questions. Together, they get a full picture of the conversation from many angles.

Sequence: [Word1] [Word2] [Word3] ... [WordN]

Each word sends queries to all words and gets back weighted information:

┌─────────────┐
│   Word1     │
│  Queries →  │
│  Attention  │
│  Weights ←  │
└─────────────┘
     ↓
┌─────────────────────────────┐
│ Weighted sum of all words'   │
│ information for Word1         │
└─────────────────────────────┘

Multi-head attention runs several of these in parallel:

Head1  Head2  Head3  ...  HeadH
  ↓      ↓      ↓          ↓
Combine all heads → Final output

Build-Up - 7 Steps

1

Foundation

Understanding sequence context importance

Concept: Words in a sentence depend on each other to make sense, so understanding context is key.

When you read a sentence, you don't just look at one word alone; you think about the words around it to understand meaning. For example, in 'The bank will close soon,' the word 'bank' could mean a river edge or a money place, and the other words help you decide which. This idea of context is the foundation for attention.

Result

You realize that to understand language, models must consider relationships between words, not just words themselves.

Understanding that words depend on each other sets the stage for why attention mechanisms are needed.

2

Foundation

Limitations of traditional sequence models

Concept: Older models like RNNs process words one by one and struggle with long-range dependencies.

Recurrent Neural Networks (RNNs) read sentences word by word in order. This means they remember previous words but can forget important information if the sentence is long. For example, in 'The cat that chased the mouse was tired,' remembering 'cat' when reading 'tired' is hard for RNNs.

Result

You see why a new method is needed to capture relationships between distant words effectively.

Knowing the weaknesses of RNNs helps appreciate the innovation of self-attention.

3

Intermediate

How self-attention works step-by-step

🤔Before reading on: do you think self-attention treats all words equally or weighs some words more? Commit to your answer.

Concept: Self-attention calculates how much each word should pay attention to every other word using learned scores.

For each word, self-attention creates three vectors: Query, Key, and Value. It compares the Query of one word with the Keys of all words to get scores. These scores are turned into weights using softmax, which decide how much each Value (word information) contributes to the final representation of the word. This lets the model focus more on important words.

Result

Each word's new representation is a weighted mix of all words, highlighting relevant context.

Understanding the Query-Key-Value mechanism reveals how models dynamically focus on different parts of the sentence.

4

Intermediate

Role of scaled dot-product attention

🤔Before reading on: do you think scaling the dot product in attention helps or is unnecessary? Commit to your answer.

Concept: Scaling the dot product by the square root of the key dimension stabilizes gradients and improves training.

The attention score is calculated by taking the dot product of Query and Key vectors. If these vectors are large, the dot product can become very big, making softmax outputs too sharp and gradients unstable. Dividing by the square root of the key size keeps values in a good range, helping the model learn better.

Result

Attention weights become more balanced, leading to more stable and effective training.

Knowing why scaling is used prevents confusion about this seemingly small but crucial detail.

5

Intermediate

Why multi-head attention improves learning

🤔Before reading on: do you think multiple attention heads learn the same or different information? Commit to your answer.

Concept: Multiple attention heads let the model look at different parts or aspects of the sentence simultaneously.

Instead of one attention calculation, multi-head attention runs several in parallel, each with its own Query, Key, and Value projections. Each head can focus on different relationships, like syntax, semantics, or position. The outputs are then combined and transformed to form a richer representation.

Result

The model captures diverse information, improving understanding and performance.

Recognizing that multiple heads specialize differently explains why multi-head attention is more powerful than single-head.

6

Advanced

Position encoding in self-attention models

🤔Before reading on: do you think self-attention alone knows word order or needs extra help? Commit to your answer.

Concept: Since self-attention treats words as a set, position encoding adds order information to the input.

Self-attention looks at all words simultaneously without inherent order. To help the model know word positions, special position encodings (like sine and cosine functions) are added to word embeddings. This lets the model distinguish between 'dog bites man' and 'man bites dog.'

Result

The model understands sequence order, which is essential for meaning.

Knowing the need for position encoding clarifies how transformers handle order without recurrence.

7

Expert

Multi-head attention internals and optimization tricks

🤔Before reading on: do you think all heads contribute equally or some can be redundant? Commit to your answer.

Concept: In practice, some attention heads may learn similar patterns or become less useful, and efficient implementations optimize computation.

Multi-head attention splits the input into smaller chunks for each head, processes them in parallel, then concatenates results. Some heads may focus on similar features, leading to redundancy. Techniques like head pruning remove less useful heads to speed up models. Also, optimized matrix multiplications and batching improve training and inference speed.

Result

Understanding these internals helps in designing efficient and effective transformer models.

Knowing that not all heads are equally important guides model compression and interpretability efforts.

Under the Hood

Self-attention computes weighted sums of input embeddings where weights come from similarity scores between queries and keys. These computations happen in parallel for all words, enabling the model to capture dependencies regardless of distance. Multi-head attention runs multiple such computations with different learned projections, allowing the model to attend to various aspects of the input simultaneously. Position encodings are added to input embeddings to provide order information since attention itself is order-agnostic.

Why designed this way?

Traditional sequence models like RNNs processed data sequentially, limiting parallelism and struggling with long-range dependencies. Self-attention was designed to allow full parallel processing and direct connections between any two words, improving efficiency and context capture. Multi-head attention was introduced to let the model learn multiple types of relationships at once, increasing expressiveness without increasing model size excessively.

Input Embeddings + Position Encoding
          ↓
┌─────────────────────────────┐
│   Linear projections to Q,K,V│
└─────────────────────────────┘
          ↓
┌─────────────────────────────┐
│   Scaled Dot-Product Attention│
│   (Q · K^T / sqrt(d_k))      │
│   → Softmax → Weighted sum V │
└─────────────────────────────┘
          ↓
┌─────────────────────────────┐
│   Repeat for each head (H)   │
└─────────────────────────────┘
          ↓
┌─────────────────────────────┐
│ Concatenate heads outputs    │
│ Linear layer to combine      │
└─────────────────────────────┘
          ↓
      Output representation

Myth Busters - 4 Common Misconceptions

Quick: Does self-attention only look at nearby words? Commit to yes or no.

Common Belief: Self-attention only focuses on nearby words like RNNs or CNNs.

Reality: Self-attention considers all words in the sequence equally, regardless of distance.

Why it matters: Believing this limits understanding of self-attention's power to capture long-range dependencies, leading to poor model design choices.

Quick: Do all attention heads learn the same information? Commit to yes or no.

Common Belief: All attention heads in multi-head attention learn the same patterns and are redundant.

Reality: Different heads specialize in different aspects of the input, capturing diverse relationships.

Why it matters: Ignoring this can cause misuse of multi-head attention and missed opportunities for model interpretability.

Quick: Is position encoding optional in transformers? Commit to yes or no.

Common Belief: Position encoding is not necessary because self-attention knows word order inherently.

Reality: Self-attention treats inputs as sets without order; position encoding is essential to provide sequence order information.

Why it matters: Without position encoding, models cannot distinguish word order, leading to poor language understanding.

Quick: Does scaling the dot product in attention have no effect? Commit to yes or no.

Common Belief: Scaling the dot product in attention is an unnecessary detail that doesn't affect training.

Reality: Scaling prevents extremely large values that cause softmax to saturate, stabilizing training and improving performance.

Why it matters: Ignoring scaling can lead to unstable training and worse model accuracy.

Expert Zone

1

Some attention heads can become redundant, and pruning them can reduce model size without much loss in accuracy.

2

The choice of position encoding (sinusoidal vs learned) affects model generalization and transfer to longer sequences.

3

Multi-head attention's parallelism enables efficient GPU utilization but requires careful memory management for large models.

When NOT to use

Self-attention and multi-head attention are less effective for very small datasets or tasks where sequence order is trivial. In such cases, simpler models like CNNs or RNNs may suffice. Also, for extremely long sequences, attention's quadratic complexity can be prohibitive; sparse or linear attention variants are better alternatives.

Production Patterns

In production, multi-head attention is used in transformer-based models like BERT and GPT for tasks such as translation, summarization, and chatbots. Techniques like head pruning, quantization, and distillation optimize these models for speed and size. Attention weights are also analyzed for interpretability to understand model decisions.

Connections

Graph Neural Networks

Both use attention mechanisms to weigh relationships between nodes or elements.

Understanding self-attention helps grasp how graph neural networks dynamically focus on important neighbors in a graph.

Human selective attention in psychology

Self-attention in models mimics how humans focus on relevant parts of information while ignoring distractions.

Knowing human attention mechanisms provides intuition for why self-attention improves model focus and understanding.

Parallel processing in computer architecture

Multi-head attention's parallel computations resemble how CPUs handle multiple tasks simultaneously.

Recognizing this parallelism clarifies why transformers are faster and more scalable than sequential models.

Common Pitfalls

#1 Ignoring position encoding in transformer inputs.

Wrong approach:

input_embeddings = word_embeddings # No position encoding added

Correct approach:

input_embeddings = word_embeddings + position_encoding

Root cause:

Misunderstanding that self-attention alone captures order leads to missing crucial sequence information.

#2 Using single-head attention instead of multi-head attention for complex tasks.

Wrong approach:

attention_output = scaled_dot_product_attention(Q, K, V) # Single head only

Correct approach:

attention_output = multi_head_attention(Q, K, V, num_heads=8)

Root cause:

Underestimating the benefit of multiple attention heads limits model expressiveness.

#3 Not scaling dot product before softmax in attention calculation.

Wrong approach:

scores = Q @ K.T # Missing division by sqrt(d_k)
weights = softmax(scores)

Correct approach:

scores = (Q @ K.T) / sqrt(d_k)
weights = softmax(scores)

Root cause:

Overlooking the scaling step causes unstable gradients and poor training.

Key Takeaways

Self-attention allows models to weigh the importance of all words in a sequence for each word, capturing context effectively.

Multi-head attention runs several self-attention processes in parallel, enabling the model to learn different types of relationships simultaneously.

Position encoding is essential to provide word order information since self-attention treats inputs as unordered sets.

Scaling the dot product in attention calculations stabilizes training and improves model performance.

Understanding these mechanisms is key to grasping how modern transformer models achieve state-of-the-art results in language tasks.