PaperlayerView exact revision on arXiv

Source-linked AI summary

FOCUS: First Order Concentrated Updating Scheme

Yizhou Liu, Ziming Liu, Jeff Gore

arXiv:2501.12243v1cs.LGcs.CLmath.OC

TL;DR

Because improving LLM pre-training is considered critical, the paper studies optimizer behavior in a hypothesized narrowing-valley landscape. It develops FOCUS by adding attraction to Signum; in GPT-2 small, it reports greater stability and twice the training-time speedup over a literature Adam baseline.

  • Problem

    Improving LLM pre-training is a critical research direction, and optimizer design requires understanding its loss landscapes to reduce training costs or reach lower losses.

    Source: §1 · p. 1

  • Method

    FOCUS augments Signum with attraction toward moving-averaged parameters to squeeze into valleys without substantially decreasing effective step size.

    Source: §1 · p. 2, §2.3 · p. 5

  • Results

    In GPT-2 small pre-training, FOCUS is more stable than Signum and Adam at similar learning rates and achieves twice the speedup over Adam from the literature in training time.

    Source: §1 · p. 1

  • Takeaways & Limitations

    The authors identify large gradient noise as an important limiting factor in LLM training and suggest testing FOCUS with limited batch sizes or low-precision training.

    Source: §6 · p. 8

    The reported speedup mainly reflects stability issues in other algorithms, which may not persist on better machines or widely used bfloat16 formats.

    Source: §3.3 · p. 6

Abstract

from arXiv · show

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.

1. Introduction

The introduction motivates better LLM pre-training and studies optimizer behavior in a hypothesized narrowing-valley landscape. It presents FOCUS, which adds attraction to Signum, and reports greater stability and a twofold training-time speedup in GPT-2 small.

1. Introduction

  • Improving pre-training is presented as a critical direction for advancing LLM capabilities, with better optimizers sought to lower training time, cost, or model size.The paper argues that understanding pre-training loss landscapes is essential for designing optimizers accordingly.Source: §1 · p. 1
  • The paper studies training dynamics using a simplified picture and toy model to generate testable hypotheses, then applies insights to real systems or updates the picture if they fail.The proposed picture is a narrowing valley in the pre-training loss landscape.Source: §1 · p. 1
  • FOCUS is more stable than Signum and Adam at similar learning rates in GPT-2 small pre-training and achieves twice the speedup over Adam from the literature in training time.Figure 1 summarizes the progression from the narrowing-valley picture to the proposed optimizer.Source: §1 · p. 1, Figure 1 · p. 2
  • When gradient noise is too large, Adam performs worse than Signum because it decreases its effective step size too much, while Signum fixes the step size.The introduction presents this toy-model result as the motivation for FOCUS.Source: §1 · p. 1
  • FOCUS adds attraction toward moving-averaged parameters to Signum, aiming to squeeze into valleys without substantially decreasing effective step size.The introduction describes the work as contributing to scientific understanding of training dynamics and practical speedup in LLM pre-training.Source: §1 · pp. 1–2

2. Methods

The section models pre-training losses as narrowing valleys, studies how noise and sharpness affect optimizer performance, and introduces FOCUS as Signum augmented with attraction toward moving-averaged parameters. Toy-landscape and MNIST results relate optimizer rankings to gradient noise and show where FOCUS improves on Signum.

2.1. Picture of loss landscapes

  • The proposed pre-training picture is a narrowing valley with heterogeneous curvatures, motivated by curvature evidence, learning-rate decay, and growing parameter norms under lower weight decay.The picture assumes training has not reached the true minimum.Source: §2.1 · pp. 2–3

2.2. Toy model

  • The toy landscape uses rotated coordinates for a valley whose orthogonal curvature increases along the valley direction, and models queried gradients with noise.Optimizer comparisons scan learning rates and select each optimizer's best performance over repeated runs.Source: §2.2 · pp. 3–4
  • Signum outperforms Adam when gradient noise is large relative to valley sharpness, while larger weight decay favors Signum and excessive decay removes FOCUS's advantage.Adam's effective step size becomes too small when noise, rather than sharpness, drives its low signal-to-noise ratio.Source: §2.2 · pp. 4–5, §2.3 · p. 6
  • On MNIST, Adam's best test loss is larger than Signum's at small batch sizes but becomes smaller after batch size increases.The authors interpret this transition as consistent with the toy-model competition between optimizers.Source: Figure 2 · p. 4, §2.2 · p. 5

2.3. FOCUS optimizer

  • FOCUS adds an attraction term based on the difference between current and moving-averaged parameters to Signum's momentum update.The attraction strength is controlled by γ and moves parameters toward the trajectory center or valley bottom.Source: §2.3 · pp. 5–6
  • FOCUS outperforms Signum with high probability in sharp toy valleys, but its advantage disappears at excessive weight decay; it can outperform Adam when noise outweighs sharpness.FOCUS also appears more stable staying deep in the valley, while its boundary against Adam is nearly the same as Signum's.Source: Figure 2 · p. 4, §2.3 · p. 6

