In the previous chapter, we built a runnable basic skeleton of the GPT model, and also pointed out its problems. In this chapter, we start to solve the first one: there is no position information.
The embedding layer can currently only tell the model what each token is, not where each token is. For the model, “the mouse loves rice” and “rice loves the mouse” are the same sentence: wherever a token appears, the same token always gets exactly the same word vector.
But word order is the soul of language. The same word at the beginning or at the end of a sentence can mean something completely different. Compare: an RNN processes tokens one by one in order, so position information is naturally contained in the processing order and needs no extra encoding. A GPT model, in contrast, processes all the tokens in parallel in one go, which is efficient, but at the cost of losing the word order. So we must add the word order information to the encoding explicitly.
Position encoding¶
How do we represent word order? The most direct idea is to give each position a number in turn: the 1st token is 1, the 2nd token is 2, and so on. We call this position encoding.
Does it look familiar? Quite like the index encoding of tokens? The difference is that the index encoding is the position of a token in the vocabulary, while the position encoding is the position of a token in the text.
Position embedding¶
Obviously, we can borrow the idea of token embedding: if we can replace the index encoding with a trainable word vector, why not the position encoding?
Position embedding replaces the position encoding with a learnable position vector. Providing both the word vector and the position vector of each token to the GPT model solves the problem that a GPT model naturally ignores word order.
import numpy as np
from src.core import *np.random.seed(42)class CharDataset(Dataset):
def __init__(self, filename, batch_size=1, context_size=64, stride=None, split=0.9):
self.filename = filename
self.context_size = context_size
self.stride = stride if stride is not None else context_size // 2
self.split = split
super().__init__(batch_size)
def load(self):
with open(self.filename, encoding="utf-8") as f:
text = f.read()
self.vocab = sorted(set(text))
self.vocab_size = len(self.vocab)
self.stoi = {ch: i for i, ch in enumerate(self.vocab)}
self.itos = {i: ch for i, ch in enumerate(self.vocab)}
self.tokens = self.encode(text)
split = int(len(self.tokens) * self.split)
self.train_data = self._pack(self.tokens[:split])
self.test_data = self._pack(self.tokens[split:])
def _pack(self, tokens):
x, y = [], []
for i in range(0, len(tokens) - self.context_size - 1, self.stride):
x.append(tokens[i: i + self.context_size])
y.append(tokens[i + 1: i + self.context_size + 1])
return x, y
def encode(self, symbols):
return [self.stoi[s] for s in symbols]
def decode(self, tokens):
return "".join(self.itos[t] for t in tokens)class GELU(Layer):
def __init__(self):
super().__init__()
self.c = np.sqrt(2.0 / np.pi)
def forward(self, x: Tensor):
tanh = np.tanh(self.c * (x.data + 0.044715 * x.data ** 3))
a = Tensor(0.5 * x.data * (1.0 + tanh))
def gradient_fn():
grad = 0.5 * (1.0 + tanh) + 0.5 * x.data * (1.0 - tanh ** 2) * self.c * (1.0 + 3.0 * 0.044715 * x.data ** 2)
x.grad += a.grad * grad
return a.attach(gradient_fn, parents=[x])class GPTLoss(Loss):
def loss(self, p: Tensor, y: Tensor):
exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
softmax = exp / np.sum(exp, axis=-1, keepdims=True)
flat_softmax = softmax.reshape(-1, softmax.shape[-1])
flat_y = y.data.reshape(-1).astype(np.int64)
n = len(flat_y)
rows = np.arange(n)
log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
ce = Tensor(0 - np.sum(log) / n)
def gradient_fn():
flat_grad = flat_softmax.copy()
flat_grad[rows, flat_y] -= 1
p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n
return ce.attach(gradient_fn, parents=[p])Model¶
GPT embedding layer¶
We add the position embedding positional_embedding to the GPT embedding layer, forming the combination “token embedding + position embedding”:
embedding(token embedding): looks up the table with the index encodings of the tokens and produces word vectors, with the shape(batch size, context size, embedding size);positional_embedding(position embedding): looks up the table with the position encodings of the tokens and produces position vectors, with the shape(context size, embedding size).
The embedding layer needs a new initialization parameter: the context size context_size, which determines the size of the weight matrix of the position embedding.
The final output is “word vector + position vector”.
Note that this is an addition of two tensors. It uses the tensor addition that we extended in our training framework, which keeps the computational graph and backpropagation going, instead of simply adding the numbers of two vectors together.
Also, the word vector (3 dimensions) and the position vector (2 dimensions) have different shapes, yet they can be added directly. This uses the broadcasting mechanism supported by the tensor: the position vector is automatically expanded to the same dimensions as the word vector, and then they are added.
class GPTEmbedding(Composite):
def __init__(self, vocab_size, context_size, embedding_size):
# Token embedding layer
self.embedding = Embedding(vocab_size, embedding_size)
# Position embedding layer
self.positional_embedding = Embedding(context_size, embedding_size)
super().__init__([self.embedding,
self.positional_embedding])
def forward(self, x: Tensor):
token = self.embedding(x)
position = self.positional_embedding(Tensor(range(x.shape[1])))
return token + positionGPT feed-forward network layer¶
class GPTFeedForward(Composite):
def __init__(self, embedding_size):
self.input = Linear(embedding_size, embedding_size * 4)
self.gelu = GELU()
self.output = Linear(embedding_size * 4, embedding_size)
super().__init__([self.input,
self.gelu,
self.output])
def forward(self, x: Tensor):
h = self.gelu(self.input(x))
return self.output(h)GPT Transformer block¶
class GPTTransformer(Composite):
def __init__(self, embedding_size):
self.feed_forward = GPTFeedForward(embedding_size)
super().__init__([self.feed_forward])
def forward(self, x: Tensor):
return self.feed_forward(x)GPT output layer¶
class GPTOutput(Composite):
def __init__(self, embedding_size, vocab_size):
self.output = Linear(embedding_size, vocab_size)
super().__init__([self.output])
def forward(self, x: Tensor):
return self.output(x)GPT model¶
Creating the GPT model also needs one extra initialization parameter: the context size context_size, used to create the embedding layer.
class GPT(Composite):
def __init__(self, vocab_size, context_size, embedding_size, blocks):
self.embedding = GPTEmbedding(vocab_size, context_size, embedding_size)
self.transformers = [GPTTransformer(embedding_size) for _ in range(blocks)]
self.output = GPTOutput(embedding_size, vocab_size)
super().__init__([self.embedding] + self.transformers + [self.output])
def forward(self, x: Tensor):
x = self.embedding(x)
for layer in self.transformers:
x = layer(x)
return self.output(x)GPT trainer¶
class GPTTrainer(Trainer):
def train(self, dataset, epochs, scheduler=None):
dataset.train()
self.model.train()
steps = 0
for epoch in range(epochs):
for i in range(len(dataset)):
if scheduler is not None:
self.optimizer.lr = scheduler.step(steps)
feature, label = dataset[i]
self.optimizer.zero_grad()
prediction = self.model(feature)
loss = self.loss_fn(prediction, label)
loss.backward()
self.optimizer.clip_grad_norm()
self.optimizer.step()
steps += 1
def test(self, dataset):
dataset.eval()
self.model.eval()
predictions = []
total_loss = 0.0
with Tensor.no_grad():
for i in range(len(dataset)):
feature, label = dataset[i]
prediction = self.model(feature)
loss = self.loss_fn(prediction, label)
predictions.append(prediction)
total_loss += float(loss.data)
loss = total_loss / len(dataset)
dataset.train()
self.model.train()
return predictions, loss
def generate(self, dataset, prompt, steps=256):
self.model.eval()
tokens = dataset.encode(prompt)
with Tensor.no_grad():
for _ in range(steps):
feature = Tensor([tokens[-dataset.context_size:]])
logits = self.model(feature)
last_logits = logits.data[0, -1]
exp = np.exp(last_logits - np.max(last_logits))
probs = exp / np.sum(exp)
token = np.random.choice(len(probs), p=probs)
tokens.append(token)
return dataset.decode(tokens)LEARNING_RATE = 0.001Batch size¶
BATCH_SIZE = 4Number of epochs¶
EPOCHS = 2Context size¶
CONTEXT_SIZE = 32Embedding size¶
EMBEDDING_SIZE = 64Number of Transformer blocks¶
BLOCKS = 2Modeling¶
When creating the model, we need to pass CONTEXT_SIZE in as an additional initialization parameter.
DATA_FILE = "../../../../data/tinyshakespeare.txt"
dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, BLOCKS)
loss_fn = GPTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = GPTTrainer(model, loss_fn, optimizer)Training¶
scheduler = WarmupCosineScheduler(LEARNING_RATE, EPOCHS * len(dataset), 100, LEARNING_RATE / 10)
trainer.train(dataset, EPOCHS, scheduler)Evaluation¶
prediction, loss = trainer.test(dataset)
print(f'prediction:\t{len(prediction)} steps, each {prediction[0].data.shape}\nloss:\t{loss}')prediction: 1742 steps, each (4, 32, 65)
loss: 2.462666325331008
Inference¶
print(trainer.generate(dataset, prompt="ROMEO:"))ROMEO:
PELUCENone iat werachald l d:
ANMyerther, itik he
TIO:
Ife.
DUKayo in, y dld s m f s pllorislds hine try be ou 'te noungeltwiomavear whe
Opld migrarthemaxhayorthe t; llleakash lu aru g d, harenout watllllicinchiowe f s my, boive
I he y a fod, gheltl I her
In this chapter, through position embedding, we give the model “word vector + position vector”, so that it can tell both “what it is” and “where it is”.
With this, we have completed the input representation part of the GPT model.
But we cannot see obvious progress in the inference result, because a more fundamental problem remains unsolved: tokens “do not talk to each other”. The prediction at each position depends only on the current token and its position, and the model has no idea what else is in the context.
In the next chapter, we will introduce the most central mechanism of the GPT model, self-attention, which lets tokens truly “see each other”.
Exercises¶
Try to count the parameters of the GPT model, including the weight matrices of the token embedding and the position embedding, and the weights and biases of the linear layers in the other modules. Which hyperparameters is the count mainly related to?