PaperlayerView exact revision on arXiv

Source-linked AI summary

Optimizer Memory Schedules for Outscaling the Overtraining Axis

Katie Everett, Shikai Qiu

arXiv:2609.04577v1cs.LG

TL;DR

Optimizer comparisons are often made at one training horizon, despite evidence that rankings and hyperparameters change with overtraining. This paper evaluates AdamW, ADANA, Muon, and SOAP across model sizes and horizons, finding that ADANA’s advantage over AdamW persists after horizon-specific tuning and grows with training, while matrix-preconditioned methods retain more constant gains.

  • Problem

    Optimizer benchmarks often compare methods at a single horizon, leaving performance and hyperparameter scaling across the overtraining axis insufficiently characterized.

  • Method

    The study compares AdamW, ADANA, Muon, and SOAP on 51M–253M models across 1×–256× overtraining, independently sweeping base learning rates and evaluating token efficiency.

  • Results

    ADANA’s scaling advantage over AdamW persists after horizon-specific fixed-memory tuning, while Muon and SOAP provide roughly constant advantages across most measured horizons.

  • Takeaways & Limitations

    Training horizon should be treated as an essential axis for optimizer evaluation and design, including the scaling of learning rate schedules, weight decay, and memory.

  • Takeaways & Limitations

    The conclusions require further validation across larger models, architectures, datasets, and longer training horizons.

Abstract

from arXiv · show

We investigate how optimizers scale across the overtraining axis and show that relative optimizer performance and optimal hyperparameters change substantially with training horizon. In particular, we study how matrix-preconditioned methods (Muon and SOAP) and a momentum-scheduled method (ADANA) scale relative to AdamW. We compare these four optimizers across models from 51M to 253M parameters and overtraining (OT) factors from 1x to 256x, sweeping the base learning rate at every setting. The preferred learning rate schedule can reverse across the overtraining axis, the best weight decay coefficient scales approximately as sqrt(OT), and longer horizons generally favor longer fixed memory. ADANA's scaling advantage over AdamW persists after tuning AdamW's fixed memory separately at each horizon. Log-time weight decay and momentum cooldown provide substantial gains for ADANA that compound as training increases. With this treatment, ADANA outscales AdamW with an exponent advantage close to that predicted by DANA theory on power-law random features. Muon and SOAP instead provide roughly constant token-efficiency advantages over AdamW across most of the measured range, although SOAP may gain further at the highest overtraining factors. ADANA begins behind both matrix-preconditioned optimizers but closes the gaps as training increases, surpassing Muon and becoming competitive with SOAP at our highest OT factors. These results establish training horizon as an essential axis for optimizer evaluation and design.

1 INTRODUCTION

The paper evaluates optimizer scaling beyond the compute-optimal horizon, where rankings and preferred hyperparameters can change substantially. Across 51M–253M models and 1×–256× overtraining, ADANA increasingly closes its gap with matrix-preconditioned methods while retaining an advantage over AdamW.

  • Optimizer rankings and preferred hyperparameters can change as training extends beyond the compute-optimal horizon.The study varies model size, overtraining factor, and optimizer treatments rather than comparing methods at one fixed horizon.
  • Muon and SOAP lead at short horizons, while ADANA closes its initial gap as overtraining increases and SOAP overtakes Muon at the longest measured horizons for 51M and 124M models.These rankings are based on lowest final validation loss after learning-rate sweeps under uniform weight decay.
  • 51M–253M models are evaluated from 1× to 256× overtraining with independently swept base learning rates and token-efficiency measurements.The token multiplier compares baseline tokens with optimizer tokens at the same loss.
  • The preferred learning-rate schedule can reverse across overtraining, and the best weight decay coefficient scales approximately as √f.The paper also finds that longer horizons generally favor longer fixed memory and that larger memory pairs with a smaller learning rate.
  • With log-time weight decay and momentum cooldown, ADANA outscales AdamW, while Muon and SOAP provide roughly constant token-efficiency advantages across most of the overtraining axis.ADANA begins behind matrix-preconditioned optimizers but gains with training, surpassing Muon and becoming competitive with SOAP.

2 RELATED WORK

