Lesson 2 of 6 · 47 min

Loss functions & optimization

Why the loss is the algebraic commitment: MSE for real targets, cross-entropy for classes and why (convex Hessian, non-saturating gradient). Then the dynamics — batch vs SGD, Adam’s per-coordinate scaling, when Adam hurts generalization, and the learning-rate regimes.

The loss is a claim about your data, not a convenience

Pick the loss and you have committed to a probabilistic story: MSE says “Gaussian noise on a real number”; cross-entropy says “Bernoulli/Categorical over classes.” Get that claim wrong and everything downstream — gradients, convexity, calibration — fights you. The modeling round probes exactly this: “why cross-entropy and not MSE for classification?” (named as a pen-and-paper item in the Priyaa FAANG writeup), “why is logistic loss convex?”, “when would you not use Adam?” The strong answer always walks from what the loss assumes → what the gradient does when you’re wrong → the production implication. This lesson gives you all three.

Start with regression, where the loss is least controversial. MSE = (1/N)·Σ(yᵢ − ŷᵢ)² is the maximum-likelihood loss under additive Gaussian noise; it penalizes the distance of residuals and is symmetric in direction. Google’s ML Crash Course makes the behavior concrete: predicting 23.1 when truth is 24 gives a per-point L2 loss of 0.81 — the square means a few large residuals dominate the total, which is exactly why MSE is sensitive to outliers and why Huber loss (quadratic near zero, linear in the tails) is the standard robust alternative. Interview angle. “MSE vs MAE?” → MAE (L1) is robust to outliers and optimizes the median; MSE optimizes the mean and punishes large misses harder. Pick by whether tail errors are catastrophic or just noise.

Now classification. Cross-entropy is the negative log-likelihood of the Bernoulli (binary) or Categorical (multi-class): minimizing it is maximum likelihood under that model. Binary form: −(1/N)·Σ[ yᵢ·log ŷᵢ + (1−yᵢ)·log(1−ŷᵢ) ]. The “why the minus sign?” follow-up has a one-line answer: likelihoods are ≤ 1, so their logs are ≤ 0; the minus flips it to a positive quantity we minimize. The elegant fact CS229 highlights: the gradient of logistic loss w.r.t. the logit simplifies to (ŷ − y)·x — proportional to the error itself, large when confidently wrong and small when right. That is precisely the gradient signal you want for classification.

Why We Don’t Use the Mean Squared Error (MSE) Loss in ClassificationDataMListic

Why cross-entropy is convex — and why that sentence scores

For the linear/logistic model, ESL derives the Hessian of the negative log-likelihood as XᵀWX with W = diag[ p(1−p) ]. Because every weight p(1−p) ≥ 0, that matrix is positive semi-definite — so the loss is convex, and if X has full column rank the minimizer is unique. That is the rigorous “why” behind the casual claim “logistic regression has a single global optimum.” Dropping “the Hessian is XᵀWX, PSD because p(1−p)≥0” is a senior-level mic-drop on the convexity question — far stronger than “it just is convex.”

The deeper reason to reject MSE on a sigmoid/softmax output is the gradient, not just the likelihood story. Pair MSE with a saturating output and, when the model is confidently wrong (predicts ~0 for a true 1), the sigmoid derivative σ(z)(1−σ(z)) is near zero — so the MSE gradient is vanishingly small exactly where you most need a big correction. Cross-entropy cancels that saturating term, leaving the clean (ŷ − y) signal. So the crisp answer to “why not MSE for classification?” is: (1) it assumes Gaussian noise on something that is a probability; (2) combined with a sigmoid it produces vanishing gradients on confident mistakes; (3) it is non-convex in the logits for that model, unlike cross-entropy.

code

1WHY MSE STALLS ON A CONFIDENTLY-WRONG SIGMOID OUTPUT23  true label y = 1, model outputs p ~ 0.001 (confidently wrong)45  MSE grad  ~  (p - y) * sigma'(z)  =  (-0.999) * [p(1-p)] ~ (-0.999)*0.001 ~ -0.0016                                        ^ saturating term ~ 0  -> almost no update78  CE  grad  ~  (p - y)              =  -0.999                  -> full-size update910  Cross-entropy removes the saturating sigma' term, so the gradient stays11  proportional to the error. This is the real reason CE beats MSE here.

code

