Combining Adam and its Inverse Counterpart to Enhance Generalization of Deep Learning Optimizers

TL;DR

Dual Adam (DualAdam), which integrates the update mechanisms of both Adam and InvAdam, ensuring convergence while enhancing generalization performance is proposed, and the diffusion theory is introduced to mathematically demonstrate InvAdam's ability to escape sharp minima.

Abstract

In the training of neural networks, adaptive moment estimation (Adam) typically converges fast but exhibits suboptimal generalization performance. A widely accepted explanation for its defect in generalization is that it often tends to converge to sharp minima. To enhance its ability to find flat minima, we propose its new variant named inverse Adam (InvAdam). The key improvement of InvAdam lies in its parameter update mechanism, which is opposite to that of Adam. Specifically, it computes element-wise multiplication of the first-order and second-order moments, while Adam computes the element-wise division of these two moments. This modification aims to increase the step size of the parameter update when the elements in the second-order moments are large and vice versa, which helps the parameter escape sharp minima and stay at flat ones. However, InvAdam's update mechanism may face challenges in convergence. To address this challenge, we propose dual Adam (DualAdam), which integrates the update mechanisms of both Adam and InvAdam, ensuring convergence while enhancing generalization performance. Additionally, we introduce the diffusion theory to mathematically demonstrate InvAdam's ability to escape sharp minima. Extensive experiments are conducted on image classification tasks and large language model (LLM) fine-tuning. The results validate that DualAdam outperforms Adam and its state-of-the-art variants in terms of generalization performance. The code is publicly available at https://github.com/LongJin-lab/DualAdam.

Figures (6)

• Figure 1: Mian idea of DualAdam. (a) The relationship between sharp minimum $\phi$ and flat one $\psi$. $\chi$ is the saddle point between $\phi$ and $\psi$; (b) The update mechanisms of Adam and InvAdam. $\boldsymbol{\hat{m}}$ and $\boldsymbol{\hat{v}}$ are the bias-corrected first- and second-order moments, respectively. Subscript $i$ represents the $i$-th element in the vector. $\Delta \boldsymbol{\theta}_i$ is the $i$-th element in parameter update in one iteration; and (c) The update mechanism of DualAdam.
• Figure 2: Optimization trajectories of InvAdam and Adam on 2-parameter loss landscapes. The red stars represent the start points, and the black circles represent the end points.
• Figure 3: Test accuracies over epochs on CIFAR-100.
• Figure 4: Comparisons of training loss, validation perplexity, and generalization gap of DualAdam and AdamW on the fine-tuning of OpenPangu-1B.
• Figure 5: Comparisons of top Hessian matrix's eigenvalues, Hessian matrix's traces, and Hessian matrix's eigenvalue densities of loss landscapes on the CIFAR-100 dataset using ResNet18.

Theorems & Definitions (1)

• Theorem 1