Gradient Descent Optimizers Interview Preparation Guide

While SGD with Momentum maintains a single exponential moving average of past gradients (first moment) and applies a uniform learning rate across all parameters, Adam tracks both the first moment and uncentered second moment (variance of squared gradients). This allows Adam to adaptively scale step sizes for individual parameters based on their historical gradient magnitudes, whereas SGD with Momentum applies a globally scaled update vector adjusted only by directional inertia.

Standard Adam coupled with L2 regularization scales the regularization penalty inversely through the historical squared gradient buffer in the denominator. This means parameters with frequent, small gradients receive disproportionately harsher regularization penalties than parameters with infrequent, large gradients. AdamW decouples weight decay entirely from the adaptive gradient calculation, applying the decay penalty directly to the weights as a simple multiplicative factor, ensuring uniform and correct L2 regularization.

During the initial phase of pre-training, randomly initialized weights produce high-variance, erratic gradients. If a large learning rate is applied immediately, these volatile gradients can cause oversized parameter updates that destabilize early representations and lead to permanent loss divergence or NaN values. A learning rate warmup gradually ramps the step size from near-zero up to the peak learning rate over several thousand steps, allowing the model to stabilize its internal feature representations safely.

Standard optimizers like Adam store additional tensor states for every trainable parameter in the model. Specifically, Adam maintains a first-moment buffer and a second-moment buffer, each stored in 32-bit floating-point precision regardless of whether the model weights are in FP16 or BF16. For a 70-billion parameter model, this requires tens of gigabytes of dedicated GPU memory solely for optimizer states, necessitating memory-sharding techniques like ZeRO or 8-bit optimizer quantization.

While Adam often converges faster in early training epochs, SGD with momentum combined with a well-tuned cosine annealing schedule frequently achieves superior generalization performance on test data for specific computer vision and image classification tasks. Additionally, SGD has a much smaller memory footprint (requiring only a single momentum buffer per parameter), making it advantageous when GPU memory is severely constrained.

Because the exponential moving averages for both first and second moments are initialized at zero, they are biased toward zero during the initial training steps, especially when decay coefficients like beta1 and beta2 are close to 1.0. The bias correction terms divide the moment estimates by (1 - beta^t), effectively scaling up the update vectors in early steps and automatically fading out to 1.0 as the step count t increases.

8-bit optimizers like those provided by BitsAndBytes compress the optimizer state buffers (such as Adam's momentum and variance tensors) from 32-bit floating-point down to 8-bit integer formats using quantile quantization. Because optimizer states require high dynamic range but can tolerate lower precision representation without destabilizing gradient accumulation, this quantization slashes optimizer state memory consumption by roughly 75% (a 4x reduction versus 32-bit) while preserving end-to-end model convergence.

Gradient accumulation is a technique used to simulate large training batch sizes when physical GPU memory cannot accommodate large micro-batches. Instead of calling optimizer.step() after every backward pass, the code accumulates computed gradients across multiple micro-batches by summing them up (while dividing the loss appropriately). The optimizer step and zero_grad calls are only triggered after the desired accumulation boundary is reached, ensuring updates reflect the aggregate batch.

Deep architectures with recurrent connections or deep self-attention layers are highly susceptible to exploding gradients during backpropagation through time or deep chains of matrix multiplications. Gradient clipping inspects the global norm of all gradient tensors and scales them down proportionally if the norm exceeds a predefined threshold (e.g., max_norm=1.0), preventing numerical overflow and sudden NaN loss spikes.

Linear scaling rules dictate that when multiplying the batch size by k, the learning rate should also be multiplied by k. However, empirical studies show that gradient noise variance grows sub-linearly with batch size in large models. Consequently, square-root scaling (scaling the learning rate by the square root of k) is frequently deployed in massive distributed pre-training runs to maintain stable convergence without requiring excessive learning rate reductions.

Standard optimizer implementations execute multiple discrete element-wise tensor operations in Python, repeatedly loading and storing parameter, gradient, and moment tensors to and from global GPU memory. Fused kernels combine moment updates, weight decay, and parameter adjustments into a single monolithic CUDA or Triton kernel executed directly in SRAM cache, drastically reducing memory bandwidth bottlenecks and accelerating step execution times.

When resuming training, both the model parameter weights and the exact optimizer state buffers (momentum tensors, variance tensors, and global step counters) must be restored from the checkpoint file. If the optimizer state buffers are omitted or initialized to zero, the moving averages lose their historical inertia, causing sudden convergence instability, step size jumps, or temporary loss spikes upon restart.