Loss Functions › Language Modeling Loss

Language-model loss functions must handle an extremely large output space. Predicting the next token from a vocabulary of 50,000 candidates is a 50,000-way classification problem applied at every position of every sequence in the training corpus. This creates both theoretical and computational challenges that shape how language models are trained.

The cross-entropy loss for language modeling is:

$$\mathcal{L}_{\text{LM}} = -\frac{1}{T} \sum_{t=1}^{T} \log p_\theta(w_t | w_1, w_2, \ldots, w_{t-1})$$

where:

• $T$: the number of tokens in the sequence
• $w_t$: the token at position $t$
• $p_\theta(w_t | w_1, \ldots, w_{t-1})$: the model's predicted probability for token $w_t$ given all preceding tokens

This is still negative log-likelihood, but applied at each time step and averaged over the entire sequence. Training a language model is the act of minimizing this sum across an enormous corpus: the model learns to predict the next word in every context it encounters.

The most common evaluation metric for language models is perplexity, which is the exponentiation of the average NLL:

$$\text{PPL} = \exp\left(\mathcal{L}_{\text{LM}}\right) = \exp\left(-\frac{1}{T} \sum_{t=1}^{T} \log p(w_t | w_{1:t-1})\right)$$

Perplexity measures how surprised the model is by the text on average. A perplexity of 10 means the model is, on average, as uncertain as if it were choosing uniformly among 10 equally probable next tokens. A perplexity of 100 means the model is effectively choosing among 100 possibilities. Lower perplexity means better predictions.

The nice property of perplexity is that it is interpretable in a way that raw NLL is not. When GPT-2 large achieves a perplexity of 35.76 on Penn Treebank, and GPT-3 achieves 20.50, the difference is immediately meaningful: GPT-3 is predicting each token from a smaller effective vocabulary of candidates.

Computing a full softmax over 50,000+ tokens at every position is computationally expensive. Several approximations exist:

Sampled softmax computes the denominator using a sample of $K$ randomly chosen negative tokens instead of the full vocabulary. At training time, this reduces the cost from $O(V)$ to $O(K)$ per prediction, where $K \ll V$. The approximation introduces bias but is effective in practice.

Noise contrastive estimation (NCE) reframes language modeling as a binary classification problem: is this word the true next word, or is it a random noise word? NCE avoids computing the full normalizing constant by treating the partition function as a parameter to be estimated.

Hierarchical softmax organizes the vocabulary into a binary tree, where each word is a leaf. The probability of a word is computed as the product of probabilities along the path from the root to the leaf. This reduces prediction cost from $O(V)$ to $O(\log V)$ per example.

Modern large language models generally use full softmax with efficient hardware implementations rather than approximations, because the cost of the softmax can be made manageable through batching and mixed-precision arithmetic. But these approximations remain important for understanding the history of language model training and for resource-constrained settings.