Projecting Encoder Hidden States to Vocabulary Logits

Welcome back. In the previous lesson, your non-causal encoder turned each corrupted token position into a context-aware vector. A masked position could use visible tokens on both sides, but the encoder did not yet say which character belongs there.

Today you’ll add the final piece of the denoiser: a small output layer that gives one score per vocabulary token, at every sequence position. You’ll check the shapes in Colab and distinguish those scores from probabilities. This lesson should take about 15–20 minutes.

From a hidden vector to token scores

Suppose the encoder returns a tensor with shape [batch, length, d_model]. In last lesson’s example, that was [2, 4, 16]: two sequences, four positions each, and a 16-number hidden vector at each position.

We want each position to have a score for every token in the vocabulary. If the vocabulary contains eight tokens, the desired output shape is [2, 4, 8]. An nn.Linear(16, 8) layer does exactly this. It reads the last dimension of each hidden vector and replaces it with eight output scores, leaving the batch and position dimensions intact.

For a hidden state $H$, the score for vocabulary token $v$ at position $\ell$ in batch item $b$ is:

The same learned weights $W$ and biases $a$ are used at every position. Predictions can nevertheless differ across positions because the encoder has produced different context-aware hidden states for them.

For a quick visual explanation of this final projection, watch these portions of Coding a ChatGPT Like Transformer From Scratch in PyTorch from StatQuest with Josh Starmer:

The PyTorch reference makes the shape rule precise:

Add the vocabulary head in Colab

Run the input-module and encoder cells from the preceding lessons first. The following cell reuses their inputs, encoder, x_t, t, and pad_id variables. It is a shape check with untrained weights, not a training step:

import torch
from torch import nn

d_model = 16
vocab_size = 8
vocab_head = nn.Linear(d_model, vocab_size)

inputs.eval()
encoder.eval()
vocab_head.eval()

with torch.no_grad():
    hidden_inputs = inputs(x_t, t)             # [2, 4, 16]
    hidden_states = encoder(
        hidden_inputs,
        mask=None,
        src_key_padding_mask=x_t.eq(pad_id),
        is_causal=False,
    )                                         # [2, 4, 16]
    logits = vocab_head(hidden_states)        # [2, 4, 8]

print("encoder output:", hidden_states.shape)
print("vocabulary logits:", logits.shape)
print("position 1, item 0:", logits[0, 1].shape)

assert hidden_states.shape == (2, 4, d_model)
assert logits.shape == (2, 4, vocab_size)
assert logits[0, 1].shape == (vocab_size,)

logits[0, 1] selects the first sequence’s masked position from last lesson. Its eight numbers are logits: relative scores for vocabulary IDs 0 through 7, not yet probabilities or predicted text.

Notice that the projection also produces scores at visible-token and [PAD] positions. That is expected: nn.Linear operates on every hidden vector it receives. The padding attention mask stops [PAD] from serving as context; it does not remove that position from the output tensor. In the next lesson, the loss will select only corrupted, non-padding positions for training.

What the scores do—and do not—tell you

You can turn one position’s logits into probabilities to inspect them:

position_logits = logits[0, 1]                 # [8]
probabilities = torch.softmax(position_logits, dim=-1)

print("candidate ID:", position_logits.argmax().item())
print("probabilities sum:", probabilities.sum().item())

argmax picks the highest-scoring vocabulary ID; softmax converts all eight scores into a distribution that sums to approximately one. Because the encoder and vocabulary head are untrained, the candidate ID has no useful linguistic meaning yet. Keep the raw logits as the model’s output; the training loss in the next lesson will use them directly.

The distinction from next-token prediction is worth keeping in view. A left-to-right model may produce logits at many positions during training, but uses them to predict subsequent tokens. Here, logits[b, ℓ] scores what token belongs at position $\ell$ given the corrupted sequence and its available context.

What to carry forward

Your denoiser now has an end-to-end shape story: token IDs enter as [batch, length], the encoder produces [batch, length, d_model], and a shared linear head produces vocabulary logits shaped [batch, length, vocab_size]. These logits exist at every position; they are not probabilities, and their existence does not mean every position should contribute to training.

Next, you’ll use the corruption mask to compute a loss only where the model was asked to recover a hidden token.