Text Preprocessing and Embeddings

From Text to Numbers: The Core Problem

A neural network does not understand words. It understands tensors of floating-point numbers. Every layer we have studied so far — linear projections, convolutions, recurrences, and attention — is defined as operations over real-valued vectors. So before a single weight can be applied to the sentence “The cat sat on the mat”, we must answer a deceptively hard question: how do we turn a string into a vector without destroying its meaning?

That question has one of the most consequential answers in modern machine learning: we use an embedding — a learnable map from discrete tokens into a continuous vector space where geometry encodes meaning. Everything that comes later in the model — the attention scores in a Transformer, the hidden state of an LSTM, the logits a language model produces — is ultimately a function of those initial vectors. Get the embedding wrong and nothing downstream can rescue the model.

The promise of embeddings: if we can arrange words in space so that semantically similar words land near each other, then linear operations in that space become semantic operations. The famous example $\text{king} - \text{man} + \text{woman} \approx \text{queen}$ is possible precisely because the learned geometry reflects meaning.

In this section we will walk the entire pipeline — from raw Unicode text to the matrix of dense vectors that the next layer consumes. We will build every piece twice: first from scratch in NumPy, so every number is verifiable; then in PyTorch with $\texttt{nn.Embedding}$ so you see how the same idea scales into a trainable layer. Along the way we will connect embeddings to the layers that will dominate the rest of the book: positional encodings, multi-head attention, Flash Attention, and the KV-cache that makes modern LLM inference feasible.

The Preprocessing Pipeline

Preprocessing is the sequence of transformations that converts raw text into a tensor of integer IDs. Different systems differ in the details but almost every modern pipeline has the same five stages:

1. Normalize — collapse surface variation that is not semantically meaningful: case, whitespace, Unicode compatibility forms, and sometimes accents. The string “The” and “the” should hit the same vocabulary entry.
2. Tokenize — split the normalized string into atomic units. Word-level splits on whitespace; character-level uses single characters; subword approaches like Byte-Pair Encoding (BPE) sit in between. The choice fundamentally shapes what the model can and cannot express.
3. Build a vocabulary — a bijection between token strings and small integers. This vocabulary is fixed at training time and shared between training and inference.
4. Index — look up each token in the vocabulary and replace it with its integer ID. The string disappears; only an array of integers remains.
5. Embed — map each integer ID to a dense, d-dimensional vector via a learnable lookup table. This is the bridge from discrete language to continuous geometry.

Visually, the data flows from a string on the left to a $T \times d$ matrix of floats on the right:

The first four stages are deterministic string manipulation. The fifth — embedding — is where the learning happens. Everything interesting in this section lives in that last arrow.

One-Hot Encoding: The Naive Baseline

The simplest way to turn an integer ID into a vector is one-hot encoding. For a vocabulary of size $V$, token $i$ becomes a vector of length $V$ with a $1$ at position $i$ and $0$ everywhere else. Formally, the one-hot vector $\mathbf{e}_i \in \mathbb{R}^V$ satisfies $(\mathbf{e}_i)_j = \delta_{ij}$, the Kronecker delta.

For our sentence with vocabulary $\{\langle \text{pad}\rangle:0,\ \langle \text{unk}\rangle:1,\ \text{the}:2,\ \text{cat}:3,\ \text{sat}:4,\ \text{on}:5,\ \text{mat}:6\}$, the six tokens become a $6 \times 7$ matrix:

Notice that both occurrences of the have identical rows — the vocabulary guarantees that identical strings produce identical vectors. This is already a property we want: a neural network should not care which copy of the it is looking at.

Why One-Hot Fails: The Curse of Orthogonality

One-hot vectors have two fatal flaws for language modelling.

1. They are orthogonal — so every pair of different words is equally dissimilar

