Understanding Gradient Descent and Optimization Landscapes

1. Problem Formulation and Objective Function

In parametric machine learning models, training consists of finding parameter values $\theta \in \mathbb{R}^d$ that minimize an empirical risk function over a dataset of $N$ observations:

$$L(\theta) = \frac{1}{N} \sum_{i=1}^{N} \ell(f(x_i; \theta), y_i)$$

Where:

• $f(x_i; \theta)$ denotes the model hypothesis parameterized by $\theta$.

• $\ell(\hat{y}, y)$ represents the per-sample loss function (e.g., squared error or cross-entropy).

• $L(\theta)$ defines the scalar loss surface over the parameter space.

2. First-Order Optimization: Standard Gradient Descent

Standard gradient descent updates parameters iteratively by moving in the direction of steepest descent, which corresponds to the negative gradient vector:

$$\nabla_{\theta} L(\theta) = \begin{bmatrix} \frac{\partial L}{\partial \theta_1} \\ \frac{\partial L}{\partial \theta_2} \\ \vdots \\ \frac{\partial L}{\partial \theta_d} \end{bmatrix}$$

The discrete update rule at step $t$ is defined as:

$$\theta_{t+1} = \theta_t - \eta \nabla_{\theta} L(\theta_t)$$

Here, $\eta > 0$ represents the learning rate (step size).

2.1 Taylor Expansion and the Choice of Learning Rate

To understand step size limits, consider a second-order Taylor expansion around the current parameter state $\theta_t$:

$$L(\theta_t + \Delta \theta) \approx L(\theta_t) + \nabla L(\theta_t)^T \Delta \theta + \frac{1}{2} \Delta \theta^T \mathbf{H} \Delta \theta$$

Where $\mathbf{H} = \nabla^2 L(\theta_t)$ is the Hessian matrix of second-order partial derivatives:

$$H_{ij} = \frac{\partial^2 L}{\partial \theta_i \partial \theta_j}$$

Substituting the gradient update $\Delta \theta = -\eta \nabla L(\theta_t)$ yields:

$$L(\theta_{t+1}) \approx L(\theta_t) - \eta \Vert{}\nabla L(\theta_t)\Vert{}^2 + \frac{\eta^2}{2} \nabla L(\theta_t)^T \mathbf{H} \nabla L(\theta_t)$$

To guarantee descent ($L(\theta_{t+1}) < L(\theta_t)$), the step size $\eta$ must satisfy:

$$0 < \eta < \frac{2}{\lambda_{\max}(\mathbf{H})}$$

Where $\lambda_{\max}(\mathbf{H})$ is the largest eigenvalue of the Hessian matrix. If $\eta > \frac{2}{\lambda_{\max}}$, the update overshoots the valley and diverges.

3. Stochastic Gradient Descent (SGD) with Momentum

Calculating $\nabla L(\theta)$ across the entire dataset $N$ becomes computationally prohibitive when $N \gg 10^5$. Stochastic Gradient Descent approximates the expectation using a mini-batch $\mathcal{B} \subset \{1, \dots, N\}$ of size $B = \vert{}\mathcal{B}\vert{}$:

$$g_t = \frac{1}{B} \sum_{i \in \mathcal{B}} \nabla_{\theta} \ell(f(x_i; \theta_t), y_i)$$

Because $g_t$ is an unbiased estimator ($\mathbb{E}[g_t] = \nabla L(\theta_t)$) with high variance, classical SGD exhibits oscillatory behavior in ill-conditioned ravines.

3.1 Polyak Momentum Formulation

Momentum dampens oscillations along high-curvature directions by accumulating an exponentially decaying moving average of past gradients:

$$v_{t+1} = \beta v_t + g_t$$

$$\theta_{t+1} = \theta_t - \eta v_{t+1}$$

Where $\beta \in [0, 1)$ acts as the momentum coefficient (commonly set to $\beta = 0.9$). Unrolling the recurrence reveals the effective velocity vector:

$$v_{t+1} = \sum_{\tau=0}^{t} \beta^{t-\tau} g_\tau$$

This produces an effective acceleration along consistent descent directions by a factor of:

$$\text{Effective Scale} = \frac{1}{1 - \beta}$$

4. Summary Comparison of First-Order Solvers

• Standard SGD

  • Memory per Parameter: $O(1)$

  • Rescaling Invariance: No

  • Hyperparameters: $\eta$

• SGD with Momentum

  • Memory per Parameter: $O(1)$ (stores velocity $v$)

  • Rescaling Invariance: No

  • Hyperparameters: $\eta, \beta$

• Adam

  • Memory per Parameter: $O(2)$ (stores first and second moments $m_t, v_t$)

  • Rescaling Invariance: Yes (coordinate-wise)

  • Hyperparameters: $\eta, \beta_1, \beta_2, \epsilon$

Follow me on LinkedIn