Position and masking

What you will be able to do

Give a transformer the order of the words, and control what each token may read. Then you can explain the two masks in every transformer, and the difference between models that read a whole text and models that write one.

NotebookDo this lesson in a notebookColab · Kaggle · GitHub

The notebook runs the shuffle test and both masks on real tensors. It trains its own tiny transformer for 8 epochs first (one to three minutes), so its scores differ a little from this page. About 35 minutes.

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. They need a free Google or Kaggle account. Stuck? See the solution notebook.

Before you startWhat to know and to have readyAbout 50 minutes

What you need to know

• The transformer block and TinyTransformer, from Transformer blocks.
• Attention weights as a table, one row for each token that reads, from Attention step by step.
• Why word order matters for some tickets, from Why word order and context matter.

What you need on your computer (only for the local route)

• The folder ticket-router with its .venv, from Set up PyTorch: Route 2, and the file tiny_transformer.pt that Transformer blocks: Train it saves. If the file is missing, train again with seed 0 and save it.

Tested with Python 3.12 and PyTorch 2.14.1 on the CPU (macOS).

A recipe on loose cards

Write each step of a recipe on a card, and drop the cards. Someone who picks them up has every step, but not the order: “bake for 20 minutes” may now come before “mix the flour”. The words are all there, but the meaning depends on where each card was.

The problem: attention does not see the order

Why word order and context matter showed a problem of the bag of words. “Damaged on arrival, not after use” and “damaged after use, not on arrival” get the same input. Does attention do better? Run the test with two made-up sentences that have the same 7 tokens in a different order. One attention layer reads each sentence, and the average of its outputs is a summary.

positions.py

import torch
from torch import nn

words = ["damaged", "on", "arrival", ",", "not", "after", "use"]
a = [0, 1, 2, 3, 4, 5, 6]          # damaged on arrival , not after use
b = [0, 5, 6, 3, 4, 1, 2]          # damaged after use , not on arrival
torch.manual_seed(0)
table = nn.Embedding(len(words), 4)                  # one vector of 4 numbers per word
attention = nn.MultiheadAttention(4, 1, batch_first=True)

def read(order):
    x = table(torch.tensor([order]))                 # (1, 7, 4)
    out, _ = attention(x, x, x)                      # (1, 7, 4)
    return out.mean(dim=1).detach()                  # one summary per sentence: (1, 4)

print(read(a))
print(read(b))
print(torch.allclose(read(a), read(b), atol=1e-6))

The output:

tensor([[-0.1657, -0.0738,  0.3466, -0.1725]])
tensor([[-0.1657, -0.0738,  0.3466, -0.1725]])
True

The two opposite sentences get the same summary. Look inside, at the output for “use”, which is at position 6 in sentence a and at position 2 in sentence b:

x_a, x_b = table(torch.tensor([a])), table(torch.tensor([b]))
out_a = attention(x_a, x_a, x_a)[0].detach()
out_b = attention(x_b, x_b, x_b)[0].detach()
print("'use' in a, position 6:", out_a[0, 6])
print("'use' in b, position 2:", out_b[0, 2])

The output:

'use' in a, position 6: tensor([-0.0194, -0.0040,  0.2147, -0.0818])
'use' in b, position 2: tensor([-0.0194, -0.0040,  0.2147, -0.0818])

The same numbers. Attention compares each query with every key, whatever their places are. So “use” reads the same set of words in both sentences, gets the same weights, and gives the same output. Shuffle the input, and attention only shuffles the output in the same way. The order never enters the calculation.

The idea: a vector for each position

Give each place in the text a number pattern of its own, and add it to the token’s vector. Then “use” at position 6 and “use” at position 2 are different inputs. This is a position embedding: a learned table with one row for each position, used exactly like the token embedding table. Its general name is positional encoding: any way of putting the position into the vectors.

torch.manual_seed(1)
position = nn.Embedding(7, 4)                        # one vector of 4 numbers per position
steps = torch.arange(7)                              # positions 0, 1, ..., 6

