Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

FOCUS: First Order Concentrated Updating Scheme

Yizhou Liu, Ziming Liu, Jeff Gore

TL;DR

The paper investigates why pre-training large language models encounters slow or unstable optimization due to gradient noise in narrowing valley loss landscapes. It proposes FOCUS, a physics-inspired optimizer that augments Signum with an attraction toward moving-average parameters to maintain larger step sizes in noisy, sharp valleys. Through a two-dimensional toy model and GPT-2 small pretraining experiments, FOCUS is shown to be more stable than Signum and faster than Adam, with a convergence analysis providing a regret bound comparable to Adam’s. The work provides both practical speedups for LLM pre-training and a framework for understanding gradient-noise effects, suggesting broader optimizer design principles for noisy, non-convex landscapes.

Abstract

Large language models (LLMs) demonstrate remarkable performance, and improving their pre-training process appears to be key to enhancing their capabilities further. Based on the documented success of Adam, learning rate decay, and weight decay, we hypothesize that the pre-training loss landscape features a narrowing valley structure. Through experiments with synthetic loss functions, we discover that when gradient query noise is high relative to the valley's sharpness, Adam's performance falls behind that of Signum because Adam reduces the effective step size too drastically. This observation led us to develop FOCUS, an optimizer that enhances Signum by incorporating attraction toward moving averaged parameters, allowing it to handle noise better while maintaining larger step sizes. In training GPT-2, FOCUS proves to be more stable than Signum and faster than Adam. These results suggest that gradient noise may be an underappreciated limiting factor in LLM training, and FOCUS offers promising solutions.

FOCUS: First Order Concentrated Updating Scheme

TL;DR

The paper investigates why pre-training large language models encounters slow or unstable optimization due to gradient noise in narrowing valley loss landscapes. It proposes FOCUS, a physics-inspired optimizer that augments Signum with an attraction toward moving-average parameters to maintain larger step sizes in noisy, sharp valleys. Through a two-dimensional toy model and GPT-2 small pretraining experiments, FOCUS is shown to be more stable than Signum and faster than Adam, with a convergence analysis providing a regret bound comparable to Adam’s. The work provides both practical speedups for LLM pre-training and a framework for understanding gradient-noise effects, suggesting broader optimizer design principles for noisy, non-convex landscapes.

Abstract

Large language models (LLMs) demonstrate remarkable performance, and improving their pre-training process appears to be key to enhancing their capabilities further. Based on the documented success of Adam, learning rate decay, and weight decay, we hypothesize that the pre-training loss landscape features a narrowing valley structure. Through experiments with synthetic loss functions, we discover that when gradient query noise is high relative to the valley's sharpness, Adam's performance falls behind that of Signum because Adam reduces the effective step size too drastically. This observation led us to develop FOCUS, an optimizer that enhances Signum by incorporating attraction toward moving averaged parameters, allowing it to handle noise better while maintaining larger step sizes. In training GPT-2, FOCUS proves to be more stable than Signum and faster than Adam. These results suggest that gradient noise may be an underappreciated limiting factor in LLM training, and FOCUS offers promising solutions.
Paper Structure (18 sections, 4 theorems, 21 equations, 8 figures, 1 algorithm)

This paper contains 18 sections, 4 theorems, 21 equations, 8 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Methods
  3. Picture of loss landscapes
  4. Toy model
  5. FOCUS optimizer
  6. Experiments
  7. Setup
  8. Stability
  9. Speedup
  10. Convergence analysis
  11. Related works
  12. Discussion
  13. Model details
  14. Toy model
  15. MLP
  16. ...and 3 more sections

Key Result

Theorem 4.2

Let $\{L_t\}_{t=1}^T$ be a sequence of convex functions with bounded gradients $\|g_t\|_\infty \leq G_\infty$. For the FOCUS optimizer with learning rate $\eta_t = \eta/\sqrt{t}$, $\beta_1, \beta_2 \in [0,1)$, $\beta_{1,t} = \beta_1 \lambda_\beta^t$, $\gamma_t = \gamma \lambda_\gamma^t$,We generaliz

Figures (8)

  • Figure 1: Insights from toy models lead to practical speedup in LLM pre-training. (a) Our philosophy that a simple and self-consistent picture provides a basis of thinking and testable hypotheses is not necessarily correct. Interactions between testing the real world and updating our picture keep us moving forward. (b) Narrowing valleys are assumed to be a key structure in LLM training loss. (c) The toy model explains that Adam can be slow with a large gradient noise (\ref{['sec:method']}). (d) By adding self-attraction to Signum, FOCUS is proposed to handle sharp landscapes with large gradient stochasticity and achieve practical speedups (\ref{['sec:exp']}).
  • Figure 2: Signum outperforms Adam when gradient stochasticity is relatively large, and FOCUS further improves Signum when the valley-like landscape is sharp. (a) Signum outperforms Adam when gradient noise is large compared to sharpness. Orange pixels refer to conditions Adam is better and yellow parts mean Signum is better. FOCUS is even better than Signum after increasing sharpness. Blue pixels refer to conditions FOCUS is the best. (b) Increasing weight decay helps Signum and FOCUS against Adam. Yet the advantage of self-attraction is lost when weight decay is too large. (c and d) Gradient noise is large when batch size is small in practice. For MNIST classification, We find that increasing batch size (decreasing noise) can lead to a transition of the optimal optimizer from Signum to Adam, showing the insights from toy models are relevant to reality. Grey dashed lines in (d) highlight $\pm 5$ %. The error bars represent standard deviations. Experiment details are in \ref{['app:models']}.
  • Figure 3: FOCUS is more stable and faster in training GPT-2 (small). (a) In float16 and with the same hyperparameters, FOCUS is slower than Signum, but Signum is unstable. The smaller learning rate of Signum leads to slower training than FOCUS yet is still unstable. (b) Similarly, in float16, FOCUS is more stable than Adam and stable Adam training is much slower. (c) We copy the optimal performance of Adam (trained in bfloat16) from liu2023sophia (black line), which is still slower than our FOCUS. (d) Decreasing the number of training steps of FOCUS (which also changes the learning rate scheduler), we find FOCUS can achieve a 2x speedup compared to Adam. More details in \ref{['app:models']}.
  • Figure 4: Training dynamics reaches steady state on the narrowing valley. The left panel is $L(u,v=0)=-cu$ showing our toy landscape can go to negative infinity ($c=0.1$ in this case). The right panel is the loss for Adam (learning rate $0.005$, $\beta_1=0.9$, $\beta_2=0.999$, weight decay $0.1$) and FOCUS (learning rate $0.005$, $\beta_1=0.9$, $\beta_2=0.9$, $\gamma=0.2$, weight decay $0.1$) on the toy landscape with $a=10$ and $c=0.1$. Both training dynamics already converge within $3000$ steps.
  • Figure 5: Pairwise comparisons between optimizers. This figure uses the same data as \ref{['fig:toy']}, a and b. We plot all the pairwise comparisons, showing more details to support the conclusion.
  • ...and 3 more figures

Theorems & Definitions (7)

  • Definition 4.1: Online Optimization Protocol
  • Theorem 4.2: Regret Bound
  • Lemma 3.1: Convex function
  • Lemma 3.2: Momentum bound
  • proof
  • Theorem 3.3: Regret Bound
  • proof