Part of Language AI Handbook
Covers neural network loss functions covering MSE, cross-entropy, label smoothing, focal loss, KL divergence, and contrastive losses for NLP tasks.
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Loss Functions
Every neural network needs a way to measure how wrong its predictions are. Without that signal, training is impossible. The loss function is that measurement, a scalar value that tells the optimizer whether the model is getting closer to correct predictions or wandering further away.
You have already seen how linear classifiers produce scores, how activation functions introduce non-linearity, and how MLPs stack layers to represent complex functions. But none of that training happens without a loss function. The loss function is the bridge between the model's output and the training signal that flows backward through backpropagation. Choosing the right loss function is not a minor implementation detail. It shapes what the model optimizes, what errors it penalizes most, and whether it even converges.
This chapter covers the most important loss functions for neural networks and NLP: regression losses, classification losses, and more specialized objectives for tasks like contrastive learning and class imbalance. For each one, you will understand why it was designed that way, what it measures, and when to choose it over alternatives.
What Makes a Good Loss Function
A loss function takes the model's predicted output and the true target , and returns a non-negative scalar indicating how far apart they are. Before diving into specific losses, it is worth spending time on the properties that make one loss function better suited to a problem than another. Three properties matter most:
- Differentiability: The optimizer needs to compute gradients of the loss with respect to every model parameter. Non-differentiable losses can stall training. In practice, some losses (like MAE) are non-differentiable at exactly one point, and this is handled using subgradients, but losses with large non-differentiable regions are problematic.
- Smoothness: Flat regions (zero gradient) prevent learning. Sharp discontinuities cause training instability. A well-designed loss function provides informative gradient signal across the full range of predictions, not just near the decision boundary.
- Alignment with task: The loss should penalize the errors that matter most for the task. A loss that treats all errors equally may produce models that fail on rare but important cases.
There is a fourth property that is less often stated explicitly but just as important: computational tractability. For language models with vocabularies of 50,000 or more tokens, computing a full softmax distribution over the entire vocabulary at every position of a long sequence is expensive. This computational pressure has driven the development of approximations like sampled softmax, noise contrastive estimation, and hierarchical softmax, all of which modify the loss to make it computable within a reasonable budget.
There is no universal best loss function. Regression problems, binary classification, multiclass classification, and sequence modeling each have losses that are better matched to their structure. The key insight is that every loss encodes assumptions about the shape of the problem: what errors are costly, how uncertainty should be handled, and what the output distribution looks like. When those assumptions match reality, training proceeds efficiently. When they do not, you get models that optimize the wrong objective.
From Likelihood to Loss
Most loss functions are not invented from scratch. They are derived from the same statistical principle: maximum likelihood estimation (MLE). The idea is simple. Given a model with parameters and a dataset of examples, we want to find the parameters that make the observed data as probable as possible under the model. Concretely, if the model defines a probability distribution over outputs given inputs , MLE finds:
Taking the logarithm converts the product to a sum and does not change the optimal solution because logarithm is a monotone function:
Negating this gives a minimization objective, which is what gradient descent needs:
This is negative log-likelihood (NLL). Cross-entropy for classification and mean squared error for Gaussian regression are both special cases of NLL under different distributional assumptions. Cross-entropy arises when you model the output as a categorical distribution. MSE arises when you model the output as a Gaussian distribution with fixed variance. Understanding this connection means you understand why these specific formulas are used, not just how to apply them.
Regression Losses
Regression losses measure the discrepancy between a continuous predicted value and a continuous target value. The two most common choices are mean squared error and mean absolute error, and their differences are significant.
Mean Squared Error
Mean squared error (MSE) is the most widely used regression loss. Given a batch of training examples with predictions and targets , MSE computes the average squared difference:
where:
- : the number of examples in the batch
- : the model's prediction for example
- : the true target value for example
- : the squared error, always non-negative
The squaring serves two purposes. First, it makes every error positive, so over-predictions and under-predictions both contribute to the loss. Second, it penalizes large errors disproportionately. An error of 2 contributes 4 to the loss, while an error of 10 contributes 100. This makes MSE sensitive to outliers.
The gradient of MSE with respect to the prediction is:
where:
- : the gradient, pointing in the direction of steepest increase of the loss
- : the residual, positive when the prediction is too high, negative when too low
The gradient is proportional to the error magnitude. Large prediction errors produce large gradients, which means the optimizer takes bigger steps to correct big mistakes. This property is exactly what makes MSE both useful and risky: useful because big errors dominate training and get fixed quickly, risky because a single outlier can dominate the loss and pull the model away from fitting the majority of the data.
There is a probabilistic interpretation of MSE that ties it back to maximum likelihood. If you assume that each target value is sampled from a Gaussian distribution centered at the model's prediction, then:
The negative log-likelihood under this Gaussian assumption is proportional to , which is exactly MSE. So MSE is the correct loss when you believe the noise in your targets is Gaussian, with equal variance everywhere. When that assumption is violated (heavy-tailed noise, heteroskedastic variance, or outliers), MSE is the wrong loss.
Mean Absolute Error
Mean absolute error (MAE) uses the absolute value instead of the square:
where:
- : the absolute difference between prediction and target, always non-negative
MAE treats all errors with equal weight per unit of magnitude. A prediction off by 2 contributes exactly twice as much as a prediction off by 1, regardless of scale. This makes MAE much more robust to outliers than MSE: an outlier with error 10 contributes 10 to the loss instead of 100.
The gradient of MAE is:
where:
- : returns if the prediction is too high, if too low, and if exact
The gradient is constant in magnitude (either or per example), which means the optimizer always takes the same-sized step regardless of error size. This is the key tradeoff: MAE is more less brittle to outliers, but converges more slowly because it does not amplify gradients for large errors. It is also technically non-differentiable at zero, though this is handled in practice by taking the subgradient (which returns 0 at exactly zero error).
The probabilistic counterpart to MAE is a Laplace distribution for the noise. If the noise follows a Laplace distribution rather than a Gaussian, negative log-likelihood gives MAE. The Laplace distribution has heavier tails than the Gaussian, meaning it assigns more probability to large errors. Choosing MAE over MSE is equivalent to assuming the data has Laplace noise rather than Gaussian noise: a decision that matters when your dataset contains outliers.
Huber Loss
Huber loss (also called smooth L1 loss) is a hybrid that behaves like MSE for small errors and MAE for large errors:
where:
- : a threshold parameter that controls the transition point between quadratic and linear behavior
- : the quadratic (MSE-like) component for small errors
- : the linear (MAE-like) component for large errors
The term ensures continuity of the loss and its gradient at the transition point . Huber loss is fully differentiable everywhere (unlike MAE), and it clips the gradient magnitude for outliers (unlike MSE). The hyperparameter controls where the transition happens. A common default is .
The gradient of Huber loss is:
For small errors, the gradient scales linearly with error (like MSE), giving strong correction signals proportional to mistake size. For large errors, the gradient is capped at , preventing outliers from dominating parameter updates. Huber loss is a principled way to get the best of both worlds: fast convergence near the optimum (quadratic behavior) with robustness to outliers (linear behavior).
In NLP, regression losses appear in tasks like predicting sentiment scores, estimating semantic similarity as a real-valued score, and quality estimation for machine translation. For any of these, Huber loss is a defensible default unless you have strong prior knowledge about the noise distribution.
Classification Losses
Classification tasks require predicting one of discrete classes. The structure of the problem is fundamentally different from regression: instead of predicting a continuous value, the model produces scores (logits) for each class, which are then converted to probabilities.
Negative Log-Likelihood
The foundation of all classification losses is maximum likelihood estimation. We want to find model parameters that maximize the probability of observing the training data. If the model outputs a probability distribution over classes, the likelihood of the true class is . We want this to be high.
Maximizing likelihood is equivalent to minimizing negative log-likelihood:
where:
- : the model's predicted probability for the true class
- : the natural logarithm
The log transformation serves two purposes. First, it converts a product of probabilities (over a dataset) into a sum, which is numerically more stable. Second, the log is a monotone transformation, so maximizing is equivalent to maximizing . Because log of a probability between 0 and 1 is negative, we negate it to get a positive loss.
Negative log-likelihood has an intuitive interpretation: it measures how surprised the model is to see the true label. If the model assigns probability 0.99 to the correct class, the loss is (barely surprised). If the model assigns probability 0.01 to the correct class, the loss is (very surprised). The loss grows without bound as the predicted probability approaches zero, which strongly penalizes confident wrong predictions.
This unbounded growth is intentional. A loss that caps at some maximum value would allow the model to be confidently wrong without incurring a proportional penalty. The logarithmic penalty means there is no "good enough" point where the model can stop improving its confidence in correct predictions, which drives the model toward sharper, more accurate probability estimates.
Cross-Entropy Loss
Cross-entropy loss is the most widely used loss for classification. It combines softmax (to convert logits to probabilities) with negative log-likelihood into a numerically stable combined operation.
Given the model's raw scores (logits) for classes, the softmax function converts them to a probability distribution:
where:
- : the raw score (logit) for class
- : the exponential of the logit, which ensures positivity
- : the normalizing constant, which ensures probabilities sum to 1
- : the predicted probability for class
Cross-entropy loss then applies negative log-likelihood to the true class:
where:
- : a one-hot indicator, equal to 1 for the true class and 0 elsewhere
- : the predicted probability for class
Because is one-hot, this sum collapses to just where is the index of the true class. The sum notation is more general, covering soft labels (non-binary values), but the one-hot version is standard for supervised classification.
Cross-entropy has roots in information theory. It measures the expected number of bits needed to encode events from distribution using a code optimized for distribution . When matches exactly, cross-entropy equals the entropy of , its minimum possible value. During training, minimizing cross-entropy pushes the model's predictions toward the true distribution of the data.
Numerical stability is a critical implementation detail. Computing softmax followed by log involves exponentials that can overflow (for large logits) or underflow (for very negative logits). The standard trick is the log-sum-exp stabilization:
where is subtracted before exponentiation to prevent overflow. PyTorch's F.cross_entropy() implements this automatically when you pass raw logits. Always pass logits directly to F.cross_entropy() rather than computing softmax separately and then applying NLL, because the combined implementation is both more numerically stable and more computationally efficient.
The gradient of cross-entropy with respect to the logits has a clean form that reveals why the loss is so well-suited to gradient descent. For the logit :
This is simply the difference between the predicted probability and the target. For the true class (), the gradient is , which is negative (pushing the logit up) unless . For incorrect classes (), the gradient is , which is positive (pushing the logit down). The optimizer directly reduces the predicted probability of incorrect classes and increases the probability of the correct class. This clean interpretation is one reason cross-entropy is so widely used: the gradients have an immediate intuitive meaning.
Binary Cross-Entropy
Binary classification is a special case where (positive or negative). Rather than using a 2-class softmax, it is standard to use a single output neuron with a sigmoid activation, which directly outputs the probability of the positive class.
Given a model output where is the sigmoid function and is the logit, binary cross-entropy (BCE) is:
where:
- : the true binary label
- : the model's predicted probability of the positive class
- : the loss contribution when (penalizes predicting low probability for a positive example)
- : the loss contribution when (penalizes predicting high probability for a negative example)
The two terms are mutually exclusive: when , the second term vanishes, and when , the first term vanishes. This loss strongly penalizes confident wrong predictions: if the model outputs for a positive example (), the loss is .
The gradient of BCE with respect to the logit (before sigmoid) has a particularly clean form:
where:
- : the sigmoid of the logit
- : the true label
This is the simplest possible gradient: just the prediction error. The sigmoid's derivative cancels with the derivative of the log, producing a numerically stable and interpretable gradient signal.
Binary cross-entropy is also the right loss for multi-label classification, where each example can belong to multiple classes simultaneously. In multi-label settings, you have output neurons, each representing an independent binary prediction. BCE is applied independently to each output, and the total loss is the sum over all binary losses. This is different from multiclass (categorical) classification, where exactly one class is correct and the outputs must sum to 1.
Categorical Cross-Entropy vs. Binary Cross-Entropy
The choice between categorical and binary cross-entropy depends on the problem structure:
- Binary cross-entropy applies when each example belongs to one of two classes or when predicting multiple independent binary attributes simultaneously (multi-label classification).
- Categorical cross-entropy applies when each example belongs to exactly one of classes (multi-class classification).
In multi-label classification, each example can have multiple positive labels (e.g., a news article tagged as both "politics" and "international"). Binary cross-entropy applies independently to each label, computing separate losses and summing them. The key distinction is that the outputs do not need to sum to 1, and each output is trained independently.
Weighted Cross-Entropy
Class imbalance is one of the most pervasive problems in NLP. Named entity recognition datasets might have 95% non-entity tokens. Sentiment datasets might have 80% neutral examples. Toxicity classifiers deal with much more benign text than harmful text. In these settings, an unweighted cross-entropy loss will cause the model to learn to predict the majority class almost all the time, achieving low loss while being useless for the rare classes of interest.
Weighted cross-entropy assigns a higher loss multiplier to underrepresented classes:
where:
- : the weight for class , typically set to the inverse frequency of class in the training set
- : the one-hot target (1 for true class, 0 otherwise)
- : the predicted probability for class
A standard heuristic for setting weights is:
where is the total number of training examples and is the number of examples in class . This ensures that each class contributes equally to the total loss in expectation, regardless of its frequency. In PyTorch, you can pass a weight tensor directly to F.cross_entropy().
Label Smoothing
Standard cross-entropy trains the model to push probabilities toward 0 or 1 for the correct and incorrect classes. This works well, but it can cause overconfidence: the model may assign extremely high probability to one class even for ambiguous examples, becoming poorly calibrated.
Label smoothing addresses this by replacing the hard one-hot labels with soft labels that reserve a small probability mass for all classes:
where:
- : the smoothing parameter, typically 0.1
- : the total number of classes
- : the index of the true class
- : the smoothed probability for the correct class (slightly less than 1)
- : the smoothed probability for each incorrect class (slightly above 0)
The cross-entropy loss with smoothed labels becomes:
where:
- : the predicted probability for the true class
- : the sum of log probabilities across all classes
Label smoothing has two effects. It discourages the model from assigning zero probability to any class, which improves calibration. It also reduces the gap in logits between the correct class and incorrect classes, making the model's representations more spread out. The original Transformer paper used , and label smoothing has since become standard in many NLP models.
To see why label smoothing improves calibration, consider what happens during training without it. The optimal model under standard cross-entropy would assign probability 1 to the true class and probability 0 to all other classes, for every training example. This means the logits for the true class would need to go to infinity, which is impossible in finite training but the model tries to approach. This drives the logit gap (difference between the top logit and the second-highest) to be very large, leading to overconfident predictions. Label smoothing provides a finite target: the optimal model assigns to the true class, which requires a finite logit gap. Training naturally converges to this finite optimal, producing better-calibrated confidence scores.
A model is well-calibrated if its predicted probabilities match the actual frequency of outcomes. A perfectly calibrated model that says "I'm 70% confident" should be correct 70% of the time. Overconfident models (assigning 99% probability to answers they get wrong) have poor calibration. Label smoothing tends to improve calibration by preventing extreme probability assignments.
Focal Loss for Class Imbalance
Many real-world classification problems have severe class imbalance: fraud detection, rare disease diagnosis, named entity recognition (where most tokens are not entities). Standard cross-entropy performs poorly in these settings because easy negative examples (confidently classified background tokens) dominate the loss, overwhelming the contribution of the rare positive class.
The problem is more subtle than just having fewer positive examples. Even after you account for class frequency, the model tends to become good at classifying the easy majority class early in training. Once the model can classify the majority class with high confidence, those examples contribute almost no loss (because ), but there are many of them. The rare, hard-to-classify minority examples contribute more loss per example, but there are so few of them that they still get drowned out.
Focal loss (introduced by Lin et al., 2017 in the context of object detection) down-weights easy examples dynamically, focusing training on hard examples:
where:
- : the model's predicted probability for the true class
- : the focusing parameter (typically 2.0)
- : the modulating factor, close to 1 for low confidence and close to 0 for high confidence
When the model is already confident about an example (), the modulating factor reduces the loss contribution by 100x. When the model is uncertain (), the factor reduces the loss by only 4x. Hard examples (where the model is wrong or uncertain) contribute more to the loss than easy examples.
Focal loss is often combined with a class-weighting factor :
where:
- : the per-class weight, which can be set to the inverse frequency of each class to further up-weight rare classes
Setting recovers standard cross-entropy. A common default is and for the positive class.
Focal loss adapts dynamically to each example based on the model's current confidence, rather than applying a fixed multiplier like class weighting. This is a powerful property: early in training when the model is uncertain about everything, focal loss behaves similarly to standard cross-entropy. As training progresses and the model masters easy examples, focal loss automatically down-weights those contributions and concentrates the learning signal on the difficult cases that remain.
KL Divergence
Kullback-Leibler (KL) divergence measures how much one probability distribution differs from a reference distribution . In neural network training, it appears when we want the model's output distribution to match a target distribution, not just predict a single class.
where:
- : the probability of event under the reference distribution
- : the probability of event under the approximate distribution
- : the log ratio, measuring how much deviates from for event
KL divergence is non-negative (by Jensen's inequality) and equals zero if and only if . Importantly, it is asymmetric: . When is the true distribution, penalizes cases where assigns low probability to events that considers likely. This is the relevant direction for model training.
The relationship between cross-entropy and KL divergence is:
where:
- : the entropy of the true distribution , a constant with respect to model parameters
- : the KL divergence between and
Since is constant during training, minimizing cross-entropy is equivalent to minimizing KL divergence. This connection explains why cross-entropy is the correct loss for maximum likelihood estimation: it minimizes the information gap between the model's distribution and the true distribution.
KL divergence appears explicitly in variational autoencoders (VAEs) as a regularization term, and in knowledge distillation, where it is used to train a student model to match a teacher model's output distribution. In distillation, the student is trained to minimize : the teacher's soft probability distribution provides richer training signal than one-hot labels, because it encodes the teacher's beliefs about how similar various classes are to each other.
The asymmetry of KL divergence has practical consequences. is large when assigns low probability to events that considers common (mode-seeking behavior). is large when spreads probability mass over regions where has low probability (mean-seeking behavior). For model training, you generally want mode-seeking behavior: learn to correctly handle the common cases well. This is why cross-entropy (which minimizes ) is preferred over the reverse direction.
Temperature Scaling
Temperature scaling modifies a model's logits before computing probabilities, controlling the sharpness of the predicted distribution. A temperature is applied as:
where:
- : the raw logit for class
- : the temperature parameter
- : the temperature-scaled probability for class
When , the standard softmax is recovered. When (low temperature), the distribution becomes sharper: the class with the highest logit gets even more probability mass. When (high temperature), the distribution becomes softer, approaching a uniform distribution as .
Temperature scaling is used in several contexts:
- Knowledge distillation: A high temperature (e.g., or ) is applied to the teacher's logits during training. This produces soft targets that reveal the relative similarity between classes. This provides richer supervision than hard one-hot labels.
- Calibration: After training, temperature can be tuned on a validation set to make the model better calibrated, without changing its predictions (only confidence scores change).
- Sampling from language models: During text generation, temperature controls the randomness of token sampling. Low temperature makes the model more deterministic, high temperature makes it more creative.
Post-training temperature scaling is notable for being a uniquely simple and effective calibration technique. Unlike platt scaling or isotonic regression, temperature scaling has only one parameter to tune, avoids overfitting on the calibration set, and does not change the model's accuracy on any task (only its confidence scores). The procedure is: freeze the model, search over values of to minimize the NLL on a held-out validation set, and apply the optimal at inference time.
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:
where:
- : the number of tokens in the sequence
- : the token at position
- : the model's predicted probability for token 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.
Perplexity
The most common evaluation metric for language models is perplexity, which is the exponentiation of the average NLL:
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.
Efficient Softmax for Large Vocabularies
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 randomly chosen negative tokens instead of the full vocabulary. At training time, this reduces the cost from to per prediction, where . 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 to 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.
Contrastive Losses
Contrastive losses are designed for representation learning tasks, where the goal is not to classify a fixed set of classes but to learn an embedding space where similar examples are close together and dissimilar examples are far apart.
The motivation comes from a fundamental limitation of cross-entropy: it requires a fixed set of predefined classes. But many real-world tasks do not fit that mold. In face verification, you want to check whether two photos show the same person, not classify the person into one of millions of identities. In semantic search, you want to retrieve documents similar to a query, not classify the query into a topic. In these settings, you need a loss that directly shapes the geometry of the embedding space, not one that categorizes examples into slots.
Contrastive learning has had enormous impact in NLP. Systems like SimCSE, DPR (Dense Passage Retrieval), and Sentence-BERT all use variants of contrastive objectives to learn sentence and passage representations that capture semantic similarity. The key insight is that you can construct training pairs (or triplets) from naturally occurring structure in data: paraphrases, question-answer pairs, adjacent sentences, or even a sentence and its dropout-augmented version.
Contrastive Loss
The original contrastive loss (Hadsell et al., 2006) operates on pairs of examples with a binary label indicating whether the pair is similar (1) or dissimilar (0):
where:
- : the Euclidean distance between the embeddings of examples and
- : 1 if the pair is similar, 0 if dissimilar
- : the margin hyperparameter, the minimum desired distance for dissimilar pairs
- : the loss for similar pairs, minimizing their distance
- : the loss for dissimilar pairs, penalizing only if they are closer than margin
For similar pairs, the loss pushes embeddings together (minimizes ). For dissimilar pairs, the loss pushes embeddings apart only when they are closer than the margin . The margin prevents the model from uselessly separating pairs that are already far apart.
Triplet Loss
Triplet loss (Schroff et al., 2015, introduced for face recognition) uses triplets of examples: an anchor , a positive example (same class or similar), and a negative example (different class or dissimilar):
where:
- : the distance between the anchor and positive embeddings
- : the distance between the anchor and negative embeddings
- : the margin hyperparameter, the minimum desired gap between positive and negative distances
- : the hinge function, applying zero loss when the constraint is already satisfied
Triplet loss ensures that the anchor is closer to its positive than to its negative by at least a margin . When the condition is already satisfied, the loss is zero and no gradient flows. This means triplet mining is critical in practice: randomly sampled triplets are mostly "easy" (already satisfying the constraint), so training needs to focus on hard triplets where the negative is closer to the anchor than the positive.
Hard negative mining is the process of selecting, for each anchor, the negative example that is currently closest to the anchor in embedding space. These hard negatives provide the most informative gradient signal because they represent the cases the model is currently getting wrong. In NLP settings, hard negatives might be selected from examples that are semantically dissimilar to the query but lexically similar, which are the cases most likely to fool an embedding model.
Triplet loss is widely used for sentence embeddings, where semantically similar sentences should be close in embedding space. In the original Sentence-BERT paper, triplet loss with hard negatives trained on sentence pairs achieved far better results than training with pairwise or pointwise losses.
InfoNCE and In-Batch Negatives
A modern extension of contrastive learning uses all other examples in the batch as negatives for each anchor, a technique that scales well and produces high-quality representations. The InfoNCE (Noise-Contrastive Estimation) loss treats the positive pair's similarity as the unnormalized score for a -way classification problem over the examples in the batch:
where:
- : the anchor embedding
- : the positive embedding
- : the embedding of the -th example in the batch (including the positive)
- : cosine similarity between two embeddings
- : a temperature hyperparameter controlling the sharpness of the distribution
InfoNCE scales naturally to large batches: with a batch of 256 examples, each anchor sees 255 in-batch negatives. This provides rich training signal without any explicit negative sampling strategy. SimCSE uses InfoNCE with dropout-augmented positives (the same sentence encoded twice with different dropout masks) and achieves state-of-the-art sentence embedding quality.
The temperature controls the strength of the InfoNCE discrimination signal. Low temperature (around 0.05 to 0.1) makes the loss more sensitive to the relative order of similarities, pushing the model to create sharp distinctions between the positive and all negatives. High temperature makes the task easier but also reduces the discrimination signal. In practice, temperature is treated as a hyperparameter to tune, or learned as a trainable parameter initialized near 0.07.
Ranking Losses
A class of losses closely related to contrastive learning is designed specifically for information retrieval and ranking problems. In ranking, the goal is not to assign correct class labels but to order items so that the most relevant items appear first.
Pairwise ranking loss (also called margin ranking loss) requires the score of the positive example to exceed the score of the negative example by at least a margin:
where and are scalar relevance scores for the positive and negative items respectively. This is structurally identical to the linear hinge loss, applied to ranking.
Listwise losses operate on the entire ranked list at once, optimizing directly for ranking metrics like NDCG or MAP. ListNet and LambdaRank are examples. These are more complex to implement but can directly optimize the ranking metric of interest rather than a surrogate.
In NLP, ranking losses appear in dense passage retrieval (DPR), where the goal is to rank relevant passages above non-relevant ones for a given query. The DPR model uses a variant of in-batch negatives similar to InfoNCE, treating each passage in the batch as a potential negative for every query in the batch.
Worked Example: Loss Surface for a 3-Class Problem
To build intuition for how different losses behave, consider a concrete scenario. A model is classifying text into three sentiment categories: negative (class 0), neutral (class 1), and positive (class 2). The model produces raw logits for each class, and we want to see how the loss changes as the model's confidence changes.
Consider a single example with true label "positive" (class 2). The model outputs logit vectors of different shapes:
- Correct and confident: logits . After softmax: approximately . Cross-entropy loss: .
- Uncertain: logits . After softmax: . Cross-entropy loss: .
- Wrong and confident: logits . After softmax: approximately . Cross-entropy loss: .
The loss spans from 0.10 to 3.0 across these three scenarios. This provides a strong gradient signal that distinguishes confident-correct from uncertain from confident-wrong. This is the key property of cross-entropy: it does not just measure whether the prediction is right or wrong, it measures the quality of the confidence estimate.
Now compare focal loss with :
- Correct and confident: .
- Uncertain: .
- Wrong and confident: .
The confident-correct loss drops from 0.10 to 0.001 (100x reduction), while the confident-wrong loss barely changes from 3.0 to 2.70. Focal loss dynamically concentrates the training signal on the hard cases while suppressing easy ones.
Loss Function Selection for NLP Tasks
Choosing the right loss function is about matching the loss to the structure of the task. Here is a practical guide for common NLP scenarios.
Text classification (spam detection, sentiment, topic): Use categorical cross-entropy for single-label, binary cross-entropy per class for multi-label. Label smoothing (0.1) is a low-cost improvement for calibration.
Named entity recognition (NER): Token-level classification. Use categorical cross-entropy with class weights or focal loss if entity classes are rare compared to non-entity tokens. Most production NER systems deal with severe label imbalance because most tokens are not entities.
Language modeling (next-token prediction): Categorical cross-entropy on the vocabulary distribution. The large vocabulary () requires efficient computation. Label smoothing is sometimes applied at the vocabulary level.
Sequence-to-sequence tasks (translation, summarization): Cross-entropy on each output token, summed or averaged across the sequence. Label smoothing is standard.
Sentence similarity and semantic search: Triplet loss or InfoNCE. SimCSE and similar approaches use a form of noise-contrastive estimation that generalizes contrastive loss to multiple in-batch negatives.
Knowledge distillation: KL divergence between student and teacher distributions, often with temperature scaling to soften teacher outputs.
Regression tasks (predicting scores, ranking): MSE for clean data, Huber loss when outliers are present.
| Task | Recommended Loss | Notes |
|---|---|---|
| Single-label classification | Categorical cross-entropy | Add label smoothing for calibration |
| Multi-label classification | Binary cross-entropy | Per-class sigmoid outputs |
| Class imbalance | Focal loss | gamma=2, alpha=inverse frequency |
| Sequence modeling | Cross-entropy | Label smoothing standard |
| Similarity learning | Triplet or InfoNCE | Hard mining critical for triplet |
| Calibration | Cross-entropy + temperature | Tune T post-training |
| Regression | MSE or Huber | Huber when outliers present |
| Information retrieval | InfoNCE or pairwise ranking | In-batch negatives scale well |
| Knowledge distillation | KL divergence | Apply temperature to teacher logits |
Code Implementation
Let us implement and compare the core loss functions in Python. We will compute each loss by hand, then verify with PyTorch's built-in implementations, and visualize how different losses behave as a function of prediction error.
Setup
We start by importing the necessary libraries and defining the losses we will compute.
import numpy as np
import torch
# Set random seed for reproducibility
np.random.seed(42)
torch.manual_seed(42)Regression Losses
Let us compute MSE, MAE, and Huber loss for a small regression example and compare their gradients.
# Predictions and targets for regression
y_pred = np.array([2.5, 0.0, 2.0, 8.0]) # model predictions
y_true = np.array([3.0, -0.5, 2.0, 7.0]) # true values
# Mean Squared Error
mse = np.mean((y_pred - y_true) ** 2)
# Mean Absolute Error
mae = np.mean(np.abs(y_pred - y_true))
# Huber Loss (delta=1.0)
delta = 1.0
residuals = np.abs(y_pred - y_true)
huber = np.where(
residuals <= delta, 0.5 * residuals**2, delta * residuals - 0.5 * delta**2
)
huber_loss = np.mean(huber)
# MSE gradients
mse_grads = 2 * (y_pred - y_true) / len(y_pred)Predictions: [2.5 0. 2. 8. ] True values: [ 3. -0.5 2. 7. ] Residuals: [-0.5 0.5 0. 1. ] MSE loss: 0.3750 MAE loss: 0.5000 Huber loss: 0.1875 MSE gradients: [-0.25 0.25 0. 0.5 ]
The MSE and Huber losses are close here because no residual exceeds the delta=1.0 threshold. Notice that the MSE gradient is proportional to the residual: the larger the error, the larger the gradient step that corrects it.
Cross-Entropy Loss
Here we compute cross-entropy manually and compare it with PyTorch's implementation, which handles numerical stability automatically.
# Raw logits for 3 classes (batch of 4 examples)
logits = torch.tensor(
[
[2.0, 1.0, 0.1], # Example 1: model prefers class 0
[0.5, 2.5, 0.3], # Example 2: model prefers class 1
[0.8, 0.2, 3.0], # Example 3: model prefers class 2
[1.5, 1.5, 0.5], # Example 4: uncertain between 0 and 1
]
)
targets = torch.tensor([0, 1, 2, 0]) # True class labels
# Manual softmax + cross-entropy
probs = F.softmax(logits, dim=1)
true_probs = probs[
torch.arange(4), targets
] # Probability assigned to true class
manual_ce = -torch.log(true_probs).mean()
# PyTorch implementation (numerically stable)
torch_ce = F.cross_entropy(logits, targets)Softmax probabilities: Example 1: [0.659 0.242 0.099], true class=0, P(true)=0.659 Example 2: [0.109 0.802 0.089], true class=1, P(true)=0.802 Example 3: [0.095 0.052 0.854], true class=2, P(true)=0.854 Example 4: [0.422 0.422 0.155], true class=0, P(true)=0.422 Manual cross-entropy loss: 0.4144 PyTorch cross-entropy loss: 0.4144
The two implementations agree. The model is most confident about examples 2 and 3 (probabilities above 0.8 for the true class), while example 4 is more uncertain.
Binary Cross-Entropy
Binary cross-entropy applies when each output is an independent binary prediction.
# Binary classification: spam detection
# Logits (raw scores before sigmoid)
logits_binary = torch.tensor([2.0, -1.5, 0.5, -2.0, 1.8])
labels_binary = torch.tensor([1.0, 0.0, 1.0, 0.0, 1.0])
# Predicted probabilities via sigmoid
probs_binary = torch.sigmoid(logits_binary)
# Manual BCE
eps = 1e-8 # numerical stability
manual_bce = -(
labels_binary * torch.log(probs_binary + eps)
+ (1 - labels_binary) * torch.log(1 - probs_binary + eps)
).mean()
# PyTorch BCE with logits (numerically stable)
torch_bce = F.binary_cross_entropy_with_logits(logits_binary, labels_binary)Binary classification predictions: Example 1: logit=2.0, P(spam)=0.881, label=1 [CORRECT] Example 2: logit=-1.5, P(spam)=0.182, label=0 [CORRECT] Example 3: logit=0.5, P(spam)=0.622, label=1 [CORRECT] Example 4: logit=-2.0, P(spam)=0.119, label=0 [CORRECT] Example 5: logit=1.8, P(spam)=0.858, label=1 [CORRECT] Manual BCE loss: 0.2165 PyTorch BCE loss: 0.2165
The model correctly classifies 4 of 5 examples. Example 3 (logit=0.5) is predicted as spam with probability 0.62 but is labeled as 1, so the model is correct. The loss reflects the combined penalty for all predictions weighted by their confidence.
Label Smoothing
# Label smoothing implementation
def cross_entropy_with_label_smoothing(logits, targets, epsilon=0.1):
"""Cross-entropy loss with label smoothing."""
K = logits.shape[1] # number of classes
# Standard cross-entropy
ce_loss = F.cross_entropy(logits, targets, reduction="none")
# Entropy of uniform distribution (penalty term)
log_probs = F.log_softmax(logits, dim=1)
uniform_loss = -log_probs.mean(dim=1) # avg log prob across all classes
# Combine: (1-eps)*CE + eps*uniform_loss
smoothed_loss = (1 - epsilon) * ce_loss + epsilon * uniform_loss
return smoothed_loss.mean()
# Compare standard CE vs label-smoothed CE
ce_standard = F.cross_entropy(logits, targets)
ce_smoothed = cross_entropy_with_label_smoothing(logits, targets, epsilon=0.1)Standard cross-entropy: 0.4144 Label-smoothed cross-entropy: 0.5235 Difference: 0.1092
Label smoothing slightly increases the loss value because it prevents the model from being completely certain about any prediction. In practice, this leads to better-calibrated models that generalize better.
Focal Loss
def focal_loss(logits, targets, gamma=2.0, alpha=None):
"""Focal loss for addressing class imbalance."""
probs = F.softmax(logits, dim=1)
# Get probabilities for true classes
batch_size = logits.shape[0]
true_probs = probs[torch.arange(batch_size), targets]
# Modulating factor
modulating_factor = (1 - true_probs) ** gamma
# Base cross-entropy
ce = F.cross_entropy(logits, targets, reduction="none")
# Focal loss
fl = modulating_factor * ce
if alpha is not None:
fl = alpha * fl
return fl.mean()
# Compare at different confidence levels
# Simulate a class imbalance scenario:
# Many easy examples (high confidence) and a few hard ones
easy_logits = torch.tensor([[3.0, 0.1, 0.1]] * 8) # 8 easy examples
hard_logits = torch.tensor([[0.9, 0.8, 0.7]] * 2) # 2 hard examples
all_logits = torch.cat([easy_logits, hard_logits])
all_targets = torch.tensor([0] * 8 + [0] * 2)
ce_all = F.cross_entropy(all_logits, all_targets)
fl_all = focal_loss(all_logits, all_targets, gamma=2.0)
# Focal loss per example
easy_probs = F.softmax(easy_logits, dim=1)[:, 0]
hard_probs = F.softmax(hard_logits, dim=1)[:, 0]Standard CE loss (all examples): 0.2839 Focal loss (gamma=2, all): 0.0811 Easy examples - avg P(true class): 0.901 Hard examples - avg P(true class): 0.367 Easy examples: CE=0.1044, Focal=0.0010, ratio=101.7x Hard examples: CE=1.0019, Focal=0.4013, ratio=2.5x
Focal loss suppresses the contribution of easy examples much more aggressively than hard ones. The easy examples have their loss reduced by a large factor, while hard examples (where the model is less confident) keep most of their loss contribution. This re-weighting forces the model to focus on the hard examples that standard cross-entropy would otherwise ignore.
Triplet Loss
def triplet_loss(anchor, positive, negative, margin=1.0):
"""Triplet loss for embedding learning."""
d_pos = torch.norm(anchor - positive, dim=1)
d_neg = torch.norm(anchor - negative, dim=1)
loss = torch.clamp(d_pos**2 - d_neg**2 + margin, min=0)
return loss.mean(), d_pos.mean().item(), d_neg.mean().item()
# Simulate embeddings in 4D space
# Anchor: average representation of the concept
# Positive: similar item (same class, different instance)
# Negative: different class
torch.manual_seed(42)
anchor = torch.randn(8, 4) # 8 triplets, 4D embeddings
positive = anchor + 0.2 * torch.randn(8, 4) # positive close to anchor
negative = anchor + 1.5 * torch.randn(8, 4) # negative further from anchor
loss_val, d_pos_mean, d_neg_mean = triplet_loss(
anchor, positive, negative, margin=1.0
)
# Also test with easy negatives (already far)
negative_easy = anchor + 3.0 * torch.randn(8, 4)
loss_easy, d_pos_easy, d_neg_easy = triplet_loss(
anchor, positive, negative_easy, margin=1.0
)Standard triplets: Avg pos distance: 0.385 Avg neg distance: 2.713 Triplet loss: 0.0102 Easy negatives (already far apart): Avg pos distance: 0.385 Avg neg distance: 4.203 Triplet loss: 0.0000
When negatives are close to the anchor (standard triplets), the model has significant loss and learns to push them apart. When negatives are already far (easy triplets), the loss is zero and no gradient flows. This zero-gradient behavior on easy triplets is why hard negative mining is essential for efficient triplet loss training.
InfoNCE Loss
def infonce_loss(anchors, positives, temperature=0.07):
"""InfoNCE loss using in-batch negatives."""
# Normalize embeddings to unit sphere (cosine similarity)
anchors_norm = F.normalize(anchors, dim=1)
positives_norm = F.normalize(positives, dim=1)
# Compute all pairwise cosine similarities: shape [batch, batch]
similarity_matrix = (
torch.matmul(anchors_norm, positives_norm.T) / temperature
)
# True pairs are on the diagonal
batch_size = anchors.shape[0]
labels = torch.arange(batch_size)
# InfoNCE is cross-entropy where the correct class is the diagonal
loss = F.cross_entropy(similarity_matrix, labels)
return loss
torch.manual_seed(0)
batch_size = 8
embedding_dim = 16
# Create anchor-positive pairs with some noise
base_embeddings = torch.randn(batch_size, embedding_dim)
anchors_infonce = base_embeddings + 0.1 * torch.randn(batch_size, embedding_dim)
positives_infonce = base_embeddings + 0.1 * torch.randn(
batch_size, embedding_dim
)
loss_high_temp = infonce_loss(
anchors_infonce, positives_infonce, temperature=1.0
)
loss_low_temp = infonce_loss(
anchors_infonce, positives_infonce, temperature=0.07
)InfoNCE loss (temperature=1.0): 1.2591 InfoNCE loss (temperature=0.07): 0.0004 Avg positive pair similarity: 0.993 Avg negative pair similarity: -0.052
Lower temperature sharpens the loss surface, penalizing cases where negative pairs have high similarity more aggressively. The key difference from triplet loss is that InfoNCE uses every other example in the batch as a negative simultaneously. This provides denser gradient signal.
Visualizing Loss Behavior







Limitations and Practical Considerations
Every loss function has failure modes, implementation pitfalls, and scenarios where it is the wrong choice. Understanding these is as important as understanding the formulas. This section covers the most common issues practitioners encounter when working with neural network losses in NLP.
When Standard Cross-Entropy Falls Short
Cross-entropy is a strong default, but it has failure modes. In imbalanced classification settings (which are more common in NLP than often appreciated: most sentences are not toxic, most tokens are not named entities, most documents are not spam), the model learns to predict the majority class. The loss gets dominated by the many easy negatives. Focal loss and class weighting both address this, but they add hyperparameters that need tuning.
Label smoothing helps with calibration but does not help with imbalance. These are two separate problems with separate solutions. Do not conflate them. A model can be well-calibrated but still fail on rare classes, and a model can handle rare classes well but still be overconfident in its probability estimates.
When dealing with severe imbalance, it is worth combining multiple approaches: class-weighted loss handles the frequency imbalance statically, while focal loss handles it dynamically based on current model confidence. Using both together can be more effective than either alone.
MSE for Probabilities is a Mistake
A common error is using MSE as a loss for classification by treating class probabilities as continuous targets. MSE does not respect the constraint that probabilities must sum to 1, and it does not penalize confident wrong predictions as sharply as cross-entropy. Cross-entropy is the principled choice for classification because it is derived from maximum likelihood under the categorical distribution. Using MSE for classification is like assuming the targets follow a Gaussian distribution when they are categorical, which produces suboptimal gradients and slow convergence.
There are edge cases where regression-style losses are used on probability-like outputs. Semantic similarity tasks sometimes predict a score between 0 and 1 representing similarity, where MSE or Huber loss is reasonable. But this is because the underlying task is a regression task (predicting a real-valued similarity score), not a classification task (predicting a discrete label).
Gradient Vanishing in BCE with Saturated Sigmoid
Binary cross-entropy applied directly to sigmoid outputs can cause gradient vanishing. When the sigmoid is nearly saturated (sigmoid output close to 0 or 1), its derivative is nearly zero, and the chain rule multiplies this small derivative into the gradient. The gradient can become so small that the optimizer makes essentially no progress.
The solution is to use F.binary_cross_entropy_with_logits() instead of computing the sigmoid first and then applying BCE. The combined function uses the numerically stable formulation:
where is the logit before sigmoid. This avoids computing the exponential of large positive numbers and maintains numerical precision across the full range of logit values. Always pass raw logits to combined loss functions rather than pre-activating them.
Triplet Loss Requires Hard Mining
Triplet loss fails silently in practice when using random triplets: most randomly sampled triplets are "easy" (already satisfying the margin constraint), so the loss is zero and no gradient flows. The model does not improve. Hard negative mining selects the hardest negative for each anchor (the negative that is currently closest to the anchor in embedding space), and semi-hard mining uses negatives that are close but still beyond the margin. This mining step is critical to making triplet loss work.
The computational cost of hard mining scales as in the number of examples: for each anchor, you need to find the hardest negative among all other examples. In practice, online hard mining within a batch (treating all non-positive examples in the batch as potential negatives and selecting the hardest) is the standard approach.
Loss Function and Evaluation Metric Alignment
The loss you train on should be related to the metric you evaluate on. Training with cross-entropy but evaluating with F1 score is common and often works well in practice. But when there is a significant mismatch, careful thought is needed. If you care about ranking (e.g., information retrieval), a pairwise or listwise ranking loss is more directly aligned with the evaluation metric than cross-entropy. If you care about precision at a specific recall threshold, focal loss with tuned and may align better with your deployment requirements.
The deeper principle here is that the loss function defines what the model considers an "error" and how much that error matters. If your deployment scenario weights certain errors (e.g., false negatives in medical diagnosis) much more heavily than others, your training loss should reflect this asymmetry. Standard symmetric losses will produce models optimized for average performance, which may not match the risk profile of your application.
Numerical Issues in Practice
Loss functions that involve logarithms or exponentials are vulnerable to numerical overflow and underflow. Key rules of thumb:
- Always use
F.cross_entropy()with raw logits, notF.nll_loss()after manually applyingF.softmax(). - Always use
F.binary_cross_entropy_with_logits(), notF.binary_cross_entropy()after applying sigmoid. - When implementing custom losses, use
torch.clamp()to keep log arguments away from zero, or use PyTorch's numerically stable built-in functions where possible. - Watch for NaN gradients: they often signal a loss computation that hits a numerical singularity. Log(0) and 0/0 are the most common causes.
Numerical stability is not an academic concern. Production models trained on large corpora with billions of parameter updates can fail silently if any loss computation produces NaN, as NaN propagates through the network and corrupts all weights. Defensive programming with stable loss implementations is essential.
Summary
Loss functions are the optimization targets that shape everything about how a neural network trains. The key takeaways are:
- MSE penalizes large errors quadratically and works well for regression with clean data. It arises as maximum likelihood under a Gaussian noise assumption. Huber loss is MSE's more robust variant that transitions to linear behavior for large errors, preventing outliers from dominating the gradient.
- Cross-entropy is derived from maximum likelihood and is the standard loss for classification. It strongly penalizes confident wrong predictions via the logarithm. The gradient is simply the prediction error: for each class. Binary cross-entropy is the two-class special case, and extends naturally to multi-label classification.
- Weighted cross-entropy addresses class imbalance by multiplying each example's loss by the inverse frequency of its class. This ensures rare classes receive proportionally more training signal.
- Label smoothing prevents overconfidence by replacing one-hot labels with soft targets. A small epsilon (0.1) is a reliable improvement for most classification tasks, reducing the logit gap and improving calibration.
- Focal loss addresses class imbalance by dynamically down-weighting easy examples through a modulating factor . The focusing parameter is a common default. Unlike class weighting, focal loss adapts to the model's current state rather than applying a fixed multiplier.
- KL divergence measures how far one distribution is from another and is the theoretical foundation linking cross-entropy to maximum likelihood estimation. It appears explicitly in distillation and variational methods.
- Temperature scaling controls the sharpness of a softmax distribution, used in distillation (high temperature), sampling (variable temperature), and calibration (post-hoc temperature search).
- Language modeling loss is cross-entropy applied at each token position over a large vocabulary, and perplexity is its exponential. This provides an interpretable measure of prediction quality.
- Contrastive and triplet losses enable representation learning without predefined class labels, forming the basis for modern sentence embedding systems. InfoNCE generalizes these to use the full batch as negatives, and is the foundation for models like SimCSE and DPR.
The next chapter covers backpropagation, which explains how the gradient of any of these loss functions flows backward through a neural network to update every parameter simultaneously. Understanding the loss is only half the picture. Backpropagation closes the loop between loss values and parameter updates.
Quiz
Ready to test your understanding? Take this quick quiz to reinforce what you've learned about loss functions.
Loss Functions Quiz
Reference
Citation details
Cite or share this article.
Continue with the full handbook
This chapter is part of Language AI Handbook. Use the handbook page to browse the complete table of contents and continue reading in sequence.
Explore Language AI HandbookStay up to date
Get articles, book updates, and news delivered to your inbox.
No spam, unsubscribe anytime.
Join the community
Sign in to remove popups, track your reading progress, and join the discussion.

Comments
No comments yet. Be the first to share your thoughts!