Token Embedding

Mathematical limitations of scalar token IDs:

1. Non-Metric Space:

Let i, j, k ∈ Z be token IDs. The distance |i - j| has no semantic interpretation:

|ID("cat") - ID("dog")| ≠ SemanticDissimilarity("cat", "dog")

1. Zero Gradient Flow:

Because integer IDs are discrete indices, the derivative with respect to an integer ID is undefined:

(∂ L) / (∂ ID_t) = undefined

We cannot adjust an integer ID by a fractional learning rate η = 0.0001.

1. Continuous Embedding Solution:

We map each discrete token i ∈ 0, ṡ, |V|-1 to a continuous vector eᵢ ∈ R^d:

eᵢ = [w_(i,1), w_(i,2), ṡ, w_(i,d)]

where each w_(i,j) ∈ R is a differentiable parameter updated via (∂ L) / (∂ w_(i,j)).

This transition from discrete symbols to continuous vector spaces is the foundational trick of modern deep learning, enabling models to interpolate between concepts, calculate semantic analogies, and perform optimization over language.

In 1957, linguist John Rupert Firth famously wrote: 'You shall know a word by the company it keeps.' Yoshua Bengio et al. (2003) operationalized this intuition in 'A Neural Probabilistic Language Model', introducing continuous word feature vectors to defeat the curse of dimensionality.

• IDs are arbitrary names; embeddings are physical coordinates.
• You cannot differentiate an integer; you can differentiate a coordinate.

Discrete integer tokens introduce fatal artificial geometric constraints:

1. The Scalar Proximity Trap: If `"cat" = 42`, `"dog" = 43`, and `"airplane" = 44`, arithmetic subtraction yields |43 - 42| = 1 and |44 - 43| = 1. The neural network is falsely forced to assume that dogs are as conceptually close to airplanes as they are to cats.
2. Orthogonality of One-Hot Vectors: Treating tokens as one-hot vectors (eᵢ ∈ R^V) eliminates false linear ordering, but makes every pair equidistant:

\|eᵢ - eⱼ\|₂ = √(2) (∀ i ≠ j)

In this one-hot regime, "king" is as distant from "queen" as it is from "refrigerator".

1. Continuous Manifold Solution: Projecting tokens into continuous dense space R^d_(model) (e.g. d = 768) allows the model to align semantic nuances along continuous geometric axes. Cosine similarity cos(θ) = (u · v) / (\|u\| \|v\|) directly reflects conceptual affinity.

Every transformer architecture, from BERT to GPT-4, begins by mapping discrete token integers into a continuous vector space (d_(model) = 768 in GPT-2 small, d_(model) = 4096 in Llama-3-8B). This projection enables subsequent linear layers to perform continuous calculus, gradient descent, and geometric attention routing.

In 2003, Yoshua Bengio et al. published "A Neural Probabilistic Language Model", introducing distributed continuous representations for words. In 2013, Tomas Mikolov et al. at Google showed with Word2Vec that continuous vector arithmetic (v_(king) - v_(man) + v_(woman) ≈ v_(queen)) naturally emerges from continuous space embeddings.

• Integers are phone numbers; vectors are GPS coordinates: A phone number tells you who to call, but not where they live; coordinates place them on the map of human thought.
• Continuous distance equals semantic difference: Closeness in vector space is closeness in meaning.

The embedding matrix W_E establishes a direct linear map between discrete token indices and continuous feature representations:

1. Matrix Dimensions & Contract:

W_E ∈ R^V × d_(model)

Where V is the vocabulary size (e.g. 50,257 in GPT-2) and d_(model) is the embedding hidden dimension (e.g. 768).

1. Row Indexing Formulation:

For an incoming discrete token integer t ∈ 0, 1, ṡ, V-1, the extracted embedding vector e_t is the exact t-th row of W_E:

e_t = W_E[t, :] ∈ R^1 × d_(model)

1. Parameter Accounting & Memory Footprint:

The total trainable parameter count in the input embedding layer is:

Params(W_E) = V × d_(model) = 50,257 × 768 = 38,597,376 parameters ( ≈ 38.60M)

At FP32 precision (4 bytes/param), this single matrix occupies 154.39 MB of VRAM; at FP16/BF16 (2 bytes/param), it occupies 77.19 MB.

Suppose V = 5 and d_(model) = 3:

W_E =

0.12 & -0.45 & 0.88 \

0.34 & 0.91 & -0.15 \

-0.78 & 0.22 & 0.05 \

0.61 & -0.11 & 0.49 \

-0.02 & 0.83 & -0.67

When processing token ID t = 2:

