"Really Great Transformer Article"
[catchy subtitle I don't actually believe and would never use because of journalistic integrity… although, I'm not a journalist… so…]
Exposing the House of Cards
sorry, i don't have a gimmick this time; i just wrote from the heart. i know you were all counting on me, and you depended on my hilarious and original jokes, but you need to move on. i'm sorry i let you down.
- I'm assuming you have a basic understanding of how transformers work (or at least that you think you do), and a pretty decent understanding of PyTorch.
- If you don't, then go read about it or something and come back.
Let's start our journey with the dataset (which is often the most fun part of training a model anyway (fight me about it)). We'll build a simple mock-BPE tokenizer without merges.
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>"
] # pretty normal subword tokenization
self.tokens = sorted(tokens[:-1], key=len, reverse=True)
self.vocabulary = {t: i for i, t in enumerate(self.tokens)} | {
"<pad>": len(self.tokens)
} # just set pad to the last token instead of trying to find it
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() # short text examples separated by newlines (you can steal this from huggingface; they won't mind)
tokenizer = Tokenizer(examples)
print(
",".join(
[
tokenizer.decode([token])
for token in tokenizer.encode("trees lose leaves in autumn")
]
)
)
Assuming that output something reasonable:
Great! Next, let's set up some hyperparameters.
cheat_help = max(
[len(example) for example in examples]
) # it's called that because having a max sequence length is for cowards, and is basically cheating (the only thing that really needs it is the position embedding, and that doesn't need it)
epochs = 50
batch_size = 8
lr = 3e-4
hidden = 512
emb_dim = 128
Okay, crunch time. Well, I mean-- the start of that. The self attention head applies a query matrix to all tokens, a key matrix to all tokens, and value matrix to all tokens. The key matrix is trained to output similar embeddings for "answers" to specific query tokens, and query tokens are trained to output embeddings that stimulate certain types of keys. (similarity is measured with euclidian distance.. just kidding, it's the scaled dot product. of course.) The array of similarity scores is then masked, softmaxed, and then passed to the value matrix which does exactly what it sounds like. It assigns value to certain tokens more than others (or at least, that's what I think it does. There's really no way of knowing).
import torch
import torch.nn as nn
class AttentionHead(nn.Module):
def __init__(self, embedding_dim, head_dim):
super().__init__()
# no bias, as it doesn't help attention enough to justify the extra cost
self.q = nn.Linear(embedding_dim, head_dim, bias=False)
self.k = nn.Linear(embedding_dim, head_dim, bias=False)
self.v = nn.Linear(embedding_dim, head_dim, bias=False)
def forward(self, embeddings, cache=None):
query = self.q(embeddings)
key = self.k(embeddings)
value = self.v(embeddings)
# if a cache is provided (see later), combine previous keys/values with the current ones along the sequence
if cache is not None:
cached_key, cached_value = cache
key = torch.cat([cached_key, key], dim=-2)
value = torch.cat([cached_value, value], dim=-2)
scores = (
query @ key.transpose(-2, -1) * (query.size(-1) ** -0.5)
) # compute scaled scores (yay dot product!)
offset = key.size(-2) - query.size(
-2
) # this is used to align the causal mask for the cache
return torch.softmax(
scores.masked_fill(
torch.triu(
torch.ones(scores.size(-2), scores.size(-1), device=scores.device),
diagonal=offset + 1,
).bool(),
-1e9, # can't be negative infinity because of floating point imprecision or something idrk
),
dim=-1,
) @ value, (key, value)
The Transformer block is a separate module for a really good reason that I haven't thought of yet. Anyway, it takes-- you guessed it-- an embedding. It gives the embedding to num_heads attention heads, and then concatenates the outputs of the heads into one really long array, which is passed to an out/projection matrix (which has a bias for some reason) which brings it back to the normal embedding shape (thank gosh, I was so scared for second there). The normalized embedding is passed to an mlp for some reason, and normalized again (in that order (no need to look at the code), because that's easier to understand, and I would never sacrifice readability for efficiency).
class TransformerBlock(nn.Module):
def __init__(self, embedding_dim, hidden_size, num_heads=2):
super().__init__()
head_dim = embedding_dim // num_heads
self.heads = nn.ModuleList(
[AttentionHead(embedding_dim, head_dim) for _ in range(num_heads)]
)
self.proj = nn.Linear(
head_dim * num_heads,
embedding_dim,
) # this is the thing that combines the different attention heads without getting rid of too much information (hopefully)
self.mlp = nn.Sequential(
nn.Linear(embedding_dim, hidden_size),
nn.GELU(),
nn.Linear(hidden_size, embedding_dim),
) # this is applied to every token separately after it goes through attention to add nonlinear transformations and whatnot (hence the dual layer mlp) to the scores (i don't really get why you wouldn't just change the input to `head_dim` and then add the embeddings together instead of using the out matrix. i assume it's because of money)
self.ln1 = nn.LayerNorm(
embedding_dim
) # external normalization is applied because transformer blocks are really unstable
self.ln2 = nn.LayerNorm(embedding_dim)
def forward(self, embeddings, cache=None):
head_outs, new_caches = [], []
for i, head in enumerate(self.heads):
head_cache = cache[i] if cache is not None else None
out, new_cache = head(self.ln1(embeddings), head_cache)
head_outs.append(out) # python list comprehension is so fast (no)
new_caches.append(new_cache)
attn_out = embeddings + self.proj(torch.cat(head_outs, dim=-1))
return (
attn_out + self.mlp(self.ln2(attn_out)),
new_caches,
)
We finally actually see more than one token at a time! In the aptly named Transformer module (trust me, I saw it turn into a car earlier), we actually do the embedding, and precict probability logits in the most underwhelming way possible.
class Transformer(nn.Module):
def __init__(self, embedding_dim, hidden_size, vocab):
super().__init__()
self.embed = nn.Embedding(
vocab, embedding_dim
) # learned embedding from one hot, pretty much. it turns a list of integers (say 20, 56) and then makes a tensor where only index 20 is 1 (the rest are zero) in the first column, and 56 for the second column (shape (2) -> (2, 1250)), and then uses a linear layer to bring the vocab size down to emb_dim while helping with some of the bias caused by tokens mapping to numbers
self.pos_embed = nn.Embedding(
cheat_help, embedding_dim
) # same idea, but this one takes each token's index and the max sequence length instead of the token id and vocab size
self.blocks = nn.Sequential(
*[TransformerBlock(embedding_dim, hidden_size) for _ in range(3)]
)
self.lm_head = nn.Linear(
embedding_dim, vocab
) # this takes the last token (which hopefully has a ton of information from those blocks), and predicts the next token (technically it works on all tokens in the batch, but we obviously already know the previous tokens, so this only applies in training)
self.lm_head.weight = (
self.embed.weight
) # weight tying to help with generalization (although if you really think about it, this can only hurt long term, because the embeddings that are mapped *to* are supposed to have a lot less specificity than the embeddings being mapped *from*)
def forward(self, tokens, cache=None):
B, T = tokens.shape
offset = (
cache[0][0][0].size(-2) if cache is not None else 0
) # these are the offsets from before
embeddings = self.embed(tokens) + self.pos_embed(
torch.arange(offset, offset + T, device=tokens.device)
.unsqueeze(0)
.expand(B, T)
) # embed the inputs so there's a mesh for the transformer blocks to shape
new_caches = []
for i, block in enumerate(self.blocks):
block_cache = cache[i] if cache is not None else None
embeddings, block_new_caches = block(
embeddings, block_cache
) # add as much contextual information as possible to each token embedding, so the lm head's job is as easy as it can be
new_caches.append(block_new_caches) # again, so fast
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
) # literally just runs the transformer to generate one token at a time by predicting on the last token, and then setting the last token to the next token and doing it again
return tokens
model = Transformer(
emb_dim, hidden, tokenizer.vocab_size
).cuda() # right now the model is trained to predict "the next most likely token" because it's trained on a dataset that conditions it to assign high probabilities to probable outputs (duh), but with rlhf or whatever, the goal shifts to be "the token that is most likely to be recieved well by humans when put together with a bunch of other tokens also generated by it"
model.requires_grad_()
print(
sum(p.numel() for p in model.parameters())
) # it's not a very big model or a very big dataset (i assume), so don't expect it to perform well
Training time! The part you probably care about least if you made it this far. Anyway, we just do gradient descent from the cross entropy loss of the model's predicted logits.
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").read().splitlines()]
to_batch = lambda encoded: torch.stack(
[
torch.tensor(example + [tokenizer.pad_id] * (cheat_help - len(example))) for example in encoded
] # pad the examples so all the examples are the same shape
).cuda() # pre-batch the tokenized examples for speed reasons
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, # don't assign any loss to predictions where the correct response is pad
).item()
model.train()
return loss # predict on a test set that the model has never "seen" before (by seen, i mean examples that the model hasn't predicted an output and had its weights updated in response to its peformance)
steps_per_epoch = len(encoded_train) // batch_size
bar = tqdm(
range(epochs * steps_per_epoch)
) # tqdm is great, as long as you don't care about being able to read your logs later
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
) # no, i'm not using a scheduler, but the distance in weight space is determined by the loss anyway (it flips if your model is bad enough, but then you're kinda screwed anyway)
model.zero_grad()
train_loss = loss.item()
if (
step % steps_per_epoch == 0
): # only check the validation loss once per epoch, because any time or data that isn't used to train the model is a waste (which is why i don't normally use a test set)
val_loss = eval_loss()
bar.set_postfix(train_loss=train_loss, validation_loss=val_loss)
Hokay, presumeably the model didn't learn langauge structure from the dataset you gave it, but I'm certain that throwing more compute at it will solve the problem. Maybe get regurgitated data from another model, but make sure it's a closed source one so you don't get the really helpful logit distribution, and instead get sparse data points that are hard for your model to learn.
They say that a small amount of good data is better than a ton of bad data, but "good" and "bad" is defined by what the model can learn, and isn't really determinable by you or me.