Loss in Pretraining / SFT

0. One shared loss

• O(1T) params, O(10k) chips, O(1 month) training -> one line of cross-entropy loss
• both pretraining and SFT use it
• diff between pretraining and SFT: loss mask.

In this blog, we will cover

• what “cross-entropy (CE)” actually means
• why it’s the same thing as negative log-likelihood (NLL)
• how KL divergence is one algebraic step away
• where normalization details quietly break your training run

1. LLM is a conditional distribution

LLM defines a distribution over the next token at every position:

• : position index, not token ID.
• : current token to generate
• : context

Prob of the whole sequence:

2. Negative log-likelihood (NLL)

We want params to produce distributions similar to the training data
-> maximum likelihood (the oldest idea in statistics):

• is the training dataset (a set of sequences)
• is the likelihood the model assigns to that entire dataset
• For a single sequence, can be calculated using the formula from Section 1.

Turn this into something an optimizer can use:

• take the log:
  • a 4k sequence is a product of rk probs
  • -> precision issues (underflows to zero fast)
  • -> is monotonic, so it keeps the same , and multiplications become additions
• flip the sign
  • optimizers minimizes
  • maximizing log-likelihood = minimizing negative log-likelihood

The per-token term is the model’s surprise at the token that actually appeared:

1.0	0.00
0.5	0.69
0.1	2.30
0.01	4.61

• nat: a unit of information/entropy when you use natural log
  • relationship btw bits and nats:
• confidently right -> zero loss
• confidently wrong -> unbounded loss

so one mislabeled example can spike your loss curve -> gradient clipping is not optional

3. Perplexity

perplexity: , where is the mean per-token NLL

• a perplexity of 8 means the model is as uncertain as choosing uniformly among 8 tokens

why

• assume LLM produces uniform distribution over tokens (true token among the )
• -> it assigns to the true token
• -> loss
• ->

PPL = the model’s effective branching factor (how many reasonable next tokens on average)

• : perfect. (vocab size): no better than random
• it’s a more intuitive unit than nats

4. Cross-entropy: same number, different story

Cross-entropy (CE): the cost of encoding samples from a true distribution using a code built for a different distribution .

• code: assigns each token a codeword of length (Shannon-optimal for ).
• cost: average codeword length. bits/nats per token.

• = entropy (Shannon’s notation). Two arguments -> cross-entropy, the two-distribution version.

• weight each codeword length by how often the token truly appears, .

• minimized only when . Any mismatch costs extra -> that extra is the KL divergence.

• In language modeling the “true distribution” at each position is empirical.

• We observed exactly one token, so is one-hot.

• Every term multiplies by zero except one:

• Cross-entropy with a one-hot target is negative log-likelihood. Not analogous —> numerically identical.
• PyTorch calls it cross_entropy. Statisticians call it NLL. Both compute the same scalar.
• Why keep both names? Cross-entropy generalizes to soft targets: label smoothing, distillation. There isn’t one-hot, and the sum no longer collapses.

5. KL divergence: what we’re actually minimizing

KL divergence : how many extra nats you pay for coding with instead of the true .

• always. Zero only when .
• not symmetric: .

Split the log to see where the decomposition comes from:

• -> cross-entropy
• -> negative entropy

So . Rearrange:

• The first term depends only on the data distribution.
• It doesn’t depend on . It has zero gradient. It’s a constant floor.

Minimizing cross-entropy is exactly minimizing — the KL divergence from the data distribution to the model.

This reframing explains several things at once.

Why loss never reaches zero.

• because language is genuinely stochastic.
  • After “the capital of China is”, the next token is nearly determined.
  • After “I think that”, it isn’t.
• That inevitable entropy is your loss floor.
• A run converging to 1.9 nats/token isn’t failing. It’s approaching the entropy of the text.

Why the model plays it safe.

• We minimize forward KL, . The expectation is taken over .
• Wherever the data has mass, the model is punished for assigning near-zero probability — infinitely, in the limit.
• Forward KL is mass-covering. It forces to spread out and cover everything the data does.
• It does this even at the cost of mass on regions the data never visits.
• This is the mathematical origin of the vague, wishy-washy tone of base models. They are structurally pushed to spread their bets rather than commit.