e₂ = W_E[2, :] = [-0.78, 0.22, 0.05]

The lookup operation extracts the 3rd row (0-indexed) with zero floating-point arithmetic.

In large models such as Llama-3-70B (V = 128,256, d_(model) = 8192), W_E contains over 1.05 billion parameters. Because looking up rows is memory-bandwidth bound rather than compute bound, modern deep learning runtimes optimize this layer using dedicated embedding kernels that pull contiguous cachelines directly from high-bandwidth memory (HBM).

The concept of viewing a lookup table as a weight matrix whose rows are adjusted by stochastic gradient descent dates back to the connectionist models of David Rumelhart and James McClelland (1986). In modern PyTorch, `torch.nn.Embedding` wraps this exact tensor operation, providing an optimized C++/CUDA backend.

• V rows of length d: Vocabulary sets the height; hidden dimension sets the width.
• Lookup is row selection: No matrix multiplies occur; W_E[i] retrieves row i in O(1) memory strides.

There is a profound duality between the algebraic definition and the hardware execution of embedding layers:

1. The Algebraic Perspective (Linear Projection):

Let token i be represented as a canonical basis vector pᵢ = [0, ṡ, 0, 1, 0, ṡ, 0] ∈ 0, 1^V, where index i is 1 and all other entries are 0. The embedding vector is the standard matrix multiplication:

eᵢ = pᵢ W_E

This proves that the embedding layer is fundamentally a linear layer without a bias term.

1. The Hardware Reality (Pointer Offset Slicing):

Multiplying a 1 × 50,257 one-hot vector by a 50,257 × 768 matrix requires 50,257 × 768 ≈ 38.6M floating-point multiplications, of which 38,596,608 are 0 × W_(k, j) = 0.

1. Execution Divergence:

Hardware treats W_E as a contiguous memory block in RAM/VRAM. Given base pointer address ptr(W_E), accessing token i simply jumps directly to offset:

Address(eᵢ) = ptr(W_E) + i × d_(model) × sizeof(dtype)

This achieves O(1) time complexity and zero floating-point operations.

Frameworks like PyTorch provide `nn.Embedding(V, d)` instead of `nn.Linear(V, d, bias=False)`. While both produce identical gradients during backpropagation, `nn.Embedding` executes thousands of times faster and uses orders of magnitude less memory bandwidth on both CPUs and GPUs.

The equivalence between one-hot selection and linear layer projections is a classical theorem in linear algebra (the i-th basis vector selects the i-th row of any matrix). In early neural NLP (circa 2000), researchers wrote one-hot matrix multiplies in software, until GPU memory limits necessitated direct array slicing kernels.

• Theory multiplies; hardware slices: Math sees linear algebra; silicon sees memory addresses.
• Never compute a zero: 50,000 zeros times a matrix is a waste of silicon power; jump directly to the target row.

#### 1. Mathematical Reduction Proof

Let p ∈ 0, 1^V be a one-hot vector where pᵢ = 1 and pₖ = 0 for all k ≠ i. Let W_E ∈ R^(V × d) be the weight matrix.

The j-th component of the output vector e = p W_E is given by the inner product:

eⱼ = ∑_(k=0)^(V-1) pₖ (W_E)_(k, j) = pᵢ (W_E)_(i, j) + ∑_(k ≠ i) 0 · (W_E)_(k, j) = 1 · (W_E)_(i, j) = (W_E)_(i, j)

Assembling across all columns j ∈ 0, 1, ṡ, d-1:

e = [ (W_E)_(i, 0), (W_E)_(i, 1), ṡ, (W_E)_(i, d-1) ] = (W_E)_(i, :)

Thus, p W_E reduces identically to row (W_E)_(i, :). \blacksquare

In row-major contiguous memory layout (C-contiguous, default in NumPy and PyTorch):

ByteOffset(i) = i × Stride = i × (d_(model) × S_(bytes))

Where S_(bytes) is the scalar element size (4 bytes for FP32, 2 bytes for FP16/BF16).

• For GPT-2 (d_(model) = 768 in FP32):

Stride = 768 × 4 = 3,072 bytes per row

For token ID i = 42:

ByteOffset(42) = 42 × 3,072 = 129,024 bytes

During the backward pass of training, the gradient (∂ L) / (∂ W_E) is extremely sparse: only rows corresponding to tokens present in the current batch receive non-zero gradient updates. PyTorch's `nn.Embedding(..., sparse=True)` exploits this by emitting `torch.sparse` tensors, slashing optimizer communication bandwidth in distributed multi-GPU training.

