🎓 A1Learner
🔍

PyTorch Positional Encoding

A positional encoding is a fixed vector that marks one position in a sequence. Attention alone cannot tell order, so the encoding is added to each word embedding.

The same word in two places

An embedding gives one word one vector, wherever the word sits. Two copies of word 5 receive two identical rows.

import torch
from torch import nn
torch.manual_seed(0)
embedding = nn.Embedding(10, 8)
pair = embedding(torch.tensor([5, 5]))
print("Two rows for the same word are equal:", torch.equal(pair[0], pair[1]))

One row from the formula

Even columns take sine of the angle and odd columns take cosine. The angle shrinks as the column pair grows, so later pairs swing more slowly.

pos = 2
dim = 8
row = []
for i in range(dim):
    angle = pos / (10000 ** ((i // 2 * 2) / dim))
    row.append(torch.sin(torch.tensor(angle)) if i % 2 == 0
               else torch.cos(torch.tensor(angle)))
print([round(v.item(), 2) for v in row])

The learned alternative

An embedding table can serve as the position table instead. Training shapes its rows, and no formula is needed.

torch.manual_seed(0)
table = nn.Embedding(8, 8)
print("Learned table shape:", tuple(table.weight.shape))
print("Rows 0 and 1 differ:", not torch.equal(table.weight[0], table.weight[1]))

Example

import torch
from torch import nn

torch.manual_seed(0)

positions = 8
dim = 8

table = torch.zeros(positions, dim)
pos = torch.arange(positions).unsqueeze(1).float()
pair = torch.arange(0, dim, 2).float()
angle = pos / torch.pow(10000.0, pair / dim)
table[:, 0::2] = torch.sin(angle)
table[:, 1::2] = torch.cos(angle)

print("Positional encodings, positions 0 to 7:")
for row in table.tolist():
    print([round(v, 2) for v in row])

embedding = nn.Embedding(10, dim)
tokens = torch.tensor([5, 5, 5])
added = embedding(tokens) + table[:3]
print("Embedding shape:", tuple(embedding(tokens).shape))
print("Added shape:", tuple(added.shape))
print("Same word at positions 0 and 1 differs after adding:",
      not torch.equal(added[0], added[1]))

Key points

  • Positional encoding: a vector of numbers that stands for one position
  • Sine and cosine: two wave functions from maths that swing between -1 and 1
  • Sinusoidal encoding: a table built from sine and cosine waves at different speeds
  • Adding the encoding gives each position its own version of the same word
  • A learned table holds one trainable row per position instead

In this example, the table holds eight positions and eight numbers per position. Even columns take sine of the angle and odd columns take cosine. The angle shrinks as the column pair grows, so later pairs swing more slowly. The prints show the grid for positions 0 to 7. Position 0 starts with 0.0 and 1.0 pairs, because sine of 0 is 0 and cosine of 0 is 1. An embedding layer then looks up the word 5 three times. Adding the first three table rows makes each copy different. The last line proves the copies differ after the adding.

Note

Sinusoidal encodings need no training, because the formula fixes every value

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why is a positional encoding added to a word embedding?