Loss Functions for Language Models

Cross-Entropy Loss The standard loss for language modeling: $$ L = -\frac{1}{N}\sum_{i=1}^{N} \log P(y_i | x_{

Where $y_i$ is the correct token and $P$ is the model's predicted probability.

How It Works

Per-Token Loss For each position, compare predicted probability distribution to ground truth:

import torch.nn.functional as F

# logits: (batch, seq_len, vocab_size)
# labels: (batch, seq_len)

loss = F.cross_entropy(
    logits.view(-1, vocab_size),
    labels.view(-1),
    ignore_index=-100  # Ignore padding tokens
)

Intuition

• If model assigns high probability to correct token → low loss
• If model assigns low probability to correct token → high loss

Special Loss Handling

Label Smoothing Prevent overconfidence by softening targets:

loss = F.cross_entropy(logits, labels, label_smoothing=0.1)

• Replaces one-hot targets with (1-ε, ε/(V-1), ...)
• Improves generalization

Focal Loss (for imbalanced data) $$ FL(p) = -(1-p)^\gamma \log(p) $$

• Down-weights easy examples
• Used when some tokens are much more common

Loss Masking

Ignore Padding

# Don't compute loss on padding tokens
loss_fct = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

Instruction Masking For instruction-tuned models, often only compute loss on the response:

[System prompt | User query | Response]
[  Masked      |   Masked   | Loss computed ]

Auxiliary Losses

Load Balancing (MoE) Encourage even expert utilization: $$ L_{aux} = \alpha \sum_{i=1}^{n} f_i \cdot P_i $$

Contrastive Loss (RLHF) Prefer chosen over rejected responses: $$ L = -\log \sigma(r_{chosen} - r_{rejected}) $$

Monitoring Loss