Creating a Transformer Trainer with a Ton of Overhead and No Substance

Noa Roggendorff

nroggendorff

Follow

Welcome back! Today, we will be upgrading our transformer with some much needed improvements to take us closer to our shared goal of world domination using the strongest LM ever fathomed. We all had that conversation, remember?

The cell creating the tokenizer and dataset is unchanged.

from collections import Counter


class Tokenizer:
    def __init__(self, corpus, vocab_size=1250, max_token_length=8):
        tokens = [
            t
            for t, _ in Counter(
                sub
                for example in corpus
                for chunk in (lambda w: [w[0], *map(lambda x: f" {x}", w[1:])])(
                    example.strip().split()
                )
                for n in [len(chunk)]
                for i in range(n)
                for j in range(i + 1, min(n, i + max_token_length) + 1)
                for sub in [chunk[i:j]]
            ).most_common(vocab_size - 1)
        ] + ["<pad>"]

        self.tokens = sorted(tokens[:-1], key=len, reverse=True)

        self.vocabulary = {t: i for i, t in enumerate(self.tokens)} | {
            "<pad>": len(self.tokens)
        }

        self.id_to_token = {v: k for k, v in self.vocabulary.items()}

        self.pad_id = self.vocabulary["<pad>"]
        self.vocab_size = len(self.vocabulary)

        self.decode = lambda ids: "".join(
            self.id_to_token[i] for i in ids if i != self.pad_id
        )

    def encode(self, text):
        ids = []
        i = 0

        while i < len(text):
            for token in self.tokens:
                if text.startswith(token, i):
                    ids.append(self.vocabulary[token])
                    i += len(token)
                    break
            else:
                ids.append(self.vocabulary.get(text[i], self.pad_id))
                i += 1

        return ids


examples = open("dataset").readlines()
tokenizer = Tokenizer(examples)

print(
    ",".join(
        [
            tokenizer.decode([token])
            for token in tokenizer.encode("trees lose leaves in autumn")
        ]
    )
)

There are some new hyperparameters, though:

cheat_help = max([len(example) for example in examples])
epochs = 50
batch_size = 8
lr = 3e-4
emb_dim = 128
hidden = 512
latent_dim = emb_dim // 4
num_heads = 2
num_blocks = 3

Our first alteration is in the attention head. We are going to implement multi-head latent attention, as the dirty peasants begging for scraps from the Deepseek table that we are. It is very cool that they published all of their findings, and I really wish more companies did stuff like that.

