Token Embeddings

Why we need embeddings (motivation before formulas)

Neural networks are built to process numbers: vectors, matrices, and tensors. But language (and many discrete domains) start as symbols:

• •Words: “apple”, “bank”
• •Subwords: “un”, “##able”
• •Characters: “a”, “b”
• •Generic tokens: IDs produced by a tokenizer

A token by itself has no natural numeric geometry. Token ID 7 isn’t “closer” to token ID 8 than to token ID 9000, yet a neural network will treat raw numbers that way.

So we introduce a representation that:

1) Is numeric (so the model can compute),

2) Has a geometry (so “similar” tokens can end up close),

3) Is learnable (so it adapts to the training data).

That representation is the token embedding.

Definition

Let the vocabulary size be V (number of distinct tokens) and let the embedding dimensionality be d.

We store an embedding matrix:

• •E ∈ ℝ^(V×d)

Each token i (an integer ID in {0, …, V−1}) is assigned the i-th row of E:

• •eᵢ = E[i] ∈ ℝᵈ

This is a lookup table: you don’t compute E[i] by multiplying all of E by something dense; you directly retrieve a row.

Intuition: embeddings as “coordinates”

You can think of embeddings as placing each token at a point in a d-dimensional space. During training, the model moves these points around to make its predictions better.

• •If two tokens behave similarly in context, training often moves their vectors closer.
• •If two tokens are used differently, their vectors drift apart.

This doesn’t guarantee a perfect semantic map, but it gives the model a flexible continuous “surface” on which to build meaning.

What embeddings are not

• •Not one-hot vectors: one-hot is sparse and high-dimensional (dimension V), and doesn’t express similarity.
• •Not fixed unless you choose them to be: E is typically learned from scratch or fine-tuned.
• •Not necessarily interpretable coordinate-by-coordinate: meaning is distributed across dimensions.

Where embeddings appear in a model

In a typical language model pipeline:

1) Text → tokenizer → token IDs (integers)

2) Token IDs → embedding lookup in E → vectors e₁, e₂, …

3) Vectors → Transformer (or other neural network) → predictions

So token embeddings are often the first learned layer of a modern NLP model.

Why lookup matters

A token ID is discrete. The model must map it to a continuous vector e ∈ ℝᵈ efficiently.

If V is large (e.g., 50k, 200k, 1M), you cannot afford to treat the input as a dense V-dimensional vector at every step. The embedding matrix lets you:

• •store V vectors once,
• •retrieve only the ones you need.

One-hot view (conceptual bridge)

Conceptually, you can represent token i as a one-hot vector xᵢ ∈ {0,1}^V with a 1 at position i.

Then embedding lookup can be written as a matrix product:

• •eᵢ = xᵢᵀ E

Let’s check shapes:

• •xᵢᵀ has shape (1×V)
• •E has shape (V×d)
• •product has shape (1×d) → a vector in ℝᵈ

This multiplication selects exactly one row of E, because all entries of xᵢ are 0 except at i.

Why implementations don’t multiply

Even though xᵢᵀ E is a nice equation, real implementations do an index operation:

• •eᵢ = E[i]

Because multiplying by a V-length one-hot vector would waste memory and compute.

Batching and sequences

In practice you have a batch of sequences:

• •Batch size: B
• •Sequence length: T
• •Token IDs: X ∈ {0,…,V−1}^(B×T)

Embedding lookup produces:

• •Embeddings: H ∈ ℝ^(B×T×d)

So each position t in each sequence b gets a vector:

• •h[b,t] = E[ X[b,t] ]

A careful note about “row vectors” vs “column vectors”

Different sources use different conventions. You might see embeddings as rows E[i] or columns. The key invariant is:

• •Each token corresponds to a length-d vector.

In this lesson we treat E as V rows, each row is eᵢ ∈ ℝᵈ.

Parameter count and memory

Embeddings can dominate parameter count.

• •Number of parameters in E: V·d

Example:

• •V = 50,000
• •d = 768

Parameters = 50,000 · 768 = 38,400,000

That’s 38.4M parameters just for token embeddings.

Basic similarity geometry

Once tokens are vectors, you can measure similarity.

Two common measures:

1) Dot product:

• •sim(a, b) = a·b = ∑ⱼ aⱼ bⱼ

2) Cosine similarity:

• •cos(a, b) = (a·b) / (‖a‖ ‖b‖)

Cosine similarity compares direction, not magnitude. Many embedding analyses use cosine similarity to focus on relational structure.

Why similarity emerges at all

Embeddings are optimized to help the model predict correct outputs. If the model benefits from treating two tokens similarly (because they appear in similar contexts), gradient descent tends to push their embeddings in similar directions.

This is not magic; it’s just shared pressure from the loss function across many training examples.

Why embeddings must be learnable

You could assign each token a random vector and never change it. The model would then have to learn everything in later layers, with no ability to shape the input representation.

Learnable embeddings let the model:

• •allocate representational capacity where needed,
• •encode frequent patterns early,
• •reduce burden on deeper layers.

Formally, E is part of the model parameters θ.

Initialization

Common initialization strategies:

• •Random normal: E[i,j] ∼ 𝒩(0, σ²)
• •Xavier/Glorot or Kaiming style scaling (depending on architecture)

A typical goal is to keep activations at reasonable scale early in training.

If σ is too large:

• •early activations can explode

If σ is too small:

• •signals and gradients can vanish

How the embedding row gets updated (the key idea)

Suppose in one training step you see token i at some position. The forward pass retrieves eᵢ = E[i]. The loss L depends on that vector through downstream computations.

Backprop computes the gradient:

• •∂L/∂E