Row-major memory addressing formulas were formalized in the earliest Fortran and C compiler specifications in the 1950s and 1970s. Modern GPU tensor cores (such as NVIDIA Hopper/Blackwell) rely on these exact linear stride equations to coalesce memory transactions across wide 128-byte cachelines.

• Index times stride equals byte address: Slicing is linear arithmetic on pointers, not floating-point matrix multiplication.
• Sparse gradients for dense rows: Forward pass reads single rows; backward pass updates only the visited rows.

Gradient flow into embedding matrix W_E:

Let loss L produce upstream gradient with respect to embedded vectors:

G = (∂ L) / (∂ H) ∈ R^(B × T × d)

Because H[b, t, :] = W_E[X[b, t], :], the gradient for row k ∈ 0, ṡ, |V|-1 of W_E is:

(∂ L) / (∂ W_E[k, :]) = ∑_(b=1)^B ∑_(t=1)^T G[b, t, :] · I(X[b, t] = k)

Key observations:

1. Sparsity: If token ID k does not appear in batch X, its gradient is strictly zero:

(∂ L) / (∂ W_E[k, :]) = vec(0)

1. Gradient Accumulation: If token ID k appears multiple times in the batch (e.g., the word 'the'), its gradients sum additively.
2. Semantic Clustering: Tokens that appear in similar predictive contexts receive similar gradient vectors over time, causing their row vectors to gravitate toward neighboring regions of R^d.

Optimizers like AdamW check for sparse gradients or update only the active rows when using `sparse=True` in `nn.Embedding`. In massive distributed models (e.g. DeepSpeed or Megatron-LM), embedding tables are sharded across GPUs using tensor parallelism.

The mathematical convergence of word vectors via gradient descent was proved by Mikolov (2013) and later analyzed by Levy & Goldberg (2014), who proved that skip-gram embedding updates with negative sampling implicitly factorize a Pointwise Mutual Information (PMI) matrix.

• Only the tokens that speak receive the gradient.
• Co-occurring words receive similar nudges, pulling them together in space.

Linear neural network layers cannot operate meaningfully on raw integer identifiers:

1. The False Linear Proximity Trap:

Assigning arbitrary scalars x ∈ 0, 1, ṡ, V-1 imposes artificial ordering and distance metrics:

|ID(airplane) - ID(dog)| = |44 - 43| = 1

|ID(dog) - ID(cat)| = |43 - 42| = 1

Arithmetic operations like addition and distance are mathematically invalid across categorical indices.

1. The High-Dimensional Vector Manifold (R^d_(model)):

By mapping each discrete token ID t to a continuous vector v_t ∈ R^(768), each semantic concept acquires independent spatial degrees of freedom.

1. The Fundamental Embedding Law:

Semantically and syntactically related words map to vectors pointing in similar directional trajectories:

cos(θ) = (v_(cat) · v_(dog)) / (\|v_(cat)\| \|v_(dog)\|) ≈ 0.98 (High Cosine Similarity)

cos(θ) = (v_(dog) · v_(airplane)) / (\|v_(dog)\| \|v_(airplane)\|) ≈ -0.12 (Orthogonal / Unrelated)

Continuous vector embeddings form the universal entry layer for all transformer language models, recommendation systems, semantic search engines (vector databases), and multimodal encoders (CLIP).

The distribution hypothesis was formulated by linguist J.R. Firth in 1957: 'You shall know a word by the company it keeps'. Mikolov et al. (2013) materialized this into continuous vector space with Word2Vec, showing that linear vector arithmetic captures semantic concepts (v_(King) - v_(Man) + v_(Woman) ≈ v_(Queen)).

• Integers are arbitrary labels; vectors are semantic geometry: Never do math on raw IDs.
• Similar words point together: Cosine similarity captures shared meaning across hundreds of dimensions.

#### 1. The Embedding Weight Matrix Layout

The token embedding matrix W_E is structured as:

W_E ∈ R^V × d_(model)

Where:

• V = 50,257 rows (vocabulary size in GPT-2).
• d_(model) = 768 columns (hidden embedding dimension).
• Row vector for token t: W_E[t, :] = [w_(t, 0), w_(t, 1), ṡ, w_(t, 767)] ∈ R^(768).

• Input Token Tensor X:

X ∈ Z^(B × T)

Contains discrete integer token IDs, where B is batch size and T is sequence length.

• Embedding Lookup Operator:

Every integer ID X_(b, t) is indexed into row X_(b, t) of W_E:

Output Tensor E = W_E[X] ∈ R^B × T × d_(model)

For example, with B = 2 and T = 4, the output shape is [2, 4, 768].

• Total Parameters:

