AlgoMaster Logo

Loss Functions and Optimization

10 min readUpdated June 1, 2026
Listen to this chapter
Unlock Audio

Every ML algorithm, from logistic regression to deep neural networks, learns by minimizing a measure of error. That measure is the loss function. The process of updating the model to reduce that error is optimization.

Together, they form the core of training. The choices you make here directly affect convergence speed, training cost, and how the model behaves in production.

What Loss Functions Measure

A loss function takes the model's prediction and the ground truth label, and outputs a single number: how wrong was the prediction? Larger values mean worse predictions. The model's entire learning process is driven by reducing this number across the training data.

Consider a spam classifier that outputs the probability an email is spam. If the model predicts 0.9 for an email that is actually spam, the loss is small, the model was confident and correct. If it predicts 0.1 for that same spam email, the loss is large. The loss function is how the model receives a learning signal: not just whether it was wrong, but how wrong, and in what direction to adjust.

This creates a feedback loop that repeats millions of times during training:

The choice of loss function determines what the model optimizes for, which shapes how it behaves in production. A recommendation system trained to minimize cross-entropy loss on click prediction optimizes for click probability. The same architecture trained with a regression loss on watch time optimizes for engagement duration. Same model, different loss function, completely different product behavior.

Common Loss Functions

Cross-Entropy for Classification

Cross-entropy is the standard loss for classification tasks. It measures the gap between the predicted probability distribution and the true label.

For binary classification (spam or not spam, fraud or legitimate), binary cross-entropy works like this: if the true label is 1 (spam) and the model predicts probability p, the loss is -log(p). If the true label is 0 (not spam), the loss is -log(1-p). The logarithm creates an asymmetric penalty structure that makes cross-entropy particularly effective.

A correct prediction with 0.95 confidence produces a loss of about 0.05. A wrong prediction with 0.95 confidence produces a loss of about 3.0. That 60x difference means the model gets punished severely for being confidently wrong, which is exactly the behavior you want. It forces the model to be well-calibrated, not just correct, but correctly confident.

This calibration property matters in production. Ad ranking systems use predicted click-through probability directly in auction pricing. If the model predicts a 5% click probability, that number feeds into bid calculations.

A model that's correct on average but wildly miscalibrated (predicting 50% when the true rate is 5%) would distort ad pricing. Cross-entropy naturally produces calibrated probabilities because the loss is minimized when the predicted probability matches the true frequency of the label.

For multi-class problems (image classification with 1,000 categories, intent detection with dozens of intents), categorical cross-entropy extends the same idea across all classes. The model outputs a probability distribution over all classes via softmax, and the loss penalizes placing low probability on the correct class.

MSE and MAE for Regression

When the target is a continuous value, you need a different notion of "wrong." Mean Squared Error (MSE) and Mean Absolute Error (MAE) are the two primary options.

MSE squares the difference between prediction and target. A prediction that's off by 10 incurs a loss of 100. A prediction that's off by 2 incurs a loss of 4. The squaring means large errors are penalized disproportionately: being off by 10 is 25 times worse than being off by 2, not 5 times worse.

This is useful when large errors are genuinely more costly. In ETA prediction, an estimate that's off by 30 minutes causes a far worse user experience than one that's off by 3 minutes, and MSE captures that asymmetry.

MAE uses the absolute difference. Off by 10 gives a loss of 10. Off by 2 gives a loss of 2. The relationship is linear. MAE is robust to outliers because a single extreme value can't dominate the total loss the way it can with MSE. House price prediction is a good example: a dataset might include a few properties sold at 10x the median price due to unique circumstances.

Under MSE, the model would distort its predictions on ordinary houses to slightly reduce its error on those outliers. Under MAE, the outliers have proportional, not outsized, influence.

Huber loss is a hybrid that behaves like MSE for small errors and MAE for large errors. It captures the "large errors matter more" property of MSE without the outlier sensitivity. Many production regression systems use Huber loss as a practical default.

Loss FunctionPenalizesOutlier SensitiveCalibrated OutputGradient BehaviorCommon Use Cases
Binary Cross-EntropyConfident wrong predictionsNoYes (probabilities)SmoothCTR prediction, spam, fraud
Categorical Cross-EntropyWrong class, high confidenceNoYes (probabilities)SmoothImage classification, intent detection
MSESquared errorYes (strongly)N/AProportional to errorETA prediction, demand forecasting
MAEAbsolute errorNoN/AConstant magnitudeHouse prices, robust regression
Huber LossMSE small errors, MAE largePartiallyN/AAdaptiveBalanced regression tasks
Contrastive / TripletSimilar items far apartNoN/ADepends on marginSearch, recommendations, similarity

