LLM Theory Foundations

• Why next-token prediction is maximum likelihood over a sequence.
• Why cross entropy is the standard training loss for language models.
• How KL divergence relates to fitting a model distribution to data.
• Why logits are trained before softmax probabilities are sampled.
• How optimization choices show up as instability, slow learning, or overfitting.

A causal LLM estimates the probability of a token sequence by multiplying conditional probabilities. In practice we add log probabilities so the objective is numerically stable and decomposes token by token.

p(x_1, ..., x_T) = product_t p(x_t | x_1, ..., x_{t-1})
log p(sequence) = sum_t log p(x_t | prefix)

The model reads tokens up to position t and is trained to predict token t + 1. A hand-written PyTorch loop usually shifts logits and labels explicitly before flattening into cross_entropy. In many Hugging Face causal-LM model classes, passing labels=input_ids delegates that shift to the model, while ignored labels such as -100 remove padding or prompt tokens from the loss.

import torch.nn.functional as F

def causal_lm_loss(logits, input_ids, ignore_index=-100):
    # logits: [batch, time, vocab], input_ids: [batch, time]
    # Predict token t+1 from the hidden state at token t.
    shift_logits = logits[:, :-1, :].contiguous()
    shift_labels = input_ids[:, 1:].contiguous()

    return F.cross_entropy(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
        ignore_index=ignore_index,
    )

1. Align once: either shift in your custom loss or let a library model do it, but do not shift both in the dataset and inside the model.
2. Mask deliberately: set labels to -100 for padding, system prompts, or examples that should condition the model without contributing loss.
3. Flatten correctly: cross_entropy expects class logits over the vocabulary and integer class-index labels, so preserve the vocab dimension as the last dimension before reshaping.
4. Test with a toy batch: make a two-sequence batch with known labels, one ignored token, and verify the averaged loss only counts the intended positions.

Sources: Hugging Face causal language modeling guide, PyTorch CrossEntropyLoss docs.

For one token with class index y, cross entropy is the negative log of the probability assigned to y. With one-hot labels, the gradient with respect to logits is the simple and important expression softmax(logits) - one_hot(y).

import torch
import torch.nn.functional as F

logits = torch.randn(2, 4, 32000)  # batch, time, vocab
targets = torch.randint(0, 32000, (2, 4))
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))