Loss Functions in Artificial Neural Networks: Types, Math & PyTorch

The loss function quantifies the error between the model's predictions and the ground truth targets, producing a scalar that the optimiser minimises via backpropagation. Every weight update in gradient descent is computed as the gradient of the loss with respect to that weight. The choice of loss function defines what the model is fundamentally optimising for.

MSE squares the residual $(y - \hat{y})^2$, so a prediction that is 10 units off contributes 100x more loss than one that is 1 unit off (not 10x). In a training batch containing a single extreme outlier, the squared error from that sample can dominate the entire gradient update, pulling all weights toward reducing that one error at the expense of correctly predicting all other samples.

Always use cross-entropy for classification. MSE applied to classification outputs (probabilities or one-hot labels) creates a non-convex loss surface with vanishing gradients for confidently wrong predictions. Cross-entropy with sigmoid/softmax activation is a proper scoring rule that penalises confident wrong predictions asymptotically (log(p) approaches infinity as p approaches 0), creating strong gradient signal throughout training.

BCE (Binary Cross-Entropy) is for single-output binary classification where the model outputs one sigmoid-activated probability. Categorical cross-entropy is for multi-class problems with C mutually exclusive classes where the model outputs C logits (softmax over the class dimension). For multi-label classification (where multiple classes can be true simultaneously), use BCE per class independently.

BCEWithLogitsLoss accepts raw logits (pre-sigmoid values) and internally computes sigmoid + binary cross-entropy using the log-sum-exp trick for numerical stability. If you manually apply sigmoid then pass to BCELoss, you risk overflow/underflow for very large or very small logit values, producing NaN losses. BCEWithLogitsLoss avoids this entirely.

Focal Loss (Lin et al., Facebook AI Research 2017) adds a modulating factor $(1-p_t)^\gamma$ to the standard cross-entropy that reduces the loss contribution of well-classified easy examples. It was designed for dense object detection (RetinaNet) where there are ~100,000 background proposals (easy negatives) versus a handful of objects. Use Focal Loss when your positive class frequency is below 5 to 10% and standard weighted BCE is not sufficient.

Yes. Implement them as plain Python functions using PyTorch tensor operations (which automatically build the autograd computational graph), or as nn.Module subclasses for stateful or parameterised losses. Verify gradients with torch.autograd.gradcheck() using double-precision inputs.

Backpropagation computes the gradient of the loss scalar with respect to every parameter by applying the chain rule backwards through the network. The loss function is the root of this chain: its gradient with respect to the final layer output initiates the backward pass. Without a differentiable loss, there is no gradient signal and no learning.

Not always. MAE is more robust to outliers but its gradient is a constant $\pm 1/N$ everywhere, which causes oscillation near the optimum (the gradient doesn't decrease as the model improves). MSE has a gradient proportional to the error magnitude, providing strong updates early in training and fine-grained updates near convergence. Huber Loss combines the advantages of both and is usually the better practical choice for noisy regression.

Label smoothing replaces hard one-hot targets with soft targets that assign small probability mass $arepsilon/(C-1)$ to each incorrect class. It prevents the model from becoming overconfident (pushing logits to $\pm\infty$), acts as a regulariser, and significantly improves model calibration. Use it when fine-tuning large transformer models (ViT, BERT, GPT), when your model is consistently overconfident on the training set, or when calibration quality matters.

KL Divergence measures distribution mismatch. In Variational Autoencoders (VAE), the KL term in the ELBO regularises the encoder's latent distribution toward the prior (typically standard Gaussian). In knowledge distillation, KL between teacher soft labels (at temperature T) and student output transfers the teacher's class similarity structure to the student. In RL (PPO), KL divergence constraints prevent policy updates that are too large.

Triplet Loss trains on anchor-positive-negative triplets: (anchor face, same-person face, different-person face). It minimises distance between anchor and positive while maximising distance to the negative by at least margin $\alpha$. The critical implementation detail is hard negative mining: naively random triplets are often trivially satisfied (negative already far from anchor). Hard negative mining selects negatives closer to the anchor, providing more informative gradients. ArcFace (Deng et al. 2019) improved further by using an angular margin in the softmax space.

InfoNCE (Noise Contrastive Estimation) maximises mutual information between query and matching key by treating all other keys in the batch as negatives. In CLIP (OpenAI), image embeddings and their paired text embeddings are the positive pairs; all other image-text combinations in the batch are negatives. Training on 400M image-text pairs with bidirectional InfoNCE creates powerful cross-modal embeddings that enable zero-shot image classification.

Options in increasing order of sophistication: (1) class_weight parameter in CrossEntropyLoss or pos_weight in BCEWithLogitsLoss (inverse frequency weighting); (2) Focal Loss (down-weights easy examples, focuses on hard misclassified ones); (3) oversampling the minority class (SMOTE, ROS) before training; (4) Dice Loss for segmentation (inherently imbalance-robust). Monitor both training loss and class-specific metrics (precision, recall per class) to assess effectiveness.

Validation loss on a held-out dataset estimates generalisation performance. If training loss decreases while validation loss plateaus or increases, the model is overfitting (memorising training data rather than learning general patterns). This is the primary signal for early stopping. Compare also training and validation accuracy or the target metric; if both training loss and metric are improving but validation metric is not, check for data leakage or distribution mismatch.