def read_in_order(order):
    x = table(torch.tensor([order])) + position(steps)   # word vector + position vector
    out, _ = attention(x, x, x)
    return out.mean(dim=1).detach()

print(read_in_order(a))
print(read_in_order(b))
print(torch.allclose(read_in_order(a), read_in_order(b), atol=1e-6))

The output:

tensor([[-0.2461, -0.0334,  0.7288, -0.4205]])
tensor([[-0.1802,  0.0059,  0.7413, -0.4107]])
False

Now the two sentences differ. These position vectors are still random, so the difference means nothing yet; training teaches the table what each position should add.

Positions in the tiny transformer

TinyTransformer has the table self.position = nn.Embedding(120, 64): 120 positions, because MAX_TOKENS cuts tickets at 120 tokens. Load the trained model from Transformer blocks, and look at its first 4 numbers for “parcel” at two places:

from ticketnet import (TinyTransformer, Vocab, labels, load_split, token_ids, evaluate, tokenize,
                       TEAMS, PAD)

train_rows, valid_rows = load_split("data/train.csv"), load_split("data/valid.csv")
vocab = Vocab([row["text"] for row in train_rows])
model = TinyTransformer(len(vocab))
model.load_state_dict(torch.load("tiny_transformer.pt"))
model.eval()
print(tuple(model.position.weight.shape), model.position.weight.numel(), "numbers")

parcel = vocab.stoi["parcel"]
print("token 'parcel'     ", model.embed.weight[parcel, :4].detach().numpy().round(2))
print("position 0         ", model.position.weight[0, :4].detach().numpy().round(2))
print("position 5         ", model.position.weight[5, :4].detach().numpy().round(2))
print("'parcel' at 0      ", (model.embed.weight[parcel, :4] + model.position.weight[0, :4]).detach().numpy().round(2))
print("'parcel' at 5      ", (model.embed.weight[parcel, :4] + model.position.weight[5, :4]).detach().numpy().round(2))

The output:

(120, 64) 7680 numbers
token 'parcel'      [-1.35 -0.68 -1.01 -0.75]
position 0          [ 0.13 -0.41 -0.04  1.04]
position 5          [ 1.76  1.14  0.25 -0.61]
'parcel' at 0       [-1.21 -1.09 -1.06  0.29]
'parcel' at 5       [ 0.42  0.46 -0.77 -1.36]

torch.load() reads the saved numbers, and load_state_dict() puts them into a new model of the same shape. Check the first number by hand: −1.35 + 0.13 = −1.22, shown as −1.21 because the page rounds each number separately. “parcel” at position 0 and at position 5 now enter the blocks as different vectors.

The shuffle test on a real ticket

Shuffle the 92 tokens of ticket T-64697, a payment ticket, and compare the trained model’s 5 scores:

import random

ticket = valid_rows[0]["text"]
tokens = tokenize(ticket)
random.seed(0)
shuffled = random.sample(tokens, len(tokens))
print(" ".join(tokens[:10]), "...")
print(" ".join(shuffled[:10]), "...")

def scores(text):
    ids = torch.tensor([vocab.encode(text)])
    with torch.no_grad():
        return [round(v, 3) for v in model(ids)[0].tolist()]

print("in order:", scores(ticket))
print("shuffled:", scores(" ".join(shuffled)))

The output:

i am writing regarding order lk - 714258 . the ...
i refund lk delivery process i receiving week so as ...
in order: [-3.104, -3.135, 10.019, -1.814, -1.48]
shuffled: [-3.283, -1.909, 10.309, -2.007, -2.905]

The scores change, so the model notices the order. Now set the position table to zero, so that the model has no positions:

saved = model.position.weight.data.clone()
model.position.weight.data.zero_()                   # remove the positions
print("in order:", scores(ticket))
print("shuffled:", scores(" ".join(shuffled)))
model.position.weight.data = saved                   # put them back

The output:

in order: [-3.571, -2.937, 10.644, -2.797, 0.005]
shuffled: [-3.571, -2.937, 10.644, -2.797, 0.005]

