The Transformer

D

LeNet works like a detective who is specialized in pictures. It collects clues using corr2d layers, and then makes decisions using linear layers.
Now we build another detective, this time specialized in text.

D

Let's call it the LittleTransformer.
It follows a similar process: collecting clues first and making decisions later.
Its new clue-collecting technique is called attention.

D

Suppose we are reading the sentence
the animal didn't cross the street because it was too tired.
When we read it, we want to know which earlier word it refers to.
Attention computes a score for every earlier word. A high score means: this one looks relevant. A low score means: probably ignore it.

D

To compare tokens, we first turn each one into a vector of numbers called an embedding.
We write one embedding as Tensor[float][[d]].
d is the length of that vector: a larger d gives each token more room to carry information.

D

A token is one unit that the model reads.
Sometimes one token is a whole word. Some other times it is part of a word or a piece of punctuation.
Different tokens may learn similar embeddings, if these tokens share similar roles.

D

Yes.
Each embedding prepares three views of itself: query, key, and value.
Query expresses what this position is looking for.
Key expresses what this position offers.
Value carries the information that can be passed along.
Attention compares queries with keys, then uses the scores to mix the values.

class AttentionParams[n, d]:
    q: LinearParams[float, n, d]
    k: LinearParams[float, n, d]
    v: LinearParams[float, n, d]

@op
def attention[l, d, n](
    x: Tensor[float][[l, d]],
    p: AttentionParams[n, d]
) -> Tensor[float][[l, n]]:
    q = linear(x, p.q)
    k = linear(x, p.k)
    v = linear(x, p.v)
    weights = (q @ k.permute(1, 0)) / n.sqrt()
    ...
    weights = softmax(weights)
    return weights @ v

k.permute(1, 0) swaps the two dimensions of k.
Then q @ k.permute(1, 0) computes one dot product for each query-key pair, giving the raw attention scores.
We keep the weights stable by dividing it with the square root of n.
softmax turns these scores into non-negative weights that add up to 1.0, so we can use them to mix the values.

D

Channels are a useful analogy, and we will return to it soon.
For our transformer, each position should focus on earlier words.*
So we add causal_mask to guide attention toward the left side of the sequence.

@op
def causal_mask[l](scores: Tensor[float][[l, l]]) -> Tensor[float][[l, l]]:
    blocked = -1e9
    return [
        [
            scores[row][col] if col <= row else blocked
            for col in iota(l)
        ]
        for row in iota(l)
    ]

causal_mask keeps the past and replaces the future with small numbers.
Then softmax gives almost zero weight to those future positions.

D

Now we revisit the analogy of channels.
There are different attentions, just like different patterns in pictures. Some tracks the grammatical structure, some focuses on subject-verb agreement, and some pays attention to nearby conexts.
A model is more effective if we seperate these attentions in different tensors, just like LeNet's output channels learn different patterns.

D

Exactly. This extra dimension is h, for multi-heads attention.
For each head, we store one set of query, key, and value weights, then merge all head outputs back together.

class MultiHeadsAttentionParams[h, d, n]:
    qw: Tensor[float][[h, d, n]]
    kw: Tensor[float][[h, d, n]]
    vw: Tensor[float][[h, d, n]]
    proj: Tensor[float][[h * n, d]]

@op
def multi_heads_attention[l, d, n, h](
    x: Tensor[float][[l, d]],
    p: MultiHeadsAttentionParams[h, d, n]
) -> Tensor[float][[l, d]]:
    q = x @ p.qw
    k = x @ p.kw
    v = x @ p.vw
    weights = (q @ k.permute(0, 2, 1)) / n.sqrt()
    weights = softmax(causal_mask(weights))
    weights = weights @ v
    merged = weights.permute(1, 0, 2).reshape([l, h * n])
    return merged @ p.proj

D

The decision-maker is called feed_forward. It consists of two linear layers.

class FeedForwardParams[d]:
    linear1: LinearParams[float, 4 * d, d]
    linear2: LinearParams[float, d, 4 * d]

@op
def feed_forward[l, d](
    x: Tensor[float][[l, d]],
    p: FeedForwardParams[d]
) -> Tensor[float][[l, d]]:
    x = larger(linear(x, p.linear1), 0)
    return linear(x, p.linear2)

The first layer expands the parameters, gives more space to untangle complex data. Then the second layer compresses the learned information to the original size.

D

Yes, 4 works well in practice.
Now we combine the clue-collector and the decision-maker into one block.

class BlockParams[n, d, h]:
    multi_heads: MultiHeadsAttentionParams[float, n, d, h]
    feed_forward: FeedForwardParams[float, d]
    ...

@op
def block[l, d, h, n](
    x: Tensor[float][[l, d]],
    p: BlockParams[n, d, h]
) -> Tensor[float][[l, d]]:
    ...
    x = x + multi_heads_attention(x, p.multi_heads)
    ...
    x = x + feed_forward(x, p.feed_forward)
    return x

D

