Show Code
from torch import nn
import torch
import mathAttention is all you need
In this experiment, we will explore the paper “Attention Is All You Need” by Vaswani et al. (2017). This paper introduces the Transformer model, which has become a foundational architecture in natural language processing and machine learning.
Show Code
import sys, os
sys.path.append(os.path.abspath(os.path.join("..", "seq2seq")))
from data_pipeline import (
train_loader, test_pairs,
english_word_to_index, english_index_to_word,
arabic_word_to_index, arabic_index_to_word,
tokenize, encode_source,
PAD_IDX, SOS, EOS, UNK,
)
SOS_IDX = arabic_word_to_index[SOS]
EOS_IDX = arabic_word_to_index[EOS]
src_batch, dec_in_batch, dec_tgt_batch, src_lengths = next(iter(train_loader))
lens = [len(tokenize(en)) for en, _ in test_pairs]
print(f"EN vocab {len(english_word_to_index)} | AR vocab {len(arabic_word_to_index)}")
print(f"PAD={PAD_IDX} SOS={SOS_IDX} EOS={EOS_IDX}")
print(f"src_batch {tuple(src_batch.shape)} | source length range {min(lens)}..{max(lens)} words")EN vocab 793 | AR vocab 1068
PAD=0 SOS=1 EOS=2
src_batch (32, 17) | source length range 8..16 wordsShow Code
# sanity check
src_batch, dec_in_batch, dec_tgt_batch, src_lengths = next(iter(train_loader))
print(
f"English vocab: {len(english_word_to_index):,} Arabic vocab: {len(arabic_word_to_index):,}")
print(f"PAD={PAD_IDX} SOS={SOS_IDX} EOS={EOS_IDX}")
print("src_batch :", tuple(src_batch.shape)) # (B, S) encoder input
# (B, T) decoder input (shifted right)
print("dec_in_batch:", tuple(dec_in_batch.shape))
print("dec_tgt_batch:", tuple(dec_tgt_batch.shape)) # (B, T) labelsEnglish vocab: 793 Arabic vocab: 1,068
PAD=0 SOS=1 EOS=2
src_batch : (32, 14)
dec_in_batch: (32, 13)
dec_tgt_batch: (32, 13)Show Code
# inspect each src , dec_in, dec_tgt triplet in the batch
batch_inx = 0
src = src_batch[batch_inx]
dec_in = dec_in_batch[batch_inx]
dec_tgt = dec_tgt_batch[batch_inx]
print("\nInspecting batch index", batch_inx)
print("src:", src)
print("dec_in:", dec_in)
print("dec_tgt:", dec_tgt)
Inspecting batch index 0
src: tensor([ 18, 12, 40, 16, 17, 401, 76, 3, 33, 2, 0, 0, 0, 0])
dec_in: tensor([ 1, 22, 56, 13, 454, 3, 26, 0, 0, 0, 0, 0, 0])
dec_tgt: tensor([ 22, 56, 13, 454, 3, 26, 2, 0, 0, 0, 0, 0, 0])Small toy example
We will build a small toy example to demonstrate the key concepts of the Transformer model.
Show Code
# 3 tokens with each having 4 dimension representation
x = torch.tensor([
[1., 2., 3., 4.], #token 1
[2., 3., 4., 5.], #token 2
[3., 4., 5., 6.] #token 3
])In transformation attention, we have 3 variables: queries, keys, and values. For simplicity, we will use the same tensor for all three variables.
Q = X
K = X
V = X
Show Code
Q = x
K = x
V = xWe first get the scores by multiplying the queries with the keys. Using the following equation, we can compute the scores: \[scores = Q.K^T\]
Show Code
for i, q_row in enumerate(Q):
for j, k_row in enumerate(K):
dot = (q_row * k_row).sum()
print(
f"Q[{i}]={q_row.tolist()} · K[{j}]={k_row.tolist()} = {dot.item()}")Q[0]=[1.0, 2.0, 3.0, 4.0] · K[0]=[1.0, 2.0, 3.0, 4.0] = 30.0
Q[0]=[1.0, 2.0, 3.0, 4.0] · K[1]=[2.0, 3.0, 4.0, 5.0] = 40.0
Q[0]=[1.0, 2.0, 3.0, 4.0] · K[2]=[3.0, 4.0, 5.0, 6.0] = 50.0
Q[1]=[2.0, 3.0, 4.0, 5.0] · K[0]=[1.0, 2.0, 3.0, 4.0] = 40.0
Q[1]=[2.0, 3.0, 4.0, 5.0] · K[1]=[2.0, 3.0, 4.0, 5.0] = 54.0
Q[1]=[2.0, 3.0, 4.0, 5.0] · K[2]=[3.0, 4.0, 5.0, 6.0] = 68.0
Q[2]=[3.0, 4.0, 5.0, 6.0] · K[0]=[1.0, 2.0, 3.0, 4.0] = 50.0
Q[2]=[3.0, 4.0, 5.0, 6.0] · K[1]=[2.0, 3.0, 4.0, 5.0] = 68.0
Q[2]=[3.0, 4.0, 5.0, 6.0] · K[2]=[3.0, 4.0, 5.0, 6.0] = 86.0Show Code
scores = Q @ K.TShow Code
scorestensor([[30., 40., 50.],
[40., 54., 68.],
[50., 68., 86.]])What does the score matrix represent?
The score matrix represents the similarity or relevance scores between each query and each key. Higher scores indicate greater similarity or relevance.
THe dot measures direction and magnitude, all vectors have the same direction, but their magnitude is different that is why we see the scores are different.
The paper also introduce a scale value \(d_k\), in our above example \(d_k=4\). We recalculate the \(scores\).
Show Code
d_k = 4
scores = (Q @ K.T) / math.sqrt(d_k)
print(scores)tensor([[15., 20., 25.],
[20., 27., 34.],
[25., 34., 43.]])Now the scores will be passed to a \(softmax()\) function to get the attention weights.
Show Code
attention_weights = torch.softmax(scores, dim=-1)
torch.set_printoptions(precision=6, sci_mode=False)
attention_weightstensor([[0.000045, 0.006693, 0.993262],
[0.000001, 0.000911, 0.999088],
[0.000000, 0.000123, 0.999877]])NOw we get the output by multiplying the attention weights with the vector \(V\).
\[ Output = \text{AttentionWeights} . V \]
Show Code
output = attention_weights @ V
outputtensor([[2.993217, 3.993217, 4.993217, 5.993217],
[2.999087, 3.999088, 4.999087, 5.999088],
[2.999877, 3.999877, 4.999876, 5.999877]])The attention weights detemine how much to take the \(V\) and how much to discard.
Try on a real example
We will take one of the training examples, and apply the above attention steps.
Show Code
src_batch[0]tensor([ 18, 12, 40, 16, 17, 401, 76, 3, 33, 2, 0, 0, 0, 0])Show Code
d_k = 64
embedding = nn.Embedding(
len(english_word_to_index),
d_k
)
x = embedding(src_batch)Show Code
x.shapetorch.Size([32, 14, 64])We need position encoding because attention attend to all words regardless of their position, so we need to attach positions to the embeddings:
Token embedding
+
Positional encoding
↓
Transformer input
Show Code
Q = x
K = x
V = x
scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
scores.shapetorch.Size([32, 14, 14])Show Code
attention_weights = torch.softmax(scores, dim=-1)
print(attention_weights.shape)
print(attention_weights[0,0:3,:])torch.Size([32, 14, 14])
tensor([[0.986650, 0.001019, 0.000264, 0.000330, 0.000424, 0.000703, 0.000936,
0.000241, 0.001348, 0.000080, 0.002001, 0.002001, 0.002001, 0.002001],
[0.000207, 0.997629, 0.000010, 0.000170, 0.000024, 0.000210, 0.000015,
0.000120, 0.000157, 0.000644, 0.000203, 0.000203, 0.000203, 0.000203],
[0.000078, 0.000015, 0.996397, 0.000358, 0.000153, 0.000675, 0.000509,
0.001119, 0.000081, 0.000104, 0.000128, 0.000128, 0.000128, 0.000128]],
grad_fn=<SliceBackward0>)Show Code
output = attention_weights @ V
output[0, 0, :]tensor([ 0.418869, 1.183535, -0.237103, -0.627570, 0.054837, -0.278081,
0.864567, -0.883654, 0.546802, -1.711450, 1.732111, 0.060963,
-0.129367, -0.221151, -0.840471, -1.218504, -0.501565, 0.290208,
-1.216797, -0.856026, 0.641887, -0.992822, 0.029074, 0.083398,
-1.554130, -0.035989, 1.602567, 0.014122, -1.196856, -0.503162,
-0.671937, 0.325050, -1.062838, -0.425071, -0.589839, -0.636615,
-0.106364, 0.059236, 0.437981, -0.183379, 1.851293, 0.366936,
-0.231619, 0.036085, 0.293636, 0.384140, -0.447928, -0.658962,
2.254594, -0.424685, 0.704885, 0.742130, -1.131663, -1.884434,
0.235231, -0.309927, -0.885611, 2.857200, -0.688647, 0.374017,
-0.959858, 0.027764, 1.983199, 0.015056], grad_fn=<SelectBackward0>)Toy Encoder Example
input–> a sequence of vectors.
Show Code
class SelfAttention(nn.Module):
def __init__(self, d_model):
super().__init__()
self.d_model = d_model
# Learnable transformation
self.W_Q = nn.Linear(d_model, d_model)
self.W_K = nn.Linear(d_model, d_model)
self.W_V = nn.Linear(d_model, d_model)
def forward(self, x):
# x: (batch, seq_len, d_model)
# 1. Create Q, K, V
Q = self.W_Q(x)
K = self.W_K(x)
V = self.W_V(x)
# 2. QK^T
scores = torch.matmul(Q, K.transpose(-2, -1))
# 3. Scale
scores = scores / math.sqrt(self.d_model)
# 4. Softmax
attention_weights = torch.softmax(scores, dim=-1)
# 5. Weighted sum of V
output = torch.matmul(attention_weights, V)
return output, attention_weightsShow Code
x = torch.tensor([
[1., 2., 3., 4.], # token 1
[2., 3., 4., 5.], # token 2
[3., 4., 5., 6.] # token 3
])
output, attention_weights = SelfAttention(4)(x)
print(output)
print(attention_weights)tensor([[ 0.152170, 2.194707, -1.071629, -3.973908],
[ 0.151867, 2.189679, -1.070900, -3.969090],
[ 0.151565, 2.184670, -1.070173, -3.964290]], grad_fn=<MmBackward0>)
tensor([[0.449563, 0.321097, 0.229341],
[0.453412, 0.320309, 0.226279],
[0.457259, 0.319499, 0.223242]], grad_fn=<SoftmaxBackward0>)Show Code
class ToyEncoder(nn.Module):
def __init__(self, d_model=64, ff_hidden=256):
super().__init__()
# first get the attention
self.self_Atten = SelfAttention(d_model)
# Layer normalize
self.norm1 = nn.LayerNorm(d_model)
# feed forward net
self.ff = nn.Sequential(
nn.Linear(d_model, ff_hidden),
nn.ReLU(),
nn.Linear(ff_hidden, d_model)
)
# Layer normalization after feed-forward
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x):
# -------------------------
# 1. Self-Attention
# -------------------------
#print("\n--- INPUT ---")
#print(x)
attention_outputs, attention_weights = self.self_Atten(x)
#print("\n--- ATTENTION WEIGHTS ---")
#print(attention_weights)
#print("\n--- ATTENTION OUTPUT ---")
#print(attention_outputs)
# residual connection +LayerNorm
x = self.norm1(x + attention_outputs)
# -------------------------
# 2. Feed Forward
# -------------------------
ff_output = self.ff(x)
# Residual connection + LayerNorm
x = self.norm2(x + ff_output)
return x, attention_weightsShow Code
x = torch.tensor([
[
[1., 2., 3., 4.],
[2., 3., 4., 5.],
[3., 4., 5., 6.]
]
])
encoder = ToyEncoder(
d_model=4,
ff_hidden=8
)
output, attention_weights = encoder(x)
print("Input:", x.shape)
print("Attention:", attention_weights.shape)
print("Output:", output.shape)Input: torch.Size([1, 3, 4])
Attention: torch.Size([1, 3, 3])
Output: torch.Size([1, 3, 4])Observe atten weights update
Show Code
SUBJECTS = ["boy", "girl", "dog", "cat", "teacher",
"doctor", "king", "queen", "bird", "fish"]
VERBS = ["chases", "loves", "sees", "eats", "reads",
"writes", "sings", "plays", "watches", "follows"]
OBJECTS = ["book", "ball", "apple", "house", "car",
"tree", "river", "mountain", "song", "letter"]
ADJECTIVES = ["big", "small", "happy", "sad",
"red", "blue", "fast", "slow", "old", "young"]
DET = ["the"]
VOCAB = DET + SUBJECTS + VERBS + OBJECTS + ADJECTIVES
print(len(VOCAB), "unique words") # 4141 unique wordsShow Code
TRAIN_SENTENCES = [
"the happy boy chases the small ball",
"the young girl reads the old book",
"the fast dog eats the red apple",
"the sad cat sees the big house",
"the old teacher writes the small letter",
"the big doctor reads the blue book",
"the young king loves the old song",
"the small queen sings the happy song",
"the fast bird watches the blue river",
"the slow fish follows the red car",
"the happy teacher plays the old song",
"the sad king watches the big mountain",
"the young doctor writes the small letter",
"the old queen loves the small tree",
"the big cat plays the fast car",
]
TEST_SENTENCES = [
"the fast queen chases the old car",
"the sad girl writes the happy letter",
"the young cat sees the small mountain",
"the big fish follows the blue book",
"the slow king plays the red apple",
"the happy doctor loves the old tree",
]Show Code
word_to_index = {w: i for i, w in enumerate(VOCAB)}
index_to_word = {i: w for w, i in word_to_index.items()}
def tokenize(sentence):
return sentence.strip().split()
def encode(sentence):
return torch.tensor([word_to_index[w] for w in tokenize(sentence)], dtype=torch.long)
train_tensor = torch.stack([encode(s) for s in TRAIN_SENTENCES]) # (15, 7)
test_tensor = torch.stack([encode(s) for s in TEST_SENTENCES]) # (6, 7)
print(train_tensor.shape, test_tensor.shape)torch.Size([15, 7]) torch.Size([6, 7])
Comments