Promoting Exploration in Memory-Augmented Adam using Critical Momenta
Adaptive gradient-based optimizers, notably Adam, have left their mark in training large-scale deep learning models, offering fast convergence and robustness to hyperparameter settings. However, they often struggle with generalization, attributed to their tendency to converge to sharp minima in the loss landscape. To address this, we propose a new memory-augmented version of Adam that encourages exploration towards flatter minima by incorporating a buffer of critical momentum terms during training. This buffer prompts the optimizer to overshoot beyond narrow minima, promoting exploration. Through comprehensive analysis in simple settings, we illustrate the efficacy of our approach in increasing exploration and bias towards flatter minima. We empirically demonstrate that it can improve model performance for image classification on ImageNet and CIFAR10/100, language modelling on Penn Treebank, and online learning tasks on TinyImageNet and 5-dataset. Our code is available at \url{https://github.com/chandar-lab/CMOptimizer}.
Code (1)
Tasks
image-classificationImage ClassificationLanguage ModellingMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
AdaMemento: Adaptive Memory-Assisted Policy Optimization for Reinforcement Learning
In sparse reward scenarios of reinforcement learning (RL), the memory mechanism provides promising shortcuts to policy optimization by reflecting on past experiences like humans. However, current memory-based RL methods …
Ensemble Learningreinforcement-learningReinforcement LearningReinforcement Learning (RL)ADAM: A Systematic Data Extraction Attack on Agent Memory via Adaptive Querying
Large Language Model (LLM) agents have achieved rapid adoption and demonstrated remarkable capabilities across a wide range of applications. To improve reasoning and task execution, modern LLM agents would incorporate me…
Do We Need Adam? Surprisingly Strong and Sparse Reinforcement Learning with SGD in LLMs
Reinforcement learning (RL), particularly RL from verifiable reward (RLVR), has become a crucial phase of training large language models (LLMs) and a key focus of current scaling efforts. However, optimization practices …
Reinforcement LearningStable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning
Effective decision-making in partially observable environments demands robust memory management. Despite their success in supervised learning, current deep-learning memory models struggle in reinforcement learning enviro…
Decision MakingManagementMeta Reinforcement Learningreinforcement-learning+1APOLLO: SGD-like Memory, AdamW-level Performance
Large language models (LLMs) are notoriously memory-intensive during training, particularly with the popular AdamW optimizer. This memory burden necessitates using more or higher-end GPUs or reducing batch sizes, limitin…
GPUQuantization