Prior work studies optimizer memory, matrix preconditioning, scaling-law forms, and changing optimizer efficiency, but comparisons across broad overtraining ranges remain limited. This paper situates ADANA, Muon, and SOAP within those lines of work.

  • Recent optimizers vary memory through multiple momentum buffers, iterate averaging, or scheduled timescales rather than keeping one fixed memory throughout training.The related methods include AdEMAMix, schedule-free optimizers, and ADANA’s scheduled memory.
  • Muon, SOAP, and Shampoo use matrix-level preconditioning, contrasting with Adam-like coordinatewise or diagonal normalization.Prior theory and experiments associate matrix preconditioning with speedups on suitable spectral modes and approximately constant compute gains.
  • Scaling-law analyses increasingly couple model size and training data through interaction exponents or shared optimizer parameters.The paper uses a shared-m functional form to improve comparison across optimizers.
  • Earlier optimizer comparisons often use one horizon, while recent evidence reports approximately 1.2×–1.4× Muon and SOAP token multipliers through 8× Chinchilla.Muon performs better through 4× in the cited study, while SOAP catches or overtakes it later in that range.

3 SCALING OPTIMIZERS ACROSS TRAINING HORIZONS

The study measures optimizer performance across overtraining factors and shows that schedules, weight decay, memory, and relative rankings must be evaluated as horizon-dependent choices. It defines outscaling through equivalent training-token requirements at matched loss.

  • The experiments compare AdamW, ADANA, Muon, and SOAP from 1× to 256× overtraining on models from 51M to 253M parameters.The overtraining factor is f = T/(20P), and increasing it proportionally increases optimizer updates in the core setup.
  • Preferred learning-rate schedules can reverse across the overtraining axis, while the best weight decay coefficient scales approximately as √f.The coefficient rules used in most experiments are c_uniform(f) = 8√f and c_log(f) = 2√f.
  • The token multiplier is the baseline’s required training tokens divided by the indicated optimizer’s tokens at the same loss.An optimizer outscales the baseline when this multiplier increases with overtraining, which can reflect a better token-decay exponent, a lower high-token loss limit, or both.
  • Muon and SOAP lead at many short and intermediate horizons, ADANA improves relative to AdamW with overtraining, and SOAP overtakes Muon at the longest measured horizons for 51M and 124M models.Muon’s token multiplier is approximately constant, ADANA’s increases, and SOAP may improve further at the highest overtraining factors.

4 HORIZON-TUNED FIXED MEMORY

Horizon-tuned fixed memory improves fixed-memory optimizers, especially at the ends of the overtraining range, but does not remove ADANA’s relative advantage over AdamW. Longer horizons generally favor longer memory, with smaller learning rates at larger memory.

  • Longer horizons generally favor longer fixed memory, although the optimal trajectories are not always monotonic.From 1× to 128× overtraining, optimal memory rises from M = 20 to M = 1280 for AdamW, M = 20 to M = 640 for Muon, and M = 40 to M = 1280 for SOAP.
  • Per-horizon memory tuning improves final validation loss most at low and high overtraining factors, while M = 50 is near-optimal at moderate horizons.At 128×, improvements are 3.3 mpt for AdamW, 3.6 mpt for Muon, and 1.3 mpt for SOAP.
  • After fixed-memory tuning, ADANA continues gaining on AdamW as overtraining increases, while Muon and SOAP retain roughly constant advantages.This indicates that ADANA’s advantage is not explained by holding AdamW’s β2 fixed across horizons.

5 LOG-TIME WEIGHT DECAY ACROSS TRAINING HORIZONS

The preferred weight-decay treatment changes with optimizer and training horizon. Log-time weight decay increasingly helps ADANA, while its effects on SOAP, AdamW, and Muon differ.

  • Weight-decay comparisons independently sweep the base learning rate for every optimizer, schedule, and overtraining factor.Token multipliers compare the tokens required by a baseline and indicated optimizer at the same loss.
  • The best weight decay coefficient scales approximately as √f, with fitted rules c(f) = 8√f and a(f) = 2√f.These rules approximate the cross-optimizer trend rather than separately optimizing every coefficient.
  • Log-time weight decay increasingly benefits ADANA as overtraining grows, modestly benefits SOAP, has little consistent effect on AdamW, and harms Muon at high OT.The Muon result may reflect a poorly matched log-time coefficient rather than the schedule itself.

6 MOMENTUM COOLDOWN FOR GROWING MEMORY

Momentum cooldown limits ADANA’s growing memory during terminal learning-rate decay. Combined with log-time weight decay, it produces the largest long-horizon gains across model sizes.

  • Momentum cooldown preserves ADANA’s growing memory through most training but shortens it during terminal learning-rate decay.It bounds the memory timescale by the learning-rate decay timescale near the end of training.
  • Momentum cooldown substantially improves ADANA at long horizons but has little consistent effect on fixed-memory AdamW.The effect is consistent across all three model sizes.
  • Log-time weight decay and momentum cooldown each improve ADANA, and their combined gains are largest at long horizons.The combined treatment outperforms ADANA with uniform weight decay across all three model sizes.
  • The combined ADANA treatment yields an approximately straight token-multiplier trend on log–log axes and is used in subsequent strongest-treatment comparisons.

