Learn AI Create the Future — with Pranay
Main site

How Models Know Word Order

Learned position embeddings

BERT and GPT-2 gave each slot its own trainable row and let the model work out what position means. It works well, up to the last row of the table.

Read these first

On this page 7
  1. Why anyone would do this instead of using waves
  2. How it works
  3. The wall at the end of the table
  4. The quieter problem
  5. Where you have seen this
  6. Remember this
  7. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

A learned position embedding gives every slot its own row of numbers, and the model trains those numbers like any other weights.

Think of numbered seats in a cinema. The seat number tells you nothing on its own. What you learn over many visits is that the front row hurts your neck and the middle is best.

Nobody told you the meaning of "row 1". You worked it out by sitting there. A learned position embedding does the same for slot 1, slot 2 and slot 511.

The model is handed blank labels and figures out what each slot is worth. That is the entire method.

Why anyone would do this instead of using waves

Waves are a guess. A clever guess, but a guess about what "position" should mean.

A learned table makes no guess. If the first token of every document is special in your data, the model discovers that and puts it in the row. If the last token matters, it discovers that too.

BERT and GPT-2 both chose this. It is one line of code, it trains alongside everything else, and on ordinary text it works at least as well as waves.

How it works

   position table (512 rows, one per slot)
   ┌──────────────────────────────┐
   │ slot   0: [ ...768 numbers ] │  ← trained
   │ slot   1: [ ...768 numbers ] │  ← trained
   │ slot   2: [ ...768 numbers ] │  ← trained
   │   ...                        │
   │ slot 511: [ ...768 numbers ] │  ← trained
   └──────────────────────────────┘

   word "cat" at slot 2   →   cat_vector + row_2   →   into the transformer

Look up the row for the slot. Add it to the word vector. That is it.

The wall at the end of the table

Now the catch, and it is a hard one. The cinema has 512 seats. There is no seat 512.

Ask for a position the table does not have and the program stops with an error. Not a bad answer. A crash.

This is why a model's context limit is often a hard number rather than a soft one. The table has a last row, and past it there is nothing to look up.

The quieter problem

There is a second issue that hurts more in practice. Rows near the end of the table are used far less during training.

Most training documents are short. Slot 5 appears in nearly every example. Slot 500 appears in very few. So the later rows are trained on far less evidence and stay closer to their random starting values.

The result is a model that behaves well early in a document and less reliably near its limit. You may have felt this as a long chat losing the thread.

Where you have seen this

  • BERT, powering a decade of search ranking.
  • GPT-2, with its famous 1024-slot limit.
  • The original vision transformer, with one learned row per image patch.
  • Any model whose context length is a flat number with no graceful degradation past it.

Remember this

  • Each slot gets a trainable row, and the model learns what each slot means.
  • The table has a fixed number of rows and asking past the end is an error.
  • Later rows get less training data, so behaviour is weakest near the limit.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Written against PyTorch 2.5.1. Runs on CPU in under a second.

The whole method, plus its two failure modes

learned_positions.py
import torch
import torch.nn as nn

torch.manual_seed(0)

MAX_POS, D_MODEL, VOCAB = 32, 16, 50

wte = nn.Embedding(VOCAB, D_MODEL)      # token embeddings, one row per token id
wpe = nn.Embedding(MAX_POS, D_MODEL)    # position embeddings, one row per slot

print("token embedding table   :", tuple(wte.weight.shape),
      f"= {wte.weight.numel():,} numbers")
print("position embedding table:", tuple(wpe.weight.shape),
      f"= {wpe.weight.numel():,} numbers")

def embed(token_ids):
    positions = torch.arange(token_ids.shape[-1])
    return wte(token_ids) + wpe(positions)      # add, do not concatenate

batch = torch.randint(0, VOCAB, (2, 12))        # two sequences of 12 tokens
print("\nembedded batch shape:", tuple(embed(batch).shape))