Without positions, the whole trained transformer gives exactly the same scores for the ticket and for word salad. The positions are the only place where order enters.

How much does the trained model use the order? With every validation ticket shuffled, it is right on 93.3%, against 94.3% in order: one point less. The bag-of-words network gets 98.0% on the shuffled tickets, because it cannot notice. On these tickets, the words decide the team, so the transformer uses order only a little. It still pays the full cost of learning it.

Other ways to encode positions

A learned table is the simplest way, and it has a limit. It has rows for 120 positions only, and positions that are rare in training stay close to random. Two other ways are common:

• Sinusoidal positions, from the first transformer of 2017: a fixed pattern of waves with different lengths for each position. A formula computes it, so nothing is learned.
• Rotary positions (RoPE): instead of adding a vector, they turn each query and key by an angle that grows with the position. So attention sees how far apart two tokens are. SmolLM2-135M, in Module 7, uses rotary positions.

Go deeper

The sinusoidal formula, and why a move by a fixed number of positions is a rotation, are worked out in The original transformer, step by step, in the optional mathematics module.

The padding mask: do not read the blank lines

The familiar situation. A form has ten answer lines, and a short answer leaves seven of them blank. Nobody reads the blank lines as part of the answer.

The problem. A batch is one tensor, so every ticket in it needs the same length. token_ids() fills the short tickets with <pad>, ID 0, up to the longest ticket in the batch. Take a short delivery ticket and put it in a batch with the longest validation ticket:

short = next(row for row in valid_rows if row["style"] == "short")
long = max(valid_rows, key=lambda row: len(vocab.encode(row["text"])))
print(short["ticket_id"], short["team"], "|", short["text"])
pair = token_ids([short, long], vocab)
print(tuple(pair.shape))
print(pair[0, :12].tolist(), "...", pair[0, -4:].tolist())
print("padding in the short ticket:", int((pair[0] == PAD).sum()))

The output:

T-64715 delivery | Ugh, need to change trampoline delivery slot. Order LK-640933
(2, 120)
[2653, 8, 1664, 2591, 618, 2620, 832, 2338, 10, 1749, 1486, 9] ... [0, 0, 0, 0]
padding in the short ticket: 107

The short ticket has 13 real tokens and 107 padding tokens. Without a mask, attention would read the padding like words. Run the model’s steps by hand, once with the mask and once without:

def team_scores(ids, use_mask):
    with torch.no_grad():
        h = model.embed(ids) + model.position(torch.arange(ids.shape[1]))
        mask = ids == PAD if use_mask else None
        h = model.blocks(h, src_key_padding_mask=mask)
        keep = (ids != PAD).unsqueeze(-1).float()
        mean = (model.norm(h) * keep).sum(1) / keep.sum(1)
        return [round(v, 3) for v in model.out(mean)[0].tolist()]

alone = token_ids([short], vocab)
print("alone, no padding:   ", team_scores(alone, use_mask=True))
print("padded, with mask:   ", team_scores(pair, use_mask=True))
print("padded, without mask:", team_scores(pair, use_mask=False))

The output:

alone, no padding:    [5.223, -1.598, -3.307, -4.811, 4.613]
padded, with mask:    [5.223, -1.598, -3.307, -4.811, 4.613]
padded, without mask: [1.351, -6.228, -1.896, 0.008, 8.922]

The teams are in the order delivery, returns, payment, warranty, account. Alone, the ticket goes to delivery (5.223), which is right. Padded, without a mask, it goes to account (8.922): 107 blank lines changed the answer.

The idea. Tell attention which keys to skip. src_key_padding_mask has the shape (batch, tokens) and is True where a token must be hidden: ids == PAD. Hidden keys get no weight, so the real tokens never read padding. With the mask, the padded ticket gets exactly the scores it gets alone. This is the padding mask. TinyTransformer uses it in every forward pass, and its average skips the padding rows too.

The causal mask: cover the answers below

The familiar situation. In a test, the answers are printed on the same sheet, below each question. A fair student covers everything below the current line with a piece of paper.

