Embeddings

Embeddings turn a list of token ids into the matrix the first block reads. The model does this with two learned tables and one addition. This page covers each table, the sum, the size of the position table and the backward pass.

Token embedding

What is it?

Token embeddingImplemented

Embeddings

tok_emb is a table with one row for each token in the vocabulary. To turn an id into a vector, the model reads that row. The whole step is NumPy indexing, tok_emb[ids], with no arithmetic.

Why does the model need it?

An id is only a row number, and the layers that follow multiply and add vectors. The table gives each token a vector of width d. Its entries are parameters, so training can change them.

Mathematics

Row t of the lookup is the table row of the t-th id #

\left(E[\text{ids}]\right)_t = E_{i_t}

What does each symbol mean?

Symbol	Meaning	In code	Shape
$E$	The token table. A parameter	`params["tok_emb"]`	(V, d)
$i_t$	The id at position t	`ids[t]`	scalar
$E[\text{ids}]$	One table row for each id. The code does not give this array a name	not named in the code	(T, d)

Tiny numerical example

Id 2 appears twice, so its row appears twice.

Looking up three idsIllustrative ExampleHand-picked numbers from test_embedding_worked_example. A real table has Xavier uniform values at first and learned values later, and it is much larger.

Step	Code	Value
Token table	`tok_emb`	`[[0.1, 0.2], [0.3, -0.1], [0.5, 0.4], [-0.2, 0.6]]`
Ids	`ids`	`[2, 0, 2]`
One row per id	`tok_emb[ids]`	`[[0.5, 0.4], [0.1, 0.2], [0.5, 0.4]]`

Tensor shapes

From ids to vectorstransformer.py:53