Instead of computing K and V from the full size embeddings each time, it now receives a pre-computed latent tensor and projects it up to K and V. This means that the total size of the cache decreases (it's stored for the whole block now), but it also means we have to do more math, and that predictions are less accurate (this is hard to measure, because nobody trains models without some form of cheating on cache, so technically, latent representations make predictions more accurate, because you can afford bigger models).

To be clear, MLA doesn't solve all yo' problems, and compute still scales linearly with input size (only the steps are smaller; the curve isn't any less steep).

import torch
import torch.nn as nn


class AttentionHead(nn.Module):
    def __init__(self, embedding_dim, head_dim, latent_dim):
        super().__init__()

        self.head_dim = head_dim

        self.q = nn.Linear(embedding_dim, head_dim, bias=False)
        self.k_up = nn.Linear(latent_dim, head_dim, bias=False)
        self.v_up = nn.Linear(latent_dim, head_dim, bias=False)

    def forward(self, embeddings, latent):
        query = self.q(embeddings)
        key = self.k_up(
            latent
        )  # projecting the latent back up to `head_dim`
        scores = query @ key.transpose(-2, -1) * (self.head_dim**-0.5)

        return torch.softmax(
            scores.masked_fill(
                torch.triu(
                    torch.ones(
                        scores.size(-2),
                        scores.size(-1),
                        device=scores.device,
                    ),
                    diagonal=key.size(-2) - query.size(-2) + 1,
                ).bool(),
                -1e9,
            ),
            dim=-1,
        ) @ self.v_up(latent)

Enabling that is the new kv_down matrix, which brings the embeddings down to an even further compressed latent size (depending on if you consider embeddings as compressed ideas, which you probably don't).

class TransformerBlock(nn.Module):
    def __init__(self, embedding_dim, hidden_size, num_heads=2, latent_dim=32):
        super().__init__()

        head_dim = embedding_dim // num_heads

        self.kv_down = nn.Linear(embedding_dim, latent_dim, bias=False)

        self.heads = nn.ModuleList(
            [
                AttentionHead(embedding_dim, head_dim, latent_dim)
                for _ in range(num_heads)
            ]
        )

        self.proj = nn.Linear(
            head_dim * num_heads,
            embedding_dim,
        )

        self.mlp = nn.Sequential(
            nn.Linear(embedding_dim, hidden_size),
            nn.GELU(),
            nn.Linear(hidden_size, embedding_dim),
        )

        self.ln1 = nn.LayerNorm(embedding_dim)
        self.ln2 = nn.LayerNorm(embedding_dim)

    def forward(self, embeddings, cache):
        normed = self.ln1(embeddings)
        latent = torch.cat([cache, self.kv_down(normed)], dim=-2)

        attn_out = embeddings + self.proj(
            torch.cat([head(normed, latent) for head in self.heads], dim=-1)
        )
        return attn_out + self.mlp(self.ln2(attn_out)), latent

Starting our trend, the Transformer module doesn't really change at all, it just stores the tensor for future use.

class Transformer(nn.Module):
    def __init__(self, embedding_dim, hidden_size, vocab, latent_dim):
        super().__init__()
        self.embed = nn.Embedding(vocab, embedding_dim)
        self.pos_embed = nn.Embedding(cheat_help, embedding_dim)

        self.latent_dim = latent_dim

        self.blocks = nn.Sequential(
            *[
                TransformerBlock(embedding_dim, hidden_size, latent_dim=self.latent_dim)
                for _ in range(num_blocks)
            ]
        )
        self.lm_head = nn.Linear(embedding_dim, vocab)
        self.lm_head.weight = self.embed.weight

    def forward(self, tokens, cache=None):
        B, T = tokens.shape
        cache = cache or [
            torch.empty(B, 0, self.latent_dim, device=tokens.device)
            for _ in self.blocks
        ]

        offset = cache[0].size(-2)

        embeddings = self.embed(tokens) + self.pos_embed(
            torch.arange(offset, offset + T, device=tokens.device)
            .unsqueeze(0)
            .expand(B, T)
        )

        new_caches = []
        for i, block in enumerate(self.blocks):
            embeddings, block_new_cache = block(embeddings, cache[i])
            new_caches.append(block_new_cache)

        return self.lm_head(embeddings), new_caches

    def generate(self, tokens, max_length=100):
        cache = None
        for _ in range(max_length - tokens.shape[1]):
            logits, cache = self(tokens if cache is None else tokens[:, -1:], cache)
            tokens = torch.cat(
                [tokens, logits[:, -1].argmax(dim=-1, keepdim=True)], dim=1
            )
        return tokens


model = Transformer(emb_dim, hidden, tokenizer.vocab_size, emb_dim // 4).cuda()
model.requires_grad_()

print(sum(p.numel() for p in model.parameters()))

Training is the same as before too, although you'll probably want to run it for longer now, because of the new parameters. to be honest, you shouldn't run this trainer at all, because it's really only for learning, and won't produce a competent model (you didn't hear it from me)

import random
from tqdm import tqdm
from torch.nn.functional import cross_entropy

encoded_train = [tokenizer.encode(e) for e in examples]
encoded_test = [tokenizer.encode(e) for e in open("test").readlines()]

to_batch = lambda encoded: torch.stack(
    [
        torch.tensor(example + [tokenizer.pad_id] * (cheat_help - len(example)))
        for example in encoded
    ]
).cuda()


def eval_loss():
    model.eval()
    with torch.no_grad():
        batch = to_batch(encoded_test)
        pred, _ = model(batch[:, :-1])

        loss = cross_entropy(
            pred.reshape(-1, pred.size(-1)),
            batch[:, 1:].reshape(-1),
            ignore_index=tokenizer.pad_id,
        ).item()

    model.train()
    return loss


steps_per_epoch = len(encoded_train) // batch_size
bar = tqdm(range(epochs * steps_per_epoch))

for step in bar:
    batch = to_batch(random.sample(encoded_train, batch_size))

    pred, _ = model(batch[:, :-1])

    loss = cross_entropy(
        pred.reshape(-1, pred.size(-1)),
        batch[:, 1:].reshape(-1),
        ignore_index=tokenizer.pad_id,
    )

    loss.backward()

    with torch.no_grad():
        for p in model.parameters():
            if p.grad is not None:
                p -= p.grad * lr

    model.zero_grad()

    train_loss = loss.item()

    if step % steps_per_epoch == 0:
        val_loss = eval_loss()

    bar.set_postfix(train_loss=train_loss, validation_loss=val_loss)

Alright, now, I want to tackle vision. I'll show you how to do image encoding, but you can also adapt this to other mediums like video or audio without too much effort. It's worth noting that this demo task will not encourage actual understanding of the image, the encoder will end up just becoming the classifier and telling the transformer what the correct answer is. There's also nothing preventing backprop from just trying to put all of the information in one of the slots, and ignoring the rest.

image_width = 32
image_height = 32
num_slots = 4

from torchvision import transforms

transform = transforms.Compose(
    [
        transforms.Resize(image_width),
        transforms.CenterCrop((image_width, image_height)),
        transforms.ToTensor(),
    ]
)

from datasets import load_dataset

dataset = load_dataset("nroggendorff/flowers")

format_prompts = lambda example: {
    "text": f"This flower is obviously {'an' if example['label'] == 'orchid' else 'a'} {example['label']}!",
}

train_set, test_set = [
    data.map(
        format_prompts,
        remove_columns="label",
    )
    for data in [dataset["train"], dataset["test"]]
]

The image encoder will take the image, and move it to embedding space for the model to echo-chamber (technical term) with the language embeddings.

class ImageEncoder(nn.Module):
    def __init__(
        self, width, height, hidden_size, embedding_dim, trunk_size, slot_count
    ):
        super().__init__()

        self.encoder = nn.Sequential(
            nn.Linear(3 * width * height, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, trunk_size),
        )  # only three-color-channel images
        self.embedders = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(trunk_size, hidden_size),
                    nn.ReLU(),
                    nn.Linear(hidden_size, embedding_dim),
                )
                for _ in range(slot_count)
            ]
        )

    def forward(self, image):
        trunk = self.encoder(image.flatten(1))
        return [embedder(trunk) for embedder in self.embedders]

class VisionTransformer(nn.Module):
    def __init__(
        self, embedding_dim, hidden_size, vocab, latent_dim, num_blocks, num_slots
    ):
        super().__init__()
        self.embed = nn.Embedding(vocab, embedding_dim)
        self.pos_embed = nn.Embedding(cheat_help, embedding_dim)
        self.latent_dim = latent_dim
        self.num_slots = num_slots

        self.encoder = ImageEncoder(
            image_width,
            image_height,
            hidden_size,
            embedding_dim,
            image_width * image_height,
            num_slots,
        )  # the idea is that the encoder will learn to put abstract information about the images in those slots

        self.blocks = nn.Sequential(
            *[
                TransformerBlock(embedding_dim, hidden_size, latent_dim=latent_dim)
                for _ in range(num_blocks)
            ]
        )
        self.lm_head = nn.Linear(embedding_dim, vocab)
        self.lm_head.weight = self.embed.weight

    def forward(self, tokens, image, cache=None):
        B, T = tokens.shape
        is_prefill = cache is None
        cache = cache or [
            torch.empty(B, 0, self.latent_dim, device=tokens.device)
            for _ in self.blocks
        ]

        offset = max(cache[0].size(-2) - num_slots, 0) # prevent gaps in the position embeddings

        embeddings = self.embed(tokens) + self.pos_embed(
            torch.arange(offset, offset + T, device=tokens.device)
            .unsqueeze(0)
            .expand(B, T)
        )

        if is_prefill:
            image_embeddings = torch.stack(self.encoder(image), dim=1) + self.pos_embed(
                torch.zeros(num_slots, device=tokens.device, dtype=torch.long)
            ).unsqueeze(0)
            embeddings = torch.cat([image_embeddings, embeddings], dim=1)

        new_caches = []
        for i, block in enumerate(self.blocks):
            embeddings, block_new_cache = block(embeddings, cache[i])
            new_caches.append(block_new_cache)

        logits = self.lm_head(embeddings)
        return (logits[:, num_slots:] if is_prefill else logits), new_caches


model = VisionTransformer(
    embedding_dim=emb_dim,
    hidden_size=hidden,
    vocab=tokenizer.vocab_size,
    latent_dim=latent_dim,
    num_blocks=num_blocks,
    num_slots=num_slots,
).cuda()
model.requires_grad_()

print(sum(p.numel() for p in model.parameters()))

Batching is more annoying now, so that's fun.

import random
from tqdm import tqdm
from torch.nn.functional import cross_entropy

encoded_train = [tokenizer.encode(e) for e in train_set["text"]]
encoded_test = [tokenizer.encode(e) for e in test_set["text"]]

train_examples = list(zip(encoded_train, train_set["image"]))
test_examples = list(zip(encoded_test, test_set["image"]))


def to_batch(examples):
    texts, images = zip(*examples)

    tokens = torch.stack(
        [
            torch.tensor(text + [tokenizer.pad_id] * (cheat_help - len(text)))
            for text in texts
        ]
    ).cuda()

    images = torch.stack(
        [img if torch.is_tensor(img) else transform(img) for img in images]
    ).cuda()

    return tokens, images


def eval_loss():
    model.eval()

    with torch.no_grad():
        tokens, images = to_batch(test_examples)

        pred, _ = model(tokens[:, :-1], images)

        loss = cross_entropy(
            pred.reshape(-1, pred.size(-1)),
            tokens[:, 1:].reshape(-1),
            ignore_index=tokenizer.pad_id,
        ).item()

    model.train()
    return loss


steps_per_epoch = len(train_examples) // batch_size
bar = tqdm(range(epochs * steps_per_epoch))

for step in bar:
    tokens, images = to_batch(random.sample(train_examples, batch_size))

    pred, _ = model(tokens[:, :-1], images)

    loss = cross_entropy(
        pred.reshape(-1, pred.size(-1)),
        tokens[:, 1:].reshape(-1),
        ignore_index=tokenizer.pad_id,
    )

    loss.backward()

    with torch.no_grad():
        for p in model.parameters():
            if p.grad is not None:
                p -= p.grad * lr

    model.zero_grad()

    train_loss = loss.item()

    if step % steps_per_epoch == 0:
        val_loss = eval_loss()

    bar.set_postfix(train_loss=train_loss, validation_loss=val_loss)

Now, it's time for my least favorite change to the architecture, Mixture of Experts. The reason I dislike it so much is that you have to use an auxiliary loss to discourage the router from giving all tasks to one expert, and that means the architecture isn't elegant. It gets the job done, though, so I can't complain.

MoE involves replacing the transformer block's single MLP with several smaller ones, and using a router/gate to pick one or many to predict on different types of prompts. Weirdly, the experts aren't pretrained on different datasets or anything; what specialties the different experts have is entirely up to the router.

class MoE(nn.Module):
    def __init__(self, embedding_dim, hidden_size, num_experts=4, top_k=2):
        super().__init__()

        self.top_k = top_k
        self.gate = nn.Linear(embedding_dim, num_experts, bias=False)

        self.experts = nn.ModuleList(
            [
                nn.Sequential(
                    nn.Linear(embedding_dim, hidden_size),
                    nn.GELU(),
                    nn.Linear(hidden_size, embedding_dim),
                )
                for _ in range(num_experts)
            ]
        )

    def forward(self, x):
        shape = x.shape
        x = x.reshape(-1, shape[-1])

        probs = torch.softmax(self.gate(x), dim=-1)
        top_probs, top_indices = probs.topk(self.top_k, dim=-1)
        top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)

        out = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            token_idx, k_idx = (top_indices == e).nonzero(as_tuple=True)
            if token_idx.numel():
                out[token_idx] += top_probs[token_idx, k_idx, None] * expert(
                    x[token_idx]
                )

        dispatch_mask = torch.zeros_like(probs).scatter_(-1, top_indices, 1.0)
        self.aux_loss = (
            len(self.experts) * (dispatch_mask.mean(dim=0) * probs.mean(dim=0)).sum()
        )

        return out.reshape(shape)

class TransformerBlock(nn.Module):
    def __init__(self, embedding_dim, hidden_size, num_heads=2, latent_dim=32):
        super().__init__()

        head_dim = embedding_dim // num_heads

        self.kv_down = nn.Linear(embedding_dim, latent_dim, bias=False)

        self.heads = nn.ModuleList(
            [
                AttentionHead(embedding_dim, head_dim, latent_dim)
                for _ in range(num_heads)
            ]
        )

        self.proj = nn.Linear(
            head_dim * num_heads,
            embedding_dim,
        )

        self.mlp = MoE(embedding_dim, hidden_size)

        self.ln1 = nn.LayerNorm(embedding_dim)
        self.ln2 = nn.LayerNorm(embedding_dim)

    def forward(self, embeddings, cache):
        normed = self.ln1(embeddings)
        latent = torch.cat([cache, self.kv_down(normed)], dim=-2)

        attn_out = embeddings + self.proj(
            torch.cat([head(normed, latent) for head in self.heads], dim=-1)
        )
        mlp_out = self.mlp(self.ln2(attn_out))
        return attn_out + mlp_out, latent, self.mlp.aux_loss

The Transformer has basically no change again. Just bookkeeping stuff is happening.

class Transformer(nn.Module):
    def __init__(self, embedding_dim, hidden_size, vocab, latent_dim):
        super().__init__()
        self.embed = nn.Embedding(vocab, embedding_dim)
        self.pos_embed = nn.Embedding(cheat_help, embedding_dim)

        self.latent_dim = latent_dim

        self.blocks = nn.Sequential(
            *[
                TransformerBlock(embedding_dim, hidden_size, latent_dim=self.latent_dim)
                for _ in range(num_blocks)
            ]
        )
        self.lm_head = nn.Linear(embedding_dim, vocab)
        self.lm_head.weight = self.embed.weight

    def forward(self, tokens, cache=None):
        B, T = tokens.shape
        cache = cache or [
            torch.empty(B, 0, self.latent_dim, device=tokens.device)
            for _ in self.blocks
        ]

        offset = cache[0].size(-2)

        embeddings = self.embed(tokens) + self.pos_embed(
            torch.arange(offset, offset + T, device=tokens.device)
            .unsqueeze(0)
            .expand(B, T)
        )

        new_caches = []
        aux_losses = []
        for i, block in enumerate(self.blocks):
            embeddings, block_new_cache, aux_loss = block(embeddings, cache[i])
            new_caches.append(block_new_cache)
            aux_losses.append(aux_loss)

        return self.lm_head(embeddings), new_caches, sum(aux_losses)

    def generate(self, tokens, max_length=100):
        cache = None
        for _ in range(max_length - tokens.shape[1]):
            logits, cache, _ = self(tokens if cache is None else tokens[:, -1:], cache)
            tokens = torch.cat(
                [tokens, logits[:, -1].argmax(dim=-1, keepdim=True)], dim=1
            )
        return tokens


model = Transformer(emb_dim, hidden, tokenizer.vocab_size, emb_dim // 4).cuda()
model.requires_grad_()

print(sum(p.numel() for p in model.parameters()))

Training is largely the same, too.

import random
from tqdm import tqdm
from torch.nn.functional import cross_entropy

encoded_train = [tokenizer.encode(e) for e in examples]
encoded_test = [tokenizer.encode(e) for e in open("test").readlines()]

to_batch = lambda encoded: torch.stack(
    [
        torch.tensor(example + [tokenizer.pad_id] * (cheat_help - len(example)))
        for example in encoded
    ]
).cuda()


def eval_loss():
    model.eval()
    with torch.no_grad():
        batch = to_batch(encoded_test)
        pred, _, aux_loss = model(batch[:, :-1]) # here it is: the bane of my existence

        loss = (
            cross_entropy(
                pred.reshape(-1, pred.size(-1)),
                batch[:, 1:].reshape(-1),
                ignore_index=tokenizer.pad_id,
            )
            + 0.01 * aux_loss
        ).item()

    model.train()
    return loss


steps_per_epoch = len(encoded_train) // batch_size
bar = tqdm(range(epochs * steps_per_epoch))

for step in bar:
    batch = to_batch(random.sample(encoded_train, batch_size))

    pred, _, aux_loss = model(batch[:, :-1])

    loss = (
        cross_entropy(
            pred.reshape(-1, pred.size(-1)),
            batch[:, 1:].reshape(-1),
            ignore_index=tokenizer.pad_id,
        )
        + 0.01 * aux_loss
    )

    loss.backward()

    with torch.no_grad():
        for p in model.parameters():
            if p.grad is not None:
                p -= p.grad * lr

    model.zero_grad()

    train_loss = loss.item()

    if step % steps_per_epoch == 0:
        val_loss = eval_loss()

    bar.set_postfix(train_loss=train_loss, validation_loss=val_loss)

Huzzah! Now, take this knowledge I have bequethed unto you and go forth. Use it to create whatever it is that you desire. Let it become more intelligent than you.