How does the model turn tokens into vectors and back?

The embedding layer is a lookup table: the token id addresses a row of a large matrix, and that row is the vector entering the first block. The unembedding runs the other way and is not a table — it compares the vector leaving the last block against every row, one by one, and returns one number per vocabulary token. In many models both ends use the same matrix.

These are the model’s two borders. Before the first one there are no vectors, only token ids, and the path up to that point is in from text to tokens. After the second one there are no vectors either, only a score per vocabulary token, and what you do with that score is in logits and logprobs. Everything between the two borders is the subject of how an LLM works inside.

One naming collision is worth clearing up first. “Embedding” in vector search means something else: one vector per sentence or document, produced by a model trained so that texts with similar meaning land close together, a retrieval embedding. Here it is one vector per vocabulary entry, learned as a side effect of training the model to predict the next word. The two spaces are not interchangeable, and swapping one for the other is a common and expensive mistake.

Both ends, one matrix

The matrix has one row per vocabulary entry and one column per model dimension. In Qwen3-8B that is 151,936 rows of 4,096 numbers each, 622 million weights. It is one of the few parts of an LLM you can describe completely in a sentence.

Textbooks present the input side as a matrix multiplication: the token becomes a vector with a 1 at its position and zeros everywhere else, and that vector multiplies the matrix. The arithmetic works, and the result is exactly the corresponding row — the zeros cancel everything else. No implementation does this. It would be 151,936 multiplications to read a row whose address you already know. In practice it is a gather: take the row, done.

One subtlety from the 2017 paper survives in most models today. After the row is read, its values are multiplied by the square root of the model dimension¹. With 4,096 dimensions that is a factor of 64. The paper records the operation without explaining the reason; the usual reading is that one shared matrix needs different scales at each end, and the factor adjusts the magnitude on the input side without touching the output side.

On the way out the logic inverts. The vector leaving the last block is multiplied by every row of the matrix, and each dot product becomes a logit. One number per vocabulary token, 151,936 of them, computed from scratch for every token generated. Going in, the cost is one memory access; coming out, it is a multiplication against the entire matrix.

That asymmetry gives you the geometric intuition for what the model is doing. A dot product is largest when the final vector points along the row. Generating text, then, is pushing a vector until it points at the region of the space where the right token lives. Every block in between exists for that: moving one point around a space of a few thousand dimensions, one nudge at a time.

Reading each block’s contribution in vocabulary space

If the output is a product against the matrix, you can apply that product before the last layer and see what the model was about to say partway through. This works because the vector is never replaced along the way: it accumulates sums along the residual stream, and what the output matrix reads at the end is the total of everything written into it. The technique became known as the logit lens, and it underpins a good share of practical interpretability.

Geva and colleagues took the idea seriously in 2022 and reinterpreted the entire feed-forward layer in those terms. Treating a token’s representation as a distribution over the vocabulary, each layer becomes an additive update to that distribution, and each update decomposes into sub-updates corresponding to individual parameter vectors — many of them promoting concepts a person can name. They put it to two concrete uses: cutting GPT-2’s toxicity by roughly half and saving 20% of the compute with an early exit².

The method has a known problem, and it matters. Projecting an intermediate layer with the output matrix assumes that layer already writes in the same basis as the last one, which is not true. Belrose and colleagues measured the drift and proposed a correction: an affine probe trained per block, converting the hidden state into a distribution over the vocabulary. On models up to 20 billion parameters, the calibrated version came out more predictive, more reliable and less biased than the plain logit lens, and the prediction trajectories it produces identify malicious inputs with high accuracy³.

Dar and colleagues pushed it one step further, showing that the transformer’s own parameters can be projected into vocabulary space and interpreted there, with no forward or backward pass at all⁴. Same idea taken to its limit: the output matrix is the dictionary that translates any internal vector into words.

Tied weights

The idea of reusing one matrix at both ends dates to 2016 and arrived from two independent directions. Press and Wolf studied the topmost weight matrix of language models, showed that it is itself a valid word embedding, and recommended tying the two; the payoff was a consistent drop in perplexity, and in machine translation the model shrank to less than half its size with no loss of performance⁵. Inan, Khosravi and Socher reached the same arrangement from a theoretical framing of the loss function rather than an empirical observation⁶. The original transformer already shipped with shared weights, citing the first of the two¹.

Why it works is reasonable enough. Both matrices describe the same set of symbols; training two independent copies spends parameters rediscovering the same structure twice, and the output copy gets a training signal heavily concentrated on frequent tokens.

What changed since then is scale, and with it the decision. The Qwen3 family, released in 2025, is the cleanest illustration because it uses the same 151,936-entry vocabulary at every size. The 0.6B, 1.7B and 4B models tie their weights; the 8B and 14B do not. The arithmetic explains why. At 1,024 dimensions the matrix holds 156 million weights, and the 0.6B model has 0.44 billion parameters outside it. Tying saves 26% of the whole model. At 4,096 dimensions each matrix holds 622 million weights, and the pair accounts for 1.25 billion of the 8B model’s 8.2 billion, leaving 6.95 billion doing the actual work.

In a small model, then, tying is close to mandatory: a quarter of the parameter budget would go to a lookup table. In a large one the saving stops paying for the constraint, because the geometry that groups similar tokens on the way in is not the geometry that separates them on the way out. A direct example: going in, you want the spellings of a word to land near each other; coming out, you want to tell them apart, because only one of them is right in that context.

The cost, in memory and in arithmetic

The output matrix is the only place in the model where the vocabulary shows up in the arithmetic, and it behaves differently from everything else because of it.

