From the Embedding Table to the Logits

Goal

Build, using only the standard library, the path from a token number to a vector and from that vector back to scores over the whole vocabulary. You compute with two methods to confirm that a lookup gives the same value as a one-hot vector times a matrix, count the parameters, build weight tying, which uses the same table once more on the output side, and see in numbers where cosine and the dot product part ways. Finally you measure that multiplying the embedding by √d changes only the magnitude and leaves the direction as it is.

Why it matters

This is where people most often go wrong when talking about models. If you think of it as split, "an embedding is a lookup and the output is a matrix multiplication", you cannot see that they are the same table, and you cannot explain why the parameter counts disagree. Why a proposal to raise the vocabulary ends in a memory meeting also makes no sense until you count V times d. This lab does not call a real model. The system Python of this Pod has no numpy, torch or transformers (numpy exists only inside /opt/onnx-lab/bin/python). So numbers such as the vocabulary size or parameter count of a real model are not used here — only values measured from the table you built. Many of the judgments are floating-point comparisons. The grader checks with abs(a-b) <= atol + rtol*abs(b), but also checks integer judgments such as shapes, parameter counts and neighbor lists. So it accepts the last digit wobbling because the order of adding differs, and filters out implementations that are actually wrong. The grader does not trust the explanations you wrote down. It actually imports your module, pokes at the functions with a different table and different numbers every time, and checks them against values it computes separately.

Steps

1. In /root/work/tf-embed/embed.py, create VOCAB = 512, DIM = 64, SEED = 20260917 and make_table(vocab, dim, seed), shape(matrix) and param_count(vocab, dim). The same seed must always give the same table.
2. Add lookup(table, ids) so that it turns a list of numbers into a list of vectors. The same number gives the same vector, and a number outside the vocabulary is an IndexError.
3. Add one_hot(token_id, vocab), row_times_matrix(vec, matrix), lookup_via_one_hot(table, ids) and one_hot_mults(vocab, dim, count) and see whether a lookup and a one-hot multiplication give the same value.
4. Create logits(table, hidden) so that it produces scores over the whole vocabulary with the very table used for the input. You must not build a new matrix.
5. With tied_params(vocab, dim), untied_params(vocab, dim) and vocab_growth(vocab, dim, factor), count the number of values when tied and when kept separate.
6. Create dot, norm, cosine, stretch, nearest_by_cosine and nearest_by_dot and see how the two lists part ways when you stretch just one row.
7. Create rms(vec) and scaled_lookup(table, ids, dim) so that you measure that multiplying by √d makes the magnitude √d times.
8. Measure everything with the fixed constants and record it in /root/work/tf-embed/embed_report.json and /root/work/tf-embed/embed_report.md.

Notes

