Multi-head attention in PyTorch - Model Pipeline Trace

Model Pipeline - Multi-head attention

This pipeline shows how multi-head attention works in a transformer model. It takes input data, splits it into multiple heads to learn different parts of the data, combines the results, and improves predictions by focusing on important features.

Data Flow - 5 Stages

1Input Embeddings

64 rows x 512 features→Input token embeddings representing words in a sentence→64 rows x 512 features

[[0.1, 0.3, ..., 0.2], [0.05, 0.4, ..., 0.1], ...]

↓

2Linear Projections for Q, K, V

64 rows x 512 features→Project input into Query, Key, and Value vectors for each head→64 rows x 8 heads x 64 features

Q shape: (64, 8, 64), K shape: (64, 8, 64), V shape: (64, 8, 64)

↓

3Scaled Dot-Product Attention per Head

Q, K, V each (64 rows x 8 heads x 64 features)→Calculate attention scores, apply softmax, and weight values→64 rows x 8 heads x 64 features

Attention weights sum to 1 per query vector

↓

4Concatenate Heads

64 rows x 8 heads x 64 features→Combine all heads back into one tensor→64 rows x 512 features

Concatenated tensor shape (64, 512)

↓

5Final Linear Layer

64 rows x 512 features→Project concatenated output to final feature space→64 rows x 512 features

Output tensor shape (64, 512)

Training Trace - Epoch by Epoch

Epoch 1: ************ (1.2)
Epoch 2: ******** (0.85)
Epoch 3: ****** (0.65)
Epoch 4: **** (0.50)
Epoch 5: *** (0.40)

Epoch	Loss ↓	Accuracy ↑	Observation
1	1.2	0.45	Model starts learning, loss is high, accuracy low
2	0.85	0.62	Loss decreases, accuracy improves as attention learns
3	0.65	0.75	Model focuses better, attention heads capture useful info
4	0.50	0.82	Loss continues to drop, accuracy rises steadily
5	0.40	0.88	Model converges well, multi-head attention effective

Prediction Trace - 5 Layers

Layer 1: Input Embeddings

Layer 2: Linear Projections to Q, K, V

Layer 3: Scaled Dot-Product Attention

Layer 4: Concatenate Heads

Layer 5: Final Linear Layer

Model Quiz - 3 Questions

Why do we split the input into multiple heads in multi-head attention?

A To make the model run slower

B To reduce the input size drastically

C To learn different parts of the input data separately

D To ignore some parts of the input

Key Insight

Multi-head attention allows the model to look at different parts of the input data at once, improving its ability to understand complex patterns. Training shows steady improvement as the model learns to focus attention effectively.