Mapping Token IDs to Vocabulary Logits with Embeddings and Linear Layers

Hello! Last time, you tracked tensors by shape, dtype, and device. Now those checks become a way to read model code: integer token IDs enter an embedding layer, and floating-point vectors leave a linear layer as scores for every token in the vocabulary.

This lesson builds a small, runnable version of that path. It is not yet a useful language model—the missing ingredient is a way to combine information from earlier tokens—but it establishes the input and output layers you will encounter in Llama-style models. Allow about 40 minutes, including the short video, two focused readings, and time to run the code.

From an ID to a vector

A token ID is an index, not a numerical description of a token. If cast has ID 2 and heal has ID 3, the difference between those numbers says nothing about their meanings. An embedding layer instead maintains a table with one learnable vector per vocabulary entry. Given ID 2, it retrieves row 2.

Watch this brief demonstration before writing any code:

If the vocabulary contains $V$ tokens and each vector has $D$ values, nn.Embedding(V, D) stores a weight table of shape $(V,D)$. For an input of shape $(B,T)$—batch size $B$, sequence length $T$—the output has shape $(B,T,D)$. PyTorch replaces each ID with its corresponding row while preserving the batch and sequence axes.

Read the embedding implementation and shape example in the Hugging Face guide. Its example assumes a tokenizer defined earlier in the article; for this lesson, focus on the layer and shapes rather than trying to run that excerpt independently.

“Learnable” does not mean that a newly created embedding already encodes useful semantics. Its entries begin with initialized values and acquire useful structure through training. In the example below, we will set values by hand so that every lookup is easy to verify.

From a vector to vocabulary scores

An embedding gives us $D$ numbers at each position, but a next-token predictor needs one score per possible next token. A linear layer performs that projection:

where $\mathbf{h}$ has $D$ values, $W$ has shape $(V,D)$, and the output $\mathbf{\ell }$ has $V$ values. PyTorch expresses this as nn.Linear(D, V). It applies the same projection to the last axis of its input, so a $(B,T,D)$ tensor becomes $(B,T,V)$ without your having to flatten the batch or sequence axes.

Each output value is a logit: an unnormalized score associated with one vocabulary ID. Logits are not probabilities and need not lie between zero and one. If we use the scores at position $t$ to predict the token after position $t$, there is one candidate score for each of the $V$ tokens.

In the same Hugging Face guide, look only at the output projection; the intervening transformer blocks are for a later module.

That distinction matters: the output head has the same basic job in our small example and in a full model, but the vector it receives is much more informative after a model has processed the preceding tokens.

Build and trace the two layers

Run this on the CPU; it needs only PyTorch. The vocabulary is deliberately tiny and the weights are set by hand. Here I, cast, and heal are readable stand-ins for tokens—real tokenizers can also produce pieces of words and special tokens.

import torch
from torch import nn

vocab = ["<eos>", "I", "cast", "heal", "fireball"]
vocab_size = len(vocab)   # V = 5
model_dim = 2             # D = 2

# Two sequences, each containing three token IDs.
ids = torch.tensor([
    [1, 2, 3],            # I cast heal
    [1, 2, 4],            # I cast fireball
], dtype=torch.int64)

embedding = nn.Embedding(vocab_size, model_dim)
head = nn.Linear(model_dim, vocab_size, bias=True)

# Fixed values make the lookup and projection inspectable.
with torch.no_grad():
    embedding.weight.copy_(torch.tensor([
        [ 0.0, 0.0],       # <eos>
        [ 1.0, 0.0],       # I
        [ 0.0, 1.0],       # cast
        [ 1.0, 1.0],       # heal
        [-1.0, 1.0],       # fireball
    ]))
    head.weight.copy_(torch.tensor([
        [ 0.0, 0.0],       # score for <eos>
        [ 0.0, 0.5],       # score for I
        [ 1.0, 0.0],       # score for cast
        [ 0.0, 1.0],       # score for heal
        [-1.0, 2.0],       # score for fireball
    ]))
    head.bias.zero_()

vectors = embedding(ids)
logits = head(vectors)

print("IDs:      ", ids.shape, ids.dtype, ids.device)
print("vectors:  ", vectors.shape, vectors.dtype, vectors.device)
print("logits:   ", logits.shape, logits.dtype, logits.device)
print("at 'cast':", logits[0, 1].tolist())

assert torch.allclose(vectors[0, 1], embedding.weight[2])
assert torch.allclose(
    logits[0, 1],
    vectors[0, 1] @ head.weight.T + head.bias
)

The shape trace is the core of the lesson:

The ID tensor is torch.int64, while the embeddings, weights, and logits are floating-point tensors. The modules and IDs all remain on the CPU in this script. If you later move these layers to an accelerator, move the input IDs to that same device before calling the embedding layer.

You can verify one projection without trusting the printed tensor. At position 1 of the first sequence, the input ID is 2 (cast), so its retrieved vector is $[0,1]$. The fireball row of head.weight is $[-1,2]$, and its bias is zero. Thus its logit is

The printed scores at that position, ordered like vocab, are [0.0, 0.5, 0.0, 1.0, 2.0]. The largest belongs to fireball. This is a consequence of the hand-set weights, not evidence that the model has learned to finish a sentence. Also notice the layout of head.weight: although nn.Linear is constructed with input size 2 and output size 5, PyTorch stores its weight as (5, 2)—one row for each output score.

What this small model can—and cannot—predict

The two sequences share I cast but end differently. At the cast position, both receive exactly the same vector and exactly the same logits:

print(torch.equal(logits[0, 1], logits[1, 1]))  # True

More importantly, any occurrence of ID 2 would receive those same logits in this two-layer setup, regardless of what came before it. The linear layer acts on each position’s vector separately; neither layer has combined information from earlier positions. You could label its output as next-token scores for a training task, but it cannot yet make those scores depend on the preceding sequence. Llama-style models place context-processing blocks between the embedding and the final vocabulary projection to address that limitation.

At each position, logits[..., j] denotes the score for vocabulary ID j. Choosing the largest score is one possible way to select a token, but there is no need to select one during this forward pass. Keep the raw logits: they are the input needed for the loss calculation in the next lesson.

Takeaways

nn.Embedding(V, D) looks up a length-$D$ vector for each integer ID, changing shape $(B,T)$ to $(B,T,D)$. nn.Linear(D, V) projects each such vector to $V$ floating-point logits, changing $(B,T,D)$ to $(B,T,V)$. Neither layer, by itself, mixes information across sequence positions.

Next, you will compare next-token logits with the tokens that actually followed and calculate cross-entropy loss and perplexity. That will give us a way to measure predictions before learning how training changes the weights.