Training Deep Learning Models with Norm-Constrained LMOs
Entraînement des Modèles d'Apprentissage Profond avec des LMO à Norme Contrainte
Abstract
In this work, we study optimization methods that leverage the linear minimization oracle (LMO) over a norm-ball. We propose a new stochastic family of algorithms that uses the LMO to adapt to the geometry of the problem and, perhaps surprisingly, show that they can be applied to unconstrained problems. The resulting update rule unifies several existing optimization methods under a single framework. Furthermore, we propose an explicit choice of norm for deep architectures, which, as a side benefit, leads to the transferability of hyperparameters across model sizes. Experimentally, we demonstrate significant speedups on nanoGPT training without any reliance on Adam. The proposed method is memory-efficient, requiring only one set of model weights and one set of gradients, which can be stored in half-precision.
Dans ce travail, nous étudions des méthodes d'optimisation qui exploitent l'oracle de minimisation linéaire (LMO) sur une boule normée. Nous proposons une nouvelle famille d'algorithmes stochastiques qui utilise le lmo pour s'adapter à la géométrie du problème et, de manière surprenante, démontrons qu'ils peuvent être appliqués à des problèmes non contraints. La mise à jour qui en résulte unifie plusieurs méthodes d'optimisation existantes sous un même cadre. De plus, nous proposons un choix explicite de norme pour les architectures profondes qui, comme avantage secondaire, permet la transférabilité des hyperparamètres entre différentes tailles de modèles. Expérimentalement, nous démontrons des accélérations significatives sur l'entraînement de nanoGPT sans aucune dépendance à Adam. La méthode proposée est efficace en mémoire, ne nécessitant qu'un ensemble de poids du modèle et un ensemble de gradients, qui peuvent être stockés en demi-précision.
| Origin | Files produced by the author(s) |
|---|