Token IDs and the Embedding Lookup

~10 min · token-ids, embedding

Level 0Token

0 XP0/94 lessons0/10 achievements

0/120 XP to next level120 XP to go0% complete

After tokenization, every token has a unique integer ID drawn from the vocabulary table. The vocabulary is fixed at training time and never changes — IDs are stable for the life of the model.

The first operation inside any Transformer is embedding lookup: for each token ID, retrieve the corresponding row of the embedding matrix. The matrix has shape (vocab_size × d_model). Index 47458 returns a 4096-dimensional vector if d_model=4096. That vector is the model's first guess at the meaning of the token, which subsequent layers refine.

Two things to remember. First, the embedding matrix is learned — it is updated during training along with everything else. Second, in many modern models the input embedding matrix and the output projection (logits → vocab) share weights, called weight tying. This saves vocab × d_model parameters and ties together "how I read this token" with "how I predict this token."

Code

Embedding lookup is just a row index·python

import torch
import torch.nn as nn

vocab_size, d_model = 128_000, 4096
emb = nn.Embedding(vocab_size, d_model)

ids = torch.tensor([47458, 5614, 15592])      # 3 tokens
vectors = emb(ids)                              # (3, 4096)
# Internally: vectors = emb.weight[ids]         # advanced indexing

Weight tying — input and output share the same matrix·python

# input embedding and output head share the same matrix
class TiedTransformer(nn.Module):
    def __init__(self, vocab, d_model):
        super().__init__()
        self.emb = nn.Embedding(vocab, d_model)
        # ... transformer blocks ...
    def head(self, hidden):
        # Use embedding weight transposed as the output projection
        return hidden @ self.emb.weight.T   # (seq, vocab) logits

External links

Press & Wolf (2017) — Using the Output Embedding to Improve LMs (weight tying)

Exercise

Load any small open-weight model (e.g., gpt2 from Hugging Face). Print the shape of model.transformer.wte.weight (the input embedding) and model.lm_head.weight (the output head). Are they the same tensor (weight-tied) or independent? Test by mutating one and checking the other.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousSpecial Tokens: Scaffolding for the Model Next →Tokenizer Quirks: Why Your LLM Can't Count R's in 'strawberry'

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.