Embedding Layer

April 2026 · ML Series · Transformers

From Text to Numbers

Neural networks operate on numbers. They multiply matrices, compute gradients, and adjust weights - all purely numerical operations. They have absolutely no concept of what a “word” is. So the very first question any language model must answer is: how do we turn text into numbers?

This is where the embedding layer enters the picture. It is the very first operation in a transformer - the bridge between the symbolic world of discrete tokens and the continuous numerical world where the transformer does all its reasoning. Every single input token passes through this layer before anything else happens.

The pipeline looks like this:

For example, the sentence “The cat sat” might be tokenized into three tokens - "The", "cat", "sat" - each assigned a unique integer ID by the tokenizer’s vocabulary. The embedding layer then converts each of those integer IDs into a dense, high-dimensional vector (typically 512 to 8192 dimensions). These vectors are what the transformer’s self-attention and feed-forward layers actually consume.

Without the embedding layer, the transformer has no input. It is, quite literally, where all the magic begins.

The Lookup Table Concept

Despite its importance, the embedding layer is one of the simplest components in a transformer. It is a lookup table - nothing more. Concretely, the embedding layer is a matrix E of shape (vocab_size, d_model):

Each row of this matrix is the embedding vector for one token in the vocabulary. Row 0 corresponds to token ID 0, row 1 to token ID 1, and so on. To get the embedding for a token with ID i, you simply take row i from the matrix:

This is a critical point that many explanations get wrong: the embedding lookup is NOT a matrix multiplication. It is an array index operation - equivalent to E[42] in any programming language. The cost is O(1) per token (or more precisely, O(d_model) to copy the row), not O(vocab_size).

For a sequence of n tokens, we perform n lookups and stack the resulting vectors into a matrix of shape (n, d_model). This is the input to the rest of the transformer.

The values in this matrix are learned parameters. They start as random numbers and are updated during training through backpropagation, just like any other weight matrix in the network. Over millions of training steps, the model learns to place semantically similar tokens near each other in this embedding space.

One-Hot Encoding: The Bad Alternative

Before embeddings, the standard way to represent categorical data for neural networks was one-hot encoding. In one-hot encoding, each token is represented as a vector of length vocab_size, with a 1 at the position corresponding to the token’s ID and 0s everywhere else:

This is a terrible representation for language. Here’s why:

There is a useful mathematical insight here. If you multiply a one-hot vector by a weight matrix, the result is exactly one row of that matrix - the row at the position of the 1. So:

This means a one-hot vector multiplied by a weight matrix is mathematically equivalent to a lookup. But the one-hot matmul is O(V × d) while the lookup is O(d). The embedding layer skips the useless multiplication by zeros entirely and goes straight to indexing the row. It is the sparse shortcut for what would otherwise be a very wasteful matrix multiplication.

Vocabulary and Tokenization Connection

The embedding layer’s vocabulary size is determined entirely by the tokenizer. The tokenizer is trained (or constructed) before the language model itself, and the vocabulary is fixed at training time. The embedding matrix must have exactly one row for every token in the vocabulary.

Different models use different tokenization strategies, which produce different vocabulary sizes:

Note the massive variation - GPT-2 has 50,257 tokens while Gemma uses 256,000. Larger vocabularies can represent more words as single tokens (fewer subword splits), but the embedding matrix becomes correspondingly larger.

Special Tokens

Every vocabulary includes special tokens that have structural rather than linguistic meaning. These get their own rows in the embedding table, just like regular tokens:

These special tokens start with randomly initialized embeddings and learn useful representations during training. For instance, BERT’s [CLS] token learns to aggregate sequence-level information, making it useful for classification tasks.

Embedding Dimension

The embedding dimension d_model is one of the most important architectural choices in a transformer. It determines how much information each token vector can carry, and it ripples through the entire model: every layer’s hidden size, every attention head’s key/value/query size, and the final output projection all depend on d_model.

The tradeoff is straightforward:

• Larger d_model → more expressive embeddings, can capture finer-grained distinctions between tokens, but requires more memory, more compute, and more training data to learn well.
• Smaller d_model → faster, lighter, more suitable for edge deployment, but may collapse semantically distinct tokens into similar regions of the embedding space.

A rough rule of thumb: double the embedding dimension, and you roughly quadruple the model’s compute requirements (since compute scales as O(d²) in the attention and feed-forward layers). The embedding layer itself scales as O(V × d), which is linear in d.

Interactive: Token-to-Vector Lookup

This animation shows exactly how the embedding layer works. We take the sentence “The cat sat”, tokenize it, and look up each token’s embedding vector from the embedding matrix. Step through to watch each token get resolved to a dense vector.

The √d_model Scaling

In the original “Attention Is All You Need” paper (Vaswani et al., 2017), there is a subtle but important step right after the embedding lookup: the embedding vectors are multiplied by √d_model before adding positional encodings.