1LOSS-BY-TASK QUICK MAP (interview rapid-fire)23  task                          loss              notes4  ---------------------------   ---------------   ---------------------------------5  real-valued regression        MSE (L2)          optimizes the mean; outlier-sensitive6  robust regression             MAE / Huber       optimizes median / quadratic-then-linear7  binary classification         BCE (log loss)    Bernoulli NLL; gradient ~ (p - y)8  multi-class (exclusive)       softmax CE        one coupled distribution, sums to 19  multi-label (independent)     sigmoid + BCE     per-label independent probabilities10  heavy class imbalance         focal loss        CE * (1 - p_t)^gamma, down-weights easy

Two more loss facts the rapid-fire round loves (Yuan Meng’s MLE bank). Softmax vs sigmoid: sigmoid gives independent per-label probabilities (multi-label); softmax couples logits into one distribution that sums to 1 (mutually-exclusive multi-class). Temperature divides logits before exponentiation: T > 1 softens the distribution (more uncertainty, used in distillation), T < 1 sharpens it. And focal loss = cross-entropy × (1−p_t)^γ, which down-weights easy examples so training focuses on the hard, rare positives — the bridge to class imbalance in Lesson 5.

Batch vs SGD: the dynamics, and why SGD oscillates

There is a third mode between the extremes, and it is what everyone actually uses: mini-batch SGD (typical batch 32–512). It keeps most of SGD’s regularizing noise while amortizing the forward/backward pass over a batch the GPU can parallelize — the sweet spot of statistical efficiency and hardware throughput. Full-batch is too slow per step and overfits to sharp minima; pure single-example SGD wastes the accelerator. Interview angle. “Why mini-batch and not full-batch or single-example?” → noise for generalization plus vectorized hardware utilization; the batch size is then a knob trading gradient-estimate variance against compute.

Batch gradient descent uses the full-data gradient per step: θ ← θ − α·∇J(θ). It is stable but “has to scan the entire training set before taking a single step” (CS229). SGD takes a noisy step per example (or mini-batch) and “can start making progress right away.” The convergence contrast is exam-ready: for a convex objective, full-batch GD “always converges (assuming α not too large) to the global minimum,” whereas SGD “may never converge — the parameters keep oscillating around the minimum,” which is why you must “slowly let α decrease to zero” to anneal the noise. That oscillation is not a bug; the gradient noise is a form of implicit regularization that helps SGD find flatter, better-generalizing minima.

Interview angle. “Batch size — how does it affect training?” Larger batches give a lower-variance gradient estimate (smoother, more parallel-efficient steps) but less regularizing noise and a tendency toward sharper minima; small batches inject noise that can generalize better but train slower and noisier. This is also why naively scaling batch size requires scaling the learning rate (the linear-scaling rule) — a detail that signals you have actually run large training jobs.

Momentum: what every modern optimizer adds first

Before Adam, the first fix to plain SGD is momentum: accumulate an exponential moving average of past gradients and step in that direction, v ← βv + (1−β)g; θ ← θ − αv. The physical picture is a ball rolling downhill — it builds speed along consistent directions and damps the zig-zag across a ravine, so it crosses ill-conditioned valleys far faster than vanilla SGD and rolls through small local bumps. Nesterov momentum is a look-ahead variant that often converges a touch faster. Interview angle. “Why momentum?” → it accelerates consistent descent and cancels oscillation in high-curvature directions; it is the m_t term Adam keeps. Knowing momentum is the prerequisite for explaining Adam properly.

Adam: per-coordinate scaling, and when it backfires

Adam (Kingma & Ba) keeps two exponential moving averages per parameter: m_t of the gradient (momentum-like direction) and v_t of the squared gradient (per-coordinate scale), bias-corrects both for their zero init, and updates θ ← θ − α·m̂_t/(√v̂_t + ε). The effect: each coordinate is normalized by its own recent gradient magnitude, so sparse and ill-scaled features get sensible step sizes without manual tuning. That adaptivity is why Adam is the default for transformers and NLP, where gradients are sparse and the loss landscape is nasty — it just works out of the box.

code

1ADAM, ONE STEP (per parameter)23  m_t = b1*m_{t-1} + (1-b1)*g          # 1st moment: direction (momentum)4  v_t = b2*v_{t-1} + (1-b2)*g^2        # 2nd moment: per-coordinate scale5  m_hat = m_t / (1 - b1^t)             # bias-correct the zero init6  v_hat = v_t / (1 - b2^t)7  theta = theta - alpha * m_hat / (sqrt(v_hat) + eps)89  defaults: b1=0.9  b2=0.999  eps=1e-810  the eps in the denominator is a stabilizer -- but it also extinguishes tiny gradients.

