TT Lab
Get started
Learn Learning paths Courses

Transformers — Compute Attention By Hand

Build the Loss by Hand

Continue in TT Lab

Goal

Starting from logits, build yourself, using only the standard library, the process of turning "how well this model does" into a single number. After climbing up through a stable log-softmax, the loss at one position, cross entropy as the average over many positions, and perplexity exponentiated back, you place side by side and measure how much shifting the labels by one position and removing the padding positions change that number. Finally you convert to a base-2 logarithm and read it in bits per token as well.

Why it matters

Both training and evaluation move by looking at this one number. But in the process of producing this number there are four places where it goes silently wrong without raising an error. When you took the logarithm, whether you shifted the labels by one position, whether you removed the padding, and what the base of the logarithm is. All four are one line of code, they raise no exception even when wrong, and they usually go wrong in the direction where the number gets better. So there is no reason to suspect them. This lab does not call a real model. The system Python of this Pod has no numpy, torch or transformers. Instead it builds deterministically a row of scores over the whole vocabulary for each position and hand-writes the same computation on top of that. So statements like "model X has a perplexity of N" are not made here. Every number that comes out was measured with data you built. If the neighboring module is about how to pick one from a distribution (temperature, top-k, top-p), this one is about how to measure how wrong that distribution is. Measuring comes before picking. The grader does not trust the explanations you wrote down. It actually imports your module, pokes at the functions directly with different logits every time, and checks them against values it computes separately. The inputs change on every run, so you cannot memorize values and plug them in.

Steps

  1. In /root/work/tf-loss/loss.py, create VOCAB, PAD_ID, SEQ, dataset() and log_softmax(xs). It goes straight from logits to log probabilities without going through probabilities.
  2. Add NEG_INF and naive_log_softmax(xs) to build a version with deliberately the wrong order. It computes the probability first and then takes the log, reproducing -inf appearing at the bottom.
  3. Add token_loss(logits, target) so that it measures the loss at one position. It is the negated log probability of the correct token.
  4. Add cross_entropy(rows, targets) so that it takes the average of the loss over many positions.
  5. Create perplexity(rows, targets) and uniform_perplexity(vocab_size, length). You confirm that the perplexity of a uniform distribution equals the vocabulary size.
  6. Create shift_pairs(rows, ids), shifted_loss(rows, ids) and unshifted_loss(rows, ids) and measure side by side when the labels are shifted by one position and when they are not.
  7. Create kept_positions(targets, pad_id) and masked_cross_entropy(rows, targets, pad_id) so that they measure with the padding positions removed. You remove them in both the numerator and the denominator.
  8. Add bits_per_token(loss) and nats_per_token(bits), and record the measured values in /root/work/tf-loss/loss_report.json and /root/work/tf-loss/loss_report.md.

Notes

Straight from logits to log probabilities

In /root/work/tf-loss/loss.py, create VOCAB (at least 12), PAD_ID, SEQ (at least 16, with the last 3 or more being PAD_ID), dataset() and log_softmax(xs). dataset() returns (로짓 줄 목록, 토큰 번호 목록) (the placeholders are the list of logit rows and the list of token numbers) and uses no random numbers. log_softmax produces log probabilities straight from logits without going through probabilities.

Run mkdir -p /root/work/tf-loss. The formula is the one line x_i - (max + log sum exp(x - max)). The point is that there is no division — if you build the probability, divide and then take the log, a very small probability becomes 0.0 and the log collapses. If the second cell of log_softmax([0.0, -800.0]) is a finite value (near -800), it is correct. Build dataset() so that the row at position t gives the highest score to the token at position t+1, and put a larger score on the positions where padding is the correct answer.

Break it on purpose

Add NEG_INF and naive_log_softmax(xs). This time you build the probability first and then take the log. At positions where the probability collapsed to 0.0, math.log throws an exception, so fill them with NEG_INF yourself. Confirm that on the same input log_softmax is finite while only this one gives -inf.

NEG_INF = float("-inf"). Only the order needs to change — divide the exped values by the sum to build the probability, and take the log of that probability. In the middle values it gives the same answer as the function from the previous step, and it parts ways only at the bottom. Put in a row with a large difference like [0.0, -800.0]. You must check whether the probability is 0.0 and filter it; if you just call math.log, it ends in an exception.