7 OUTSCALING ACROSS THE OVERTRAINING AXIS

Across the overtraining axis, ADANA outscales AdamW with tuned treatments, while Muon and SOAP retain roughly constant token-efficiency advantages over most of the measured range. ADANA closes its initial gaps to both matrix-preconditioned optimizers as training horizons lengthen.

  • 7.1 ADANA OUTSCALES ADAMW: ADANA’s fitted token-decay exponent advantage over AdamW remains positive across the tested range of high-token-limit differences.The high-token-limit difference itself is less well identified.
  • 7.1 ADANA OUTSCALES ADAMW: Fitted ADANA equivalent-OT exponents range from 1.15 to 1.20, close to the DANA prediction 2 −κ = 1.15 for κ = 0.85.This holds across all three model sizes and both AdamW weight-decay baselines.
  • 7.2 MATRIX-PRECONDITIONED OPTIMIZERS PROVIDE APPROXIMATELY CONSTANT GAINS OVER ADAMW: Muon provides a 1.4×–1.7× token multiplier over AdamW and SOAP provides 1.3×–1.7× across most overtraining factors.SOAP reaches approximately 1.9× at 128× OT on the 51M model, suggesting possible further gains.
  • 7.3 ADANA OUTSCALES THE MATRIX-PRECONDITIONED OPTIMIZERS: ADANA starts behind Muon and SOAP but closes both gaps with increasing training horizon, surpassing Muon and becoming competitive with SOAP at the highest tested OT factors.These trends rely heavily on the highest tested OT factors and may not continue beyond them.

8 DISCUSSION AND LIMITATIONS

The results make training horizon a central dimension of optimizer evaluation: both rankings and hyperparameter choices can change with overtraining. Scheduled memory and matrix preconditioning show complementary scaling behavior, but transfer beyond the tested settings remains limited.

  • Optimizer rankings and preferred learning-rate, weight-decay, and memory settings can change across the training horizon.Long-horizon comparisons therefore require hyperparameter scaling rules rather than single-horizon choices.
  • Muon and SOAP provide large advantages over AdamW across most horizons, while ADANA’s relative efficiency increases with training horizon.Tuning fixed memory separately at each horizon does not explain ADANA’s scaling advantage.
  • The conclusions require validation across larger models, architectures, datasets, and longer training horizons.The paper reports preliminary evidence for robustness to another preconditioner ordering but leaves broader transfer to future work.
  • Batch dependence remains unresolved because the fixed-token experiment does not identify the mechanism or establish a general batch scaling law for ADANA.The main optimizer comparisons use a global sequence batch of 256.

9 CONCLUSION

The paper shows that optimizer rankings and optimal treatments change beyond compute-optimal horizons, making training horizon essential for optimizer evaluation and design. It finds scheduled memory can improve relative efficiency with additional training, while matrix preconditioning provides substantial but roughly constant gains.

  • Optimizer rankings and optimal treatments change as models are trained beyond compute-optimal token budgets.
  • ADANA’s advantage over AdamW persists after separately tuning AdamW’s fixed memory at each training horizon.
  • Log-time weight decay and momentum cooldown amplify ADANA’s advantage, producing equivalent-OT scaling close to the 2 −κ prediction from DANA theory.
  • Matrix-preconditioned optimizers provide large but approximately constant gains over AdamW across most of the measured range.
  • Scheduling optimizer memory is presented as a promising route to methods whose relative efficiency improves with additional training.

B.4 DATASET, TRAINING ORDER, SEEDS, AND EVALUATION

The experiments use decoder-only Transformers trained on a fixed FineWeb subset with matched data order and validation targets. Comparisons use full-set validation loss after standardized optimizer updates and decoupled weight decay.

  • Dataset and evaluation: The dataset contains 360.5 billion FineWeb training tokens, with disjoint source files reserved for validation.
  • Dataset and evaluation: All experiments process the tokenized corpus in fixed sequential order without shuffling, so matched configurations see identical tokens in identical updates.
  • Seeds and evaluation: Each configuration uses a fresh initialization with random seed 42, while matched comparisons share initialization and ordered training data.
  • Seeds and evaluation: Validation uses 1,920 held-out sequences and 3,932,160 next-token targets, reporting mean per-token cross-entropy over the complete set.
  • Optimization: Decoupled weight decay is applied as an independent parameter update rather than an L2 penalty or preconditioned gradient term.

