A language model gives logits over V tokens at each position. Softmax turns logits into probabilities; log_softmax computes log probabilities stably. The next-token loss at position t scores token t+1 using logits at t. Scoring the current token instead leaks the answer through the input alignment.
The logit at position t scores the following token, and the mask belongs to that target token.
import torch
logits = torch.tensor([[1000.0, 1000.0, 1000.0]], dtype=torch.float64)
print((-torch.log_softmax(logits, dim=-1)).tolist())With T tokens, only T-1 next-token targets exist. A target mask may exclude prompt positions and padding: include target_mask[t+1] when selecting the loss from logits[t]. Divide by the number of included target tokens. Summing losses or dividing by all positions rewards different sequence lengths.
Attention combines values using weights from query-key compatibility. For q,k of shape (T,D), scores are q@k.T divided by sqrt(D). Softmax over keys gives one distribution per query. Causal attention forbids future keys: scores strictly above the diagonal become negative infinity before softmax.
import torch
scores = torch.zeros((3, 3))
future = torch.triu(torch.ones((3, 3), dtype=torch.bool), diagonal=1)
weights = scores.masked_fill(future, float("-inf")).softmax(dim=-1)
print(weights.tolist())Zeroing future weights after softmax leaves rows with total mass below one. Mask first. The diagonal is allowed: a query can attend to its own position. This exercise is one attention head with no batch axis, dropout or positional embeddings. It is a small numerical model, not an implementation of a deployed LLM. No tokenizer package or weights are required.
Causal attention assigns no probability mass to future keys; mask the scores before softmax so allowed weights still normalize to one.
In the build, test loss with a prompt mask and unequal target probabilities. For attention, changing a future value must leave every earlier output unchanged. An all-zero score matrix has an easy hand-computed causal average.