# Train on short sequences only, then look at which rows actually moved.
loss = embed(torch.randint(0, VOCAB, (4, 16))).pow(2).mean()
loss.backward()
grad = wpe.weight.grad.abs().sum(dim=1)
print("\ngradient reaching each position row after one backward pass:")
print("  rows 0-15  (used) : total", f"{grad[:16].sum().item():.4f}")
print("  rows 16-31 (unused): total", f"{grad[16:].sum().item():.4f}")

# Ask for a position the table does not have.
try:
    wpe(torch.tensor([MAX_POS]))
except IndexError as e:
    print("\nasking for position 32:", str(e).split("\n")[0])
Output
token embedding table   : (50, 16) = 800 numbers
position embedding table: (32, 16) = 512 numbers

embedded batch shape: (2, 12, 16)

gradient reaching each position row after one backward pass:
  rows 0-15  (used) : total 1.9565
  rows 16-31 (unused): total 0.0000

asking for position 32: index out of range in self

Reading that output

The gradient on unused rows is exactly 0.0000. Not small. Zero. A row that never gets indexed receives no gradient at all, so it never learns anything.

That is the mechanism behind "the model gets worse near its context limit". It is not a subtle degradation, it is starvation. Rows indexed once per ten thousand examples learn roughly one ten-thousandth as fast as row 5.

One caveat on that zero. With weight decay, optimisers like AdamW still nudge unused rows toward zero even with no gradient, because decoupled decay applies to every parameter. So they drift, they do not learn.

index out of range in self is the entire story of a hard context limit, in one exception. There is no row 32 and no sensible default. nn.Embedding refuses rather than guessing.

Stretching the table to a longer context

Because the rows are only numbers, you can resample them the way you resize an image. Vision transformers do this routinely when the input resolution changes.

interpolate_positions.py
import torch
import torch.nn as nn

torch.manual_seed(0)
wpe = nn.Embedding(32, 16)
w = wpe.weight

# Treat the table as a 1-D signal of length 32 with 16 channels, and resample it.
longer = torch.nn.functional.interpolate(
    w.T.unsqueeze(0), size=64, mode="linear", align_corners=False
).squeeze(0).T

print("table shape:", tuple(w.shape), "->", tuple(longer.shape))
print("new row 1 == 0.75*old[0] + 0.25*old[1] ? ",
      torch.allclose(longer[1], 0.75 * w[0] + 0.25 * w[1], atol=1e-6))
print("new row 3 == 0.75*old[1] + 0.25*old[2] ? ",
      torch.allclose(longer[3], 0.75 * w[1] + 0.25 * w[2], atol=1e-6))
Output
table shape: (32, 16) -> (64, 16)
new row 1 == 0.75*old[0] + 0.25*old[1] ?  True
new row 3 == 0.75*old[1] + 0.25*old[2] ?  True

Every new row is a weighted blend of two old rows. Nothing new was learned. You have doubled the number of slots and halved the resolution between them.

This works acceptably for images, where neighbouring patches are genuinely similar. It works badly for text without a finetuning pass afterwards, because the model was trained expecting one row per token, not a blend.

Common mistakes

Concatenating instead of adding. Position embeddings are added to token embeddings, not appended. Concatenating changes the width of every downstream matrix and burns capacity on channels that only ever carry position.

Initialising the table too large. GPT-2 initialises positions with a standard deviation of about 0.01, an order of magnitude smaller than its token embeddings. A position table at the same scale as the token table swamps word identity early in training.

Forgetting to offset positions when using a KV cache. During generation you feed one token at a time, so torch.arange(seq_len) yields [0] on every step. You need torch.arange(past_len, past_len + new_len). This bug produces text that is fluent and progressively confused, which makes it hard to spot.

Padding on the left without adjusting positions. With left padding, real token 0 sits at index 5. Unless you build position_ids from the attention mask, every sequence in the batch gets a different, wrong alignment.

Try it yourself