• Execution contract: the grader imports /root/work/tf-embed/embed.py as a Python module and uses VOCAB, DIM, SEED, make_table, shape, param_count, lookup, one_hot, row_times_matrix, lookup_via_one_hot, one_hot_mults, logits, tied_params, untied_params, vocab_growth, dot, norm, cosine, stretch, nearest_by_cosine, nearest_by_dot, rms and scaled_lookup directly. It does not run it as a script, so if __name__ == "__main__" is not needed.
• make_table(vocab, dim, seed) creates one random.Random(seed) and fills it row by row, starting from cell 0 of token 0. For each cell it draws random() once and subtracts 0.5. That way the grader can build the same table separately and match the values.
• shape(matrix) returns (줄 수, 칸 수) (the placeholders are the number of rows and the number of cells), and raises a ValueError if the number of cells differs from row to row. An empty table is (0, 0).
• lookup(table, ids) raises an IndexError for a number outside the vocabulary. A negative number is also outside the vocabulary — if you leave Python's negative indexing alone, it counts from the back and silently hands you the wrong row.
• row_times_matrix(vec, matrix) produces a vector of length d from a row vector of length V and a V×d table. It is out[c] = sum(vec[r] * matrix[r][c] for r in range(V)).
• one_hot_mults(vocab, dim, count) is the number of multiplications. It is V times d per token, and 0 with a lookup. Do not measure time; count this number.
• logits(table, hidden) is a list of length V. It is out[t] = sum(table[t][c] * hidden[c] for c in range(d)), and it is a ValueError if the length of the hidden vector differs from the width of the table. Using the table as it is is weight tying.
• vocab_growth(vocab, dim, factor) is a dictionary with the keys vocab, bigger_vocab, dim, tied, bigger_tied, untied, bigger_untied and saved. saved is the value when kept separate minus the value when tied.
• cosine(a, b) is 0.0 if a vector of length 0 comes in. rms([]) is also 0.0.
• stretch(table, token_id, factor) returns a new table with only that row stretched by factor. Do not modify the table you received in place.
• nearest_by_cosine and nearest_by_dot return the indices of the k items in descending score order. They leave out the item itself, and in case of a tie the smaller number comes first.
• scaled_lookup(table, ids, dim) is the result of lookup multiplied by math.sqrt(dim).
• Step 8 builds the table with VOCAB = 512, DIM = 64, SEED = 20260917 and probes token 137. The neighbors use k = 5, the row to stretch is the fifth cosine neighbor (the fourth cell counting from 0), and the stretch multiplier is 4.0. The multiplier for raising the vocabulary is 2.
• This Pod has no internet. pip install does not work and import numpy does not work in the system Python either. math and random are enough.
• Official documents: Attention Is All You Need · PyTorch — MultiheadAttention · NumPy — matmul · Python — math
• Common mistakes: filling the table cell by cell so the seed order goes off, letting negative numbers pass through, flipping rows and columns in the one-hot multiplication, building a new matrix for the logits, not counting the separate case as double, not dividing by length in cosine, stretch modifying the original table, and multiplying by d instead of √d.

The shape of the table and the parameter count

In /root/work/tf-embed/embed.py, create VOCAB = 512, DIM = 64, SEED = 20260917 and make_table(vocab, dim, seed), shape(matrix) and param_count(vocab, dim). make_table fills the table row by row from cell 0 of token 0 with a single random.Random(seed), and for each cell subtracts 0.5 from random().

The table is a list of lists. If you create rng = random.Random(seed) only once and draw dim values per row, the row-by-row order is preserved. If you go cell by cell, a different table comes out even with the same seed. Have shape check the number of cells in each row and raise a ValueError if they differ — a table with a mismatched shape silently produces strange values later. param_count is a product, not a sum.

Pull out rows by number

Add lookup(table, ids). It turns a list of numbers into a list of vectors. If the same number appears twice, exactly the same vector must come out twice, and a number outside the vocabulary (including negative ones) is an IndexError.

Pulling out the row is all there is. But table[-1] silently gives the last row in Python, so you have to check 0 <= token_id < len(table) yourself. If you copy the row you pulled out with list() and return it, there is no risk of the caller touching the original. That the same number gives the same vector is not a bug but a property — an embedding has no context, and the context comes later from attention.

A lookup is the same value as a one-hot multiplication

Add one_hot(token_id, vocab), row_times_matrix(vec, matrix), lookup_via_one_hot(table, ids) and one_hot_mults(vocab, dim, count). The value computed by the one-hot multiplication must equal lookup, and one_hot_mults returns the number of multiplications, V times d per token.

One cell of row_times_matrix is sum(vec[r] * matrix[r][c] for r in range(V)). If you flip rows and columns, even the length goes off, so check with shape before looping. A one-hot has only one position at 1.0, so when you multiply and add, only that row survives — the values must be exactly the same to be correct. Do not try to measure time; count the multiplications. A lookup is 0.

Produce scores with the same table

Create logits(table, hidden). It turns one hidden vector into scores over the whole vocabulary. You must use the very table used for the input, as it is (weight tying). It is a ValueError if the length of the hidden vector differs from the width of the table.