The dot product of any two distinct one-hot vectors is zero: $\mathbf{e}_i \cdot \mathbf{e}_j = \delta_{ij}$. That means the model has no way to learn that cat and kitten are closer in meaning than cat and photosynthesis. Every word is equidistant from every other word. Geometry carries zero semantic information.

2. They are extremely high-dimensional and sparse

A vocabulary of $V = 128{,}000$ (LLaMA 3) means every word becomes a 128,000-dimensional vector. Storing a 1,000-token sequence at float32 would take $1000 \times 128000 \times 4 \text{ bytes} \approx 512\,\text{MB}$ — and the first layer would need to multiply that by a $128000 \times d$ weight matrix, almost all of which gets multiplied by zeros. This is wasteful in both memory and compute.

3. The number of parameters explodes linearly in vocab size

Any dense layer fed by a one-hot input has $V \times d$ weights. A word-level model for English can easily reach $V > 10^5$ — already in the tens of millions of parameters just to consume the input.

Embeddings: Dense Vectors with Meaning

An embedding is a function $\text{Embed}: \{0, 1, \ldots, V-1\} \to \mathbb{R}^d$ that assigns every token ID a dense vector of dimension $d \ll V$. The function is represented as a lookup table: a matrix $\mathbf{E} \in \mathbb{R}^{V \times d}$ whose $i$-th row is the vector for token $i$. The lookup is simply row indexing:

$\text{Embed}(i) = \mathbf{E}[i, :] \in \mathbb{R}^d$

Three design decisions give embeddings their power:

1. Every row is a separate learnable parameter. During training, gradient descent can move each token’s vector independently, so the model discovers its own geometry.
2. The dimension $d$ is small. Typical values: 128 (word2vec), 300 (GloVe), 768 (BERT-base), 4096 (LLaMA-2-7B), 12288 (GPT-3). Much smaller than any real vocabulary.
3. Rows can be close. Nothing stops cat and kitten from pointing in nearly the same direction — and training actively encourages it when their contexts are similar. That is the Distributional Hypothesis made concrete: you shall know a word by the company it keeps.

The Mathematics of Embedding Lookup

Here is a piece of trivia with serious consequences: embedding lookup is a matrix multiplication in disguise.

Let $\mathbf{O} \in \{0,1\}^{T \times V}$ be the one-hot matrix for a sequence of $T$ tokens, and $\mathbf{E} \in \mathbb{R}^{V \times d}$ the embedding matrix. Then the embedded sequence is:

$\mathbf{X} = \mathbf{O}\mathbf{E} \in \mathbb{R}^{T \times d}$

Because row $i$ of $\mathbf{O}$ contains a single $1$ at column $\text{id}_i$ and zeros elsewhere, the inner product $\mathbf{O}[i,:] \cdot \mathbf{E}[:,j] = \mathbf{E}[\text{id}_i, j]$. The full product collapses to row selection:

$\mathbf{X}[i, :] = \mathbf{E}[\text{id}_i, :]$

So $\mathbf{O}\mathbf{E}$ and the fancy index $\mathbf{E}[\text{ids}]$ compute the exact same thing. The matmul form costs $O(T \cdot V \cdot d)$; the lookup form costs $O(T \cdot d)$. For modern vocabularies with $V > 10^5$ the speedup is five orders of magnitude per token.

The interactive below makes the identity concrete. Click any row of $\mathbf{O}$ to see exactly which row of $\mathbf{E}$ is “picked up”. Toggle between the full matmul and the direct lookup and notice that the output matrix is bit-identical.

The sparse gradient

The matmul view also explains the gradient. If loss $L$ depends on $\mathbf{X} = \mathbf{O}\mathbf{E}$, then by the chain rule:

$\frac{\partial L}{\partial \mathbf{E}} = \mathbf{O}^{\top} \frac{\partial L}{\partial \mathbf{X}}$