The senior twist — and the real answer to “when would you not use Adam?” — is that Adam can generalize worse than SGD on vision. The Adam-vs-SGD study (OPT-ML 2021) reports that “on image classification problems, [Adam’s] generalization performance is significantly worse than SGD,” attributing it to loss flattening: adaptive scaling pushes weights/logits up, gradient magnitudes shrink toward zero, and the optimizer gets stuck in a region with inferior test error because it has no signal to escape. SGD, being slower and noisier, “ends up going to solutions with superior generalization.” So the playbook is: Adam to prototype fast and for transformers; SGD+momentum (with a schedule) when chasing state-of-the-art accuracy on dense vision benchmarks. AdamW (decoupled weight decay) narrows but does not erase the gap.

Why not just use Newton’s method (the true second-order step that rescales by the Hessian)? Because computing and inverting a Hessian over millions of parameters is O(n²)–O(n³) — completely intractable at deep-learning scale. Adam is best read as a cheap diagonal approximation to that second-order rescaling: v_t estimates per-coordinate curvature with one extra moving average, no matrix inversion. Interview angle. “Why don’t we use second-order optimizers for deep nets?” → the Hessian is too big to form or invert; adaptive methods (Adam, RMSProp) buy most of the benefit — coordinate-wise step sizing — at first-order cost. That framing shows you know what Adam approximates and why the exact thing is off the table.

Adam Optimization Algorithm (C2W2L08)DeepLearning.AI

Adam to move fast and forgive your tuning; SGD+momentum to land the best vision number. The optimizer is not a default you set once — it is a choice you revisit when generalization, not speed, becomes the goal.

Learning rate: the cheapest hyperparameter to reason about

The learning rate α is the single step size that decides whether any of the above converges. CS229’s anchor for convex problems: GD “always converges, assuming α is not too large.” The three regimes are worth memorizing as symptoms: too large → loss diverges/explodes (steps overshoot the basin; fix with a smaller α or gradient clipping); too small → loss stalls (near-zero progress; fix with a larger α or warm-up); just right → fast, monotone-ish decrease, ideally paired with a decay schedule. For deep nets the production recipe is warm-up for a few hundred steps to dodge early instability, then cosine or step decay so the noisy SGD tail can settle into a flat minimum.

code

1LEARNING-RATE REGIMES (read the loss curve)23  alpha too large   loss spikes / NaNs / oscillates wildly   -> shrink alpha, clip grads4  alpha too small   loss crawls, barely moves                -> raise alpha, add warm-up5  alpha ~ right     smooth, fast decrease to a plateau        -> add a decay schedule6  flat at HIGH loss both train AND val stuck high            -> UNDERFIT: add capacity first78  Trap: if train+val both plateau high, the fix is usually MORE MODEL, not LESS alpha.

Interview angle. A favorite trap: “your loss plateaued — lower the learning rate?” Only if it plateaued after descending and you want a finer landing. If both train and validation are stuck high from the start, that is underfitting (Lesson 1) — you need capacity or better features, and dropping α further just makes it crawl. Diagnosing “is this an optimization plateau or a capacity ceiling?” before touching α is the senior reflex.

Interview prep

Loss-and-optimization questions reward the “assumption → gradient → production” walk and a couple of memorized facts (the Hessian, the Adam-vs-SGD generalization result). Lead with intuition, not LaTeX — Chip Huyen’s rubric says strong answers “prioritize intuitive explanations over mathematical equations.” Have these ready.

01“Why cross-entropy not MSE for classification?” → CE is the Bernoulli/Categorical NLL; MSE+sigmoid gives vanishing gradients on confident mistakes and is non-convex in the logits.
02“Why is logistic loss convex?” → Hessian is XᵀWX with W=diag[p(1−p)]⪰0, so PSD; full-rank X ⇒ unique minimum.
03“Why the minus sign in cross-entropy?” → likelihoods ≤ 1 ⇒ logs ≤ 0; the minus makes it a positive quantity to minimize.
04“Batch GD vs SGD?” → batch is stable but scans all data per step; SGD progresses immediately but oscillates, so anneal the learning rate.
05“What does Adam adapt?” → per-coordinate step size via v_t (EMA of g²) plus momentum via m_t; great for sparse/transformer gradients.
06“When would you NOT use Adam?” → dense vision when chasing SOTA — Adam can generalize worse (loss flattening); use SGD+momentum with a schedule.
07“Learning-rate too large vs too small?” → too large diverges/oscillates; too small stalls; pair the right α with warm-up + decay.
08“Softmax vs sigmoid; temperature?” → softmax = one coupled distribution (mutually exclusive); sigmoid = independent labels; temperature softens (T>1) or sharpens (T<1) logits.