Build a nn.Embedding(64, 16) and train it on sequences of length 16 only. Then print the row norms for all 64 rows. Rows 16 upward will sit near their initialisation while the first 16 have moved. That plot is the clearest picture of why the tail of a context window is unreliable.

What to learn next

Researcher — Mathematics and papers.

Definition

A learned absolute position embedding is a matrix $E_{pos} \in \mathbb{R}^{L \times d}$ for maximum length $L$ and width $d$. The input to layer zero is

$$ h^{(0)}m = E{tok}[x_m] + E_{pos}[m] $$

with both matrices trained by the same objective. Parameter cost is $Ld$: for BERT-base, $512 \times 768 = 393{,}216$; for GPT-2 small, $1024 \times 768 = 786{,}432$. Under 1 percent of total parameters in both cases.

Effect on the attention logit

Expanding the pre-softmax score for query position $m$ and key position $n$:

$$ \frac{1}{\sqrt{d_k}} (e_m + p_m)^\top W_Q W_K^\top (e_n + p_n) $$

$$ = \frac{1}{\sqrt{d_k}}\Big[ \underbrace{e_m^\top A e_n}_{\text{content-content}} + \underbrace{e_m^\top A p_n}_{\text{content-position}} + \underbrace{p_m^\top A e_n}_{\text{position-content}} + \underbrace{p_m^\top A p_n}_{\text{position-position}} \Big] $$

with $A = W_Q W_K^\top$. Ke et al. (2021), Rethinking Positional Encoding in Language Pre-training (TUPE), argue the two cross terms are harmful: they force one shared projection to serve two heterogeneous signals, and the position-position term is the only one carrying clean positional structure. TUPE untangles them with separate projections and reports faster pre-training convergence on GLUE.

Learned versus sinusoidal

Vaswani et al. (2017) report the two are within noise on WMT En-De. Wang et al. (2021), On Position Embeddings in BERT, probe what the learned tables contain and find BERT's rows do develop translation-like structure, but imperfectly, and the tables are noticeably noisier in their tail.

The practical differences are not about quality within the trained range:

LearnedSinusoidal
Parameters$Ld$0
Beyond $L$IndexErrordefined, low quality
Tail rowsundertrainedidentical treatment at every position
Adapts to datayesno
Extensioninterpolate plus finetunechange nothing

The undertrained-tail effect

If document lengths follow distribution $p(\ell)$, position $m$ receives gradient from a fraction $\Pr[\ell > m]$ of examples. Web-scale corpora are heavy-tailed but overwhelmingly short-dominated, so $\Pr[\ell > m]$ falls steeply. Row $m$ near $L$ can see orders of magnitude fewer updates than row 0.

This has two visible consequences. Quality degrades near the stated context limit rather than at it. And packing multiple documents into one training sequence — standard practice for throughput — is partly a fix for this, since it keeps every row occupied.

Extension by interpolation

For a table $E_{pos} \in \mathbb{R}^{L \times d}$ and target length $L' > L$, resample along the position axis:

$$ \tilde{E}[m'] = \sum_{m} w(m', m)\, E_{pos}[m], \qquad \sum_m w(m', m) = 1 $$

with $w$ from linear or bicubic interpolation. Dosovitskiy et al. (2021) use 2-D bicubic interpolation of ViT position embeddings for fine-tuning at higher resolution, and it is standard there. Touvron et al. (2021), DeiT, confirm it. For text, interpolation without a subsequent finetuning pass degrades quality substantially, which is one reason RoPE-based extension displaced it.

Current status

Learned absolute embeddings dominated 2018 to 2021 and are now uncommon in new decoder-only language models. RoPE is the default there. They remain in use where the sequence length is bounded and known: ViT patch grids, encoder-only classifiers, and models with a fixed input geometry.

Kazemnejad et al. (2023) place absolute position embeddings last among five schemes for downstream length generalisation, behind T5's relative bias, ALiBi, RoPE and NoPE.

Papers

What to learn next