Skip to content
LLMs from Scratch

Chapter 2 · Bonus material

Lookup Is a Matrix Multiply

An embedding layer and a linear layer are the same operation. Watch a one-hot vector pick a row, and see why indexing is the shortcut.

Adapted from the original notebook by Sebastian Raschka · the Python shown is the book's; the demos run a faithful JavaScript port in your browser

Bonus · Lookup is a matrix multiply

ids
2
3
1

weight table · 4 × 5

0
0.34-0.18-0.30-0.591.58
1
1.301.28-0.20-0.16-0.40
2
0.70-1.81-1.160.33-0.63
3
-2.84-0.78-1.41-0.410.80

a learned weight matrix — one row per ID

01 / 05

An embedding table is a learned weight matrix— 4 rows (one per ID in a 4-word vocabulary) × 5 dimensions, trained like any other weight. That's the clue: a weight matrix can be reached by a matrix multiply.

The real code

The demo above runs a TypeScript port; here is the Python it mirrors, from the bonus notebook. First, the embedding lookup:

embeddings-and-linear-layers.ipynb
import torch
 
torch.manual_seed(123)
embedding = torch.nn.Embedding(num_idx, out_dim)   # a 4 × 5 learned table
 
idx = torch.tensor([2, 3, 1])
embedding(idx)                                      # grabs rows 2, 3, 1

Now the same result as a matrix multiply on a one-hot encoding — an nn.Linear with no bias, its weights tied to the embedding table:

embeddings-and-linear-layers.ipynb
onehot = torch.nn.functional.one_hot(idx)          # each id → a 1-hot row
 
linear = torch.nn.Linear(num_idx, out_dim, bias=False)
linear.weight = torch.nn.Parameter(embedding.weight.T)
 
linear(onehot.float()) == embedding(idx)           # identical numbers

Two things to remember

  1. An embedding layer is a matrix multiply on a one-hot vector. The two paths return the same numbers — the bonus proves it on the seed-123 table.
  2. The lookup is the optimization. A one-hot times a big matrix is mostly multiply-by-zero; indexing the row skips all of it. That's why real models use nn.Embedding over 50,257 rows instead of nn.Linear.

Adapted from Build a Large Language Model (From Scratch) by Sebastian Raschka — original notebook (Apache 2.0). The weight matrix and results on this page are torch's real seed-123 output, verified number-for-number against the Python.