But importantly:

• •Only the rows of E corresponding to tokens that appeared in the batch get nonzero gradients.

This matches the lookup behavior: you only “touch” the embeddings you used.

A small derivation: gradient for a simple downstream linear layer

Consider a toy model:

1) Lookup: eᵢ = E[i]

2) Linear: z = Weᵢ + b

3) Loss L depends on z

Let’s compute ∂L/∂eᵢ.

We use chain rule:

• •∂L/∂eᵢ = (∂z/∂eᵢ)ᵀ (∂L/∂z)

But z = Weᵢ + b, so:

• •∂z/∂eᵢ = W

Thus:

• •∂L/∂eᵢ = Wᵀ (∂L/∂z)

And since eᵢ is the i-th row of E, the gradient for E[i] is exactly ∂L/∂eᵢ.

Finally, a gradient descent update (learning rate η):

• •E[i] ← E[i] − η (∂L/∂E[i])

Every time token i appears, its vector is nudged to reduce loss.

Frequency effects

Tokens that appear more often get updated more often.

This can be good (more data) but also can cause imbalance:

• •very frequent tokens may dominate training dynamics
• •rare tokens may remain undertrained

Some tokenization strategies (subwords) help reduce the number of truly rare tokens by composing words out of more frequent pieces.

Embedding tying (brief but important)

In many language models, the input embedding matrix E is tied (shared) with the output projection matrix used to predict token logits.

If output logits are computed as:

• •logits = H Eᵀ

then E serves two roles:

• •mapping token IDs to vectors (input)
• •mapping hidden states back to vocabulary scores (output)

This reduces parameters and often improves performance, but it couples constraints: the same geometry must serve both input and output.

Pretrained embeddings vs learned from scratch

You can initialize E using pretrained vectors (word2vec/GloVe) or from a pretrained transformer.

Comparing options:

Even when using pretrained embeddings, fine-tuning (updating E) is common.

Why d matters

The embedding dimensionality d is the length of each token vector eᵢ ∈ ℝᵈ.

Choosing d is a trade-off between:

• •capacity: how many distinctions the model can encode at the token level
• •generalization: whether the model learns robust patterns vs memorizing
• •compute/memory: E has V·d parameters, and embeddings flow through the whole network

Capacity intuition

With larger d, each token has more degrees of freedom.

• •Small d (e.g., 32, 64): forces compression; may be insufficient for complex tasks
• •Medium d (e.g., 256, 512, 768): common in transformers
• •Large d (e.g., 1024, 2048+): higher capacity but expensive

But “bigger d” isn’t automatically better. If you don’t have enough data or model structure to use it, the extra dimensions can become noisy.

Memory and compute scaling

Parameters in E scale linearly in d:

• •params(E) = V·d

The embedding activations for a batch scale as:

• •activation size ≈ B·T·d

So increasing d increases both parameter memory and the size of tensors passed through attention/MLP blocks.

Matching d to model hidden size

In transformers, token embeddings are usually produced in the same dimension as the model’s hidden size (often called d_model).

That way, you can add other vectors (like positional encodings) and pass embeddings directly into attention blocks without extra projections.

A geometric view: dot products and norms

Many downstream operations depend on dot products.

For random vectors with independent components, dot products tend to grow with d unless normalized. This is one reason initialization and normalization layers matter.

If a, b have typical component scale σ, then expected magnitude:

• •‖a‖ ≈ √d · σ

So as d grows, norms grow unless σ shrinks. Good initialization tries to keep these scales stable.

Practical heuristics

Common rules of thumb (not laws):

• •For small vocabularies and simple tasks, smaller d can work well.
• •For large language models, d is typically tied to architecture size.
• •If V is extremely large (e.g., millions), you may need smaller d or specialized techniques (adaptive embeddings, hashing, etc.).

The key is: d is a design knob controlling representational bandwidth at the input.

Where embeddings sit in the Transformer pipeline

A simplified Transformer input step:

1) Token IDs X ∈ {0,…,V−1}^(B×T)

2) Token embeddings: H_tok[b,t] = E[X[b,t]]

3) Add positional information: H₀ = H_tok + P

4) Pass through stacked attention + MLP blocks

The crucial point: attention doesn’t operate on token IDs; it operates on vectors.

Positional information and why it’s separate

Token embeddings alone do not encode order. The tokens “dog bites man” vs “man bites dog” would be the same multiset of embeddings.

Transformers typically add a positional encoding/embedding P ∈ ℝ^(T×d) (or learned position embeddings) so each position t has its own vector pₜ.

Then:

• •h₀,ₜ = e_{xₜ} + pₜ

This simple addition works because both are in ℝᵈ.

Masking connection (why you’ll need it soon)

Once you have embeddings and pass them into attention, the model computes attention scores between positions.

But you often must prevent attention to:

• •padding tokens (so the model doesn’t treat padding as data)
• •future tokens in autoregressive models (causal masking)

Masking operates on attention score matrices, not on E directly—but embeddings are what make attention possible in the first place.

Embeddings beyond language

The same idea applies broadly:

• •Recommender systems: user ID → embedding, item ID → embedding
• •Graphs: node ID → embedding
• •Multimodal: discretized image patches or code tokens → embeddings

Whenever you have a discrete symbol set, an embedding matrix is a standard first tool.

What to watch for when you implement

A minimal checklist:

• •Ensure your token IDs are in range [0, V−1]
• •Decide how to handle padding tokens (often reserve an ID for PAD)
• •Decide whether PAD’s embedding is learned or fixed to 0
• •Confirm embedding dimension d matches the model’s expected hidden size

These decisions will directly affect training stability and correctness.