N_(params) = V × d_(model) = 50,257 × 768 = 38,597,376 parameters

• FP32 Storage (4 bytes per param):

38,597,376 × 4 = 154,389,504 bytes ≈ 154.39 MB (147.24 MiB)

• FP16 / BF16 Storage (2 bytes per param):

38,597,376 × 2 = 77,194,752 bytes ≈ 77.19 MB (73.62 MiB)

Because embedding parameters scale directly with vocabulary size, models with large vocabularies (e.g., Gemma with V = 256,000 or LLaMA-3 with V = 128,256) allocate up to 25% of their total parameter budget to embedding matrices alone.

Lookup tables for categorical features have existed in database engineering for decades. In deep learning, early language models tied input embeddings to output unembedding matrices (Weight Tying, Press & Wolf, 2017) to halve this substantial memory footprint.

• [B, T] → [B, T, d]: The fundamental tensor expansion of the embedding layer.
• V × d_(model): Multiply vocabulary by embedding dimension to find the matrix size.

#### 1. The Mathematical Equivalence Theorem

Let W_E ∈ R^(V × d) be an embedding weight matrix whose j-th row is denoted wⱼ^T:

W_E = w₀^T \ w₁^T \ … \ w_(V-1)^T

Let eᵢ ∈ R^(1 × V) be a one-hot canonical basis vector where entry i is 1 and all other entries are 0:

eᵢ = [0, ṡ, 0, 1 (index i), 0, ṡ, 0]

By the definition of matrix multiplication:

v = eᵢ · W_E = ∑_(j=0)^(V-1) e_(i, j) · wⱼ^T

Expanding the sum across all indices j:

v = 0 · w₀^T + ṡ + 1 · wᵢ^T + ṡ + 0 · w_(V-1)^T = wᵢ^T

The matrix multiplication product eᵢ · W_E evaluates strictly and identically to the i-th row of W_E.

1. Conceptual / Theoretical (Linear Algebra):

An embedding layer is simply a standard dense linear layer without a bias vector:

y = x · W_E

Operating on one-hot categorical input distributions.

1. Computational / Engineering (Systems):

Executing an explicit 50,000-dimensional matrix multiplication on millions of zeros is grossly wasteful. Hardware engines exploit the mathematical equivalence to implement y = W_E[i] directly.

This mathematical identity provides the theoretical bridge allowing backpropagation and autodiff engines (PyTorch `autograd`) to treat embedding layers as fully differentiable linear operators.

The mathematical equivalence of table lookups and one-hot matrix multiplication dates to early categorical representation in generalized linear models. It was popularized in deep learning by Yoshua Bengio et al. (2003) in A Neural Probabilistic Language Model.

• eᵢ · W = W[i, :]: The fundamental one-hot projection identity.
• Linear in theory, lookup in practice: The dual identity that unifies NLP with linear algebra.

#### 1. Naive GEMM Matrix Multiply (The Inefficient Way)

Computing eᵢ · W_E via standard General Matrix Multiply (GEMM):

• Input shape: [1, 50257]
• Weight shape: [50257, 768]
• Required FLOPs:

FLOPs = 2 × V × d_(model) = 2 × 50,257 × 768 = 77,194,752 FLOPs per token

• Computational Waste:

Of the 50,257 entries in eᵢ, exactly 50,256 are 0.0. Thus, 99.998\% of GPU tensor core cycles are completely wasted multiplying and adding zeros.

In `torch.nn.Embedding`, the forward pass computes a direct virtual memory pointer offset:

Memory Address = Base Address + (Token ID × d_(model) × sizeof(float))

• Worked Numerical Trace for Token ID = 42:

Offset = 42 × 768 × 4 bytes = 129,024 bytes = 0x1F800

Target Address = 0x7FFF0000 + 0x1F800 = 0x7FFF1F800

• Execution:

The GPU memory controller copies 768 contiguous floats (3,072 bytes) from VRAM into the register file via a single memory transaction.

• FLOPs Consumed: EXACTLY ZERO.

| Metric | Naive GEMM Multiply | Direct Pointer Indexing |

| :--- | :--- | :--- |

| FLOPs per Token | 77.2 × 10⁶ FLOPs | 0 FLOPs |

| Memory Bandwidth | Choked by full matrix scan | Single 3 KB contiguous read |

| Execution Latency | High GPU kernel overhead | ≈ 1,000× faster |

This memory addressing design powers `nn.Embedding` in PyTorch and `tf.keras.layers.Embedding`. In backward passes, PyTorch uses sparse or sorted scatter-add kernels to accumulate gradients into specific rows without updating unaffected embedding weights.