The problem. The next lesson trains a model to predict the next token at every position. If position 3 can read position 4, it reads the answer instead of predicting it. Then the model fails as soon as it must write new text, where the next token does not exist yet.

The idea. Hide, for every token, all the tokens after it. In the attention table, that is every cell to the right of the diagonal: a triangle.

The causal mask on six tokens: a token may read itself and the tokens before it, never the tokens after it.

Text description of the diagram

The diagram is a grid of 6 rows and 6 columns, for the tokens “my”, “parcel”, “arrived”, “broken”, a comma and “please”. Each row is a token that reads, and each column is a token it may read. The cells on the diagonal and below it are blue: a token may read itself and every token before it. The cells above the diagonal are grey and marked “hidden”: these are the later tokens, the answers the model must predict. A note says that row “broken” may read “my”, “parcel”, “arrived” and “broken”, but not the comma or “please”. 15 of the 36 cells are hidden.

This is the causal mask. “Causal” because the past can affect a token and the future cannot. PyTorch makes it for you:

causal = nn.Transformer.generate_square_subsequent_mask(6)
print(causal)

The output:

tensor([[0., -inf, -inf, -inf, -inf, -inf],
        [0., 0., -inf, -inf, -inf, -inf],
        [0., 0., 0., -inf, -inf, -inf],
        [0., 0., 0., 0., -inf, -inf],
        [0., 0., 0., 0., 0., -inf],
        [0., 0., 0., 0., 0., 0.]])

The mask is added to the scores before softmax. A score of minus infinity becomes a weight of exactly 0, because e to the power of minus infinity is 0. Here it is on one attention head, with 6 random token vectors:

torch.manual_seed(0)
x = torch.randn(1, 6, 8)                              # 6 tokens, 8 numbers each
head = nn.MultiheadAttention(8, 1, batch_first=True)
_, weights = head(x, x, x, attn_mask=causal)
print(weights[0].detach().numpy().round(2))

The output:

[[1.   0.   0.   0.   0.   0.  ]
 [0.6  0.4  0.   0.   0.   0.  ]
 [0.38 0.2  0.42 0.   0.   0.  ]
 [0.29 0.23 0.22 0.26 0.   0.  ]
 [0.22 0.17 0.21 0.18 0.22 0.  ]
 [0.12 0.16 0.26 0.15 0.15 0.16]]

The upper triangle is all 0. The first token can read only itself, so its weight is 1. Every row still adds up to 1, spread over the tokens it may read.

Try both masks on the trained tiny transformer. Predict first, then pick a mask.

Explore

Which cells does each mask hide?

Hidden cells

0 of 36

Visible cells

36

"my" reads

6 tokens

No cell is hidden. Every token reads every token, also the 2 padding tokens: "late" gives 0.37 of its attention to padding. 0 of 36 cells are hidden; each row's visible weights add up to 1.

Rows: the token that reads. Columns: the token it reads.

	my	parcel	is	late	pad	pad
my	0.10	0.13	0.26	0.14	0.23	0.14
parcel	0.20	0.11	0.28	0.15	0.14	0.11
is	0.04	0.46	0.08	0.17	0.11	0.13
late	0.13	0.24	0.09	0.16	0.32	0.05
pad	0.24	0.04	0.23	0.12	0.15	0.22
pad	0.21	0.02	0.14	0.14	0.13	0.36

Real attention weights of the trained tiny transformer (block 1, the average of its 4 heads) for "my parcel is late" and 2 padding tokens, computed with each mask.

Text description of the explorer

The explorer shows a table of 6 rows and 6 columns for the tokens “my”, “parcel”, “is”, “late” and two padding tokens. Rows are the token that reads; columns are the token it reads. One row of buttons chooses the mask. No mask: no cell is hidden, and “late” gives 0.37 of its attention to the two padding tokens. Padding mask: the two padding columns are hidden, 12 of 36 cells. The weights of the real tokens grow to fill the gap, and “late” now gives 0.46 to “parcel”. Causal mask: the 15 cells above the diagonal are hidden. “my” reads only itself, with weight 1.00, and “parcel” gives 0.51 to “my” and 0.49 to itself. Both masks: 18 cells are hidden. The real tokens read as with the causal mask, and only the padding rows change. Each row’s visible weights add up to 1.

