Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Training Deep Learning Models with Norm-Constrained LMOs

Thomas Pethick, Wanyun Xie, Kimon Antonakopoulos, Zhenyu Zhu, Antonio Silveti-Falls, Volkan Cevher

TL;DR

The paper introduces a norm-aware optimization framework for training deep networks using an LMO over norm-balls. It unifies several non-Euclidean methods under the stochastic Conditional Gradient umbrella, proposes Constrained and Unconstrained Scion variants, and provides theoretical convergence guarantees with width-invariant hyperparameters. The main practical payoff is Scion’s memory efficiency and transferability of learning rates across model sizes, demonstrated by substantial speedups on NanoGPT and successful scaling to GPT-3B and Vision tasks. The work also introduces ALMOND as a related LMO-based method for large-batch training and develops thorough analyses to support its empirical effectiveness.

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 using our algorithm, Scion, 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. The code is available at https://github.com/LIONS-EPFL/scion .

Training Deep Learning Models with Norm-Constrained LMOs

TL;DR

The paper introduces a norm-aware optimization framework for training deep networks using an LMO over norm-balls. It unifies several non-Euclidean methods under the stochastic Conditional Gradient umbrella, proposes Constrained and Unconstrained Scion variants, and provides theoretical convergence guarantees with width-invariant hyperparameters. The main practical payoff is Scion’s memory efficiency and transferability of learning rates across model sizes, demonstrated by substantial speedups on NanoGPT and successful scaling to GPT-3B and Vision tasks. The work also introduces ALMOND as a related LMO-based method for large-batch training and develops thorough analyses to support its empirical effectiveness.

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 using our algorithm, Scion, 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. The code is available at https://github.com/LIONS-EPFL/scion .

Paper Structure

This paper contains 61 sections, 11 theorems, 126 equations, 12 figures, 11 tables, 4 algorithms.

Table of Contents

  1. Introduction
  2. Contributions
  3. Preliminaries
  4. Relationship between steepest descent and \ref{['eq:uSCG']}
  5. LMO
  6. Sharp operator
  7. Steepest descent
  8. Preliminaries
  9. Method
  10. Input radius scaling
  11. Alternative norm choices
  12. Boundary initialization
  13. Semi-orthogonal
  14. Column-wise normalized Gaussian
  15. Row-wise normalized Gaussian
  16. ...and 46 more sections

Key Result

lemma 2

lem:transfer:lmo Consider an $L$-layer neural network with input dimension $d_0$, hidden layer widths $\{d_1, d_2, \dots, d_{L-1}\}$, and output dimension $d_L$. Training samples $(\mathbf{z}, \mathbf{y})$ are drawn from some distribution $\mathcal{P}$, where $\mathbf{z} \in \mathbb{R}^{d_0}$ and $\

Figures (12)

  • Figure 1: Coordinate check at initialization. Preactivations are not constant with the spectral scaling $\sqrt{\tfrac{d_\mathrm{out}}{d_\mathrm{in}}}$, when $d_\mathrm{in}>d_\mathrm{out}$.
  • Figure 2: Performance on NanoGPT with between 64M and 1B parameters. The optimal learning rate of Scion is invariant to width.
  • Figure 3: Batch size sensitivity on NanoGPT (124M). The generalization of Scion is less sensitive to larger batches.
  • Figure 4: Scion leads to 30% fewer epochs for ViT on ImageNet and $>$40% wallclock speedup due to a larger critical batch size.
  • Figure 5: Hyperparameter transfer for all three last layer choices on MLP.
  • ...and 7 more figures

Theorems & Definitions (27)

  • Remark 4
  • lemma 2: Spectral $\operatorname{lmo}$ for the gradient with respect to $\mathbf{W}^{(\ell)}$
  • proof
  • lemma 3: Width-invariance of the maximal update learning rate
  • proof
  • Remark 5
  • Remark 6
  • Remark 8
  • lemma 4: \ref{['eq:uSCG']} template inequality
  • proof
  • ...and 17 more