Direct pointer addressing has been foundational to computer systems since the von Neumann architecture. Incorporating it into automatic differentiation graphs allowed deep learning frameworks to handle massive vocabularies without grinding GPU compute clusters to a halt.

• Pointer math costs zero FLOPs: Base address plus index times row stride.
• Don't multiply fifty thousand zeros: Index the memory pointer directly.

The volume of a d-dimensional hypersphere of radius r is:

V_d(r) = (π^(d/2)) / (Γ(d/2 + 1)) r^d

Consider the fraction of volume contained in an inner core of radius 1 - ε within a unit hypersphere (r = 1):

(V_d(1 - ε)) / (V_d(1)) = ((1 - ε)^d) / (1^d) = (1 - ε)^d

For d = 768 (GPT-2) and a tiny crust thickness ε = 0.01 (1\% of radius):

(1 - 0.01)^(768) = (0.99)^(768) ≈ 0.00045

Thus, 1 - 0.00045 = 99.955\% of the hypersphere's entire volume resides in the outer 1\% spherical shell.

This geometric shell concentration governs vector databases (Pinecone, Milvus, FAISS) and token embeddings. Approximate nearest neighbor algorithms (HNSW) exploit this phenomenon by searching along surface graphs rather than partitioning hyper-volumes.

Richard Bellman coined the term 'Curse of Dimensionality' in 1957 in his book 'Dynamic Programming', observing that the computational cost of grid-based optimization grows exponentially with dimension.

• High-dimensional oranges are all peel and no fruit.
• In d-dimensions, volume migrates to the surface: (1 - epsilon)^d vanishes.

Let u, v ~ N(0, I_d) be two independent random vectors in R^d.

The inner product is the sum of d independent random variables:

u · v = ∑_(i=1)^d uᵢ vᵢ

By independence:

E[uᵢ vᵢ] = E[uᵢ] E[vᵢ] = 0

Var(uᵢ vᵢ) = E[uᵢ²] E[vᵢ²] = 1 × 1 = 1

Thus, E[u · v] = 0 and Var(u · v) = d.

Normalizing by norms \|u\|, \|v\| ≈ √(d):

cos(θ) = (u · v) / (\|u\| \|v\|) ≈ (N(0, d)) / (d) = N(0, 1/d)

For d = 768, the standard deviation is (1) / (√(768)) ≈ 0.0361.

By Chebyshev's inequality, 99.7\% of all pairs have |cos θ| ≤ 3 × 0.0361 ≈ 0.108 (83.8^∘ ≤ θ ≤ 96.2^∘).

This near-orthogonality allows modern LLMs to store thousands of distinct factual concepts, syntactic rules, and memory associations in the same layer without interference, known as vector superposition.

The Johnson-Lindenstrauss Lemma (1984) established that high-dimensional points can be projected into lower dimensions while preserving pairwise Euclidean distances, relying directly on random vector orthogonality.

• Random high-dimensional vectors are perpendicular strangers.
• Cosine similarity has standard deviation 1/sqrt(d); as d grows, dot product vanishes to zero.

An educational lookup-table mapping diagram in a balanced square 1:1 format on warm cream parchment (#fcf0db) with a refined 6px sepia border (#e5d5be).

Top section: Title in navy serif typography: "Token Embeddings: Mapping Tokens to Continuous Space". Subtitle in warm slate: "How a one-hot vector selects an embedding row without computing a full matrix product."

Middle section: Side-by-side architecture diagram:

Left: Discrete Token Stream ["quantum", "entanglement", "state"] mapped to integer token IDs [18420, 29114, 1184].

Middle: Embedding Weight Matrix W_E ∈ R^(V × D) (Vocabulary V = 50,257, Embedding Dimension D = 768).

Right: 768-dimensional continuous dense embedding vectors e_t ∈ R^D.

Show the equivalence: W_E^T · one\_hot(i) is computationally implemented as a direct table lookup: W_E[i, :].

Bottom section: Mathematical insight: "Embedding layers are unconstrained linear lookup tables where semantically related tokens cluster into adjacent angular neighborhoods."

Balanced square composition, distinct row-selection highlighting, crisp typography.

Every transformation preserves structural invariants while mapping coordinate spaces. When operating in high dimensions, numerical precision and computational efficiency require adhering strictly to canonical mathematical laws.

Formal derivation and peer-reviewed analysis:

Radford et al. (2019), Language Models are Unsupervised Multitask Learners, OpenAI ↗

"Token Embedding Vectors: Projecting Discrete Tokens into Continuous Semantic Space: Invariance, symmetry, and structure govern mathematical truth across continuous and discrete domains."