Two details are worth seeing. The padding rows still get weights: the model computes them, and TinyTransformer throws them away when it averages the real tokens. And with padding at the end, the causal mask already hides the padding from the real tokens, so “both” changes only the padding rows.

Encoders and decoders

The masks decide what kind of model you have. A model that reads a whole text at once, with only a padding mask, is an encoder. A model with a causal mask, where every token reads only the past, is a decoder: it can write text one token at a time.

	Encoder (BERT-like)	Decoder (GPT-like)
Each token reads	every token, before and after	itself and the tokens before it
Masks	padding	causal (and padding in batches)
Good for	reading: classify, search, find names	writing: continue a text, answer, translate
In this course	TinyTransformer (nn.TransformerEncoder)	SmolLM2-135M in Module 7

BERT (2018) is the classic encoder, and the GPT models are the classic decoders. For routing tickets, an encoder fits: the whole ticket is there, and “not” at the end can change the meaning of “delivery” at the start. To write a reply, a model must work like a decoder, because the next word of the reply does not exist yet.

When to use what, and what it costs

• Positions are needed whenever order can change the meaning. A learned table is simple but stops at its maximum length; Larkfield’s stops at 120 tokens. Rotary positions handle longer texts better.
• The padding mask is needed whenever a batch holds texts of different lengths. That is almost always. It costs nothing to compute; forgetting it changes answers without an error.
• The causal mask is needed in any model that learns to predict the next token. It also limits the model: a token can never use what comes after it, even when the whole text is there.

Guided practice: masks on real tensors

1. Open the notebook from the box at the top of this page, and run the setup cells. The training cell ends after 8 epochs with a validation accuracy of about 92%.
2. In section 1, run the shuffle test with one attention layer. You see True: the same summary for the two sentences.
3. Add the position table and run it again. You see False.
4. In section 2, set the trained model’s position table to zero, and compare the scores of T-64697 in order and shuffled. You see the same 5 scores twice.
5. In section 3, run the short ticket alone, padded with the mask and padded without it. The first two lines are equal.
6. In the Your turn cell, make causal for 6 tokens and count its hidden cells in hidden_cells. Run the check cell. It prints Check passed.

Show an example answer to step 6

causal = nn.Transformer.generate_square_subsequent_mask(6)
hidden_cells = int(torch.isinf(causal).sum())
print(hidden_cells)

15

torch.isinf() is True where a value is infinite. The upper triangle of a 6 × 6 table has 5 + 4 + 3 + 2 + 1 = 15 cells.

Your turn: four tokens

Change one condition: a text of 4 tokens. Before you run it, draw the causal mask on paper and count its hidden cells.

Show the answer

four = nn.Transformer.generate_square_subsequent_mask(4)
print(four)
print("hidden cells:", int(torch.isinf(four).sum()))

tensor([[0., -inf, -inf, -inf],
        [0., 0., -inf, -inf],
        [0., 0., 0., -inf],
        [0., 0., 0., 0.]])
hidden cells: 6

3 + 2 + 1 = 6 cells, the triangle above the diagonal. For n tokens, the causal mask hides n × (n − 1) ÷ 2 cells, almost half of the table.

Common mistake

A padding mask turned the wrong way

Symptom: Tomás writes the mask as “the tokens to keep”, ids != PAD. Nothing stops with an error, but the validation accuracy falls, and some tickets get nan (not a number) as their scores:

ids_valid = token_ids(valid_rows, vocab)
y_valid = labels(valid_rows)

def scores_with(mask):
    with torch.no_grad():
        h = model.embed(ids_valid) + model.position(torch.arange(ids_valid.shape[1]))
        h = model.blocks(h, src_key_padding_mask=mask)
        keep = (ids_valid != PAD).unsqueeze(-1).float()
        return model.out((model.norm(h) * keep).sum(1) / keep.sum(1))