3. Experiments

The experiments evaluate FOCUS on GPT-2, comparing its stability and speed with Signum and Adam. FOCUS achieves practical speedup over Adam, including an approximately twofold speedup under a reduced-step schedule.

3.1. Setup

  • GPT-2 (125 million parameters) was trained on OpenWebText; Adam's comparison baseline was its hyperparameter-scanned optimum, while FOCUS was not systematically scanned.The Adam baseline was reported by Liu et al. (2023a); FOCUS tests used hyperparameters around Adam's optimal choice.Source: §3.1 · p. 6

3.2. Stability

  • In float16, FOCUS was more stable than Adam; Signum remained unstable at a learning rate of 10^-4, while Adam stabilized at 10^-4 but trained much slower and had higher final validation loss than FOCUS.Source: §3.2 · p. 6

3.3. Speedup

  • With fewer steps and a faster-decaying scheduler, FOCUS approximately achieved 2x speedup over Adam: 5×10^4 FOCUS steps matched 10^5 Adam steps, and 2×10^4 matched 5×10^4.The final learning rate remained the same as the step count decreased.Source: §3.3 · pp. 6–7, Figure 3 · p. 7

4. Convergence analysis

The convergence analysis studies FOCUS in online convex optimization using regret. It concludes that FOCUS achieves the same convergence rate as Adam, while noting limits to what the worst-case analysis reveals about optimizer design.

4. Convergence analysis

  • The analysis considers a sequence of convex losses revealed over time, with the optimizer observing the gradient at its current parameters and performance measured by regret.The feasible set and gradients are assumed bounded.Source: §4 · pp. 7–8
  • FOCUS achieves the same convergence rate as Adam on the same online convex optimization problem.The analysis states that its regret has the same dependence on T as Adam.Source: §4 · pp. 7–8
  • Although decaying β1,t and γt help bound terms in the regret analysis, the worst-case bound does not capture the advantages of momentum or attraction.The paper notes those terms may themselves be small, so decay may not be necessary; the analysis offers limited guidance for designing algorithms.Source: §4 · p. 8

5. Related works

Related work situates FOCUS among studies of valley-shaped loss landscapes, optimizer behavior, and parameter averaging. The paper distinguishes its use of parameter EMA, which changes training dynamics directly, from EMA used to improve model predictions.

Valley-like loss landscapes

Optimizers

  • FOCUS is presented as physics-inspired, unlike mostly mathematics-inspired optimizers such as Lion (Chen et al., 2024), MARS (Yuan et al., 2024), and Sophia (Liu et al., 2023a).The related-work discussion also notes prior findings that Adam is not always better than Signum (Zhao et al., 2024).Source: §5 · p. 8

Weight averaging

  • Earlier work used trajectory EMA to improve prediction stability and generalizability, whereas FOCUS uses parameter EMA to modify the optimizer and training dynamics directly (Grill et al., 2020; He et al., 2020; Morales Brotons et al., 2024).The optimizer is generalized to time-dependent decay and force strength.Source: §5 · p. 8

6. Discussion

The discussion frames FOCUS through a free-energy picture in which attraction favors smaller volumes as entropy falls in a narrowing valley. It identifies gradient noise and the batch size needed to reduce it as open issues for LLM training.

6. Discussion

  • In the paper's physics picture, free energy is F = E − ΘS, where E is loss, Θ is temperature related to batch size and learning rate, and S is entropy.The optimizer is described as seeking a minimum of free energy.Source: §6 · p. 8
  • Attraction adds an energy term favoring small volumes, countering the entropy decrease as optimization moves through a narrowing valley.The paper describes the optimizer as becoming stuck when gradient and entropic forces balance.Source: §6 · p. 8
  • The batch size needed to make LLM gradient noise small remains an open question, and the authors call for further study of FOCUS and their physical picture.They suggest trying FOCUS with limited batch sizes and low-precision training; the observed speedup is interpreted as indicating that gradient noise may limit LLM training.Source: §6 · p. 8

A.1. Toy model

The toy-model experiments explore optimizer behavior across landscape sharpness, gradient noise, learning rates, and weight decay. They use repeated runs and track loss and valley-position outcomes.

A.1. Toy model

  • The two-dimensional landscape uses a 30-degree rotation and sets sharpness a = 10 and parameter c = 0.1.The landscape coordinates are rotated by θ = π/6 radians.Source: §A.1 · p. 11
  • The Figure 2a sweep varies landscape sharpness across 96 values and gradient-noise standard deviation across 100 values.Learning rates span 20 logarithmically spaced values from 10^-3 to 1.Source: §A.1 · p. 11
  • The Figure 2b phase diagram varies weight decay across 96 values, alongside 100 noise levels and 20 learning rates.Weight decay ranges from 0 to 0.5.Source: §A.1 · p. 12
  • Each configuration was tested in 50 independent replicates; initial positions were randomly initialized with standard deviation 10−4, and optimization proceeded for 1,000 steps, tracking historical minimum loss and valley-axis projection at rotation angle π/6.Source: §A.1 · p. 12