Because $\mathbf{O}^{\top}$ is almost entirely zeros, almost every entry of the gradient $\partial L / \partial \mathbf{E}$ is also zero. Only the rows corresponding to token IDs that actually appeared in the batch receive updates. Frameworks exploit this by storing the embedding gradient in a sparse format and by providing specialized optimizer variants (e.g. $\texttt{torch.optim.SparseAdam}$).

Building It From Scratch in NumPy

Let us implement the whole pipeline — normalize, tokenize, build vocab, index, and embed — using nothing but NumPy. Every line is annotated with the exact values it produces. Click any line on the right to see the execution state on the left.

all zeros — 42 floats, all 0.0

       <pad> <unk>  the  cat  sat   on  mat
row0    0    0    1    0    0    0    0
row1    0    0    0    1    0    0    0
row2    0    0    0    0    1    0    0
row3    0    0    0    0    0    1    0
row4    0    0    1    0    0    0    0
row5    0    0    0    0    0    0    1

         d0    d1    d2    d3
<pad>   0.00  0.00  0.00  0.00
<unk>   0.10  0.20  0.30  0.40
the     0.20  0.10  0.05  0.30
cat     0.80  0.60  0.20  0.10
sat     0.40  0.70  0.90  0.20
on      0.30  0.30  0.10  0.50
mat     0.70  0.50  0.30  0.20

       d0    d1    d2    d3
row0  0.20  0.10  0.05  0.30   ← the
row1  0.80  0.60  0.20  0.10   ← cat
row2  0.40  0.70  0.90  0.20   ← sat
row3  0.30  0.30  0.10  0.50   ← on
row4  0.20  0.10  0.05  0.30   ← the (identical to row0)
row5  0.70  0.50  0.30  0.20   ← mat

       d0    d1    d2    d3
row0  0.20  0.10  0.05  0.30   ← E[2]
row1  0.80  0.60  0.20  0.10   ← E[3]
row2  0.40  0.70  0.90  0.20   ← E[4]
row3  0.30  0.30  0.10  0.50   ← E[5]
row4  0.20  0.10  0.05  0.30   ← E[2] again
row5  0.70  0.50  0.30  0.20   ← E[6]

[[0.2  0.1  0.05 0.3 ]
 [0.8  0.6  0.2  0.1 ]
 [0.4  0.7  0.9  0.2 ]
 [0.3  0.3  0.1  0.5 ]
 [0.2  0.1  0.05 0.3 ]
 [0.7  0.5  0.3  0.2 ]]

1import numpy as np
2
3# --- 1. Raw input text ---
4text = "The cat sat on the mat"
5
6# --- 2. Normalize (lowercase + strip whitespace) ---
7normalized = text.lower().strip()
8
9# --- 3. Tokenize by whitespace ---
10tokens = normalized.split()
11# tokens = ['the', 'cat', 'sat', 'on', 'the', 'mat']
12
13# --- 4. Build vocabulary with special symbols ---
14vocab = {'<pad>': 0, '<unk>': 1}
15for tok in tokens:
16    if tok not in vocab:
17        vocab[tok] = len(vocab)
18
19# --- 5. Convert tokens to integer IDs ---
20input_ids = np.array([vocab.get(t, vocab['<unk>']) for t in tokens])
21
22# --- 6. One-hot encode each token (the naive approach) ---
23vocab_size = len(vocab)
24one_hot = np.zeros((len(input_ids), vocab_size))
25for i, idx in enumerate(input_ids):
26    one_hot[i, idx] = 1.0
27
28# --- 7. Define an embedding matrix (normally learned) ---
29d_embed = 4
30E = np.array([
31    [0.00, 0.00, 0.00, 0.00],   # <pad>
32    [0.10, 0.20, 0.30, 0.40],   # <unk>
33    [0.20, 0.10, 0.05, 0.30],   # the
34    [0.80, 0.60, 0.20, 0.10],   # cat
35    [0.40, 0.70, 0.90, 0.20],   # sat
36    [0.30, 0.30, 0.10, 0.50],   # on
37    [0.70, 0.50, 0.30, 0.20],   # mat
38])
39
40# --- 8. Way A: one-hot @ E   (O(T*V*d) — wasteful) ---
41emb_matmul = one_hot @ E
42
43# --- 9. Way B: direct row-lookup   (O(T*d) — what real frameworks do) ---
44emb_lookup = E[input_ids]
45
46print("input_ids:", input_ids)
47print("embedded shape:", emb_lookup.shape)
48print(emb_lookup)

