Language modeling review

Masking

3 common types of mask:

• Causal mask: current tokens cannot attend to future tokens. Built into decoder attention.
• Attention/padding mask: direct attention to ignore padding tokens (variable-length sequences batched together). Shape (B, L)
• Completion/loss mask: 1 on completion tokens; the loss is multipled by it before summing. Prompt tokens condition the prediction but contribute no loss.

Decoding

Generation samples/searches over . Many methods: greedy (argmax), beam search, temperature sampling (, high temp smooths out peaky probabilities), truncated sampling (top-k tokens, top-p i.e., smallest set with cumulative prob > p), lookahead search (MCTS & friends etc., score simulated future continuations often with a reward/value model before committing - decode time cousin of inference time reasoning).

Pretraining -> mid-training -> SFT

The boundary at the end of pretraining is fuzzy and very important to post-training:

• Pretraining: next-token cross-entropy on massive raw web/code data; every token contributes to the loss.
• Mid-training: the annealing / high-quality-data phase at the very end of pretraining, before any instruction data. Quality ramps up here. Could be considered part of post-training, but often isn’t discussed. Continued pretraining goes here.
• SFT (supervised fine-tuning): still cross-entropy, but now on. instruction–response pairs with a completion mask: only the response tokens contribute to the loss, the prompt merely conditions. The model learns to respond, not just continue text.

Modern post-training pipelines start from a mid-trained checkpoint.

KL divergence

• Asymmetric:
• Non-negative:
• The central objective of alignment: the RLHF KL penalty keeps the policy close to its SFT checkpoint, and the DPO loss is derived from a KL-regularized reward-maximization objective.

People try to approximate this in a lot of RLHF loss functions.

In next-token prediction, minimizing cross-entropy is equivalent to minimizing the KL-divergence to the truth fixed by the data. In LM training, the label is one-hot so and the cross-entropy = KL = NLL. In RLHF the KL term is added on purpose to keep the policy near its reference.

Estimation of KL divergence using Monte Carlo sampling

Core idea is that Monte-Carlo sampling is a way to estimate a distribution by sampling from it many times and taking the sample mean of those random trials.

For an autoregressive LM, the probability of a completion is the product of its conditional next-token probabilities:

Therefore, the log probability ratio for a sampled sequence becomes a sum over its tokens:

Thus, KL contributions are summed over tokens within each sequence, while the Monte Carlo estimator is averaged over sampled sequences:

Because the token contributions are summed, longer completions can accumulate a larger KL penalty. Some implementations divide by completion length when a per-token rather than per-sequence penalty is desired.

• The integrand is just a log of probability ratio on each rollout you already generated
• The naive average () is unbiased but high variance and can go negative on a finite sample
• RLHF uses the low-variance, always- estimator , with @ApproximatingKLDivergence

Reverse KL vs forward KL

Notice how the KL formula above measures the divergence from the current policy to the reference. This is called reverse KL, commonly used in RLHF.

Reverse KL divergence is used instead of standard (forward) KL to achieve mode-seeking behavior. While forward KL forces models to cover all possible outcomes (averaging, mean-seeking), reverse KL penalizes the model for outputting wrong answers, forcing it to focus intensely on the most precise, highest-quality modes (mode-seeking).

TL;DR: forward KL smudges the model’s coverage, reverse KL sharpens peaky modes.

Sigmoid & pairwise likelihood

A preference (“response is better than ”) is a binary outcome, so a score difference is squeezed through a sigmoid into a probability:

Three things to recognize:

• The sigmoid converts a score difference into a binary probability.
• Only relative reward differences matter: shifting both rewards by a constant changes nothing.
• is binary negative log-likelihood.

The Bradley-Terry model uses this formulation to train a reward model. DPO reuses the same shape.

The goal of RL and why for LLMs

The goal of RL is to maximize the expected reward under the policy.

A human preference (or a verifier) is a single scalar at the end of a response:

• Sparse: one label per whole response, not per token
• Delayed: the quality of a token depends on the whole completion
• No token-level target, and ordinary gradients can’t pass back through the discrete sampled token

So there is no cross-entropy target for “preferred.” Policy gradients optimize the expected evaluator score anyway, first by learning that signal as a reward model, then by optimizing it. That is the whole motivation for RLHF.