If you take the dot product with the hidden vector for each row, you get a list of length V. If you build and use a new matrix, the values come out completely different — tying means reusing that very table. There is one good way to check. Put the embedding of some token into hidden as it is. The dot product with itself is the square of its length, so that token's score becomes the largest.

How much does tying reduce

Create tied_params(vocab, dim), untied_params(vocab, dim) and vocab_growth(vocab, dim, factor). vocab_growth is a dictionary with the keys vocab, bigger_vocab, dim, tied, bigger_tied, untied, bigger_untied and saved, and saved is the value when kept separate minus the value when tied.

When tied, the table is one set; when kept separate, the same shape is two sets. When you raise the vocabulary by factor, both values grow by that multiple even though you did not touch the width — that is the cost of the vocabulary size. Everything here is an integer judgment, so do not compute with floating point and round.

Length changes the order

Create dot(a, b), norm(vec), cosine(a, b), stretch(table, token_id, factor), nearest_by_cosine(table, token_id, k) and nearest_by_dot(table, token_id, k). The neighbors are the indices of the k items in descending score order, leaving out the item itself. In case of a tie, the smaller number comes first.

cosine is the dot product divided by the product of the two lengths, and if a length is 0 there is no direction, so it is 0.0. stretch must build a new table — if you modify the original in place, every later judgment goes off. For sorting, key=lambda item: (-점수, 번호) (the placeholders are the score and the index) can express the tie rule in one line. Stretch one row and look at the two lists side by side. The cosine list stays the same, while in the dot-product list the stretched row comes to the front.

Multiplying by √d changes only the magnitude

Create rms(vec) and scaled_lookup(table, ids, dim). rms is the root mean square and is 0.0 for an empty vector. scaled_lookup is the result of lookup multiplied by math.sqrt(dim).

You are multiplying every cell by the same number, so the direction does not change at all — if you measure the cosine, it is the same as before multiplying. Only the magnitude changes, and it is exactly math.sqrt(dim) times. If you multiply by dim as it is, the magnitude becomes d times and is a completely different value. For rms, divide the dot product by the number of cells and then take the square root.

Leave what you measured as a record

Build the table with VOCAB = 512, DIM = 64, SEED = 20260917, probe token 137 and measure everything. The neighbors use k = 5, the row to stretch is the fifth cosine neighbor, the multiplier is 4.0, and the multiplier for raising the vocabulary is 2. Write vocab, dim, seed, probe, table_shape, params, tied_params, untied_params, saved, bigger_vocab, bigger_tied, bigger_untied, lookup_mults, one_hot_mults, max_abs_diff, repeat_same, logit_len, logit_argmax, logit_argmax_is_self, cos_neighbors, dot_neighbors, neighbors_differ, stretch_target, stretch_factor, cos_after_stretch, dot_after_stretch, cos_unchanged_by_scale, rms_plain, rms_scaled and rms_ratio in /root/work/tf-embed/embed_report.json, and write /root/work/tf-embed/embed_report.md in the five sections ## 무엇을 쟀나 ## 조회와 원-핫 곱은 같은 연산이다 ## 가중치를 묶으면 무엇이 줄어드나 ## 내적과 코사인이 갈리는 자리 ## √d 를 곱하면 무엇이 달라지나 (the Korean headings mean "What was measured", "A lookup and a one-hot multiplication are the same operation", "What shrinks when you tie the weights", "Where the dot product and cosine part ways" and "What changes when you multiply by √d").

Do not write the numbers by hand; fill them in with values obtained by actually running your own code. one_hot_mults is the value for looking up the single token 137, and lookup_mults is 0. max_abs_diff is the largest absolute value among the differences between the values of lookup and lookup_via_one_hot. repeat_same is whether the two vectors are the same when you look up the same number twice. logit_argmax is the index with the largest score when you put the embedding of token 137 into the hidden vector as it is. cos_unchanged_by_scale is whether the cosine before and after multiplying by √d is the same. rms_ratio is the magnitude after multiplying divided by the magnitude before.