Self-attention mechanism in PyTorch - Model Pipeline Trace

Model Pipeline - Self-attention mechanism

The self-attention mechanism helps a model look at all parts of a sentence to understand the importance of each word when making predictions. It compares each word to every other word to decide what to focus on.

Data Flow - 5 Stages

1Input Embeddings

1 sentence x 5 words x 8 features→Convert words into vectors representing their meaning→1 sentence x 5 words x 8 features

[[0.1, 0.3, ..., 0.2], [0.0, 0.5, ..., 0.1], ...]

↓

2Linear Projections (Q, K, V)

1 x 5 x 8→Create Query, Key, and Value vectors by multiplying embeddings with weight matrices→1 x 5 x 8 for each Q, K, V

Q: [[0.2, 0.1, ...], ...], K: [[0.3, 0.0, ...], ...], V: [[0.5, 0.2, ...], ...]

↓

3Attention Scores

Q: 1 x 5 x 8, K: 1 x 5 x 8→Calculate scores by dot product of Q and K transpose, then scale→1 x 5 x 5

[[1.2, 0.5, 0.3, 0.7, 0.9], ...]

↓

4Softmax on Scores

1 x 5 x 5→Convert scores to probabilities that sum to 1 for each word→1 x 5 x 5

[[0.4, 0.1, 0.1, 0.2, 0.2], ...]

↓

5Weighted Sum of Values

Attention weights: 1 x 5 x 5, V: 1 x 5 x 8→Multiply attention weights by V vectors and sum to get output→1 x 5 x 8

[[0.3, 0.2, ..., 0.4], ...]

Training Trace - Epoch by Epoch

Loss
1.2 |****
0.9 |***
0.7 |**
0.5 |*
0.4 |

Epoch	Loss ↓	Accuracy ↑	Observation
1	1.2	0.45	Model starts learning, loss is high, accuracy low
2	0.9	0.60	Loss decreases, accuracy improves
3	0.7	0.72	Model learns important word relations
4	0.5	0.80	Better focus on relevant words
5	0.4	0.85	Training converges with good accuracy

Prediction Trace - 5 Layers

Layer 1: Input Embeddings

Layer 2: Linear Projections to Q, K, V

Layer 3: Attention Scores Calculation

Layer 4: Softmax on Scores

Layer 5: Weighted Sum of Values

Model Quiz - 3 Questions

What does the softmax step do in self-attention?

A Turns scores into probabilities that sum to 1

B Multiplies queries and keys

C Creates word embeddings

D Calculates loss during training

Key Insight

Self-attention lets the model weigh how important each word is compared to others in a sentence, helping it understand context better. Training shows loss going down and accuracy going up, meaning the model learns to focus on the right words.