D LEARNING RATE SCHEDULE SELECTION

Learning-rate schedule preference depends on overtraining horizon: linear decay is favored at 1× OT, while cosine decay is favored at 8× and 32× OT. The crossover persists across weight-decay coefficients and motivates cosine decay for the paper’s high-OT experiments.

  • Schedule construction: Both schedules use the same linear warmup before decaying from the peak learning rate to zero.
  • Schedule crossover: Linear decay performs better at 1× OT, whereas cosine decay performs better at 8× and 32× OT for all four optimizers.
  • Schedule crossover: The crossover remains after jointly selecting the best learning rate and weight decay coefficient for each schedule, optimizer, and OT factor.
  • Schedule crossover: Cosine decay’s high-OT advantage appears across multiple weight decay coefficients rather than arising from one favorable coefficient.
  • Protocol choice: Cosine decay to zero is used as the primary schedule because the remaining experiments focus mainly on high OT factors, not because it is claimed universally optimal.

E.4 WD COEFFICIENT SCALING RULES ACROSS OT FACTORS

The best weight decay coefficient generally increases with overtraining, and fitted exponents near one-half motivate square-root scaling rules. These rules substitute for per-horizon joint tuning when computation is limited, while direct tuning remains preferred when feasible.

  • Empirical scaling: The best weight decay coefficient generally increases with OT factor for both uniform and log-time schedules, so fixing the 1× coefficient can incur increasing regret.
  • Empirical scaling: The independently fitted exponents are 0.527 for uniform WD and 0.473 for log-time WD.
  • Scaling rule: Both fitted exponents are close to 1/2, motivating a square-root exponent for both schedules.
  • Scaling rule: The square-root prescriptions are used when repeating the complete learning-rate and weight-decay grid is computationally impractical, while direct joint tuning is used when feasible.
  • Memory horizon: Longer overtraining factors shift the low-regret region toward longer fixed memory.

F.2 HORIZON-TUNED MEMORY RESULTS

Longer horizons generally favor longer fixed memory, but memory and learning rate must be tuned jointly. Per-horizon tuning improves fixed-memory optimizers, while scheduled memory can avoid the compromise imposed by one fixed memory throughout training.

  • Longer horizons favor longer memory: From 1× to 128× OT, optimal memory increases from M = 20 to M = 640 for Muon, M = 20 to M = 1280 for AdamW, and M = 40 to M = 1280 for SOAP.The trajectories are not always monotonic.
  • Memory and learning rate coupling: The minimizing base learning rate generally decreases as effective memory grows, with the relationship depending on optimizer and horizon.Longer memory is useful at long horizons but is typically paired with a smaller learning rate.
  • Scheduled versus fixed memory: ADANA’s scheduled memory can use shorter memory early and lengthen it during training, avoiding one fixed memory–learning-rate compromise.This provides one plausible explanation for why per-horizon tuned constant memory does not fully match scheduled memory.
  • Per-horizon tuning: At 128× OT, per-horizon tuning improves validation loss by 3.3 mpt for AdamW, 3.6 mpt for Muon, and 1.3 mpt for SOAP.The fixed-memory comparison uses M = 50, corresponding to β2 = 0.98 for AdamW and SOAP and β = 0.98 for Muon.
  • Interpretation and limitations: Per-horizon tuning is an oracle comparison limited to the 51M model and one memory coefficient per optimizer.The study selects the best fixed memory and learning rate separately at each horizon rather than deriving a transferable memory prescription.

H.3 STRONGEST TREATMENT COMPARISONS

Using each optimizer’s strongest sufficiently covered treatment, the comparisons show that ADANA becomes increasingly competitive as overtraining grows. Its relative token efficiency rises against AdamW, Muon, and SOAP, while its agreement with the DANA reference remains theoretically qualified.

  • Strongest treatments: ADANA’s relative token efficiency increases with overtraining factor against log-time-WD AdamW, uniform-WD Muon, and log-time-WD SOAP.ADANA begins behind at short horizons, then catches or overtakes each comparator as training increases.
  • Strongest treatments: Figure 23 compares ADANA with log-time weight decay and momentum cooldown against AdamW and SOAP using log-time weight decay and Muon using uniform weight decay.The top, middle, and bottom rows show equivalent OT, token multipliers, and floor–decay profiles.
  • Floor–decay interpretation: Across the profiled functional forms, the data favor an improved ADANA token-decay exponent, while high-token-limit differences remain less well identified.The profiles compare tradeoffs between token-decay exponents and high-token validation-loss limits.
  • DANA PLRF comparison: The fitted equivalent-OT slopes with log-time weight decay and momentum cooldown are 1.20, 1.16, and 1.15 for the 51M, 124M, and 253M models.These values are close to the PLRF prediction 2 −κ = 1.15.
  • DANA PLRF comparison: The 2 −κ reference is theoretically motivated rather than a direct prediction for these transformer experiments.DANA theory compares DANA with SGD on PLRF, whereas these experiments compare adaptively preconditioned ADANA with AdamW.