ids(T,)(16,)
tok_emb[ids](T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

The table itself. V depends on the corpustransformer.py:21

rows(V,)
tok_emb(V, d)(V, 32)

V depends on the corpus. It is not a value in config.py. d = emb_dim = 32.

Actual implementation

init_params creates the table, and the lookup is the first term of embed_forward, shown under Embedding sum.

transformer.pyinit_paramsGitHub

21        "tok_emb": xavier_uniform(V, d, rng), # lookup table with one row per vocabulary token

Code and math connection

Equation	Role	Lines	Code
Row t of the lookup is the table row of the t-th id	Computes it	transformer.py:53	`return params["tok_emb"][ids] + params["pos_emb"][: len(ids)]`
Row t of the lookup is the table row of the t-th id	Tests it	tests/test_transformer.py:34	`X0 = embed_forward(params, np.array([2, 0, 2]))`

Where is it used?

Source

transformer.py, lines 52-53, embed_forward

PDF reference

Section 4.5, Eq. 39, p. 10

Used by

Called by: transformer.forward

Tested by: tests.test_transformer.test_embedding_worked_example

Tests it: tests/test_transformer.py, line 34 (test_embedding_worked_example), for Row t of the lookup is the table row of the t-th id

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Row t of the lookup is the table row of the t-th id	Computes it	`transformer.py`	`embed_forward`	53	GitHub: transformer.py, line 53

Positional embedding

What is it?

Positional embeddingImplemented

Embeddings

pos_emb is a second learned table, with one row for each position. init_params creates it with xavier_uniform, like any other matrix, and training changes it. It is not computed from a fixed formula. A sequence of length T uses the first T rows.

Why does the model need it?

The token table gives an id the same row wherever it appears. Adding a row chosen by position makes the vectors differ, so the same token at two positions enters the first block as two different vectors.

Note

In general, attention compares vectors by their contents, so two identical vectors get the same query, key and value. The causal mask limits which earlier rows each row can read, but it does not say how far back a row is. Position information gives the model a way to tell these cases apart. This is a general explanation. Nothing in this repository measures that effect.

Mathematics

The first T rows of the position table #

\left(P[:T]\right)_t = P_t, \quad t = 0, \ldots, T - 1

What does each symbol mean?

Symbol	Meaning	In code	Shape
$P$	The position table. A parameter	`params["pos_emb"]`	(Tmax, d)
$T$	The length of this sequence	`len(ids)`	scalar
$P[:T]$	The rows this sequence uses. The code does not give this array a name	not named in the code	(T, d)

Tiny numerical example

Slicing the position tableIllustrative ExampleHand-picked numbers from test_embedding_worked_example. The rows depend only on how long the sequence is, not on which ids it holds.

Step	Code	Value
Position table	`pos_emb`	`[[0, 0.1], [0.1, 0], [0.2, 0.1]]`
T = 3	`pos_emb[: 3]`	`[[0, 0.1], [0.1, 0], [0.2, 0.1]]`
T = 2	`pos_emb[: 2]`	`[[0, 0.1], [0.1, 0]]`

Tensor shapes

The table is cut to the length of the sequencetransformer.py:53

pos_emb(Tmax, d)(32, 32)
pos_emb[:T](T, d)(16, 32)

Tmax = context_len = 32. d = emb_dim = 32. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Actual implementation

The positions are learned because the code makes them so. This line puts pos_emb in the same parameter dictionary as every other matrix, and embed_backward gives it a gradient:

transformer.pyinit_paramsGitHub

22        "pos_emb": xavier_uniform(cfg.context_len, d, rng) # second table with one row per position, weighted average of values

The comment on that line says "weighted average of values". That does not describe the code, which draws a plain table with one row per position.

Code and math connection

Equation	Role	Lines	Code
The first T rows of the position table	Computes it	transformer.py:53	`return params["tok_emb"][ids] + params["pos_emb"][: len(ids)]`
The first T rows of the position table	Tests it	tests/test_transformer.py:34	`X0 = embed_forward(params, np.array([2, 0, 2]))`

Where is it used?

Source

transformer.py, lines 52-53, embed_forward

PDF reference

Section 4.5, Eq. 39, p. 10

Used by

Called by: transformer.forward

Tested by: tests.test_transformer.test_embedding_worked_example

Tests it: tests/test_transformer.py, line 34 (test_embedding_worked_example), for The first T rows of the position table

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The first T rows of the position table	Computes it	`transformer.py`	`embed_forward`	53	GitHub: transformer.py, line 53

Embedding sum

What is it?

Embedding sumImplemented

Embeddings

embed_forward adds the two lookups, row by row. The result, X0, is the first thing the blocks read.

Why does the model need it?

Every block expects a matrix of shape (T, d). This is where a list of ids becomes that matrix, and where the model gets both what each token is and where it stands.

Mathematics

One row of the sum #

X^{(0)}_t = E_{i_t} + P_t

The whole sum, as the code writes it #

X^{(0)} = E[\text{ids}] + P[:T]

What does each symbol mean?

Symbol	Meaning	In code	Shape
$X^{(0)}$	The input of the first block. The code calls it X in forward, and dX0 names its gradient	`X`	(T, d)
$E_{i_t}$	The token row for the id at position t	`params["tok_emb"][ids][t]`	(d,)
$P_t$	The position row for position t	`params["pos_emb"][t]`	(d,)

Tiny numerical example

Pick ids and a length. Choose the same id twice to see two things. The rows of X0 differ, because their positions differ. And the token gradient for that id collects both rows of dX0, which the backward section explains.

Pick ids, and follow them both waysIllustrative ExampleThe tables and dX0 are the hand-picked numbers of test_embedding_worked_example and test_embedding_backward_sums_repeated_ids. They are not learned values, and dX0 is chosen by hand, not computed from a loss. The page starts with the id repeated, as those tests do.

Tids[0]ids[1]ids[2]

Id	Row of tok_emb	Selected at
0	`[0.1, 0.2]`	t = 1
1	`[0.3, -0.1]`	not used
2	`[0.5, 0.4]`	t = 0, t = 2
3	`[-0.2, 0.6]`	not used

tok_emb[ids], shape (3, 2)
0.5	0.4
0.1	0.2
0.5	0.4

pos_emb[:T], shape (3, 2)
0	0.1
0.1	0
0.2	0.1

X0, their sum, shape (3, 2)
0.5	0.5
0.2	0.2
0.7	0.5

Id 2 is at positions 0 and 2, so row 2 of the token gradient is dX0[0] + dX0[2].

dX0, shape (3, 2)
1	0
0	1
2	2

grads["tok_emb"], np.add.at, shape (4, 2)
0	1
0	0
3	2
0	0

grads["pos_emb"], shape (3, 2)
1	0
0	1
2	2

Wrong: grads[ids] += dX0, shape (4, 2)
0	1
0	0
2	2
0	0

Tensor shapes

Both terms have the same shape, and so does their sumtransformer.py:53

tok_emb[ids](T, d)(16, 32)
pos_emb[:T](T, d)(16, 32)
X0(T, d)(16, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32.

Actual implementation

transformer.pyembed_forwardGitHub

52def embed_forward(params, ids):53    return params["tok_emb"][ids] + params["pos_emb"][: len(ids)]

Code and math connection

Equation	Role	Lines	Code
One row of the sum	Computes it	transformer.py:53	`return params["tok_emb"][ids] + params["pos_emb"][: len(ids)]`
The whole sum, as the code writes it	Computes it	transformer.py:53	`return params["tok_emb"][ids] + params["pos_emb"][: len(ids)]`
The whole sum, as the code writes it	Calls it	transformer.py:158	`X = embed_forward(params, ids)`
The whole sum, as the code writes it	Tests it	tests/test_transformer.py:34	`X0 = embed_forward(params, np.array([2, 0, 2]))`

Where is it used?

Source

transformer.py, lines 52-53, embed_forward

PDF reference

Section 4.5, Eq. 39, p. 10

Used by

Called by: transformer.forward

Tested by: tests.test_transformer.test_embedding_worked_example

Calls it: transformer.py, line 158 (forward), for The whole sum, as the code writes it
Tests it: tests/test_transformer.py, line 34 (test_embedding_worked_example), for The whole sum, as the code writes it

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
One row of the sum	Computes it	`transformer.py`	`embed_forward`	53	GitHub: transformer.py, line 53
The whole sum, as the code writes it	Computes it	`transformer.py`	`embed_forward`	53	GitHub: transformer.py, line 53

Why the position table is longer than a window

The table has 32 rows, set by context_len. Training windows have 16 tokens, set by seq_len. These are separate settings, and they do different jobs.

context_len is a size in the model. It decides the shape of pos_emb, and it is a field of ModelConfig (transformer.py:10). seq_len is not a field of ModelConfig, so the model does not know how long a window is. All it checks is that a sequence fits the table:

transformer.pyforwardGitHub

156    assert len(ids) <= cfg.context_len, "sequence longer than the position table"

config.py also asserts, on import, that seq_len is not above context_len.

seq_len is how the callers cut their sequences. Training uses it to cut the windows (train.py:44-45), and talk.py passes it as the generation window (talk.py:24-25). generate then keeps only that many of the latest ids:

generate.pygenerateGitHub

32        logits, _ = forward(params, np.array(ids[-window:]), cfg)

So the table has to be at least as long as the longest sequence, and it can be longer. embed_forward slices the table to the length of its input, so a sequence of any length up to context_len works. A short prompt relies on this. Why this project sets the two to different values is not available from the implementation.

What this means at the values in config.py. No sequence goes past seq_len, so rows 16 to 31 of pos_emb are never read. They are still created, saved in checkpoints and updated. They get a gradient of 0 from embed_backward, but Adam.step applies weight decay whatever the gradient is, so they shrink a little at each step. The table below shows this on a small model:

Which rows of pos_emb change in one stepIllustrative ExampleStored results of one real loss_and_grads and one real Adam step on the hand-picked tiny model of tests/test_gradcheck.py, which has T = 5 and 8 rows. This is not the trained model.

Row of pos_emb	Non-zero gradient	Changed, weight_decay = 0	Changed, weight_decay = `0.01`
0	yes	yes	yes
1	yes	yes	yes
2	yes	yes	yes
3	yes	yes	yes
4	yes	yes	yes
5	no	no	yes
6	no	no	yes
7	no	no	yes

Actual Model Output

The shapes of the two tables in a checkpoint, and of one real result of embed_forward.

Embedding shapes in the checkpointActual Model OutputCommit 66da45b. Checkpoint model.npz sha256 74de2ec88a63. Seed 42, from config.py, since a checkpoint does not record its own. Read from checkpoints/model, an ignored local file, trained on data/corpus.csv (git-ignored, and not at the pinned commit), with training settings that by inference differ from config.py.The checkpoint holds 38 tensors and 63,871 parameters. X0 is the real result of embed_forward on one window of 16 tokens.

Value	Shape
`tok_emb`	`(575, 32)`
`pos_emb`	`(32, 32)`
`X0 = tok_emb[ids] + pos_emb[:T]`	`(16, 32)`

Where these numbers came from

Repository commit	`66da45bea72176fa5d2f8dba2ea1c7e597e30dd3`
`model.npz` sha256	`74de2ec88a63ceb131a27e03cc63070c6d8c4174cc78bc189195e9356e84550e`
`model.json` sha256	`f93002949162068d573aab9b38990b12aec7f32b468ec9eaaef986075df7d3b4`
`model.tok.json` sha256	`bf874ab4e4de2b61cb936f38635a7c589c5c0d29879947a5b45d13be2456b17d`
Corpus `data/corpus.csv` sha256	`c1bec80f274294e6ce989fcc63fd7022a206f415e7d3df85e82c410e56dc4f4f`
Tokenizer rebuilt from this corpus equals the checkpoint's	Yes
Model shape, from the checkpoint	`{"context_len":32,"emb_dim":32,"ff_dim":128,"ln_eps":0.00001,"num_heads":2,"num_layers":2,"vocab_size":575}`
`config.py`, when this was exported	`{"adam_eps":1e-8,"batch_size":16,"beta1":0.9,"beta2":0.999,"context_len":32,"corpus_column":"text","corpus_path":"data/corpus.csv","emb_dim":32,"epochs":20,"ff_dim":128,"grad_clip":1,"head_dim":16,"ln_eps":0.00001,"lr_max":0.001,"lr_min":0.0001,"max_new_tokens":50,"max_rows":2000,"num_heads":2,"num_layers":2,"num_merges":500,"seed":42,"seq_len":16,"stride":1,"temperature":0.8,"top_k":0,"top_p":0.9,"val_fraction":0.1,"warmup_steps":50,"weight_decay":0.01}`
Seed	`42`
Environment	`Python 3.14.4, NumPy 2.5.2, pandas 3.1.0`

Embedding backward

What is it?

Embedding backwardImplemented

Embeddings

embed_backward takes dX0, the gradient of the loss on the first block's input, and fills the gradients of both tables. Each row of dX0 is added to the token row of its id and to the position row of its position. Nothing is returned, because the backward pass ends here.

Why does the model need it?

tok_emb and pos_emb are parameters, so the optimizer needs a gradient for each of them. The ids are integers and have no gradient, so these two tables are the last place a gradient can go.

Mathematics

Each row of the sum depends on one token row, with a derivative of 1. So the chain rule keeps only the positions that hold that id.

The chain rule for one entry of the token table Derived from implementation#

\frac{\partial L}{\partial E_{vj}} = \sum_{t} dX^{(0)}_{tj}\, \frac{\partial X^{(0)}_{tj}}{\partial E_{vj}}, \qquad \frac{\partial X^{(0)}_{tj}}{\partial E_{vj}} = \begin{cases} 1 & i_t = v \\ 0 & \text{otherwise} \end{cases}

Gradient accumulation per token row #

dE_v = \sum_{t\,:\, i_t = v} dX^{(0)}_t

Each position has its own row, so there is nothing to accumulate in the position table.

Gradient on the position table #

dP_t = \begin{cases} dX^{(0)}_t & t < T \\ 0 & t \geq T \end{cases}

Why repeated ids must accumulate. If an id appears at two positions, the loss depends on its table row twice, once through each position. The derivative is the sum of both effects, so both rows of dX0 belong in the gradient. Dropping either one gives a wrong gradient.

np.add.at does one addition for each element of ids, so it gets this right. The shorter g[ids] += dX0 does not. NumPy reads each selected row once, adds, and writes back, so for a repeated id the last write replaces the earlier ones. The code uses np.add.at for the token table. The position table can use a plain slice because its rows are never repeated.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$dX^{(0)}$	Gradient of the loss on the first block's input. loss_and_grads calls it dX, and the function's parameter is dX0	`dX0`	(T, d)
$dE$	Gradient on the token table	`grads["tok_emb"]`	(V, d)
$dP$	Gradient on the position table	`grads["pos_emb"]`	(Tmax, d)
$v$	A row of the token table	not named in the code	scalar
$i_t$	The id at position t	`ids[t]`	scalar

Tiny numerical example

A repeated id accumulatesIllustrative ExampleHand-picked numbers from test_embedding_backward_sums_repeated_ids. dX0 is chosen by hand, not computed from a loss.

Step	Code	Value
Ids	`ids`	`[2, 0, 2]`
Gradient on the sum	`dX0`	`[[1, 0], [0, 1], [2, 2]]`
Token table gradient	`np.add.at(grads["tok_emb"], ids, dX0)`	`[[0, 1], [0, 0], [3, 2], [0, 0]]`
Position table gradient	`grads["pos_emb"][: len(ids)] += dX0`	`[[1, 0], [0, 1], [2, 2]]`
Not in the code, for contrast	`g[ids] += dX0`	`[[0, 1], [0, 0], [2, 2], [0, 0]]`

One row: grads["tok_emb"][2] = dX0[0] + dX0[2] = [1, 0] + [2, 2] = [3, 2]

Checking the backward numericallyIllustrative ExampleA test-time idea, not part of the pipeline: gradcheck.numerical_grad is only called by tests. It nudges each number up and down by h = 1e-5 and measures how (X0 * dX0).sum(), whose gradient on X0 is exactly dX0, changes. Tests compare the central differences here with that real function.

Gradient	From the backward formula	From central differences	Largest gap
`grads["tok_emb"]`	`[0.0000, 1.0000][0.0000, 0.0000][3.0000, 2.0000][0.0000, 0.0000]`	`[0.0000, 1.0000][0.0000, 0.0000][3.0000, 2.0000][0.0000, 0.0000]`	`2.0e-11`
`grads["pos_emb"]`	`[1.0000, 0.0000][0.0000, 1.0000][2.0000, 2.0000]`	`[1.0000, 0.0000][0.0000, 1.0000][2.0000, 2.0000]`	`1.3e-11`

Tensor shapes

Into the token table. Rows of ids not in this window stay 0transformer.py:56

dX0(T, d)(16, 32)
grads["tok_emb"](V, d)(V, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32. V depends on the corpus. It is not a value in config.py.

Into the position table. Rows from T onward stay 0transformer.py:57

dX0(T, d)(16, 32)
grads["pos_emb"](Tmax, d)(32, 32)

T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. d = emb_dim = 32. Tmax = context_len = 32.

Actual implementation

transformer.pyembed_backwardGitHub

55def embed_backward(dX0, ids, grads):56    np.add.at(grads["tok_emb"], ids, dX0) # grads["tok_emb"][ids[t]] += dX0[t]57    grads["pos_emb"][: len(ids)] += dX0

Code and math connection

Equation	Role	Lines	Code
The chain rule for one entry of the token table	Computes it	transformer.py:56	`np.add.at(grads["tok_emb"], ids, dX0) # grads["tok_emb"][ids[t]] += dX0[t]`
Gradient accumulation per token row	Computes it	transformer.py:56	`np.add.at(grads["tok_emb"], ids, dX0) # grads["tok_emb"][ids[t]] += dX0[t]`
Gradient accumulation per token row	Calls it	transformer.py:188	`embed_backward(dX, ids, grads)`
Gradient accumulation per token row	Tests it	tests/test_transformer.py:42	`assert np.allclose(grads["tok_emb"][2], [3, 2]) # rows 0 and 2 both landed here`
Gradient on the position table	Computes it	transformer.py:57	`grads["pos_emb"][: len(ids)] += dX0`
Gradient on the position table	Tests it	tests/test_transformer.py:44	`assert np.allclose(grads["pos_emb"], dX0)`

Where is it used?

It is the last step of loss_and_grads, after every block has run backward (transformer.py:188).

Source

transformer.py, lines 55-57, embed_backward

PDF reference

Section 4.5, Eq. 40, p. 10

Used by

Called by: transformer.loss_and_grads

Tested by: tests.test_transformer.test_embedding_backward_sums_repeated_ids

Calls it: transformer.py, line 188 (loss_and_grads), for Gradient accumulation per token row
Tests it: tests/test_transformer.py, line 42 (test_embedding_backward_sums_repeated_ids), for Gradient accumulation per token row
Tests it: tests/test_transformer.py, line 44 (test_embedding_backward_sums_repeated_ids), for Gradient on the position table

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
The chain rule for one entry of the token table	Computes it	`transformer.py`	`embed_backward`	56	GitHub: transformer.py, line 56
Gradient accumulation per token row	Computes it	`transformer.py`	`embed_backward`	56	GitHub: transformer.py, line 56
Gradient on the position table	Computes it	`transformer.py`	`embed_backward`	57	GitHub: transformer.py, line 57

How this reaches a parameter update. embed_backward fills grads directly. batch_loss_and_grads averages the gradients over the batch (train.py:61), clip_grad_norm may scale them (train.py:62), and opt.step uses them (train.py:66). Both tables are matrices, so both also get weight decay. A token row with no id in the batch has a gradient of 0 for that step. It can still move, through what Adam holds from earlier steps and through weight decay.

Row of pos_emb	Non-zero gradient	Changed, weight_decay = 0	Changed, weight_decay = `0.01`
0	yes	yes	yes
1	yes	yes	yes
2	yes	yes	yes
3	yes	yes	yes
4	yes	yes	yes
5	no	no	yes
6	no	no	yes
7	no	no	yes

Row of pos_emb	Non-zero gradient	Changed, weight_decay = 0	Changed, weight_decay = `0.01`
0	yes	yes	yes
1	yes	yes	yes
2	yes	yes	yes
3	yes	yes	yes
4	yes	yes	yes
5	no	no	yes
6	no	no	yes
7	no	no	yes

Row of pos_emb	Non-zero gradient	Changed, weight_decay = 0	Changed, weight_decay = `0.01`
0	yes	yes	yes
1	yes	yes	yes
2	yes	yes	yes
3	yes	yes	yes
4	yes	yes	yes
5	no	no	yes
6	no	no	yes
7	no	no	yes