Why? The reason comes down to scale matching:

The math works out cleanly: if each component of the embedding is drawn from N(0, 1/√d_model), then after scaling by √d_model, the L2 norm of the embedding vector is approximately √d_model, giving it a magnitude comparable to the positional encoding.

Now both signals have comparable magnitude, so neither drowns out the other when added together. The model can learn to use both the identity (what token) and position (where in the sequence) information.

Weight Tying (Shared Embeddings)

One of the most elegant tricks in transformer design is weight tying (also called shared embeddings). The idea is simple but powerful: use the same weight matrix for both the input embedding and the output prediction layer.

The Setup

In a language model, two matrices have suspiciously similar shapes:

Notice that W = E^T would work perfectly - the output projection is just the transpose of the embedding matrix. Weight tying does exactly this: it sets the output projection (the “LM head”) to be the transpose of the input embedding matrix.

Why It Works

Weight tying makes deep semantic sense. Consider what each matrix does:

• Input embedding E: maps token "cat" to a vector v_cat
• Output projection E^T: computes the dot product of the hidden state with v_cat. If the hidden state is close to v_cat, the logit for "cat" is high.

This creates a beautiful symmetry: to predict "cat", the model needs to produce a hidden state that is close (in dot-product space) to the input embedding of "cat". The embedding space serves double duty as both the input representation and the output target space.

Benefits

Models That Use Weight Tying

• GPT-2 - ties input and output embeddings
• T5 - ties encoder input, decoder input, and decoder output embeddings (three-way sharing)
• ALBERT - ties embeddings and also shares parameters across all transformer layers
• BLOOM - ties input and output embeddings
• Gemma - ties input and output embeddings

Some models (like LLaMA) do not use weight tying. When the model is large enough, the parameter savings are relatively small (< 3% of total parameters), and having separate matrices can provide additional expressivity.

Embedding Space Geometry

Once a model is trained, the embedding matrix is no longer random noise. It has learned a rich geometric structure that captures linguistic relationships. Let’s explore what this space looks like.

Semantic Clustering

Similar words end up near each other in embedding space. If you project the high-dimensional embeddings into 2D (using t-SNE or UMAP), you’ll see clear clusters:

• Royalty cluster: king, queen, prince, princess, monarch, royal
• Animal cluster: cat, dog, bird, fish, horse, elephant
• Number cluster: one, two, three, four, five, six
• Punctuation cluster: . , ! ? ; :

Within each cluster, finer-grained structure exists. “King” and “queen” are closer to each other than either is to “prince” or “princess,” reflecting their shared semantic role as rulers rather than offspring.

The Famous Analogy

The most celebrated property of embedding spaces is the ability to encode analogies as vector arithmetic:

This works because the embedding space encodes relational differences as consistent vector offsets. The vector from “man” to “king” is approximately the same as the vector from “woman” to “queen” - both represent the concept of “royalty.”

Subword Structure

Since modern tokenizers use subword units (BPE, WordPiece), the embedding layer must also handle morphological structure. The token "un" learns an embedding that, when combined with "happy" through the transformer layers, produces the meaning of “unhappy.” Similarly, "##ing" (a WordPiece suffix) learns a representation that modifies the base verb into its progressive form.

Cosine Similarity

The standard way to measure similarity in embedding space is cosine similarity:

Typical values for well-trained embeddings:

• Synonyms: cos(“happy”, “joyful”) ≈ 0.75–0.85
• Related: cos(“doctor”, “hospital”) ≈ 0.55–0.65
• Unrelated: cos(“cat”, “democracy”) ≈ 0.05–0.15
• Antonyms: cos(“hot”, “cold”) ≈ 0.40–0.55 (surprisingly similar, because they share context)

Initialization Strategies

How the embedding matrix is initialized before training has a measurable impact on convergence speed and final model quality. The values must be random (to break symmetry) but carefully scaled.

Random Normal Initialization

The most common strategy. Each element of the embedding matrix is drawn independently from a normal distribution:

The 1/√d_model factor keeps the variance of the output manageable. For d_model = 768, the standard deviation is 1/√768 ≈ 0.036. This means initial embedding values are very small numbers, typically between -0.1 and 0.1.

Truncated Normal

Like normal initialization, but values beyond 2σ are re-sampled. This prevents extreme outlier values in the initial embedding:

Used by BERT and many Google-origin models. The truncation prevents individual embedding dimensions from having extreme initial values that could cause gradient instabilities.

Xavier / Glorot Initialization

Occasionally used, drawing from a uniform distribution:

Less common for embeddings specifically, but sometimes used in frameworks that apply Xavier initialization to all weight matrices uniformly.

Pretrained Initialization

For fine-tuning tasks, embeddings can be initialized from pretrained word vectors like Word2Vec or GloVe. This gives the model a head start with semantically meaningful embeddings. However, this approach has largely fallen out of favor since modern models are pretrained end-to-end on massive corpora.