The loss at one position

Add token_loss(logits, target). It is the negated log probability the model gave to the correct token. If you gave the correct one probability 1, it is 0, and it grows as the probability shrinks.

It is one line — -log_softmax(logits)[target]. If you forget to flip the sign, all the values become negative and the statement "the loss goes down" is reversed. You must not use the logits themselves. What you gave to the other tokens is not counted separately — because the sum is 1, the correct one's share is the rest's share.

The average over many positions

Add cross_entropy(rows, targets). It produces token_loss at each position and returns the average. If targets is empty, it is 0.0.

It is an average, not a sum. If you measure by the sum, a long sentence is always a bad sentence and you cannot compare texts of different lengths. Pair them with zip(rows, targets), add, and divide by len(targets). What you put in the denominator here becomes a problem again in step 7.

The scale of perplexity

Create perplexity(rows, targets) and uniform_perplexity(vocab_size, length). The first is exp(평균 손실) (the Korean words mean mean loss), and the second builds logit rows with all scores equal and measures the perplexity. Confirm that the result equals vocab_size.

It is the one line math.exp(cross_entropy(rows, targets)). It is easy to confuse with taking exp at each position and averaging, but for a uniform distribution the two values happen to be equal, so that test cannot tell them apart. uniform_perplexity builds rows of the form [[0.0] * vocab_size] * length and you can put in any correct answer — all the probabilities are 1/V, so the loss is log V, and exponentiating gives V. This value is the starting point of the perplexity scale.

The labels are shifted by one position

Create shift_pairs(rows, ids), shifted_loss(rows, ids) and unshifted_loss(rows, ids). shift_pairs pairs them as (rows[:-1], ids[1:]) so that the logits at position t predict the token at position t+1. The other two are the cross entropy of the shifted version and the unshifted version.

The last row has no next token to predict, and the first token has no row that predicted it. So the logits drop one from the back and the tokens drop one from the front. If you shift in the opposite direction (rows[1:], ids[:-1]), it is handing over what it has already seen as the answer, so the number strangely gets good. On your data the shifted loss must be clearly smaller than the unshifted one — because dataset() put the score on the token at position t+1.

Measure with the padding removed

Create kept_positions(targets, pad_id) and masked_cross_entropy(rows, targets, pad_id). The first is the list of numbers of the positions whose correct answer is not padding, and the second is the average obtained by adding only those positions and dividing by the number of those positions. If there are no positions left, it is 0.0.

There are two places to remove — the adding side and the dividing side. If you miss the dividing side and divide by the total length, the value silently gets smaller, and since the direction is always toward getting better, there is no reason to suspect it. If you keep kept_positions separate, you can see with your own eyes what was left. If you use this function on the shifted pair, it must be larger than the value measured with the padding included — because padding is an easy position to predict.

Place the four numbers side by side

Add bits_per_token(loss) and nats_per_token(bits), and write vocab_size, pad_id, seq_len, pad_count, kept, dropped, unshifted_loss, shifted_loss, masked_loss, unshifted_perplexity, shifted_perplexity, masked_perplexity, bits_per_token, uniform_perplexity, probe_gap, naive_is_inf and stable_logprob in /root/work/tf-loss/loss_report.json, and write /root/work/tf-loss/loss_report.md in the four sections ## 무엇을 쟀나 ## 한 칸 어긋나면 ## 패딩을 빼면 ## 비트로 재면 (the Korean headings mean "What was measured", "If it is off by one position", "If you remove the padding" and "Measured in bits").

Do not write the numbers by hand; fill them in with values obtained by actually running your own code. masked_loss is the value of using masked_cross_entropy on the pair shifted with shift_pairs, and shifted_loss is the value of measuring the same pair without masking. bits_per_token is computed based on masked_loss — check for yourself whether 2 ** 그 값 (the Korean words mean that value) equals masked_perplexity. uniform_perplexity is uniform_perplexity(VOCAB, 8). probe_gap is fixed at 800, naive_is_inf is whether naive_log_softmax([0.0, -800.0])[1] is -inf, and stable_logprob is the second cell that log_softmax gave on the same input.