Loss Functions › Code Implementation
Language AI HandbookAll chapters Loss Functions › Code ImplementationOne section · Sections in this chapter13 / 17 ↓
Reading tools Definitions & inline links
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
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.