Follow-ups dig into mechanism: “derive the logistic gradient” (it collapses to (ŷ−y)x), “why does SGD generalize better than full-batch sometimes?” (gradient noise ≈ implicit regularization toward flat minima), “focal loss — what problem?” (down-weights easy examples for imbalance), and the business reframe “the cost of a false negative is 100× a false positive — change the loss” (weighted/cost-sensitive cross-entropy, or directly optimize expected cost). Connecting the loss to business cost — not just statistics — is what Dan Lee’s Amazon product-return prompt is testing.

paperAdam: A Method for Stochastic OptimizationKingma & Ba paperAdam vs. SGD: Closing the generalization gap on image classificationOPT-ML 2021 paperCS229 main notes — GLMs, convexity, SGD (Stanford)Stanford CS229

Checkpoint

A junior trains a binary classifier with a sigmoid output and MSE loss. It learns the easy examples but barely improves on the confident mistakes. What is the cleanest explanation and fix?

AMSE + sigmoid gives a vanishing gradient on confident mistakes (the σ′ term → 0); switch to cross-entropy, whose gradient stays proportional to (ŷ − y)BThe learning rate is too low — raise it 10×CThe model needs more layers to represent the hard examples

Checkpoint

You are fine-tuning a transformer for text classification on a single GPU and want fast, stable convergence with minimal tuning. Which optimizer choice is most defensible, and why?

ASGD with a fixed learning rate, because it always generalizes better than AdamBFull-batch gradient descent, since it is guaranteed to reach the global minimumCAdam (or AdamW), because its per-coordinate scaling handles sparse, ill-conditioned transformer gradients out of the box

Checkpoint

An interviewer asks you to justify, rigorously, why logistic regression has a unique optimum. Best answer?

ABecause gradient descent with a small learning rate always finds itBThe negative log-likelihood has Hessian XᵀWX with W=diag[p(1−p)]⪰0, so it is convex; with full-rank X the minimizer is uniqueCBecause the sigmoid is monotonic, the loss can only have one minimum

Checkpoint

Training loss diverges to NaN within the first few hundred steps of a new run. Most likely cause and first action?

AThe model is underfitting — add capacityBThe data is leaking — audit the featuresCLearning rate too large (steps overshoot the basin) — lower α and/or add warm-up and gradient clipping

Checkpoint

A vision team reports their Adam-trained ResNet trains to lower training loss than the SGD baseline but tests ~1.5% worse. What is the principled read?

AAdam is simply buggy for CNNs — never use itBA known Adam-vs-SGD generalization gap on vision (loss flattening); switch the final run to SGD+momentum with a schedule, or AdamW, to recover test accuracyCThe test set is mislabeled — Adam cannot hurt generalization

Could you justify cross-entropy over MSE mechanically, write the Adam update, and name when SGD beats Adam — with the learning-rate regimes?

New to itGetting thereConfident

Takeaways

The loss encodes a noise model: MSE = Gaussian (real targets), cross-entropy = Bernoulli/Categorical (classes).
Cross-entropy beats MSE for classification: convex (Hessian XᵀWX ⪰ 0), and its gradient (ŷ−y) does not vanish on confident mistakes.
Batch GD is stable but slow per step; SGD oscillates (anneal α) and its noise acts as implicit regularization.
Adam adapts per-coordinate step sizes (default for transformers); on dense vision it can generalize worse — use SGD+momentum for SOTA.
Learning rate: too large diverges, too small stalls; warm-up then decay, and rule out underfitting before lowering α.

Next: regularization & calibration — L1/L2/elastic-net, dropout, early stopping, and why probabilities miscalibrate.

Sources

Free to read · better with Enzo

Learn it with Enzo

Save your progress, answer the checkpoints, and let Enzo quiz you on what you just read.

Continue with Enzo Sign up free