I BATCH SIZE AND OPTIMIZER UPDATE COUNT

At fixed token count, batch size changes both the number of optimizer updates and gradient noise, making optimizer scaling difficult to identify. The experiments show distinct batch-dependent performance and memory trends, while broader comparisons remain limited by unresolved hyperparameter and batch-size definitions.

  • Update count: S = T/(B × L) approximately, so increasing batch size reduces the number of sequential optimizer updates at fixed token budget.Here T is the token budget, B the batch size in sequences, and L the sequence length.
  • Update count: Increasing B simultaneously reduces optimizer updates and stochastic variation, so fixed-token comparisons cannot isolate batch effects from update-count effects.An optimizer may appear to scale poorly with batch because it benefits less from larger batches or more from the additional updates available at smaller batches.
  • Measured batch scaling: At 8× OT, ADANA has lower loss at B = 32 but worsens from B = 64, whereas AdamW stays near its B = 32 performance through approximately B = 256.Both optimizers lose efficiency at the largest measured batch sizes.
  • Hyperparameter scaling: For AdamW, optimal M2 = 1/(1 − β2) decreases approximately as 1/B, while ADANA’s optimal g3 increases approximately linearly with B and its jointly optimized learning rate decreases.The experiments jointly optimize AdamW’s learning rate and β2, or ADANA’s learning rate and g3, with the other listed settings fixed.
  • Interpretation: ADANA’s batch-size advantage may depend on sequential updates and noisy gradients, because increasing B reduces updates and may reduce the high-noise regime where multi-timescale methods gain most.The proposed mechanism is motivated by DANA analyses and accelerated-SGD-style results rather than established as a causal explanation here.
  • Open limitations: Relative batch scaling remains uncertain because joint tuning of ADANA’s g3, κ, and δ, AdamW’s coupled moment timescales, and batch-conditioned weight decay was not fully explored.The relevant effective batch size is also ambiguous because tokens within a sequence are correlated; matched (B, L) comparisons could distinguish tokens per update from independent information.

J.3 LADANA RESULTS AND LIMITATIONS

LADANA preserves much of ADANA’s outscaling under uniform weight decay, but its long-horizon behavior becomes sensitive to denominator scheduling and the weight-decay treatment. Momentum cooldown addresses stale long-memory effects during terminal decay and provides substantial, complementary gains for ADANA, while several stability questions remain unresolved.

  • LADANA results: LADANA with log-time second-moment memory performs similarly to ADANA through 128× OT on 51M models and slightly underperforms at 256×.The optimizers also remain close across shorter horizons on the 124M and 253M models, whereas fixed-v LADANA falls behind at long horizons on 51M.
  • LADANA limitations: Under log-time weight decay, LADANA matches ADANA through 64× OT but deteriorates sharply at 128× without a LADANA-specific weight-decay coefficient sweep.The 128× result is therefore treated as an unresolved interaction among coefficient choice, preconditioner order, and possible instability rather than evidence of inherent instability.
  • Momentum cooldown: Momentum cooldown shortens ADANA’s memory during sufficiently rapid learning-rate decay, concentrating the momentum kernel on more recent updates.Before the learning-rate and momentum timescales cross, cooldown leaves the original log-time memory rule unchanged; the principal treatment applies it to both moment states.
  • Momentum cooldown: Cooling AdamW’s first moment improves validation loss by 2.30 millipoints at 8× OT, while cooling only its second moment produces unstable losses from 15.2 to 42.9.Cooling both moments restores stable training but adds no gain beyond first-moment cooling alone.
  • Momentum cooldown: For ADANA, memory-window cooldown improves validation loss by 11.47 millipoints, exceeding the 0.92-millipoint gain from tapering only the long-memory contribution.Combining tapering with window cooldown improves loss by 6.13 millipoints, so shortening the memory window is the strongest tested intervention.
  • Implications and open questions: Momentum cooldown benefits ADANA substantially across all three model sizes, and combining it with log-time weight decay yields the lowest validation losses at the longest measured horizons.The paper treats stale memory as the motivation for cooldown, not as an empirically established explanation of its gains, and leaves analogous Muon or SOAP rules and stability mechanisms open.
Loading 2609.04577v1…