Embeddings and rotary positions

Self-attention is permutation-equivariant if you only feed token embeddings: shuffling tokens shuffles outputs the same way. Language cares about order, so you inject positional information. The other half of the first layer is a lookup: token id t becomes row E[t].

The big idea

Vocabulary size V, model width d. An embedding matrix E ∈ R^{V × d} maps token id t to row E[t]. That is V · d parameters. For Llama-3-class V ≈ 128k and d = 4096, roughly 0.5B parameters before any layers.

Many models tie this table with the LM head (W_U = Eᵀ) to save parameters and couple input and output geometry. Others untie them for quality. If W_U = Eᵀ, the score for vocabulary item v given residual h is proportional to h · E[v]. Sampling high-probability tokens means "find embedding vectors near the current residual." Untying costs another V d parameters — material at V=128k, d=4096 (~0.5B). Labs disagree. Read the config.

Walking through position recipes

Sinusoidal (original transformer): fixed sin/cos functions of position and dimension. No learned params, extrapolates in theory, largely abandoned for LLMs.

Learned absolute (GPT-2): a second table P ∈ R^{n_ctx × d} added to token embeddings. Simple. Does not extrapolate past n_ctx cleanly.

RoPE won for decoder-only LLMs. Instead of adding a position vector, RoPE rotates pairs of coordinates in q and k by an angle θ · position. No extra parameters in the basic form. Works with causal LMs. Extends, with scaling tricks, to longer context.

Treat a 2D slice (x₁, x₂) as a complex number and multiply by e^{i m θ}. Queries at position m and keys at n pick up relative phase e^{i (m−n) θ} in their dot product. Stack many frequencies across the head dimension.

Details and long-context RoPE scaling live in modern architectures and the RoFormer paper.

A clock picture: each position is an angle. The angle between two clocks is the relative distance. The picture breaks at very long context, where the rotations wrap in ways the model was not trained for, which is why scaling tricks exist.

End-to-end for one new decode token: id t → row E[t] (d,). RoPE will rotate Q/K derived from this vector using position index n. Then the stack of blocks runs, and the LM head produces logits over V.

Common mistakes

• Forgetting positions entirely. Attention then cannot tell "dog bites man" from "man bites dog".
• Assuming learned absolute positions extrapolate. They stop at n_ctx.
• Dropping RoPE caches when you export or quantize. Generation then looks like a sampler bug.

Check yourself

1. Llama-3-class V ≈ 128k, d = 4096. About how many embedding parameters is that? Answer: About 128000 × 4096 ≈ 0.52B.
2. What does RoPE make the attention score depend on, that a plain dot product of token embeddings would miss? Answer: Relative distance m − n via a rotation of Q and K.

Key takeaways

• Token embeddings are a V × d table. Tying them to the LM head saves V d parameters.
• Attention needs a position signal. RoPE injects relative position by rotating Q and K.
• Learned absolute positions are simple and do not extrapolate cleanly.
• Forgetting norms or RoPE caches at export breaks generation in confusing ways.

Further learning

Next: the transformer block.

RoFormer: Enhanced Transformer with Rotary Position EmbeddingPosition encoding recently has shown effective in the transformer architecture. It enables valuable supervision for dependency modeling between elements at different positions of the sequence. In this paper, we first …arXiv.org