Tutorial 03: Embeddings and Tokenization¶
Before a transformer can process text, it must convert human-readable tokens (numbers representing pieces of words) into dense vectors that the mathematical layers can operate on. This tutorial covers tokenization, the embedding lookup, and how tied weights connect embeddings to the output layer.
1. Tokenization: Text to Numbers¶
Transformers do not read characters or words — they read token IDs. A tokenizer splits text into subword pieces and maps each piece to an integer.
Example with a GPT-style tokenizer:
Each token ID is an index into a vocabulary of ~30,000-150,000 entries. The tokenizer is trained separately from the model — by inference time, it is fixed.
Our engine does not include a tokenizer. Callers pass raw token IDs. This is intentional — tokenization is a pre-processing concern, not an inference concern. The engine receives:
2. The Embedding Matrix¶
The embedding layer is a lookup table. Given a vocabulary of size V and an
embedding dimension of D, the embedding is a weight matrix of shape (V, D):
Each row i of this matrix is the learned vector representation of token i.
When the model sees token ID 5, it looks up row 5 of the embedding matrix.
The Math¶
This is not a matrix multiply — it is a gather operation. Given a batch of token IDs, we select the corresponding rows:
import mlx.core as mx
vocab_size, hidden_size = 100, 8
# The embedding weight matrix — in practice, this comes from model weights
W = mx.random.normal((vocab_size, hidden_size))
# Input: a batch of 2 sequences, each of length 3
token_ids = mx.array([[5, 12, 3],
[7, 1, 42]]) # shape: (2, 3)
# Lookup: select rows 5, 12, 3 for the first sequence, etc.
embeddings = W[token_ids, :] # shape: (2, 3, 8)
This is exactly how our engine implements it:
# core/embedding.py
class Embedding:
def __init__(self, vocab_size: int, embedding_dim: int, weight: mx.array):
self.vocab_size = vocab_size
self.embedding_dim = embedding_dim
self.weight = weight # shape: (vocab_size, embedding_dim)
def __call__(self, x: mx.array) -> mx.array:
# x contains token IDs, shape: (batch, seq_len)
# Returns dense vectors, shape: (batch, seq_len, embedding_dim)
return self.weight[x, :]
The magic is in self.weight[x, :]. MLX's indexing syntax handles the batch
dimension automatically — you do not need to loop over sequences.
Why This Works¶
MLX indexing with an integer array performs a gather:
W = mx.array([[10, 11], # token 0
[20, 21], # token 1
[30, 31], # token 2
[40, 41]]) # token 3
ids = mx.array([3, 0, 2]) # select tokens 3, 0, 2
result = W[ids, :]
# result = [[40, 41], ← row 3
# [10, 11], ← row 0
# [30, 31]] ← row 2
Each row is replaced by the embedding vector for the corresponding token.
3. Embedding as a Linear Layer¶
There is a beautiful connection: the embedding matrix can also be used as the
output projection. At the end of the transformer, after processing through
all layers, you have a hidden vector of size D. To predict the next token, you
need a vector of size V (one score per token in the vocabulary).
The most efficient way to do this is to multiply the hidden vector by the transpose of the embedding matrix:
This means token i's embedding vector, when dotted with the final hidden
state, gives the score for predicting token i next.
class Embedding:
# ... existing code ...
def as_linear(self, x: mx.array) -> mx.array:
"""Use the embedding matrix as an output projection (logit head).
This exploits the fact that embedding tokens and projecting to vocabulary
are symmetric operations — they share the same weight matrix, just used
in different directions.
Args:
x: Hidden state from the final transformer layer.
Shape: (batch, seq_len, hidden_size)
Returns:
Logits over vocabulary. Shape: (batch, seq_len, vocab_size)
"""
return linear(x, self.weight)
This is called weight tying — the embedding weights and the output projection weights are the same matrix. Not all models do this (Qwen3 ties them by default; Llama 3 does not), but it is very common because:
- It halves the memory cost of the embedding+head (one matrix instead of two).
- It provides a natural regularization (the embedding and prediction must be consistent).
- Empirically, it rarely hurts performance.
4. Building the Full Model's Embedding Layer¶
In our engine, the model's __init__ method sets up the embedding:
# core/model.py — inside Qwen3ForCausalLM.__init__
# The weight store loads the embedding from safetensors
embed_weight = weights.linear("model.embed_tokens")
# Check if the embedding is quantized (4-bit weights stored as scales+packed)
if isinstance(embed_weight, QuantizedWeights):
self.embedding = QuantizedEmbedding(
vocab_size=config.vocab_size,
embedding_dim=config.hidden_size,
weight=embed_weight,
)
else:
self.embedding = Embedding(
vocab_size=config.vocab_size,
embedding_dim=config.hidden_size,
weight=embed_weight,
)
And in the forward pass:
# core/model.py — inside Qwen3ForCausalLM.__call__
def __call__(self, inputs, offset=0, cache=None, logits_to_keep=None):
# Step 1: Convert token IDs to dense vectors
h = self.embedding(inputs) # (batch, seq_len) → (batch, seq_len, hidden_size)
# ... process through transformer layers ...
# Final step: project to vocabulary logits
if self.w_lm_head is not None:
# Untied embeddings — separate output projection
return _linear(h, self.w_lm_head)
else:
# Tied embeddings — reuse the embedding matrix
return self.embedding.as_linear(h)
5. Tensors in Practice: The Shape Flow¶
Let's trace the actual tensor shapes through the embedding stage:
# Input
tokens = mx.array([[15496, 11, 995, 0]]) # "Hello, world!\n"
# tokens shape: (1, 4) — batch of 1, sequence length 4
# Embedding lookup
h = model.embedding(tokens)
# h shape: (1, 4, 4096) — 4096 is the hidden_size
# Each token now has a 4096-dimensional vector representation.
# These vectors will be refined by each transformer layer.
The shape (batch, seq_len, hidden_size) is the universal format for sequence
data in transformers. Every layer in the model preserves this shape (or a
transposed variant of it).
6. Summary¶
| Concept | What It Does | Shape |
|---|---|---|
| Tokenization | Text → integer IDs | string → (seq_len,) |
| Embedding lookup | Token IDs → dense vectors | (batch, seq) → (batch, seq, hidden) |
| Embedding as linear | Hidden vectors → vocab logits | (batch, seq, hidden) → (batch, seq, vocab) |
| Weight tying | Same matrix for both | (vocab_size, hidden_size) shared |
The embedding layer is the gateway between discrete symbols and continuous vector space. Everything that follows — attention, feed-forward networks, normalization — operates in this continuous space.