Two lines deserve extra attention. Line 41 does the lookup the slow way — $\texttt{one\_hot @ E}$ — to prove the mathematical identity. Line 44 does it the fast way — $\texttt{E[input\_ids]}$ — which is what every production framework actually runs. Both produce the same $6 \times 4$ matrix. Verify it yourself: rows 0 and 4 are bit-for-bit identical because both correspond to the token the.

The Same Thing in PyTorch: nn.Embedding

The NumPy implementation is crystal-clear but it cannot be trained. PyTorch’s $\texttt{nn.Embedding}$ takes the same idea, wraps the matrix in a $\texttt{nn.Parameter}$, and plugs it into autograd — so every row becomes a learnable parameter of the model.

Three features of $\texttt{nn.Embedding}$ are worth calling out before the code:

• Sparse gradient: only rows whose indices appeared in the forward pass receive gradient updates.
• $\texttt{padding\_idx}$: the row at this index is initialized to zeros and its gradient is permanently masked to zero — so pad positions stay dead forever.
• $\texttt{max\_norm}$: optionally renormalizes each looked-up row to a fixed L2 norm on the fly — useful when you want to prevent a runaway-embedding pathology.

Below we run a full forward pass, define a contrived squared-error loss, and call $\texttt{loss.backward()}$. The annotations walk through the exact gradient computed for each row — including the critical observation that row 2 (the word “the”) accumulates contributions from both positions where it appears.

tensor([[0.2000, 0.1000, 0.0500, 0.3000],
        [0.8000, 0.6000, 0.2000, 0.1000],
        [0.4000, 0.7000, 0.9000, 0.2000],
        [0.3000, 0.3000, 0.1000, 0.5000],
        [0.2000, 0.1000, 0.0500, 0.3000],
        [0.7000, 0.5000, 0.3000, 0.2000]],
       grad_fn=<EmbeddingBackward0>)

1import torch
2import torch.nn as nn
3
4# Same vocab as before — 7 tokens total
5vocab_size = 7
6d_embed = 4
7input_ids = torch.tensor([2, 3, 4, 5, 2, 6])
8
9# nn.Embedding: a learnable lookup table
10embedding = nn.Embedding(
11    num_embeddings=vocab_size,
12    embedding_dim=d_embed,
13    padding_idx=0,
14)
15
16# Pin the weights to our hand-designed matrix for reproducibility
17with torch.no_grad():
18    embedding.weight.copy_(torch.tensor([
19        [0.00, 0.00, 0.00, 0.00],
20        [0.10, 0.20, 0.30, 0.40],
21        [0.20, 0.10, 0.05, 0.30],
22        [0.80, 0.60, 0.20, 0.10],
23        [0.40, 0.70, 0.90, 0.20],
24        [0.30, 0.30, 0.10, 0.50],
25        [0.70, 0.50, 0.30, 0.20],
26    ]))
27
28# --- Forward: lookup by ID ---
29embedded = embedding(input_ids)   # shape: (6, 4)
30
31# --- Training signal: push each token toward a target vector ---
32target = torch.tensor([
33    [0.0, 0.0, 0.0, 0.0],   # want "the"  ~ 0
34    [1.0, 1.0, 1.0, 1.0],   # "cat" ~ 1
35    [1.0, 1.0, 1.0, 1.0],   # "sat" ~ 1
36    [0.5, 0.5, 0.5, 0.5],   # "on"  ~ 0.5
37    [0.0, 0.0, 0.0, 0.0],   # "the" ~ 0  (same target as row 0)
38    [1.0, 1.0, 1.0, 1.0],   # "mat" ~ 1
39])
40
41loss = ((embedded - target) ** 2).mean()
42loss.backward()
43
44# Only rows for tokens that appeared get non-zero gradients
45print("grad for row 2 ('the'):", embedding.weight.grad[2])
46print("grad for row 3 ('cat'):", embedding.weight.grad[3])
47print("grad for row 0 ('<pad>'):", embedding.weight.grad[0])

