Embedding Layer¶
The Embedding layer is a learnable lookup table that maps integer category indices to dense continuous vectors. It serves as an optimized, implicit linear transformation for discrete categorical variables — such as product IDs, word tokens, or user IDs — in neural networks.
Mathematical Foundation¶
Problem Statement¶
Let $\mathcal{V} = \{0, 1, \ldots, N-1\}$ denote a discrete vocabulary of $N$ unique categories. The goal is to learn a continuous vector representation mapping:
$$ \phi : \mathcal{V} \to \mathbb{R}^d $$
that assigns each category $v \in \mathcal{V}$ a $d$-dimensional vector called its embedding. The embedding dimension $d$ is a hyperparameter chosen by the practitioner (typically $d \ll N$).
Weight Matrix¶
The mapping is parameterized by a single trainable weight matrix:
$$ \mathbf{W} \in \mathbb{R}^{N \times d} $$
where row $i$, denoted $\mathbf{W}_{i, :}$, is the embedding vector for category $i$. It is initialized from a Kaiming-uniform distribution:
$$ W_{ij} \sim \mathcal{U}\!\left(-\sqrt{\frac{1}{N}},\ \sqrt{\frac{1}{N}}\right) $$
which keeps the expected norm of the initial embedding vectors stable across vocabulary sizes $N$.
Equivalence with One-Hot Encoding + Linear¶
An Embedding layer computes exactly the same function as a bias-free Linear layer applied to one-hot encoded inputs. Recall from Categorical Encoding that OneHotEncoder maps a sequence of $L$ category indices $\mathbf{i} = (i_1, \ldots, i_L)$ to the binary indicator matrix $\mathbf{E} \in \{0,1\}^{L \times N}$, whose $k$-th row is the standard basis vector $\mathbf{e}_{i_k}^\top$.
Feeding $\mathbf{E}$ into a bias-free Linear layer with weights $\mathbf{W} \in \mathbb{R}^{N \times d}$ gives $\mathbf{Y} = \mathbf{E}\mathbf{W}$, whose $k$-th row is
$$ \mathbf{y}_k = \mathbf{e}_{i_k}^\top \mathbf{W} = \sum_{j=0}^{N-1} (\mathbf{e}_{i_k})_j \, \mathbf{W}_{j, :} = 1 \cdot \mathbf{W}_{i_k, :} + \sum_{j \neq i_k} 0 \cdot \mathbf{W}_{j, :} = \mathbf{W}_{i_k, :} $$
Since every term but one is multiplied by zero, the product simply extracts row $i_k$ of $\mathbf{W}$. Embedding performs that extraction directly, never materializing $\mathbf{E}$:
| Feature / Metric | OneHotEncoder + Linear(N, d) |
sorix.nn.Embedding(N, d) |
|---|---|---|
| Input Representation | Binary matrix $\mathbf{E} \in \{0, 1\}^{L \times N}$ | Integer tensor $\mathbf{i} \in \mathbb{Z}^L$ |
| Input Memory | $O(L \cdot N)$ elements | $O(L)$ integers |
| Forward Pass Time | $O(L \cdot N \cdot d)$ (dense matmul) | $O(L \cdot d)$ (row lookup) |
| FLOPs Efficiency | $N-1$ zero products per sample | No zero products |
| Gradient | Dense, over all $N$ rows | Sparse, only the rows used |
| Preprocessing | Explicit one-hot transform | Integer IDs (e.g. from LabelEncoder) |
For a realistic vocabulary ($N = 50{,}000$) the one-hot route is prohibitive on both axes, which is why Embedding exists as a dedicated layer.
Forward Pass¶
Given a tensor of integer indices $\mathbf{i}$, the forward pass is a row selection:
$$ \text{Embedding}(\mathbf{i}) = \mathbf{W}[\mathbf{i}], \qquad \mathbf{y}_k = \mathbf{W}_{i_k, :} $$
Indices of any shape are accepted: an input of shape $(*)$ produces an output of shape $(*, d)$. For a batch $\mathbf{I} \in \mathbb{Z}^{B \times L}$ the output has shape $(B, L, d)$, where element $(b, l)$ is row $\mathbf{W}_{\mathbf{I}_{b,l}, :}$.
Indices must lie in $[0, N)$. Out-of-range values — including the $-1$ that encoders often emit for unseen categories — raise an IndexError, because reading the wrong row would also send the gradient to the wrong row in the backward pass.
Backward Pass¶
Let $\mathcal{L}$ be a scalar loss and $\partial\mathcal{L}/\partial\mathbf{Y} \in \mathbb{R}^{L \times d}$ the gradient arriving from downstream layers, where $\partial\mathcal{L}/\partial\mathbf{y}_k$ is the gradient w.r.t. the $k$-th output row.
Because row $\mathbf{W}_{v,:}$ contributes to the output once for every position $k$ where $i_k = v$, the chain rule sums those contributions:
$$ \frac{\partial \mathcal{L}}{\partial \mathbf{W}_{v, :}} = \sum_{\{k \,\mid\, i_k = v\}} \frac{\partial \mathcal{L}}{\partial \mathbf{y}_k} \qquad v \in \{0, \ldots, N-1\} $$
Rows for categories absent from the batch receive $\mathbf{0}$, so the gradient is sparse: only the looked-up rows are updated. This is implemented as a scatter-add, which accumulates (rather than overwrites) on duplicate indices:
grad_W = xp.zeros_like(W.data)
xp.add.at(grad_W, indices, out.grad.data) # xp is numpy or cupy
The gradient w.r.t. the input indices is undefined — they are discrete — so the backward pass stops here. Embedding is therefore always a leaf of the graph on its input side.
Parameters & Device¶
The layer exposes its single trainable parameter $\mathbf{W}$ through parameters(), so it works with every sorix optimizer (Adam, SGD, RMSprop, ...). With device='cuda' both the lookup and the scatter-add run through CuPy on the GPU.
Multi-Field Categorical Inputs¶
When a sample has several categorical fields, each field gets its own Embedding and the resulting vectors are concatenated along the feature axis:
$$ \mathbf{z} = \left[\phi_{\text{cat}}(v_{\text{cat}}) \;\parallel\; \phi_{\text{brand}}(v_{\text{brand}})\right] \in \mathbb{R}^{d_{\text{cat}} + d_{\text{brand}}} $$
This combined vector is then processed by downstream Linear and activation layers.
# Uncomment the next line and run this cell to install sorix
#!pip install 'sorix @ git+https://github.com/Mitchell-Mirano/sorix.git@main'
import numpy as np
import sorix
from sorix import tensor
from sorix.nn import Embedding
Instantiating the Layer¶
The constructor requires two arguments: the vocabulary size $N$ (num_embeddings) and the embedding dimension $d$ (embedding_dim).
# N=10 categories, d=4 embedding dimensions
emb = Embedding(num_embeddings=10, embedding_dim=4)
print(f"W shape: {emb.W.shape}")
print(f"W requires_grad: {emb.W.requires_grad}")
W shape: sorix.Size([10, 4]) W requires_grad: True
Forward Pass — Row Selection¶
Passing integer indices retrieves the corresponding rows of $\mathbf{W}$.
# Look up rows 2, 0, 5 from W
indices = np.array([2, 0, 5])
out = emb(indices)
print(f"indices shape : {indices.shape}")
print(f"output shape : {out.shape}")
print()
# Verify: output rows equal the corresponding weight rows
match = np.allclose(out.data, emb.W.data[indices])
print(f"output matches W rows? {match}")
indices shape : (3,) output shape : sorix.Size([3, 4]) output matches W rows? True
out
tensor([[ 0.19533172, -0.06835324, -0.08989462, 0.09925262],
[-0.19922554, -0.10904585, 0.08168502, 0.07237582],
[-0.04870158, 0.04121277, 0.11020941, 0.01153292]], requires_grad=True)
Batched Lookup¶
For a batch of $B$ sequences each of length $L$, pass a 2-D index array. The output shape is $(B, L, d)$.
# Batch of 2 sentences, each with 5 word tokens
token_ids = np.array([[1, 5, 3, 0, 7],
[2, 4, 1, 6, 8]])
out_batch = emb(token_ids)
print(f"index batch shape : {token_ids.shape}")
print(f"embedding out shape: {out_batch.shape}")
index batch shape : (2, 5) embedding out shape: sorix.Size([2, 5, 4])
Index Validation¶
Indices outside $[0, N)$ are rejected. This matters most for $-1$, which would otherwise select the last row (Python's negative indexing) and silently accumulate its gradient there.
emb_guard = Embedding(num_embeddings=5, embedding_dim=3)
for bad in [np.array([0, -1, 2]), np.array([0, 5])]:
try:
emb_guard(bad)
except IndexError as e:
print(f"{str(bad):<12} -> IndexError: {e}")
[ 0 -1 2] -> IndexError: Embedding indices must be in [0, 5), but got values in [-1, 2]. Note that unseen categories encoded as -1 must be mapped to a dedicated index instead. [0 5] -> IndexError: Embedding indices must be in [0, 5), but got values in [0, 5]. Note that unseen categories encoded as -1 must be mapped to a dedicated index instead.
Backward Pass — Scatter-Add Gradient¶
When an index appears multiple times, its gradient contributions are summed via scatter-add. Below we verify this property: index 0 is looked up twice, so its row in W.grad receives twice the gradient of index 2 (looked up once).
sorix.manual_seed(0)
emb2 = Embedding(num_embeddings=5, embedding_dim=4)
# index 0 appears twice, index 2 once
idx = np.array([0, 2, 0])
out2 = emb2(idx)
# scalar loss = mean of all outputs
loss = out2.mean()
loss.backward()
print(f"W.grad[0] (used twice): {emb2.W.grad.data[0]}")
print(f"W.grad[2] (used once) : {emb2.W.grad.data[2]}")
print(f"W.grad[1] (not used) : {emb2.W.grad.data[1]}")
print()
print(f"grad[0] == 2 * grad[2]? {np.allclose(emb2.W.grad.data[0], 2 * emb2.W.grad.data[2])}")
W.grad[0] (used twice): [0.16666667 0.16666667 0.16666667 0.16666667] W.grad[2] (used once) : [0.08333334 0.08333334 0.08333334 0.08333334] W.grad[1] (not used) : [0. 0. 0. 0.] grad[0] == 2 * grad[2]? True
End-to-End Training Loop¶
The Embedding layer integrates transparently with sorix optimizers via parameters().
Here we train embeddings to reconstruct a target vector for each of three categories.
sorix.manual_seed(42)
emb3 = Embedding(num_embeddings=5, embedding_dim=4)
optimizer = sorix.optim.Adam(emb3.parameters(), lr=1e-2)
# Target: each of the 3 categories maps to the all-ones vector
target = tensor(np.ones((3, 4), dtype=np.float32))
idx_train = np.array([0, 1, 2])
for epoch in range(50):
optimizer.zero_grad()
out = emb3(idx_train)
loss = ((out - target) ** 2).mean()
loss.backward()
optimizer.step()
if epoch % 10 == 0 or epoch == 49:
print(f"Epoch {epoch:3d} | loss = {float(loss.data):.6f}")
Epoch 0 | loss = 1.061576 Epoch 10 | loss = 0.875063 Epoch 20 | loss = 0.711495 Epoch 30 | loss = 0.571944 Epoch 40 | loss = 0.455625 Epoch 49 | loss = 0.369055
Multi-field Categorical Inputs¶
When a sample has multiple categorical fields, look up each embedding independently and concatenate along the feature axis before passing to downstream Linear layers.
word_emb = Embedding(num_embeddings=10000, embedding_dim=32)
cat_emb = Embedding(num_embeddings=50, embedding_dim=8)
word_ids = np.array([3, 14, 52, 7])
cat_ids = np.array([5, 12, 11, 23])
w = word_emb(word_ids) # (4, 32)
c = cat_emb(cat_ids) # (4, 8)
combined = sorix.cat([w, c], dim=1) # (4, 40)
print(f"word embeddings : {w.shape}")
print(f"cat embeddings : {c.shape}")
print(f"combined shape : {combined.shape}")
word embeddings : sorix.Size([4, 32]) cat embeddings : sorix.Size([4, 8]) combined shape : sorix.Size([4, 40])
Inference Mode (no_grad)¶
During evaluation, wrapping the forward pass in sorix.no_grad() disables gradient tracking, reducing memory usage.
emb4 = Embedding(num_embeddings=10, embedding_dim=4)
with sorix.no_grad():
out_eval = emb4(np.array([1, 3, 5]))
print(f"out.requires_grad during inference: {out_eval.requires_grad}")
out.requires_grad during inference: False