Memory. In Qwen3-8B, 622 million weights at 16 bits take 1.24 GB. With the two matrices untied, 2.5 GB before a single block. It is not what dominates the bill for a model that size, but it is a number routinely dropped from quick estimates.

Compute per token. The output pass is 622 million multiply-accumulates, about 1.24 GFLOP. A full forward pass of the 8B model runs around 16 GFLOP, so the output head alone accounts for something close to 8% of the cost of generating a token. For a model with the same vocabulary and a tenth of the parameters, that same head becomes a much larger slice.

The logits tensor. This is the one that bites in practice. It holds 151,936 numbers per position; at 32 bits, 594 KB. A single position is irrelevant. But if you materialise the logits for the whole prompt — which is what happens in training, and in any code asking for logprobs over the full sequence — a 4,096-token context becomes a 2.5 GB tensor, for one request, before batching. That is why training libraries compute the loss in chunks and why most APIs return logprobs only for generated tokens and only for the top of the distribution.

Where this fails

The matrix rank caps what can be said. The logits across all possible contexts form a matrix whose rank cannot exceed the model dimension, because it is the product of two factors that wide. Yang and colleagues formalised this in 2017 as the softmax bottleneck: if the true distribution of language requires a higher rank, no amount of training fixes it, and the output is limited by construction. They proposed a mixture of softmaxes and improved the state of the art on Penn Treebank and WikiText-2, to perplexities of 47.69 and 40.68⁷. The scope is recurrent models in 2017 and today’s dimensions are far larger; the structural argument still stands.

Some words cannot win. Demeter, Kimmel and Downey showed a more concrete consequence of the same mechanism. Because a logit is a dot product and the probability comes from a softmax over those products, words whose vector sits in the interior of the convex hull of the embedding space have their probability bounded above by the words on the hull⁸. They can never be the most likely prediction, in any context. This is not undertraining, it is geometry.

The trip out is not the inverse of the trip in. It is tempting to read the figure’s symmetry as “undoing the embedding”, and it is not. Going in is exact addressing; coming out is a comparison against every row that returns a distribution. Nothing guarantees that running an embedding vector through the output recovers the token it came from, and two rows sitting close together does not mean the model confuses those tokens — only that the vectors point in similar directions.

The embedding row knows nothing about context. At the point the matrix is read, the vector for “bank” is the same in the sentence about money and the one about rivers. That is by construction, and it is what separates this vector from what retrieval calls an embedding. Treating an LLM’s matrix rows as semantic word representations applies to them an expectation inherited from word2vec and its line of work, which trained vectors for exactly that purpose and measured quality on syntactic and semantic similarity tests⁹. An LLM’s matrix was optimised for something else.

Changing the vocabulary breaks both ends. Adding tokens means resizing the matrix, and in a model with tied weights the new row lands on the input and the output at once. Initialising new rows with noise usually degrades the model immediately; what works is starting from the mean of the existing rows, or from the mean of the subtokens that used to spell the term.

What to do with this

1. Compare models by non-embedding parameters. That is the number predicting capability. Two “0.6B” models with different vocabularies can differ by 30% in what actually computes.
2. Do not use an LLM’s matrix as a search embedding. For retrieval, use a model trained for it. The distinction is in what an embedding is.
3. Ask for logprobs only on what you will use. The full tensor is hundreds of KB per position, and nearly every API returns just the top of the distribution by default.
4. Count the matrix separately when sizing memory for a small model. Below about 2 billion parameters it is a large slice, and it disappears from estimates made by rule of thumb.
5. When expanding the vocabulary, initialise from what exists. The mean of the old subtoken rows, not noise. And check whether the model ties its weights before touching either end.
6. Distrust the logit lens on intermediate layers. The raw projection is biased; if the conclusion matters, use a probe calibrated per layer.

Footnotes

1. Vaswani et al. (2017) describe, in section 3.4, sharing the same weight matrix between the two embedding layers and the pre-softmax linear transformation, with the input-side weights multiplied by the square root of the model dimension. ↩ ↩²

2. Geva et al. (2022) treated a token’s representation as a distribution over the vocabulary and showed each feed-forward layer to be an additive update decomposable into interpretable sub-updates; they applied the result to cut roughly half of GPT-2’s toxicity and save 20% of the compute with an early exit. ↩

3. Belrose et al. (2023) trained an affine probe per block to convert hidden states into a distribution over the vocabulary and showed it to be more predictive, more reliable and less biased than the logit lens, on models up to 20 billion parameters. ↩

4. Dar et al. (2022) proposed interpreting transformer parameters by projecting them into vocabulary space, with no forward or backward pass, and applied it to aligning parameters across models sharing a vocabulary. ↩

5. Press and Wolf (2016) showed that the topmost weight matrix of a language model is a valid word embedding and that tying it to the input matrix reduces perplexity; in machine translation the model drops to less than half its size with no loss of performance. ↩

6. Inan et al. (2016) reached the same sharing from a theoretical framing of the loss function, arguing that training against one-hot targets wastes both information and parameters. ↩

7. Yang et al. (2017) framed language modelling as matrix factorisation and showed that a softmax over distributed embeddings caps the rank of the logit matrix, proposing a mixture of softmaxes that brought perplexity to 47.69 on Penn Treebank and 40.68 on WikiText-2. ↩

8. Demeter et al. (2020) showed, through numerical, theoretical and empirical analysis, that words in the interior of the convex hull of the embedding space have their probability bounded by the words on the hull, which structurally prevents them from being the most likely prediction. ↩

9. Mikolov et al. (2013) proposed two architectures for learning continuous word vectors from very large corpora, with large quality gains at much lower computational cost, measured on a test set of syntactic and semantic word similarities. ↩