TLDR: This research introduces a two-phase training algorithm for deep neural networks that leverages the changing convexity of loss functions. It starts with the Adam optimizer for initial non-convex regions and switches to the Conjugate Gradient (CG) method once the loss function becomes convex near the optimum. The transition point is detected by monitoring the gradient norm. Experiments on various Vision Transformer and VGG5 architectures demonstrate that this adaptive approach significantly improves convergence speed and accuracy compared to using Adam alone.
Training deep neural networks is a fundamental task in machine learning, aiming to minimize a loss function that measures how well a model fits its training data. The efficiency of this process heavily relies on the characteristics of the loss function, particularly whether it is convex or non-convex.
Traditionally, many deep learning models use optimization methods like Adam, which are well-suited for non-convex loss functions. This is because loss functions in deep learning often start in complex, non-convex regions with multiple potential minima. However, a crucial insight is that around any local minimum, the loss function inherently becomes convex. In these convex environments, second-order optimization methods, such as the Conjugate Gradient (CG) algorithm, offer guaranteed superlinear convergence, meaning they can find the minimum much faster.
Researchers Tomas Hrycej, Bernhard Bermeitinger, Massimo Pavone, G¨otz-Henrik Wiegand, and Siegfried Handschuh have proposed an innovative two-phase training algorithm that capitalizes on this property. Their core hypothesis is that loss functions in real-world tasks typically transition from an initial non-convex state to a convex state as they approach the optimal solution.
The Two-Phase Approach
The proposed algorithm works in two distinct phases:
- Phase 1 (Non-convex region): In the initial stages of training, when the model parameters are far from the optimum and the loss function is likely non-convex, the algorithm uses Adam. Adam is effective at navigating these complex landscapes.
- Phase 2 (Convex region): As the training progresses and the model approaches a local minimum, the loss function is hypothesized to become convex. At this ‘swap point’, the algorithm switches to the Conjugate Gradient (CG) method, which then efficiently guides the model to the minimum with accelerated convergence.
The key to this approach is accurately identifying the ‘swap point’ – the moment the loss function transitions from non-convex to convex. The algorithm detects this by observing the gradient norm’s dependence on the loss. Initially, as the loss decreases, the gradient norm might increase (indicating a non-convex region). When the gradient norm starts to systematically decrease, it signals entry into a convex region, prompting the switch to CG.
Also Read:
- Mini-Batch Adam’s Implicit Bias Diverges from Full-Batch Behavior
- Optimizing Federated Learning in Wireless Networks: A New Approach to Bias and Variance
Empirical Validation
To test their hypothesis and the effectiveness of the two-phase algorithm, the researchers conducted experiments using various model architectures, including reduced variants of the Vision Transformer (ViT) and the convolutional network VGG5. These models were trained on popular datasets like CIFAR-10, CIFAR-100, and MNIST.
The results consistently supported their hypothesis. All tested models exhibited the predicted pattern: the gradient norm initially increased with decreasing loss (non-convex phase) and then decreased after a turning point (convex phase). The two-phase Adam+CG algorithm consistently outperformed traditional Adam-only training in terms of both convergence speed and final accuracy on the training set. Even considering the additional computational steps of CG’s line search, the benefits remained substantial.
This research suggests that by adapting the optimization strategy to the local convexity of the loss function, significant improvements in deep neural network training can be achieved. While the experiments were conducted on relatively smaller models due to resource constraints, the consistent patterns observed across diverse architectures indicate a promising direction for optimizing larger, more complex models in the future. The full research paper can be found here.