Segment Embeddings and Token Type

Some transformer architectures add additional embeddings on top of the token and position embeddings. The most notable example is BERT’s segment embeddings (also called token type embeddings).

BERT’s Three-Embedding Sum

BERT processes pairs of sentences for tasks like natural language inference and question answering. It needs a way to distinguish which sentence each token belongs to. The solution is elegant: add a third embedding that encodes the sentence membership.

The segment embedding table is tiny - just 2 rows × d_model values (one for sentence A, one for sentence B). But it gives the model crucial information about sentence boundaries.

GPT Models Don’t Need Segments

Decoder-only models like GPT process one continuous sequence at a time. There is no concept of “sentence A” vs “sentence B” - everything is just a stream of tokens. So GPT models only use token embeddings + positional embeddings (no segment embedding).

Other Embedding Types

Various models add other specialized embeddings:

• Language embeddings (XLM, mBERT): indicate the language of the input (English=0, French=1, etc.)
• Role embeddings (some dialogue models): distinguish user turns from assistant turns
• Modality embeddings (multimodal transformers): distinguish image tokens from text tokens

All of these follow the same pattern: a small lookup table whose output is added element-wise to the main token + position embedding.

Parameter Count

The embedding layer’s parameter count is straightforward to compute:

But what’s surprising is how large this number can be relative to the total model:

Two clear patterns emerge:

• In smaller models, the embedding layer is a massive fraction of total parameters - one-third of GPT-2’s parameters are just the embedding table. Weight tying is especially valuable here.
• In larger models, the transformer layers dominate (they scale as O(d²) per layer), so the embedding fraction shrinks below 2%. Weight tying saves less proportionally, which is why models like LLaMA skip it.
• Large vocabularies like Gemma’s 256K tokens can push the embedding fraction back up, even for large models.

Implementation in Pseudo-Code

Here is a clean implementation of the embedding layer with all the features we’ve discussed:

// ─── Embedding Layer ───────────────────────────────────

class TokenEmbedding:
    // Parameters
    E : matrix of shape (vocab_size, d_model)   // the embedding table
    scale : float = sqrt(d_model)               // optional scaling factor

    function __init__(vocab_size, d_model, use_scaling=True):
        // Initialize from N(0, 1/sqrt(d_model))
        E = random_normal(0, 1/sqrt(d_model), shape=(vocab_size, d_model))
        scale = sqrt(d_model) if use_scaling else 1.0

    function forward(token_ids):
        // token_ids: integer tensor of shape (batch, seq_len)
        // Returns: float tensor of shape (batch, seq_len, d_model)

        embeddings = E[token_ids]    // lookup, NOT matmul
        embeddings = embeddings * scale

        return embeddings

And here’s how it fits into the full transformer input pipeline:

// ─── Full Transformer Input Pipeline ──────────────────

class TransformerInput:
    token_emb  : TokenEmbedding(vocab_size, d_model)
    pos_emb    : PositionalEncoding(max_seq_len, d_model)
    seg_emb    : SegmentEmbedding(2, d_model)   // BERT only
    dropout    : Dropout(rate=0.1)

    function forward(token_ids, segment_ids=None):
        // Step 1: Token embeddings (lookup + scale)
        x = token_emb.forward(token_ids)

        // Step 2: Add positional encoding
        x = x + pos_emb(positions)

        // Step 3: Add segment embedding (if BERT-style)
        if segment_ids is not None:
            x = x + seg_emb[segment_ids]

        // Step 4: Dropout for regularization
        x = dropout(x)

        return x   // shape: (batch, seq_len, d_model)

And for weight-tied output projection:

// ─── Weight-Tied Output Projection ────────────────────

class LMHead:
    shared_emb : reference to TokenEmbedding.E   // same matrix!

    function forward(hidden_states):
        // hidden_states: (batch, seq_len, d_model)
        // shared_emb.T:  (d_model, vocab_size)
        logits = hidden_states @ shared_emb.T
        return logits   // shape: (batch, seq_len, vocab_size)

PyTorch One-Liner

In practice, PyTorch’s nn.Embedding handles all of this:

// PyTorch implementation
embedding = nn.Embedding(num_embeddings=50257, embedding_dim=768)
x = embedding(token_ids)   // shape: (batch, seq_len, 768)

// Weight tying
lm_head = nn.Linear(768, 50257, bias=False)
lm_head.weight = embedding.weight   // shared!

Summary

The embedding layer is deceptively simple - it is just a matrix lookup - but it is the foundation on which the entire transformer is built. Here are the key takeaways:

The embedding layer converts the discrete, symbolic world of language into the continuous, numerical world that neural networks can process. Every breakthrough in modern NLP - from BERT to GPT-4 - begins with this simple table lookup.