Technical Summary: OCP-GN – A Scalable Second-order Optimizer for Stochastic Optimization

Problem Statement

The training of large-scale neural networks typically relies on general-purpose first-order optimizers such as Stochastic Gradient Descent (SGD) and Adaptive Moment Estimation (Adam). The authors argue that these existing algorithms often fail to effectively improve both the convergence speed and the generalization performance of neural network training. They posit that neural network training can be modeled as the evolution of a high-dimensional dynamical system within a parameter space along a loss landscape. The core challenge addressed is how to systematically guide this evolutionary trajectory to converge to the optimal solution region more quickly, stably, and robustly—a problem the authors identify as critical yet underexplored.

Methodology

The paper proposes OCP-GN, a novel second-order optimization algorithm derived from the Optimal Control Principle (OCP). The methodology bridges the rigorous mathematical framework of OCP with the practical constraints of large-scale stochastic optimization through three primary components:

1. Gauss-Newton-Bartlett (GNB) Estimator

To avoid the prohibitive computational cost of calculating the full Hessian matrix ($O(d^2)$), the authors employ a Gauss-Newton-Bartlett estimator to approximate the diagonal of the Hessian.

• Formulation: For a Mean Squared Error (MSE) loss function, the Hessian is approximated by the outer product of the Jacobian.
• Bartlett's Identity: By resampling the model's predictive distribution to obtain synthetic labels $\hat{y} \sim N(\ell(x), \sigma^2)$, the authors utilize Bartlett's second identity to show that the quantity $\nabla\Psi \odot \nabla\Psi$ serves as an unbiased estimator for the diagonal of the Gauss-Newton matrix.
• Efficiency: This approach reduces the computational complexity of the Hessian approximation to $O(d)$, making it scalable for high-dimensional parameter spaces.

2. Closed-Form OCP Update Rule

The algorithm inherits the update rule from the OCP method but adapts it for stochastic settings:

• Original OCP: The update involves an inner iteration loop $\phi_l(x_k)$ defined by a recurrence relation involving the gradient and Hessian.
• Closed-Form Derivation: The authors derive a closed-form expression for the update term $\phi_l(x_k)$ to separate the learning rate parameter matrix $M$ from the update magnitude.
• Stochastic Adaptation: The full gradient and Hessian are replaced with bias-corrected stochastic estimates ($\hat{g}_k$ and $\hat{H}_k$) using exponential moving averages (similar to Adam), where $H_k$ is approximated as $g_k^2$.

3. Stabilization and Regularization Mechanisms

To ensure stability in large-scale stochastic optimization, the authors introduce several technical enhancements:

• Clipping: Lower and upper bounds are applied to the diagonal Hessian estimates ($\hat{H}_k$) to prevent division by zero and ensure positive definiteness. Additionally, the term $\kappa_k = I - M\hat{H}_k$ is clipped to ensure the update direction strictly follows the gradient correction.
• Inner Loop Control: Instead of a fixed number of inner iterations, the number of inner loop iterations ($b$) is treated as a tunable hyperparameter.
• Weight Decay: A weight decay term is incorporated to prevent overfitting and enhance generalization.

The final update rule combines these elements into a single step involving a clipped, scaled gradient correction term.

Key Contributions

1. Novel Optimizer Design: The introduction of OCP-GN, which adapts the Optimal Control Principle for deep learning by integrating a GNB estimator for Hessian approximation.
2. Scalability: Achieving a computational complexity of $O(d)$ for the Hessian approximation, enabling the application of second-order information to large-scale neural network training.
3. Theoretical Analysis: The paper provides a convergence analysis under specific assumptions (Lipschitz continuity, Polyak–Łojasiewicz condition), proving that the algorithm converges to the optimum at a linear rate.
4. Stability Enhancements: The proposal of specific clipping thresholds and bias-correction mechanisms to ensure robust performance in stochastic environments.

Experimental Results

The authors evaluated OCP-GN on image classification tasks using CIFAR-10 and CIFAR-100 datasets across ResNet-34 and Vision Transformer (ViT) architectures. The method was compared against the baseline optimizer AdamW.

• Performance: OCP-GN demonstrated significant superiority over AdamW in both convergence speed and test performance.
• Specific Metrics:
  • CIFAR-10 (ViT): OCP-GN achieved 87.50% accuracy vs. AdamW's 78.39%.
  • CIFAR-10 (ResNet-34): OCP-GN achieved 94.84% vs. AdamW's 93.67%.
  • CIFAR-100 (ViT): OCP-GN achieved 62.70% vs. AdamW's 51.96%.
  • CIFAR-100 (ResNet-34): OCP-GN achieved 74.22% vs. AdamW's 72.64%.

Significance and Claims

The paper claims that OCP-GN offers a rigorous and powerful mathematical framework for designing optimizers suitable for neural network training. By systematically guiding the evolutionary trajectory of the dynamical system (parameter space), the proposed method addresses the limitations of existing optimizers regarding convergence speed and generalization.

The authors assert that the integration of the Optimal Control Principle with efficient Hessian estimation provides a scalable solution that exhibits strong robustness and significant superiority in convergence performance and generalization capability across various model architectures and dataset scales. The work suggests that second-order information, when properly approximated and stabilized, can be effectively leveraged in stochastic optimization without incurring prohibitive computational costs.