Shifted Targets Turn Logits into Next-Token Loss

Align each logits row with the token one step to its right, then calculate stable cross-entropy and average only valid targets. Diagnose double shifts, wrong axes, padding, and future leakage.

Each logits row must be paired with the token immediately after its visible context. For the frozen window:

model input ID	model input token	target ID	target token
0	`<bos>`	1	`A`
1	`A`	2	`B`
2	`B`	3	`C`
3	`C`	7	`<eos>`

The input tensor is [0,1,2,3]; the target tensor is [1,2,3,7]. The row at input B predicts C. It must not receive C through preprocessing or an unmasked future route.

Calculate One Stable Cross-Entropy

For one logits row $z:(V)$ and target ID $y$ ,

\ell=-\log\frac{e^{z_y}}{\sum_j e^{z_j}} =-z_y+\log\sum_j e^{z_j}.

Calculate log-sum-exp stably with $m=\max_j z_j$ :

\log\sum_j e^{z_j}=m+\log\sum_j e^{z_j-m}.

For logits $[1,2,0]$ and target index 1, $m=2$ and

\ell=-2+2+\log(e^{-1}+1+e^{-2})\approx0.407606.

The class indices here use zero-based implementation indexing; “target index 1” selects the second logit.

Average Only Valid Targets

If a batch uses padding or truncated regions, calculate the mean over valid target positions:

L=\frac{\sum_{b,t}M_{b,t}\ell_{b,t}} {\sum_{b,t}M_{b,t}},

where $M$ is 1 for included targets. Dividing by $BT$ despite excluded targets changes the reported scale. A batch with no valid targets must be rejected or handled explicitly to avoid division by zero.

Common Shift Errors

pairing each logits row with its current input token;
shifting twice in data preparation and again in the loss function;
including a target outside the vocabulary;
averaging padded targets;
applying softmax across time instead of vocabulary;
exposing a future target through an incorrect causal mask.

Q1. Align one frozen target

Which target ID is paired with the logits row whose input token is C in the frozen window?

Your answer

Compute it first, then check your number.

Hint

Read the final row of the alignment table.

Solution

The target ID is 7.