Computing Token Confidence Scores from Vocabulary Distributions

Hello! Last time, you created a fully masked canvas and used the denoiser to obtain a vocabulary distribution at every position. The canvas stayed masked: the distributions were proposals, not generated text.

Today you will turn those distributions into token confidence scores. This gives you a way to compare the model’s strongest proposals across positions. We will compute the scores, but leave the decision about how many positions to reveal for the next lesson.

What does “confidence” measure here?

For a masked position $i$, suppose the denoiser assigns these probabilities to three possible characters:

The model’s leading proposal is a. A common confidence score for that proposal is simply its probability, 0.60:

The two outputs answer different questions: ${\hat{v}}_i$ says which token?; $c_i$ says how much probability did the model assign to that token? Both are computed separately for every sequence position.

For a quick visual refresher on interpreting softmax outputs, watch these portions of “Neural Networks Part 5: ArgMax and SoftMax” by StatQuest with Josh Starmer:

Now compare positions, rather than characters within one position:

Position 2 has the highest-confidence proposal, even though it appears last in the sequence. Confidence-based decoding can use this ranking; it need not reveal text strictly from left to right. The score is still a heuristic, not a measured probability that the character is correct. In particular, several confident-looking characters do not automatically form a coherent sentence together.

Compute scores from your Colab tensors

Reuse probabilities and x from the previous lesson. Their shapes are $[B,L,V]$ and $[B,L]$, respectively. The previous cell already set the probabilities of PAD and MASK to zero before normalizing over the vocabulary.

# probabilities: [batch, length, vocab]
# x:             [batch, length]

top_confidence, top_token_id = probabilities.max(dim=-1)
# Both tensors have shape [batch, length].

# Only positions still awaiting generation should receive a usable score.
eligible = x.eq(mask_id)
confidence = top_confidence.masked_fill(~eligible, -torch.inf)

print("proposed token IDs:", top_token_id[0].tolist())
print("confidence scores:", confidence[0].tolist())
print("still-masked positions:", eligible[0].tolist())

On your initial fully masked canvas, every position is eligible. Later, x might also contain a fixed prompt prefix, already revealed characters, or padding. Those positions should not compete with masked positions for the next reveal. Setting their scores to -inf excludes them from a highest-score comparison without changing the underlying probability distributions.

There is one useful distinction to keep in mind. The code scores the top candidate at each position. If you instead choose a different candidate—for example, b from the first row of our table—its confidence is its own probability, 0.30, not the top candidate’s 0.60. For any tensor of selected token IDs shaped $[B,L]$, you can look up their probabilities with:

selected_confidence = probabilities.gather(
    dim=-1,
    index=selected_token_ids.unsqueeze(-1),
).squeeze(-1)

You do not need selected_token_ids for today’s greedy proposals; the max call already returned their scores. The lookup becomes useful if a later sampler selects tokens other than the most probable ones.

Takeaway

Token confidence in this lesson is the denoiser’s probability for a proposed token at a particular position. probabilities.max(dim=-1) computes the strongest candidate and its score across all positions at once; an eligibility mask keeps fixed or non-generation positions out of comparisons. These scores help rank proposals, but they are not guarantees of correctness, and computing them does not change the canvas.

Next, you will use a reveal schedule to work out how many positions may be committed at each denoising step.