wrong = scores_with(ids_valid != PAD)                 # True on the real tokens: they are hidden
right = scores_with(ids_valid == PAD)                 # True on the padding: it is hidden
print("wrong mask: accuracy", round((wrong.argmax(1) == y_valid).float().mean().item(), 4),
      "| tickets with nan scores:", int(torch.isnan(wrong).any(1).sum()))
print("right mask: accuracy", round((right.argmax(1) == y_valid).float().mean().item(), 4))
first = int(torch.isnan(wrong).any(1).nonzero()[0])
print(valid_rows[first]["ticket_id"], len(vocab.encode(valid_rows[first]["text"])), "tokens:", wrong[first])

wrong mask: accuracy 0.155 | tickets with nan scores: 23
right mask: accuracy 0.9433
T-64707 120 tokens: tensor([nan, nan, nan, nan, nan])

Cause: in PyTorch, True in a mask means hide. Tomás’s mask hides the real words, so each ticket reads only its padding: 15.5%, worse than chance. A ticket with 120 tokens has no padding at all, so every key is hidden. Softmax then divides 0 by 0, which gives nan.

Fix: write the mask as the tokens to hide, ids == PAD, as ticketnet.py does. Test it with one ticket alone and the same ticket padded: the scores must be equal.

What it means for a real AI system: each library has its own rule for masks. In Hugging Face transformers, the attention_mask is the opposite: 1 means keep. Read the documentation of each mask, and keep a test that compares a ticket alone with the same ticket in a padded batch.

See it end to endTrain a tiny language modelCase study

The tiny language model is a real decoder trained with this causal mask.

Knowledge check

Check your understanding

Without position vectors, why does attention give the same summary for two sentences with the same words in a different order? What do position vectors change?

Show the answer

Attention compares each token’s query with every key and averages the values, and nothing in that calculation depends on where a token is. So each word gets the same output in both sentences, only in another place. An average over the tokens is then exactly the same (True in the lesson). The trained tiny transformer, with its position table set to zero, gave identical scores for a ticket and its shuffled words. Position vectors add a different learned pattern for each place. So the same word at two places becomes two different inputs, and the order can affect the result.

Check your understanding

Predict: a short ticket is padded with 107 padding tokens in a batch, and the model has no padding mask. Can its answer change? What does the mask do, and what is the shape of src_key_padding_mask for a batch of 2 tickets of 120 tokens?

Show the answer

Yes. Without a mask, attention reads the padding tokens like words, and they change every real token’s output. In the lesson, short delivery ticket T-64715 went to account (8.922) when padded without a mask, but to delivery (5.223) alone. The padding mask is True where a token is padding, so no token gives it any weight. With the mask, the padded ticket gets exactly the scores it gets alone. The mask has one value for each token of each ticket: shape (2, 120).

Check your understanding

Larkfield wants two models: one that routes a finished ticket, and one that writes the first words of a reply. Which should be an encoder and which a decoder, and which masks does each need?

Show the answer

The router should be an encoder, like TinyTransformer: the whole ticket is there, so every token may read every other token, before and after it. It needs only the padding mask. The reply writer must be a decoder: it writes one token at a time, and each token can depend only on the tokens before it. It needs the causal mask, and the padding mask when it trains on padded batches.

Recap

• Attention alone ignores word order: shuffled words give a shuffled output and the same summary.
• A position embedding adds a learned vector for each place. TinyTransformer has 120 of them; sinusoidal and rotary positions are other ways.
• The padding mask hides padding tokens. Without it, a short ticket’s answer can change with its batch.
• The causal mask hides every later token, the triangle above the diagonal, so a model cannot read the answer it must predict.
• An encoder (BERT-like) reads everything; a decoder (GPT-like) reads only the past and can write text.

Next: how a decoder learns from plain text with no labels, and why Larkfield cannot do it at full scale.

Next: Training objectives