Optimization Mathematics for Deep Learning
Vanilla gradient descent is only the starting point. This lesson builds the intermediate math behind modern neural-network training: stochastic updates, momentum, Adam with full equations, schedules, regularization, pathological gradients, and a practical recipe you can use tomorrow.
Beyond Vanilla Gradient Descent
In Lesson 3 you met the basic update \(\theta \leftarrow \theta - \eta\nabla J(\theta)\). Deep networks break the assumptions behind that simple rule: the loss is non-convex, gradients are noisy, curvatures differ wildly across coordinates, and depth amplifies numerical pathologies. Optimization for deep learning is the set of algorithms and tricks that make training reliable despite all of that.
Convexity & Non-Convex Loss Landscapes
A function \(f\) is convex if for any \(x,y\) and \(\alpha\in[0,1]\):
\[f\big(\alpha x + (1-\alpha)y\big) \le \alpha f(x) + (1-\alpha)f(y)\]
Equivalently (when twice differentiable): the Hessian \(H = \nabla^2 f\) is positive semi-definite everywhere.
Convex problems have a unique global minimum that local methods can reach. Deep-network loss surfaces are non-convex: saddles, flat plateaus, and many near-equivalent minima. Empirically, SGD still finds solutions that generalize well — the landscape is “benign enough,” not geometrically simple.
What you actually face in practice
- Saddle points — \(\nabla J \approx 0\) but some eigenvalues of \(H\) are positive and some negative
- Sharp vs flat minima — flat basins often correlate with better generalization
- Plateaus — tiny gradients for long stretches (e.g., saturated sigmoids)
\(f(x)=x^2\) is convex (\(f''=2>0\)). \(f(x)=x^4 - x^2\) has two local minima at \(x=\pm 1/\sqrt{2}\) and a local maximum at \(0\) — a tiny non-convex toy of what networks face in high dimension.
Do not expect a unique “true” minimum. Many parameter vectors achieve nearly the same training loss; SGD’s noise and implicit biases often land in basins that transfer better to unseen data than a sharp full-batch minimum would.
SGD & Mini-Batches
Stochastic gradient descent (SGD) replaces the full gradient with an estimate from one sample or a small batch \(B\):
\[g_t = \dfrac{1}{|B|}\sum_{i\in B}\nabla_\theta L\big(f(x^{(i)};\theta), y^{(i)}\big)\]
\[\theta_{t+1} = \theta_t - \eta\, g_t\]
Why mini-batches win
| Batch size | Gradient quality | Steps / epoch | Typical use |
|---|---|---|---|
| 1 (true SGD) | Very noisy | Many | Online / streaming |
| 32–256 | Balanced | Moderate | Default for many nets |
| Full batch | Exact | One | Small convex problems |
Noise is not only a bug: mild stochasticity can help escape sharp saddles and improve generalization. Very large batches often need a larger learning rate (linear scaling rule) and more careful warmup.
Frameworks report “steps,” not “epochs,” for optimizers. When you change batch size, reinterpret \(\eta\) and schedule milestones in steps, or keep the same number of epochs and scale \(\eta\) carefully.
Linear scaling rule (intuition)
If you multiply batch size by \(k\) and keep the same number of epochs, you take \(k\times\) fewer steps. A common heuristic is \(\eta \leftarrow k\eta\) (with warmup), so the expected update magnitude per epoch stays comparable. It is a starting guess, not a theorem — always validate.
Dataset size \(N=50{,}000\), batch \(64\) → about \(781\) steps/epoch. Switching to batch \(256\) (\(k=4\)) yields \(\approx 195\) steps/epoch. Try \(\eta_{\mathrm{new}} = 4\eta_{\mathrm{old}}\) after a short warmup, then check that training loss curves (vs epochs) still look healthy.
Momentum & Nesterov
Momentum accumulates a velocity vector so updates continue in consistently useful directions and damp oscillations in high-curvature axes.
\[v_{t+1} = \mu\, v_t - \eta\, g_t\]
\[\theta_{t+1} = \theta_t + v_{t+1}\]
Typical \(\mu \in [0.9, 0.99]\). Equivalently: \(v \leftarrow \mu v + g\), then \(\theta \leftarrow \theta - \eta v\).
Nesterov Accelerated Gradient (NAG)
Nesterov evaluates the gradient at a lookahead point \(\theta + \mu v\), which often reduces overshoot:
\[v_{t+1} = \mu\, v_t - \eta\, \nabla J(\theta_t + \mu v_t)\]
\[\theta_{t+1} = \theta_t + v_{t+1}\]
Suppose successive gradients alternate \(+10\) and \(-9\) on one axis. Plain SGD chatters. With \(\mu=0.9\), the velocity averages the cancelling components and progresses along the consistent direction.
AdaGrad, RMSProp & Adam
Adaptive methods give each parameter its own effective learning rate based on gradient history — crucial when sparse features or layers have very different scales.
AdaGrad
\[G_t = G_{t-1} + g_t \odot g_t\]
\[\theta_{t+1} = \theta_t - \dfrac{\eta}{\sqrt{G_t}+\varepsilon}\, g_t\]
Good for sparse gradients, but \(G_t\) only grows — the rate can shrink toward zero too early.
RMSProp
\[v_t = \rho\, v_{t-1} + (1-\rho)\, g_t \odot g_t\]
\[\theta_{t+1} = \theta_t - \dfrac{\eta}{\sqrt{v_t}+\varepsilon}\, g_t\]
Exponential moving average of squared gradients keeps the denominator fresh.
Adam (full equations)
Adam (Adaptive Moment Estimation) tracks first and second moments with bias correction:
\[m_t = \beta_1 m_{t-1} + (1-\beta_1)\, g_t\]
\[v_t = \beta_2 v_{t-1} + (1-\beta_2)\, g_t \odot g_t\]
\[\hat{m}_t = \dfrac{m_t}{1-\beta_1^t}, \qquad \hat{v}_t = \dfrac{v_t}{1-\beta_2^t}\]
\[\theta_{t+1} = \theta_t - \eta\, \dfrac{\hat{m}_t}{\sqrt{\hat{v}_t}+\varepsilon}\]
Defaults: \(\beta_1=0.9\), \(\beta_2=0.999\), \(\varepsilon=10^{-8}\), often \(\eta=10^{-3}\).
Most modern training defaults to Adam or AdamW. For large-scale vision or language models, carefully tuned SGD+momentum or AdamW with cosine schedules remain competitive. Always verify whether your library’s “weight_decay” is L2-in-loss or decoupled (AdamW).
At \(t=1\), \(g_1=0.2\), \(\beta_1=0.9\), \(\beta_2=0.999\): \(m_1=0.02\), \(v_1=4\cdot10^{-5}\). Bias-corrected \(\hat{m}_1=0.2\), \(\hat{v}_1=0.04\). Update magnitude \(\approx \eta\cdot 0.2/\sqrt{0.04}=\eta\) — without correction, early steps would be far too small.
Quick comparison
| Method | Uses \(m_t\) | Uses \(v_t\) | Bias correction | Notes |
|---|---|---|---|---|
| SGD | — | — | — | Needs careful \(\eta\) |
| Momentum | velocity | — | — | Fixed \(\eta\) |
| AdaGrad | — | cumulative \(g^2\) | — | Can stall late |
| RMSProp | — | EMA of \(g^2\) | — | Fresh scale |
| Adam | EMA of \(g\) | EMA of \(g^2\) | yes | Strong default |
| AdamW | same | same | yes | Decoupled decay |
Learning Rate Schedules
A fixed \(\eta\) is rarely optimal for the whole run. Schedules shrink (or briefly grow) the step size as training progresses.
| Schedule | Form (sketch) | When |
|---|---|---|
| Step decay | \(\eta \leftarrow \gamma\eta\) every \(K\) epochs | Classic CNNs |
| Exponential | \(\eta_t = \eta_0 e^{-kt}\) | Smooth decay |
| Cosine annealing | \(\eta_t = \eta_{\min} + \tfrac{1}{2}(\eta_{\max}-\eta_{\min})(1+\cos\pi t/T)\) | Modern default |
| Linear warmup | Ramp \(\eta\) from near 0 for first \(W\) steps | Transformers, large batches |
| 1-cycle / CLR | Rise then fall within a cycle | Fast training |
\[\eta_t = \begin{cases} \eta_{\max}\,\dfrac{t}{W} & t < W \\[6pt] \eta_{\min} + \dfrac{1}{2}(\eta_{\max}-\eta_{\min})\big(1+\cos\dfrac{\pi(t-W)}{T-W}\big) & t \ge W \end{cases}\]
If loss spikes after a decay, the drop was too aggressive. If validation plateaus for many epochs with a high \(\eta\), decay earlier. Learning-rate finder sweeps (LR range test) give a good initial \(\eta_{\max}\).
L1 / L2 Regularization & Weight Decay
Regularization penalizes complexity so the optimizer prefers simpler weights that often generalize better.
\[J(\theta) = \dfrac{1}{m}\sum_{i=1}^{m} L^{(i)} + \lambda R(\theta)\]
L2: \(R(\theta)=\tfrac{1}{2}\|\theta\|_2^2\). L1: \(R(\theta)=\|\theta\|_1\).
SGD equivalence of L2 and weight decay
\[\theta \leftarrow (1 - \eta\lambda)\,\theta - \eta\, g\]
For plain SGD this matches adding \(\lambda\theta\) to the gradient of L2. With Adam, prefer AdamW: apply \((1-\eta\lambda)\theta\) separately from the adaptive \(\hat{m}/\sqrt{\hat{v}}\) step.
In PyTorch, weight_decay in Adam is historically coupled; use torch.optim.AdamW for decoupled decay — it often improves validation accuracy at the same \(\lambda\).
Vanishing & Exploding Gradients
Backpropagation multiplies Jacobians layer by layer. If those factors have singular values \(\ll 1\), gradients vanish; if \(\gg 1\), they explode.
\[\dfrac{\partial L}{\partial \theta^{(\ell)}} = \dfrac{\partial L}{\partial h^{(L)}} \dfrac{\partial h^{(L)}}{\partial h^{(L-1)}} \cdots \dfrac{\partial h^{(\ell+1)}}{\partial h^{(\ell)}} \dfrac{\partial h^{(\ell)}}{\partial \theta^{(\ell)}}\]
Mitigations
- Activation choice — ReLU / GELU avoid the flat tails of sigmoid/tanh
- Initialization — Xavier / He scale variance with fan-in/fan-out
- Residual connections — identity paths keep gradients flowing
- Gradient clipping — if \(\|g\| > c\), set \(g \leftarrow c\, g/\|g\|\)
- Normalization layers — Batch/Layer Norm stabilize activations
Clip at \(c=1\): if \(\|g\|=5\), replace \(g\) with \(g/5\). Training continues with a bounded step — standard for RNNs and unstable Transformer runs.
Hessian & Second-Order Intuition
The Hessian \(H=\nabla^2 J\) describes local curvature. Eigenvalues tell you which directions are steep or flat; Newton’s method uses \(H^{-1}\) to rescale steps:
\[\theta_{t+1} = \theta_t - H^{-1}\nabla J(\theta_t)\]
Full Hessians are \(O(d^2)\) in memory for \(d\) parameters — impossible for modern nets. Still, the intuition matters:
- Adaptive methods approximate a diagonal preconditioner related to curvature
- Large positive eigenvalues → need small steps along those axes
- Near-zero eigenvalues → flat directions; progress is slow without preconditioning
- Negative eigenvalues → local saddles / maxima along those axes
K-FAC, Shampoo, and Sophia are practical second-order / quasi-second-order methods for large models. For most engineers, understanding that Adam is a cheap diagonal preconditioner is enough to tune \(\beta_2\) and \(\varepsilon\) wisely.
Early Stopping & Generalization
Optimization drives training loss down; generalization asks whether validation/test loss follows. Early stopping monitors a held-out metric and stops (or restores the best checkpoint) when it stops improving.
Stop if validation loss has not improved by more than \(\delta\) for \(P\) consecutive evaluations. Restore \(\theta\) from the best checkpoint.
What the curves usually look like
- Training loss keeps falling
- Validation loss falls, then rises (overfitting) or plateaus
- The “best” model for deployment is often not the final iterate
Patience \(P=10\), evaluate each epoch. Best val loss \(0.42\) at epoch 47; epochs 48–57 never beat \(0.42-\delta\). Restore epoch-47 weights and stop — even if training loss at epoch 57 is lower.
Batch Normalization — Math Briefly
Batch Norm re-centers and re-scales activations per mini-batch, then learns an affine transform. It stabilizes the distribution of inputs to each layer and often allows higher learning rates.
\[\mu_B = \dfrac{1}{m}\sum_{i=1}^{m} x_i, \qquad \sigma_B^2 = \dfrac{1}{m}\sum_{i=1}^{m}(x_i-\mu_B)^2\]
\[\hat{x}_i = \dfrac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \varepsilon}}, \qquad y_i = \gamma\hat{x}_i + \beta\]
At inference, replace \(\mu_B,\sigma_B^2\) with running averages collected during training.
Layer Norm / RMSNorm dominate Transformers (batch-size-independent). Batch Norm remains common in CNNs. Mixing train/eval modes incorrectly (e.g., BN in train mode at test time with batch size 1) silently destroys accuracy.
Initialization & Numerical Stability
Even the best optimizer fails if initial activations explode or die. Variance-preserving initializations keep forward and backward signals healthy.
\[W_{ij} \sim \mathcal{N}\!\left(0, \dfrac{2}{n_{\mathrm{in}}}\right) \quad\text{or}\quad \mathrm{Uniform}\!\left(-\sqrt{\tfrac{6}{n_{\mathrm{in}}}}, \sqrt{\tfrac{6}{n_{\mathrm{in}}}}\right)\]
Xavier/Glorot uses \(2/(n_{\mathrm{in}}+n_{\mathrm{out}})\) — better matched to tanh/sigmoid.
Stability checklist
- Watch for NaN/Inf in loss — lower \(\eta\), clip gradients, check labels
- Log gradient norms per layer — sudden spikes precede explosions
- Prefer mixed precision with loss scaling when using FP16/BF16
- Keep \(\varepsilon\) in Adam / Batch Norm from being pathologically tiny on low-precision hardware
Layer with \(n_{\mathrm{in}}=512\), ReLU: He std \(=\sqrt{2/512}\approx 0.0625\). Initializing with std \(0.5\) would inflate activations and likely explode gradients within a few steps.
Practical Training Recipe
A battle-tested starting recipe for many intermediate deep-learning problems:
| Choice | Default starting point |
|---|---|
| Optimizer | AdamW (\(\beta_1=0.9\), \(\beta_2=0.999\)) |
| Learning rate | \(10^{-3}\) (CNNs) or \(10^{-4}\)–\(5\cdot10^{-4}\) (Transformers) |
| Schedule | Linear warmup + cosine decay |
| Weight decay | \(0.01\) (decoupled); often \(0\) on biases / Norm scales |
| Batch size | Largest power of 2 that fits; scale \(\eta\) if you change it a lot |
| Regularization | Augmentation + light dropout; early stopping on val |
| Gradients | Clip global norm at 1.0 if unstable |
| Checkpointing | Save best validation metric, not only last epoch |
Debugging training that “won’t learn”
- Loss stuck at chance level → check labels, loss reduction, and learning rate (try \(10\times\) up and down)
- Loss \(\to\) NaN → lower \(\eta\), enable clipping, verify input normalization
- Train loss falls, val never does → more aug / weight decay / earlier stop; check leakage
- GPU underused → raise batch size or use mixed precision; profile data loading
Shuffle on, correct train/eval modes for Dropout/BatchNorm, labels in the right index range, loss averaged not summed inconsistently, and a baseline that beats a constant predictor before you trust fancy schedules.
Next lesson — matrix calculus & backpropagation — explains how \(g_t\) is actually computed. This lesson told you what to do with \(g_t\); that one shows where \(g_t\) comes from.
Summary
Optimization for deep learning upgrades vanilla gradient descent into a toolbox for noisy, non-convex, high-dimensional training.
\[m_t=\beta_1 m_{t-1}+(1-\beta_1)g_t,\quad v_t=\beta_2 v_{t-1}+(1-\beta_2)g_t^2\]
\[\hat{m}_t=\dfrac{m_t}{1-\beta_1^t},\quad \hat{v}_t=\dfrac{v_t}{1-\beta_2^t},\quad \theta\leftarrow\theta-\eta\dfrac{\hat{m}_t}{\sqrt{\hat{v}_t}+\varepsilon}\]
Revisit Lesson 4 for why cross-entropy is the right loss for many classifiers, then continue to Lesson 6 to derive gradients systematically. With both, you can read optimizer papers and training code with confidence.
Train with AdamW + warmup/cosine, regularize with weight decay and early stopping, and treat exploding/vanishing gradients as first-class numerical design constraints — not afterthoughts.