We also want x to stay well scaled, since practical models often stack many blocks in sequence.

class LayerNormParams[d]:
    weight: Tensor[float][[d]]
    bias: Tensor[float][[d]]

@op
def layer_norm[l, d](
    x: Tensor[float][[l, d]],
    p: LayerNormParams[float, d]
) -> Tensor[float][[l, d]]:
    eps = 1e-5
    means = x.avg(1).reshape([l, 1])
    variances = ((x - means) ** 2).avg(1).reshape([l, 1])
    return ((x - means) / (variances + eps).sqrt()) * p.weight + p.bias

layer_norm stabilizes x. Let's insert it before collecting clues or making decisions.

D

Good. Now we can draft predict.
Our LittleTransformer speaks 65 distinct tokens, and each embedding has length 32.
The input is a sequence of token IDs.
The output is a table of scores, one row per position and one column per token in the vocabulary.

vocab_size = 65
dim = 32

class ModelParams[l, n, h]:
    ...
    block1: BlockParams[n, dim, h]
    block2: BlockParams[n, dim, h]
    layer_norm: LayerNormParams[dim]
    linear: LinearParams[vocab_size, dim]

class LittleTransformer(Model):
    def predict[l, n, h](
        indices: Tensor[int][[l]],
        p: ModelParams[l, n, h]
    ) -> Tensor[float][[l, vocab_size]]:
        ...
        x = block(x, p.block1)
        x = block(x, p.block2)
        x = layer_norm(x, p.layer_norm)
        logits = linear(x, p.linear)
        return logits
    ...

D

Exactly. The next step is to turn token IDs into embeddings with emb.
Each entry in indices points to one row of the embedding table.
emb gathers the embeddings stored in the corresponding indices.

@op
def emb[n, v, d](
    indices: Tensor[int][[n]],
    embedding: Tensor[float][[v, d]]
) -> Tensor[float][[n, d]]:
    return [
        embedding[idx]
        for idx in indices
    ]

Next we give two embeddings to predict.

D

One for meanings and one for positions. These two embeddings are complementary. Here we simply add them together.

class ModelParams[l, n, h]:
    token_emb: Tensor[float][[vocab_size, dim]]
    position_emb: Tensor[float][[l, dim]]
    block1: BlockParams[float, n, dim, h]
    block2: BlockParams[float, n, dim, h]
    layer_norm: LayerNormParams[float, dim]
    linear: LinearParams[float, vocab_size, dim]

class LittleTransformer(Model):
    def predict[l, n, h](
        indices: Tensor[int][[l]],
        p: ModelParams[l, n, h]
    ) -> Tensor[float][[l, vocab_size]]:
        token_emb = emb(indices, p.token_emb)
        position_emb = emb(iota(l), p.position_emb)
        x = token_emb + position_emb
        x = block(x, p.block1)
        x = block(x, p.block2)
        x = layer_norm(x, p.layer_norm)
        logits = linear(x, p.linear)
        return logits
    ...

Now predict is complete. Let's define loss. What's the type of ys_pred?

D

Right, here's our definition of loss. For each input, ys_pred contains the scores for all tokens, while ys has the index of the correct one. Then cross_entropy tells how wrong the predictions are.

class LittleTransformer(Model):
    ...
    def loss[b, l](
        ys_pred: Tensor[float][[b, l, vocab_size]],
        ys: Tensor[int][[b, l]]
    ) -> float:
        ys_pred = ys_pred.reshape([b * l, vocab_size])
        ys = ys.reshape([b * l])
        return cross_entropy(ys_pred, ys)
    ...

D

We use our old RMS update to stabilize gradients. Now our LittleTransformer is complete!

class LittleTransformer(Model):
    def predict[l, n, h](
        indices: Tensor[float][[l]],
        p: ModelParams[l, n, h]
    ) -> Tensor[float][[l, vocab_size]]:
        token_emb = emb(indices, p.token_emb)
        position_emb = emb(iota(l), p.position_emb)
        x = token_emb + position_emb
        x = block(x, p.block1)
        x = block(x, p.block2)
        x = layer_norm(x, p.layer_norm)
        logits = linear(x, p.linear)
        return logits

    def loss[b, l](
        ys_pred: Tensor[float][[b, l, vocab_size]],
        ys: Tensor[int][[b, l]]
    ) -> float:
        ys_pred = ys_pred.reshape([b * l, vocab_size])
        ys = ys.reshape([b * l])
        return cross_entropy(ys_pred, ys)

    def update(
        self: ,
        P: Tuple[float, float, float],
        g: float
    ) -> Tuple[float, float, float]:
        r = smooth(beta, P[2], g ** 2)
        alpha1 = 0.01 / (sqrt(r) + epsilon)
        v = smooth(mu, P[1], g)
        return (P[0] - (alpha1 * v), v, r)

    def inflate(self: , p: float) -> Tuple[float, float, float]:
        return (p, 0, 0)

    def deflate(self: , P: Tuple[float, float, float]) -> float:
        return P[0]