Masked token Experiment

We will design a smalll experiment that uses the masked token task to train the model, then observe how attention updates.

from torch import nn
import torch
import math

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_weights

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_weights

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",
]

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")   # 41


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)

41 unique words
torch.Size([15, 7]) torch.Size([6, 7])