Embeddings and position¶
A Transformer cannot consume text. It consumes integer token IDs, looks up a vector for each ID, and gives the model a way to distinguish where those vectors occur.
Those are separate jobs:
- token embeddings answer "which vocabulary item is this?";
- position mechanisms answer "where is it, or how far apart are two positions?";
- attention and FFNs turn those initial vectors into contextual states.
From an ID to a vector¶
Let the vocabulary contain V tokens and the model width be d. The embedding table is a learned matrix
For token ID \(t\), lookup returns row \(E_t\). This is an indexed read, not a one-hot matrix multiplication in an efficient implementation, although it is mathematically equivalent to multiplying a one-hot vector by \(E\).
import torch
from torch import nn
vocab_size = 32_000
d_model = 512
embedding = nn.Embedding(vocab_size, d_model)
token_ids = torch.tensor([[41, 9, 9, 2]]) # (batch=1, sequence=4)
x = embedding(token_ids)
assert x.shape == (1, 4, 512)
Two occurrences of token ID 9 begin with the same embedding row. They become different hidden states after position information and surrounding context are applied.
Embeddings are learned coordinates, not dictionary definitions¶
An embedding dimension does not normally correspond to a named human feature such as "plural" or "France." Training adjusts the coordinates so the full network can reduce next-token loss. Relations can be distributed across many dimensions and transformed at every layer.
Similarly, an embedding row is not the model's complete knowledge about its token. It is only the starting state. Facts and behaviors depend on the entire stack of learned parameters.
The output side: hidden state to vocabulary logits¶
After the final block and normalization, a language-model head maps hidden state \(h_t \in \mathbb{R}^d\) to V logits:
Some architectures tie \(W_{\text{vocab}}\) to the input embedding matrix. The original Transformer shared embedding and pre-softmax weights and scaled input embeddings by \(\sqrt{d}\) (Section 3.4). Weight tying is a model configuration choice, not a law of Transformers.
Why order must be introduced¶
Attention without a position mechanism is permutation-equivariant: reordering the input reorders the outputs, but the operation has no inherent concept that one token came first. Language needs order:
The vocabulary items are the same; their relations are not.
flowchart LR
tokenA[Token embedding]
positionA[Position information]
initialState[Position-aware state]
attention[Attention]
contextualState[Contextual state]
tokenA --> initialState
positionA --> initialState
initialState --> attention
attention --> contextualState There are several ways to inject order. Do not collapse them into the generic word "positional embedding."
Absolute sinusoidal positions¶
The original Transformer adds a deterministic vector to each token embedding. For position \(p\) and coordinate pair indexed by \(i\):
Different coordinate pairs rotate at different frequencies. Nearby positions have related patterns, while the collection of frequencies can distinguish a wide range of positions (Section 3.5).
import torch
def sinusoidal_positions(length: int, d_model: int) -> torch.Tensor:
if d_model % 2:
raise ValueError("This compact implementation requires even d_model")
position = torch.arange(length, dtype=torch.float32)[:, None]
coordinate = torch.arange(0, d_model, 2, dtype=torch.float32)
inv_frequency = 1.0 / (10_000 ** (coordinate / d_model))
angles = position * inv_frequency[None, :]
table = torch.empty(length, d_model)
table[:, 0::2] = angles.sin()
table[:, 1::2] = angles.cos()
return table
pe = sinusoidal_positions(length=128, d_model=512)
assert pe.shape == (128, 512)
Learned absolute position embeddings use a trainable lookup table instead. Both approaches attach an absolute-position signal to the stream before or within the stack.
RoPE: rotate queries and keys¶
Rotary position embedding (RoPE) takes a different route. It rotates pairs of coordinates in queries and keys according to position before their dot product is computed (RoFormer).
For one 2D coordinate pair and angular frequency \(\theta\):
The position-aware query and key are
Their dot product contains a relative-position identity:
The absolute rotations combine so the attention compatibility depends on the relative offset \(s-p\). The full method applies a bank of frequencies across the head dimension.
import torch
def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
"""Rotate (x0, x1) to (-x1, x0) in every adjacent pair."""
even = x[..., 0::2]
odd = x[..., 1::2]
return torch.stack((-odd, even), dim=-1).flatten(-2)
def apply_rope(
q: torch.Tensor,
k: torch.Tensor,
positions: torch.Tensor,
base: float = 10_000.0,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Apply RoPE to q and k shaped (batch, heads, sequence, head_dim)."""
head_dim = q.shape[-1]
if q.shape != k.shape or head_dim % 2:
raise ValueError("q and k must match and have an even head dimension")
pair_coordinates = torch.arange(0, head_dim, 2, device=q.device)
inv_frequency = 1.0 / (base ** (pair_coordinates.float() / head_dim))
angles = positions.float()[:, None] * inv_frequency[None, :]
cos = angles.cos().repeat_interleave(2, dim=-1)[None, None, :, :]
sin = angles.sin().repeat_interleave(2, dim=-1)[None, None, :, :]
return q * cos + rotate_pairs(q) * sin, k * cos + rotate_pairs(k) * sin
q = torch.randn(2, 8, 16, 64)
k = torch.randn_like(q)
q_rotated, k_rotated = apply_rope(q, k, torch.arange(16))
assert q_rotated.shape == q.shape == k_rotated.shape
For a real released implementation, follow Meta's Llama 3 precompute_freqs_cis and apply_rotary_emb. That implementation represents each coordinate pair as a complex number; complex multiplication performs the rotation.
ALiBi: bias attention by distance¶
Attention with Linear Biases (ALiBi) does not add position vectors or rotate queries and keys. It adds a head-specific penalty proportional to query-key distance directly to attention logits (Press et al.). Conceptually:
The causal condition \(s \le p\) still comes from the mask. The linear term says how a head's score changes with distance.
flowchart TD
positionQuestion{Where does position enter?}
additive["Absolute vector added to token state"]
rotary["Rotation applied to queries and keys"]
bias["Distance bias added to attention logits"]
positionQuestion -->|"Sinusoidal or learned"| additive
positionQuestion -->|RoPE| rotary
positionQuestion -->|ALiBi| bias Position limits: configured length is not proven quality¶
A checkpoint's context limit involves more than whether code can allocate a larger tensor. It can depend on:
- lengths and packing used during training;
- the positional method and its frequencies or biases;
- attention masks and cache implementation;
- memory available for activations and KV cache;
- empirical quality at the target distance.
Changing max_position_embeddings can make a runtime accept longer inputs without proving that the checkpoint will use the extra region well. Techniques that rescale or interpolate RoPE likewise need model-specific evaluation.
A debugging checklist¶
When position handling is wrong, look for:
- off-by-one cache positions - new tokens reuse or skip a RoPE index;
- padding treated as content - masks allow attention into pad positions;
- left-padding position IDs - IDs do not match the runtime's convention;
- double application - position information is added or rotated twice;
- dimension mismatch - RoPE is applied over the wrong head subdimension;
- unsupported extension - the runtime accepts a length the checkpoint was not validated to handle.
Check your understanding¶
- Why do two identical token IDs start with equal token vectors but later acquire different hidden states?
- Where does RoPE operate: token embeddings, Q/K, V, or logits?
- Which identity makes relative offset appear in a RoPE dot product?
- Why is increasing a configuration length not evidence of long-context quality?
- How does ALiBi's insertion point differ from RoPE's?
Next: attention from scratch.