Contrastive Loss for Embeddings

The previous chapter covered how embeddings represent entities as dense vectors where similar items end up close together. But how does the training process actually push similar items closer and dissimilar items apart? That's the job of contrastive loss functions.

Unlike cross-entropy or MSE, contrastive losses don't predict a label. They operate on relationships between pairs or groups of examples.

Triplet loss works with three items at a time: an anchor (the reference item), a positive (an item that should be similar to the anchor), and a negative (an item that should be dissimilar). The loss penalizes cases where the negative is closer to the anchor than the positive, with a margin buffer. If the anchor is a user's recently watched video, the positive might be another video they watched in the same session, and the negative a random video from the catalog.

InfoNCE (used in models like CLIP and SimCLR) generalizes this by treating all other items in the same mini-batch as negatives. Given a batch of 4,096 pairs, each positive pair is contrasted against 4,095 negatives simultaneously. This is computationally efficient because the negatives come for free from the batch, but it requires large batch sizes to work well. That's a direct system design trade-off: larger batches need more GPU memory.

The choice of negative sampling strategy has major impact on embedding quality. Random negatives (picking any item from the catalog) are easy to compute but produce mediocre embeddings because the model learns to distinguish obviously different items without developing fine-grained discrimination. Hard negatives (items that are similar but not matches, like two different action movies) force the model to learn subtler distinctions, producing better embeddings at the cost of a more complex mining pipeline.

Gradient Descent: How Models Learn

The loss function tells the model how wrong it is. Gradient descent is the algorithm that uses that signal to improve.

The idea is straightforward: compute the gradient of the loss with respect to each model parameter. The gradient points in the direction of steepest increase, so moving in the opposite direction reduces the loss. Take a step in that direction, recompute, repeat.

The learning rate controls the step size, and getting it right is one of the most consequential decisions in training. Too large, and the model overshoots the minimum, with the loss oscillating wildly or diverging entirely. Too small, and convergence takes forever, burning GPU-hours without meaningful progress. A typical starting learning rate for Adam is 1e-3 to 3e-4. Getting this wrong by a factor of 10 can mean the difference between a model that converges in 2 hours and one that doesn't converge at all.

Computing the gradient over the full dataset before each update (batch gradient descent) is mathematically clean but impractical. With 100 million training examples, a single gradient computation takes minutes. Mini-batch gradient descent computes the gradient on a small subset (typically 32 to 512 examples), producing a noisy but useful estimate of the true gradient. The noise actually helps: it prevents the model from getting stuck in sharp, narrow minima that generalize poorly. Training on mini-batches also enables parallelism across GPU cores, which is why batch size affects training throughput directly.

Optimizers: SGD vs Adam

Vanilla gradient descent treats all parameters equally, applying the same learning rate to every weight in the model. Production optimizers are more sophisticated.

SGD with momentum adds a velocity term that accumulates gradient history. Instead of reacting only to the current gradient, the optimizer builds up speed in directions where the gradient consistently points, and dampens oscillations in noisy directions. Think of a ball rolling downhill: it accelerates through valleys and resists sudden changes in direction. Momentum values of 0.9 are standard.

Adam (Adaptive Moment Estimation) goes further by maintaining per-parameter learning rates. It tracks both the mean and variance of past gradients for each parameter. Parameters that receive sparse, infrequent gradient updates get larger effective learning rates. Parameters with frequent, consistent gradients get smaller ones. This adaptivity makes Adam much less sensitive to the initial learning rate: a default of 3e-4 works well across a wide range of architectures.

Adam dominates in practice because it reduces the engineering time spent on hyperparameter tuning. A model that converges in 10 epochs with Adam versus 50 epochs with poorly-tuned SGD means 5x fewer GPU-hours. At $2-3 per GPU-hour on cloud infrastructure, that's the difference between $500 and $2,500 per training run. Multiply by daily retraining and the cost difference compounds.

SGD with momentum can match or exceed Adam's final model quality on certain tasks, particularly large-scale image classification (ResNet, EfficientNet), where careful scheduling and tuning are feasible. But that tuning requires expertise and experimentation time that Adam doesn't need.

