Stochastic Gradient Descent Variants

Stochastic gradient descent (SGD) variants are optimization algorithms that improve upon basic SGD by adjusting update steps using momentum, adaptive learning rates, or second-order information. They are essential for training modern machine learning models, including deep neural networks.

Stochastic gradient descent (SGD) is an iterative optimization method that approximates gradient descent by using a randomly selected subset of data to estimate the gradient. Variants of SGD have been developed to address its limitations, such as slow convergence and sensitivity to learning rate. These variants include momentum, Nesterov accelerated gradient, AdaGrad, RMSProp, and modern adaptive methods like Adam, which are widely used in Machine learning and Deep learning to train models efficiently.

The core idea of SGD dates back to the Robbins–Monro algorithm of the 1950s, which introduced stochastic approximation for root-finding. In machine learning, SGD minimizes an objective function that is typically a sum of per-example loss functions. The basic update rule is \( w := w - \eta \nabla Q_i(w) \), where \( \eta \) is the learning rate and \( Q_i \) is the loss for the \( i \)-th sample. While simple, this update can be slow to converge and may oscillate, especially in ravines of the loss landscape. Variants address these issues by modifying the update direction, the learning rate, or both.

Momentum

Momentum is a technique that accelerates SGD by accumulating a velocity vector in the direction of persistent gradients. Introduced by Boris Polyak in 1964, momentum mimics physical inertia: the update at step \( t \) is \( v_t = \mu v_{t-1} - \eta \nabla Q_i(w_t) \) and \( w_{t+1} = w_t + v_t \), where \( \mu \) is the momentum coefficient (often 0.9). This helps the optimizer move faster along consistent directions and dampens oscillations in high-curvature regions. Momentum is particularly effective for training deep networks, as it smooths the noisy gradient estimates.

Nesterov Accelerated Gradient

Nesterov accelerated gradient (NAG) is a variant that adds a lookahead step. Proposed by Yurii Nesterov in 1983, NAG computes the gradient at the projected position \( w_t + \mu v_{t-1} \) rather than at the current position. The update becomes \( v_t = \mu v_{t-1} - \eta \nabla Q_i(w_t + \mu v_{t-1}) \) and \( w_{t+1} = w_t + v_t \). This correction reduces overshooting and provides a more accurate estimate of the future gradient, leading to faster convergence in convex settings. NAG is often used in training neural networks and has been incorporated into many libraries.

AdaGrad

AdaGrad, introduced by John Duchi, Elad Hazan, and Yoram Singer in 2011, adapts the learning rate per parameter based on the historical sum of squared gradients. For each parameter \( w_j \), the update is \( w_j := w_j - \frac{\eta}{\sqrt{G_{j,j} + \epsilon}} \nabla Q_i(w_j) \), where \( G_{j,j} \) accumulates the squared gradients and \( \epsilon \) is a small constant for numerical stability. AdaGrad works well for sparse data, as it gives larger updates to infrequent features. However, the accumulation of squared gradients causes the learning rate to shrink over time, which can halt training prematurely.

RMSProp

RMSProp, proposed by Geoffrey Hinton in his lecture notes in 2012, addresses AdaGrad's diminishing learning rate by using an exponentially decaying average of squared gradients. The update maintains a moving average \( E[g^2]_t = \rho E[g^2]_{t-1} + (1-\rho) g_t^2 \), where \( \rho \) is the decay rate (typically 0.9). The parameter update is \( w := w - \frac{\eta}{\sqrt{E[g^2]_t + \epsilon}} g_t \). RMSProp is effective in non-convex settings and has been widely used in training recurrent networks and deep learning models.

Adam

Adam (Adaptive Moment Estimation), introduced by Diederik Kingma and Jimmy Ba in 2015, combines momentum and RMSProp. It maintains both a first moment (mean) and a second moment (variance) of gradients, with bias correction for early steps. The update is \( m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t \), \( v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 \), and \( \hat{m}_t = m_t / (1-\beta_1^t) \), \( \hat{v}_t = v_t / (1-\beta_2^t) \). The parameter update is \( w := w - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon} \). Adam has become a default optimizer for many Deep learning tasks due to its robustness and fast convergence. Variants like AdamW, which decouples weight decay, and AMSGrad, which addresses convergence issues, have also been developed.

Modern Adaptive Methods

Beyond Adam, several adaptive methods have been proposed. AdaBelief (2020) adjusts the step size based on the belief in the current gradient direction. RAdam (Rectified Adam) introduces a rectifier to stabilize the early training phase. Lion (Evolved Sign Momentum), discovered by Google Brain in 2023, uses sign operations to reduce memory usage and has shown competitive performance. These methods are often used in training large language models and other large-scale systems, where efficiency and stability are critical.

Practical Considerations

Choosing the right SGD variant depends on the problem. For convex problems, NAG often provides theoretical guarantees. For deep networks, Adam or RMSProp are common starting points. Learning rate scheduling, such as warmup and decay, is often combined with these optimizers. Mini-batch size also affects performance; larger batches provide smoother gradients but require more memory. In distributed training, variants like LARS (Layer-wise Adaptive Rate Scaling) and LAMB (Layer-wise Adaptive Moments) are used to scale to large batches, as seen in systems like AWS Trainium and Google Cloud.

Impact on Machine Learning

SGD variants have been instrumental in the success of modern Artificial intelligence. They enable training of deep networks with millions of parameters on massive datasets, as done by organizations like OpenAI, Google DeepMind, and Anthropic. The choice of optimizer can significantly affect model accuracy and training speed. Research continues to refine these methods, with new variants emerging regularly. Understanding their properties is essential for practitioners in Machine learning and related fields.