Gradient accumulation is a feature, not a bug. Because the same row is selected twice (once per occurrence of the), its gradient is the sum of the per-position gradients. That is how a token learns from every context it appears in — the more contexts, the stronger the learning signal. This is also why embeddings for frequent words (like the) typically converge much faster than embeddings for rare words.

The panel below traces exactly which rows of $\mathbf{E}$ receive gradient for this batch. Rows that never appeared stay dark. Toggle $\texttt{padding\_idx}$ off and watch row 0 light up; toggle it on again and the lock reinstates.

Visualizing the Embedding Space

Numbers are hard to feel. Geometry is easy to feel. The visualization below shows a toy embedding space with $d = 3$ (so we can actually draw it). Each token is a coloured point in 3-space, and tokens in the same semantic cluster — royalty, people, animals, fruits, activities — were placed near each other. Rotate, zoom, and pan to see the clusters from every angle.

What you are seeing is a learned structure. Nothing in the training objective said “put royalty on one side and fruit on the other”. The distributional hypothesis — that words with similar contexts should have similar vectors — is enough. Once the loss has been minimized, similar words inevitably land near each other because pushing them apart would increase the loss.

Two geometric facts drop out of the training that seem magical at first:

• Direction encodes semantic relations. $\vec{\text{king}} - \vec{\text{man}} + \vec{\text{woman}} \approx \vec{\text{queen}}$ — the male→female direction is (approximately) a translation in the embedding space.
• Distance encodes semantic similarity. Cosine similarity between embedding vectors is one of the most widely used proxies for semantic similarity in information retrieval, semantic search, and RAG systems.

How Embeddings Learn Meaning

Where does the geometry come from? It is not hand-designed — it is a byproduct of training the model on its primary task. Whether the task is classifying sentiment, predicting the next token, or answering a question, the gradient signal has to pass through the embedding layer on its way back to the weights. Each such pass nudges the rows of $\mathbf{E}$ in whatever direction reduces the loss.

Concretely, for a language-modelling objective with cross-entropy loss$L = -\log p(w_t \mid w_{<t})$, the gradient with respect to the embedding of an input word is:

$\frac{\partial L}{\partial \mathbf{E}[\text{id}_i]} = \frac{\partial L}{\partial \mathbf{X}[i,:]}$

If two words tend to appear in the same contexts, the upstream gradient flowing into their embedding rows is, on average, similar. Over many updates their vectors drift in the same direction — and they end up geometrically close. This is exactly the distributional hypothesis expressed as an equation.

The same idea, pre-trained on huge text corpora, produced the classic static embeddings — word2vec (skip-gram and CBOW, Mikolov et al., 2013) and GloVe (Pennington et al., 2014). Today’s LLMs learn their embeddings jointly with the rest of the Transformer, so word2vec-style pre-training is largely historical. But the lesson is the same: geometry emerges from prediction.

Embeddings in Modern Transformers

The embedding layer is the first layer in every decoder-only LLM on the planet. Everything downstream — self-attention, Flash Attention, KV-cache, RoPE — is a transformation of the matrix the embedding layer produces. Understanding how embeddings connect to those mechanisms is the difference between reading the diagrams and understanding them.

Token embedding + positional encoding