A.2. MLP

The MLP experiments use a five-layer, 128-unit-per-layer network on MNIST to compare optimizer configurations over learning rates and batch sizes. Training and evaluation track losses and test accuracy.

A.2. MLP

  • The MNIST model has five hidden layers of 128 units, ReLU activations, and batch normalization before each linear layer.Inputs are flattened 784-pixel images, and the output has 10 classes.Source: §A.2 · p. 13
  • The experiments compare four optimizer configurations, including FOCUS variants with γ = 0.2 and γ = 0.4.The listed FOCUS configurations use β1 = 0.9 and β2 = 0.99.Source: §A.2 · p. 13
  • Learning rates span 20 logarithmically spaced values, while batch sizes range over 12 powers of two from 2^1 to 2^12.Each configuration has three independent replicates and runs for 400 steps.Source: §A.2 · p. 13
  • The experiments record training loss, test loss, and test accuracy at each step using cross-entropy as the optimization objective.The dataset includes 60,000 training images and 10,000 test images.Source: §A.2 · p. 13

A.3. GPT-2

The GPT-2 experiments train a 125-million-parameter model on OpenWebText and compare Signum, Adam, and FOCUS configurations. Training uses distributed data parallelism across eight V100 GPUs.

A.3. GPT-2

  • The experiments implement GPT-2 small with 125 million parameters, 12 layers, 12 heads, and embedding dimension 768.The implementation disables bias and dropout during pre-training.Source: §A.3 · p. 14
  • Training uses OpenWebText with a 9-billion-token training set and a 4.4-million-token validation set.Text is tokenized with the standard GPT-2 tokenizer.Source: §A.3 · p. 14
  • Training runs across eight NVIDIA V100 GPUs with gradient accumulation to maintain an effective batch size of 480, using float16 precision.The effective batch size is 12 examples per GPU across five accumulation steps and eight GPUs.Source: §A.3 · p. 14
  • The appendix also compares Adam and FOCUS on the toy narrowing valley, where both training dynamics reach steady state within 3,000 steps.That comparison uses a = 10 and c = 0.1.Source: Figure 4 · p. 15

B. Supplementary results

The supplementary results discuss how the narrowing-valley picture bounds optimizer progress and interpret additional optimizer comparisons. They also note that FOCUS's small runtime overhead may reflect its implementation.

B. Supplementary results

  • The constructed landscape is an unrealistic local approximation because its loss can decrease to negative infinity.The authors distinguish this from real cases, where the valley minimum is 0.Source: B. Supplementary results · p. 15
  • The narrowing-valley picture implies that finite learning rates can leave optimizers far from the true valley minimum, creating a speed-versus-depth trade-off.The authors connect this training-dynamics account to the possibility that neural scaling laws reflect training dynamics rather than model limits.Source: B. Supplementary results · p. 15
  • In pairwise toy-landscape comparisons, FOCUS may be more stable in mixed regions without reaching significantly deeper in the valley.The authors interpret the checkered pattern as stochasticity rather than sensitivity to hyperparameters in deterministic training.Source: B. Supplementary results · p. 15
  • FOCUS takes slightly longer for the same number of steps, but the authors hypothesize that the small overhead reflects unsophisticated code with room for improvement.They state that the algorithm is no more complex than Adam and requires no more memory.Source: B. Supplementary results · pp. 15–16
  • Figure 5 presents detailed pairwise optimizer comparisons using the same data as Figure 2a and 2b.The figure provides additional detail supporting the stated conclusions.Source: Figure 5 · p. 16

C. Convergence proof

The convergence appendix establishes a regret bound for FOCUS under convex losses and bounded gradients, using a momentum bound and termwise estimates.

C. Convergence proof

  • FOCUS momentum coordinates remain bounded by the gradient infinity norm when gradients are bounded and momentum coefficients lie in [0, 1).The proof establishes the result by induction from zero initial momentum.Source: C. Convergence proof · p. 16
  • For convex loss sequences with bounded gradients, FOCUS has a regret bound under the stated learning-rate, decay, and feasible-region diameter conditions.The theorem assumes η_t = η/√t, λβ and λγ in [0, 1), and bounded feasible-region distances; the supplied passage does not include the bound's formula.Source: C. Convergence proof · p. 17
  • The proof bounds the regret expression by estimating its four terms using convexity and momentum and gradient assumptions.It combines the four term estimates to obtain Equation (10).Source: C. Convergence proof · pp. 16–19

Related papers

Loading 2501.12243v1…