Adaptive gradient-based optimizers, particularly Adam, have left their mark in training large-scale deep learning models. The strength of such optimizers is that they exhibit fast convergence while being more robust to hyperparameter choice. However, they often generalize worse than non-adaptive methods. Recent studies have tied this performance gap to flat minima selection: adaptive methods tend to find solutions in sharper basins of the loss landscape, which in turn hurts generalization. To overcome this issue, we propose a new memory-augmented version of Adam that promotes exploration towards flatter minima by using a buffer of critical momentum terms during training. Intuitively, the use of the buffer makes the optimizer overshoot outside the basin of attraction if it is not wide enough. We empirically show that our method improves the performance of several variants of Adam on standard supervised language modelling and image classification tasks.

翻译：自适应梯度优化器，尤其是Adam，已在训练大规模深度学习模型中留下了深刻印记。这类优化器的优势在于收敛速度快，且对超参数选择更具鲁棒性。然而，它们通常比非自适应方法泛化性能更差。近期研究将这一性能差距归因于平坦极小值选择：自适应方法倾向于在损失景观中更尖锐的盆地中寻找解，从而损害泛化能力。为解决这一问题，我们提出了一种新的记忆增强型Adam版本，通过使用关键动量项缓冲区在训练中促进向更平坦极小值的探索。直观上，若盆地宽度不足，缓冲区的使用会使优化器超出吸引域。实验表明，我们的方法在标准监督语言建模和图像分类任务上提升了多种Adam变体的性能。