A vanilla self-attention layer is permutation-invariant: shuffle the input tokens and the output shuffles the same way. For language that is disastrous — “dog bites man” and “man bites dog” must produce different outputs. The standard fix is to add a positional encoding $\mathbf{P} \in \mathbb{R}^{T \times d}$ to the token embedding before feeding the sum into the first attention block:

$\mathbf{X}_0 = \mathbf{E}[\text{ids}] + \mathbf{P}$

In the original Transformer of Vaswani et al. (2017), $\mathbf{P}$ is the deterministic sinusoidal function; in GPT-2 it is a learnable parameter of the same shape as$\mathbf{E}[\text{ids}]$. Modern LLMs (LLaMA, Gemma, Mistral) drop additive positions entirely in favour of RoPE (Rotary Position Embedding), which rotates the query and key vectors inside attention itself. Either way, the token embedding from our $\texttt{nn.Embedding}$ is the starting material.

The $\sqrt{d_{\text{model}}}$ scaling trick

You will often see embeddings multiplied by $\sqrt{d_{\text{model}}}$ before being added to positional encodings: $\mathbf{X}_0 = \mathbf{E}[\text{ids}] \cdot \sqrt{d_{\text{model}}} + \mathbf{P}$. This is in the original Transformer paper. The motivation: if embeddings are initialized with variance $1/d_{\text{model}}$ (Xavier init), scaling by $\sqrt{d_{\text{model}}}$ restores unit variance so the embedding does not get drowned out by the positional encoding at initialization.

Weight tying with the output projection

A language model needs to do two conversions: tokens → vectors at the input, and vectors → token probabilities at the output. The input uses $\mathbf{E} \in \mathbb{R}^{V \times d}$; the output uses a linear layer $\mathbf{W}_{\text{out}} \in \mathbb{R}^{d \times V}$. Modern LLMs almost universally tie these: $\mathbf{W}_{\text{out}} = \mathbf{E}^{\top}$. This cuts parameter count roughly in half for the embedding/unembedding pair (40M → 20M for GPT-2) and often improves quality because every gradient step updates both roles of each row coherently.

Embeddings, multi-head attention, and Flash Attention

The matrix $\mathbf{X}_0 \in \mathbb{R}^{T \times d}$ coming out of the embedding + position stage feeds directly into multi-head attention. The queries, keys, and values are linear projections of this matrix:

$\mathbf{Q} = \mathbf{X}_0 \mathbf{W}_Q,\quad \mathbf{K} = \mathbf{X}_0 \mathbf{W}_K,\quad \mathbf{V} = \mathbf{X}_0 \mathbf{W}_V$

So the embedding is literally the ink that attention paints with. This matters for two modern optimizations:

• Flash Attention (Dao et al., 2022) computes $\text{softmax}(\mathbf{Q}\mathbf{K}^{\top}/\sqrt{d_k})\mathbf{V}$ without ever materializing the full $T \times T$ attention matrix. Its memory cost is $O(T \cdot d)$ instead of $O(T^2)$. The $T \cdot d$ factor is set by the embedding dimension — making $d$ too large directly raises Flash Attention’s on-chip memory requirement.
• Multi-head attention splits $d$ across $h$ heads, giving each head $d_k = d / h$ dimensions. The embedding dimension is therefore constrained to be divisible by the number of heads. LLaMA-2-7B uses $d = 4096$ and $h = 32$, giving $d_k = 128$ per head.

The KV-cache and the embedding’s role at inference

At autoregressive inference, a model generates one token at a time. Naive implementations re-run attention over the full sequence on every step — $O(T^2 \cdot d)$ per step. The KV-cache fixes this: it stores every layer’s $\mathbf{K}$ and $\mathbf{V}$ from previous steps so only the new token’s single row must be projected and attended. The cache lives in GPU memory and its size is:

$\text{KV-cache bytes} = 2 \times L \times T \times d \times \text{bytes-per-float}$

