Loss Functions and Training Objectives - The Art of Teaching AI
A model’s loss function defines what “good” means during training. It quantifies the gap between predicted outputs and ground truth, and the optimizer (e.g., gradient descent) works to minimize this gap. Choosing the right loss function is as important as choosing the right architecture—it directly shapes what the model learns.
Loss Functions for Classification
Cross-Entropy Loss
The standard choice for multi-class classification. For each example, the model outputs a probability distribution over classes (via softmax). Cross-entropy measures how surprised the model is by the true label—if the model assigns high probability to the correct class, loss is low; if it assigns low probability, loss is high.
$$L = -\sum_{i=1}^{C} y_i \log(\hat{p}_i)$$
where $y_i$ is 1 if class $i$ is correct and 0 otherwise, and $\hat{p}_i$ is the model’s predicted probability for class $i$.
Cross-entropy has desirable properties: it penalizes confident wrong predictions more severely than uncertain wrong predictions (which is usually desired), and it provides strong gradients for learning even when predictions are very wrong.
Binary Cross-Entropy
For binary classification (only two classes), cross-entropy simplifies:
$$L = -[y \log(\hat{p}) + (1-y) \log(1-\hat{p})]$$
Often paired with a sigmoid activation (outputting a probability between 0 and 1) rather than softmax.
Focal Loss
Standard cross-entropy treats all examples equally. In heavily imbalanced datasets (one class is 99% of the data), the model achieves high accuracy by predicting the majority class and ignoring minorities. Focal loss down-weights easy examples (where the model is already confident) and up-weights hard examples (where the model is uncertain about a minority class):
$$L = -\alpha_t (1 - \hat{p})^{\gamma} \log(\hat{p})$$
The factor $(1 - \hat{p})^{\gamma}$ is small when the model is confident, reducing its contribution to loss. When the model is uncertain, this factor is large, amplifying loss. This focuses training on hard examples and helps with class imbalance.
Loss Functions for Regression
Mean Squared Error (MSE)
The standard choice for regression, penalizing the squared difference between prediction and target:
$$L = \frac{1}{n} \sum_{i=1}^{n} (\hat{y}_i - y_i)^2$$
MSE is sensitive to outliers (the squared term amplifies large errors), making it ideal when extreme predictions are particularly bad. However, a few outliers can dominate training and pull the model’s attention away from typical examples.
Mean Absolute Error (MAE)
Penalizes absolute differences, less sensitive to outliers:
$$L = \frac{1}{n} \sum_{i=1}^{n} |\hat{y}_i - y_i|$$
MAE is more robust but provides smaller gradients when errors are small, which can slow learning in the early phases when the model is far from the target.
Huber Loss
Combines MSE and MAE: for small errors, it behaves like MSE (smooth gradient); for large errors, it behaves like MAE (bounded sensitivity). This balances learning speed with outlier robustness:
$$L = \begin{cases} \frac{1}{2}(y - \hat{y})^2 & \text{if } |y - \hat{y}| \leq \delta \ \delta(|y - \hat{y}| - \frac{\delta}{2}) & \text{otherwise} \end{cases}$$
where $\delta$ is a tunable parameter controlling the transition point.
Loss Functions for Ranking and Metric Learning
Triplet Loss
In recommendation systems or face recognition, the goal isn’t to predict a label but to learn embeddings where similar items cluster together. Triplet loss takes three examples: an anchor, a positive (similar to anchor), and a negative (dissimilar). It encourages the anchor and positive to be close in embedding space and anchor and negative to be far apart:
$$L = \max(0, d(a, p) - d(a, n) + m)$$
where $d$ is Euclidean distance and $m$ is a margin (gap to enforce). The $\max(0, \cdot)$ ensures the constraint is satisfied; if it already is, loss is zero and the model isn’t penalized.
Contrastive Loss
Similar idea: encourage similar pairs to be close and dissimilar pairs to be far apart. Contrastive loss is used in self-supervised learning where unlabeled data is paired (e.g., augmented views of the same image) to provide signal without manual labels.
Regularization and Modified Losses
L1 and L2 Regularization
To prevent overfitting, add a penalty for model complexity (magnitude of weights):
$$L_{total} = L_{task} + \lambda \sum w^2 \quad \text{(L2)}$$ $$L_{total} = L_{task} + \lambda \sum |w| \quad \text{(L1)}$$
L2 encourages small weights uniformly. L1 encourages sparsity—some weights go exactly to zero, removing features from the model. This acts as feature selection, useful when you want interpretability or when features are high-dimensional.
Label Smoothing
Instead of training the model to output one-hot labels (1 for correct, 0 for incorrect), train it to assign 90% probability to the correct label and 10% distributed across others. This prevents overconfidence and often improves generalization:
$$y_{smooth} = (1 - \alpha) y_{true} + \alpha \frac{1}{C}$$
where $\alpha$ is the smoothing strength and $C$ is the number of classes.
Custom and Task-Specific Losses
Weighted Loss Functions
Assign different weights to different examples. Examples from underrepresented groups in a dataset can be up-weighted so the model learns to handle them well despite being less frequent. Hard examples can be up-weighted to focus training on challenging cases.
$$L = \frac{1}{n} \sum w_i l(y_i, \hat{y}_i)$$
Domain-Specific Losses
In medical imaging, False Negatives (missing a tumor) are catastrophic; False Positives (false alarm) are inconvenient. A weighted loss can penalize FN more:
$$L = w_{FN} \cdot \text{FN_cost} + w_{FP} \cdot \text{FP_cost}$$
In recommendation systems, different metrics (e.g., ranking correctness, diversity, novelty) might be combined in a custom loss.
Loss Function Selection Checklist
-
Does the loss reflect the goal? If optimizing for accuracy, cross-entropy works; if optimizing for ranking, use ranking losses.
-
Are there class imbalances or example-wise importance differences? Use weighted or focal losses.
-
Are there outliers? Use robust losses like Huber or MAE rather than MSE.
-
Do you care about calibration? Cross-entropy naturally produces well-calibrated probabilities; other objectives may not.
-
Is interpretability important? Simpler losses (MSE, cross-entropy) are easier to understand than complex combinations.
-
Can you compute gradients? The optimizer needs to compute loss gradients; non-differentiable losses require special handling.
Evolution During Training
A common practice is to start with one loss function and switch during training. For example:
- Warm-up phase: use a simpler, more stable loss to get the model into a reasonable regime
- Main phase: switch to the final loss to optimize the specific objective
- Fine-tuning phase: use a different loss (e.g., focal loss) to focus on hard examples
The loss function isn’t fixed—thoughtful practitioners adjust it as training progresses to balance stability, learning speed, and the final objective.
Understanding loss functions is understanding what you’re asking the model to do. A well-chosen loss makes learning easier and faster; a poorly chosen one causes training to stall or the model to optimize the wrong objective entirely.