Contrast with reverse KL.

• takes its expectation over the model.
• It’s mode-seeking. can collapse onto one high-probability region and ignore the rest.
• You can’t optimize it directly in pretraining. You’d need at arbitrary model samples, and is just a dataset.
• But this asymmetry is what RL-based post-training exploits. It’s part of why RLHF produces sharper, more opinionated models than more SFT does.

Where KL shows up as an explicit term.

• RLHF/DPO: a KL penalty against a frozen reference policy appears directly, . It keeps the policy close to the reference, so chasing reward doesn’t drift the model into gibberish.
• Distillation: match a teacher’s full soft distribution instead of a one-hot label. The sum in no longer collapses. You’re minimizing KL to a genuinely dense target.
• Pretraining and SFT are the special case. The target is one-hot, and the KL interpretation stays implicit.

6. Pretraining loss

With that groundwork, the pretraining objective is simple:

• Every position predicts its successor.
• A 4096-token sequence yields 4095 supervised prediction problems.
• The crucial engineering fact: all of them are computed in a single forward pass.
• Teacher forcing feeds the ground-truth prefix, not the model’s own samples.
• The causal mask ensures position can’t see position or beyond.
• No sequential unrolling at training time.

This density is the whole reason next-token prediction scales.

• Every token is a label.
• A batch of 4M tokens is 4M supervised examples.
• Zero human annotation.

7. SFT loss

SFT changes the data, not the math.

• Training examples are now (prompt, response) pairs.
• We want the model to learn .
• Not to get better at generating prompts.

So we mask:

is the set of response positions. Same per-token term, same everything, while a narrower index set

	1	2	3	4	5	6
token	<user>	1+1?	<asst>	It's	2	<eos>
segment	prompt	prompt	prompt	response	response	response
mask	0	0	0	1	1	1

• Positions 1–3 still participate fully in the forward pass. They’re context, visible through attention.
• They just contribute no gradient.
• Implemented by setting those labels to the ignore index, or by multiplying per-token losses by a float mask.

7.1 The normalization trap

How you average across a batch is where SFT implementations actually differ. Not a rounding error.

Scheme	Formula	Effect
Token-level		Every token weighted equally; long responses dominate
Sequence-level		Every example weighted equally; short responses amplified

• Token-level is the more common default. It generally gives smoother gradients.
• But it has a sharp edge under gradient accumulation and data parallelism.
• Normalize by the global token count across all micro-batches and all ranks. Not per-micro-batch.

The bug, concretely:

• Micro-batch A has 100 response tokens. B has 900.
• Divide each by its own local count before summing, and A’s tokens get weighted 9× more than B’s.
• Your effective loss becomes a function of how examples shuffled into micro-batches.

Why it’s easy to miss:

• This shipped in widely-used training libraries. Hugging Face Trainer carried a version of it for a long time.
• Training doesn’t crash. It just converges somewhat worse.
• You have no baseline to notice against.

The fix: accumulate the unnormalized sum and the token count separately. All-reduce both. Divide once.

7.2 Sequence packing

SFT pipelines concatenate multiple examples into fixed-length sequences to avoid wasting compute on padding. Two things must hold:

1. Attention must not cross example boundaries.
   • Use a block-diagonal mask, or FlashAttention’s variable-length API.
2. The last token of each example must not predict the first token of the next.
   • The shift-by-one that produces labels doesn’t know about your boundaries.
   • Mask those junction positions explicitly.

Getting packing subtly wrong is a common source of “the loss looks fine but the model is worse” outcomes.

8. Summary

	Pretraining	SFT
Objective
Supervised positions	all	response only
Data	raw corpus	(prompt, response) pairs
Target distribution	one-hot	one-hot
Implicit divergence	forward KL to data	forward KL to demonstrations
Scale	trillions of tokens	– examples

One expression, three names:

• negative log-likelihood if you came from statistics
• cross-entropy if you came from information theory
• forward KL to the empirical distribution if you want to know what it’s doing to your model

The one-hot target is what makes all three the same. Introduce label smoothing, distillation, or an RLHF penalty term, and they pull apart. That’s when the distinctions start to matter.