where $L$ is the number of layers, $T$ the context length, and $d$ the model dimension — the same $d$ set by the embedding layer. At $L=32,\ T=8192,\ d=4096$, float16, the KV-cache per sequence is about 2 GB. This is why LLM inference is often memory-bound, and why techniques like Grouped-Query Attention, Multi-Query Attention, and KV-cache quantization exist: they all attack the $T \times d$ term that the embedding layer fixes.

Subword Tokenization and the Scaling Story

Our toy pipeline used whitespace tokenization. Real LLMs use subword tokenization — usually Byte-Pair Encoding (BPE) or its byte-level variant. The motivation is practical:

1. Word-level vocabularies are either enormous (to cover every possible word form) or suffer massive out-of-vocabulary rates. Morphologically rich languages (Finnish, Turkish) are catastrophic for word tokenizers.
2. Character-level tokenizers have tiny vocabularies (~256 bytes) but produce very long sequences, which is expensive for attention because of the $T^2$ cost.
3. Subword (BPE, WordPiece, SentencePiece) hits the sweet spot: a vocabulary of a few tens of thousands of subword units that can represent any string by concatenation. Common words get their own token; rare words decompose into meaningful pieces.

Watch the algorithm run on a small toy corpus. Each step merges the most-frequent adjacent pair; after ten merges common substrings like low, est, and er have been absorbed into single subword tokens.

Notice the trend. Modern LLMs have kept $d_{\text{model}}$ roughly fixed (~4k) while growing the vocabulary — the embedding layer keeps getting heavier. That is because more vocabulary means fewer tokens per sentence (better compression → shorter sequences → cheaper attention) and better coverage of rare languages and code. The tradeoff: more parameters and bigger gradients on the embedding layer.

Summary: The Gateway to Every Language Model

Embeddings are the layer that turns language into geometry. Everything a language model does — attention, generation, classification — is a transformation of the matrix produced by the embedding lookup. Mastering this one operation unlocks the rest of the stack.

Concepts to carry forward

1. Preprocessing is a pipeline: normalize → tokenize → build vocab → index → embed. The first four stages are deterministic; the fifth is learned.
2. One-hot is a stepping stone, not a solution. Its orthogonality makes it semantically blind and its sparsity makes it computationally wasteful.
3. Embedding lookup equals one-hot matmul. $\mathbf{O}\mathbf{E}$ and $\mathbf{E}[\text{ids}]$ compute the same thing; only cost differs.
4. Gradients are sparse. Only rows for tokens that appeared in the forward pass receive updates — the reason embedding layers are cheap to train despite their size.
5. Padding and unknown tokens are special. $\texttt{padding\_idx}$ freezes the pad row; $\langle \text{unk}\rangle$ catches out-of-vocab inputs.
6. Embeddings set $d$. The embedding dimension becomes the $d_{\text{model}}$ of attention, the width of every hidden state, and the constant in Flash Attention memory and KV-cache size.

What is next

In the next section we will take the embedded matrix $\mathbf{X}_0 \in \mathbb{R}^{T \times d}$ we built here and feed it into an LSTM classifier. You will see the embedding layer sitting as the first module of a real, trainable sentiment analyzer — the gateway through which every token in your dataset passes on its way to a prediction.

References

• Mikolov, T., Chen, K., Corrado, G. & Dean, J. (2013). Efficient Estimation of Word Representations in Vector Space. arXiv:1301.3781.
• Pennington, J., Socher, R. & Manning, C. D. (2014). GloVe: Global Vectors for Word Representation. EMNLP 2014.
• Vaswani, A. et al. (2017). Attention Is All You Need. NeurIPS 2017 / arXiv:1706.03762.
• Sennrich, R., Haddow, B. & Birch, A. (2016). Neural Machine Translation of Rare Words with Subword Units. ACL 2016 / arXiv:1508.07909.
• Dao, T., Fu, D. Y., Ermon, S., Rudra, A. & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022 / arXiv:2205.14135.