Concepts -- 2.5 Transformers

← Concepts

This page describes how transformers work. Based on demo D5.

In D5 (diagram below):

• Loop 1
  • Input the letter “h” (T1).
  • Infer. Result = “e”.
• Loop 2
  • Input the letters “h”, “e” (T1,T2).
  • Infer. Result = “l”.
• ……
• Loop 8
  • Input the letters “h”, “e”, “l”, “l”, “o”, “ “, “w”, “o”, (T1…T8).
  • Infer. Result = “r”.
• Loop 9
  • Input the letters “e”, “l”, “l”, “o”, “ “, “w”, “o”, “r” (T1…T8).
  • …………

Init

class TinyTransformer(nn.Module):
    def __init__(self):
        super().__init__()
        self.token_embed = nn.Embedding(vocab_size, embed_dim)
        self.pos_embed = nn.Embedding(block_size, embed_dim)
        self.q = nn.Linear(embed_dim, embed_dim)
        self.k = nn.Linear(embed_dim, embed_dim)
        self.v = nn.Linear(embed_dim, embed_dim)
        self.ffn = nn.Sequential(
            nn.Linear(embed_dim, 64),
            nn.ReLU(),
            nn.Linear(64, embed_dim), )
        self.out = nn.Linear(embed_dim, vocab_size)

Forward

    def forward(self, idx):
        B, T = idx.shape
        token_vecs = self.token_embed(idx)
        positions = torch.arange(T, device=device)
        pos_vecs = self.pos_embed(positions)
        x = token_vecs + pos_vecs
        Q = self.q(x)
        K = self.k(x)
        V = self.v(x)
        scores = Q @ K.transpose(-2, -1)
        scores = scores / math.sqrt(embed_dim)
        mask = torch.tril(torch.ones(T, T, device=device))
        scores = scores.masked_fill(mask == 0, float("-inf"))
        weights = F.softmax(scores, dim=-1)
        context = weights @ V
        x = x + context
        x = x + self.ffn(x)
        logits = self.out(x)
        return logits

Training

model = TinyTransformer().to(device)
model = TinyTransformer().to(device)

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(num_epochs):
    X, Y = get_batch()
    logits = model(X)
    B, T, C = logits.shape
    loss = loss_fn(
        logits.reshape(B * T, C),
        Y.reshape(B * T),
    )
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Generate text

model.eval()
idx = torch.tensor([[stoi["h"]]], dtype=torch.long).to(device)
with torch.no_grad():
    for _ in range(40):
        idx_cond = idx[:, -block_size:]
        logits = model(idx_cond)
        last_logits = logits[:, -1, :]
        probs = F.softmax(last_logits, dim=-1)
        next_id = torch.multinomial(probs, num_samples=1)
        idx = torch.cat([idx, next_id], dim=1)
generated = "".join(itos[i] for i in idx[0].tolist())

TF demo D5 overview

26.0624 (v1 26.0623)