PropertySGDSGD + MomentumAdam
Learning rateSingle, globalSingle, globalPer-parameter, adaptive
Tuning effortHighMediumLow
Convergence speedSlow without tuningModerateFast out of the box
Final model qualityBest with careful tuningGoodGood (sometimes slightly worse)
Memory per parameterNone1 extra value (velocity)2 extra values (mean + variance)
Best forVision models with tuning budgetGeneral useDefault choice, NLP, embeddings

Learning Rate Schedules

A fixed learning rate is rarely optimal across the full training run. Early in training, larger steps help the model explore the loss surface quickly. Later, smaller steps allow fine-tuning near the optimum without overshooting.

Step decay reduces the learning rate by a fixed factor at predefined epochs. Divide by 10 at epoch 30 and again at epoch 60 is a common recipe for image classification. It's simple and effective, but requires knowing roughly how many epochs training will take.

Cosine annealing smoothly decreases the learning rate following a cosine curve from the initial value down to near zero. This avoids the abrupt drops of step decay and has become the standard for transformer training. The smooth decay gives the model time to settle into a good minimum without sudden disruptions.

Warmup starts with a very small learning rate and ramps up over the first few hundred or thousand steps before transitioning to the main schedule. Random initial weights can produce wild gradients early in training, and a large learning rate amplifies that instability. Warmup lets the model stabilize before applying the full learning rate. BERT, GPT, and most transformer architectures use warmup as standard practice.

ScheduleBehaviorCommon UseKey ParametersTypical Setting
ConstantFlat throughoutQuick experimentsLR1e-3 to 1e-4
Step decayDrops at fixed intervalsImage classification (ResNet)LR, step size, decay factorDivide by 10 every 30 epochs
Cosine annealingSmooth curve to near-zeroTransformer trainingInitial LR, total steps1e-3 down to 1e-6
Warmup + decayRamp up, then decreaseLarge models (BERT, GPT)Warmup steps, peak LRWarmup 10% of total steps
One-cycleRise to peak, then decayFast convergenceMax LR10x default, then cosine

Why This Matters for System Design

All of these choices, loss function, optimizer, learning rate schedule, translate directly into production concerns: how long training takes, how much it costs, and whether it succeeds at all.

Training Cost

Optimizer choice affects how many epochs the model needs to converge, which maps directly to GPU-hours and dollars. Consider a concrete example: training a recommendation model with 50 million training examples and 200-dimensional embeddings, retrained daily.

  • With Adam + cosine schedule + batch size 512: converges in 3 epochs, about 2 hours on 4 V100 GPUs. Cost per run: ~$24.
  • With SGD + fixed LR + batch size 64: converges in 15 epochs, about 8 hours on 4 V100 GPUs. Cost per run: ~$96.
  • Annual cost of daily retraining: $8,760 vs $35,040.

The optimizer choice alone creates a 4x cost difference. Loss function choice matters too: contrastive loss with in-batch negatives can train embedding models in hours using large batches, while triplet loss with randomly sampled negatives might need days to reach equivalent quality.

Convergence and Stability

In production, training failures are expensive. If a training run diverges at hour 15 of a 20-hour job, those 15 GPU-hours are wasted. Several mechanisms reduce this risk:

  • Gradient clipping caps the gradient norm to prevent exploding gradients. This is standard practice for transformer training and costs essentially nothing to apply.
  • Warmup prevents early instability from random initial weights.
  • Adam's adaptive learning rates are inherently more stable than a fixed global learning rate.

A robust optimization setup (Adam + warmup + gradient clipping + cosine decay) doesn't guarantee convergence, but it dramatically reduces the probability of a wasted training run compared to vanilla SGD with a fixed learning rate.

DecisionFaster / Cheaper OptionSlower / Costlier OptionTypical Impact
OptimizerAdam (fewer epochs)Untuned SGD (more epochs)2-5x training time difference
Batch sizeLarger (better GPU utilization)Smaller (underutilized GPUs)2-3x throughput difference
LR scheduleWarmup + cosine (stable, fast)Fixed LR (risk of divergence)Prevents wasted runs
Loss functionTask-appropriate choiceMismatched lossMay never converge
Gradient clippingApplied (prevents divergence)Not applied (risk of NaN losses)Prevents total training failures

Quiz

Loss Functions and Optimization Quiz

10 quizzes