No More Adam: Learning Rate Scaling at Initialization is All You Need | Summary

This article explains the key points of No More Adam: Learning Rate Scaling at Initialization is All You Need.

• 2024-12-16 (arXiv)
• Xu, Minghao, Xiang, Lichuan, Cai, Xu, Wen, Hongkai.
• Department of Computer Science, University of Warwick, Coventry, UK, Collov Labs
• Paper
• Github

Read this article in Korean

Summary

• AdamW, an optimizer that rescales parameter updates using gradient histories, incurs substantial memory and computation costs by maintaining a squared-gradient history for every parameter. In the reported ImageNet-1K image-classification search with ViT-S/16, a Vision Transformer that processes image patches, SGD-SaI reaches 72.92% top-1 accuracy versus AdamW’s 73.04%, while separate optimizer profiling shows approximately half the optimizer-state memory.
• SGD-SaI combines stochastic gradient descent with momentum (SGDM), which smooths updates using previous gradients, with Scaling at Initialization (SaI), which fixes a learning-rate multiplier for each parameter block using the first batch. Its gradient signal-to-noise ratio (g-SNR) is the block’s gradient norm divided by its gradient standard deviation, with numerical stabilization; that standard deviation measures variation across gradient entries within the block.
• On GPT-2-Small pretraining, SGD-SaI reports training and validation losses of 2.98 and 3.09, versus 3.11 and 3.12 for both AdamW and Adam-mini, an optimizer that shares adaptive scales within parameter blocks. Its optimizer update takes 9.31 ms versus 13.42 ms for AdamW and 27.21 ms for Adam-mini; these measurements concern optimizer updates rather than complete training iterations.
• The experiments also cover convolutional neural networks (CNNs), which process images with learned convolution filters, and low-rank adaptation (LoRA), which fine-tunes models through low-rank parameter updates. Evidence for larger models is limited to profiling, early convergence is slower than the Adam baselines, and the diffusion-model comparison, involving an image generator, provides qualitative examples without quantitative evaluation.

1. Introduction

The paper attributes part of Adam’s usefulness to local gains, multipliers that adjust learning rates across parameters with different gradient patterns. Maintaining first- and second-moment tensors, running averages of gradients and squared gradients, makes these adjustments expensive as models grow; SGD-SaI asks whether an initial gradient measurement can supply useful fixed gains without maintaining the second moment.

• Figure 1 profiles optimizer-state memory and optimizer-step time on a single NVIDIA A100-80GB; it does not measure end-to-end training speed.
• Figure 2 distinguishes elementwise adaptive gains in Adam(W), meaning Adam and AdamW, partitionwise adaptive gains in Adam-mini, uniform gains in SGD, and fixed blockwise gains in SGD-SaI.

The profiles show increasing optimizer-state costs with model size and lower measured costs for SGD-SaI than the plotted adaptive optimizers. The runtime axis measures the optimizer step, so the plot does not establish total training acceleration.

SGD-SaI assigns different gains to default parameter blocks once, whereas Adam(W) and Adam-mini repeatedly recompute gains. The diagram also shows the finer partitioning used by Adam-mini.

2. Related Work

The related work connects adaptive optimizers, memory-reduction methods, and gradient measurements at initialization. Adafactor shares second-moment information across dimensions, low-bit optimizers compress stored states, and Adam-mini shares second moments within selected blocks; SGD-SaI removes the dynamically updated second moment.

• Pruning at Initialization, which selects a sparse network before training, and zero-cost neural architecture search, which scores untrained architectures from inexpensive measurements, motivate using initial gradient structure.
• These connections motivate the method but do not independently establish that a first-batch statistic remains adequate throughout training.

3. Problem Setting

The parameter vector has d entries divided into B architecture-defined blocks, and optimization updates parameters by subtracting a global learning rate times an update direction. Momentum SGD uses a running gradient average; Adam additionally maintains a moving average of squared gradients and uses its bias-corrected value, adjusted for initialization, to rescale updates.

• The stated objective is to replace the second-moment-dependent local gain without storing or updating a full second-moment tensor.
• The distinction concerns optimizer states: removing one parameter-sized state tensor does not halve total training memory, which also includes parameters, gradients, and intermediate model activations.

4. Methods

The method develops a blockwise gradient statistic, examines its distribution across parameter types and training steps, and then freezes normalized blockwise scales at initialization. Its justification combines observed stability in the analyzed model with a local argument about small parameter updates.

4.1. Memory Efficient Local Gain: g-SNR

For block i, the paper defines the gradient norm as G_norm^(i) = sqrt(sum_j (g_j^(i))²) and the within-block variance as G_var^(i) = (1/d_i) sum_j (g_j^(i) − ḡ^(i))², where d_i is the number of entries and ḡ^(i) is their mean. It computes G_snr^(i) = G_norm^(i)/sqrt(G_var^(i) + ε), where ε is a small numerical-stability constant, and normalizes each block by the maximum g-SNR across blocks, yielding a multiplier between 0 and 1.

• The effective block learning rate is the global learning rate multiplied by this normalized statistic.
• The variance is calculated across gradient entries within a parameter block; the formula does not directly estimate stochastic gradient variance across independently sampled batches.
• Storing one scale per block requires O(B) storage, proportional to the number of blocks, while the retained momentum tensor still requires O(d) storage, proportional to the parameter count.

4.2. Statistics Analysis for g-SNR

On ViT-S/16 trained on ImageNet-1K, the paper observes different g-SNR values across parameter types and relatively stable values over the plotted training trajectory. Similar structures in different layers generally have similar values, while bias parameters have consistently low values.

• Figure 3 shows the distribution across parameter types; Figure 4 traces shallow, middle, and deep blocks and several parameters outside those blocks.
• The accompanying first-order Taylor argument, a local approximation of gradient changes, assumes sufficiently small parameter updates to motivate small changes in gradients and their statistics. It does not provide a global bound establishing the adequacy of first-batch scales across arbitrary training trajectories.
• The first block’s norm1 weights, belonging to its first normalization layer, are a reported exception to otherwise similar values across corresponding layer structures.

Weight blocks span markedly different g-SNR ranges, while bias values cluster in a narrow low range. This supports considering block-specific scaling in the analyzed model, without demonstrating that the resulting scales are optimal.

The traced statistics remain nearly flat within each parameter type while differing across types. These trajectories provide the paper’s empirical motivation for freezing scales at initialization.

The accuracy surfaces show that SGD-SaI retains useful performance over more of the selected learning-rate and weight-decay grid. Its ResNet18 peak is slightly higher than SGDM’s, while its ViT peak remains slightly below AdamW’s.

4.3. Proposed Methods Detail: SGD-SaI

SGD-SaI computes and normalizes g-SNR from the initial batch, stores those scales, and applies them to momentum updates for the rest of training. It uses PyTorch’s default parameter partitioning and decoupled weight decay, which applies regularization directly to parameters so the decay term does not enter the gradient statistic.

• Algorithm 1 initializes momentum with the first gradient, then uses m_t^(i) = μm_(t−1)^(i) + (1−μ)g_t^(i) on subsequent steps, where μ is the momentum coefficient.
• The update first applies weight decay and then subtracts the global learning rate times the stored normalized g-SNR times momentum.
• The algorithm appears as Figure 6; Section 4.3’s reference to Algorithm 6 is inconsistent with the displayed SGD-SaI algorithm.

The displayed Algorithm 1 places g-SNR calculation inside the first-step branch and reuses the resulting scales thereafter. Weight decay is applied separately from the scaled momentum update.

5. Experiments

The evaluation covers large language model (LLM) pretraining, image-classification pretraining, parameter-efficient fine-tuning (PEFT), which updates a restricted parameter set, and convolutional image classification. Large-model efficiency measurements report optimizer-state memory and optimizer-step time separately from model-quality experiments.

• Full pretraining comparisons use GPT-2-Small on OpenWebText, a text corpus, and ViT-S/16 on ImageNet-1K.
• Larger GPT-2 and Llama-2 models are profiled without complete pretraining runs.
• Fine-tuning includes GPT-2 on E2E, a natural-language generation dataset, and a diffusion image generator trained to capture visual concepts.

5.1. LLM Pre-train

GPT-2-Small has 125M parameters and is pretrained on OpenWebText against AdamW and Adam-mini using settings inherited from the cited Adam-mini study. Figure 8 reports lower final training and validation losses for SGD-SaI despite slower initial convergence, alongside approximately half AdamW’s optimizer-state memory.

• Training loss is 2.98 for SGD-SaI and 3.11 for both baselines; validation loss is 3.09 versus 3.12.
• Optimizer updates take 9.31 ms for SGD-SaI, 13.42 ms for AdamW, and 27.21 ms for Adam-mini.
• The section does not provide a complete training configuration or uncertainty estimates for these GPT-2-Small loss comparisons.

Table 1 profiles GPT2-1.5B and Llama2-7B in full FP32 precision, meaning 32-bit floating-point values, with context length 1024 and batch size 1 on a single NVIDIA A100-80GB. GPT2-1.5B uses 5.93 GB of optimizer states with SGD-SaI versus 11.86 GB with AdamW, and optimizer steps take 68.0 ± 21.0 ms versus 138.0 ± 6.0 ms.

• For Llama2-7B, SGD-SaI reports 25.15 GB and 180.0 ± 30.0 ms; Adam-mini reports 27.21 GB and 421.0 ± 22.0 ms.
• AdamW, Adam, and Prodigy, an Adam variant that automatically adjusts step size, encounter out-of-memory errors in the Llama2-7B configuration.
• The abstract claims a 25.15 GB saving over AdamW for Llama2-7B, but Table 1 lists 49.48 GB versus 25.15 GB, a difference of 24.33 GB, consistent with Figure 1.

5.2. ViT Pre-train

For ViT-S/16 on ImageNet-1K, Table 2 reports peak top-1 validation accuracy, the fraction of examples whose highest-scoring prediction is correct, averaged over three seeds: SGD-SaI achieves 72.92 ± 0.07%, SGDM 63.80 ± 0.35%, AdamW 73.04 ± 0.31%, Adam-mini 72.29 ± 0.43%, and Prodigy 73.24 ± 0.21%. SGD-SaI has a higher grid-average accuracy than AdamW, while AdamW retains a slightly higher peak.

• Across the reported hyperparameter trials, average accuracy is 57.55 ± 18.46% for SGD-SaI, 37.21 ± 35.43% for AdamW, and 36.65 ± 35.39% for Adam-mini. These grid-average results describe sensitivity across settings and should be distinguished from seed variation at the best setting.
• The grid uses learning rates {0.1, 0.01, 0.001, 0.0001} and weight decays {0.01, 0.001, 0.0001}; grid-average performance depends on these choices.
• Prodigy is excluded from the grid-average comparison because the paper reports convergence failures when its hyperparameters are modified.

Figure 7 shows slower early convergence for SGD-SaI and comparable final image-classification performance. Table 3 separately profiles optimizer efficiency: for ViT-H/14, SGD-SaI uses 2.42 GB and takes 54.0 ± 13.0 ms per optimizer step, versus AdamW’s 4.86 GB and 124.0 ± 4.0 ms.

• For ViT-S/16, SGD-SaI reports 0.08 GB and 12.4 ± 0.2 ms, compared with AdamW’s 0.17 GB and 45.0 ± 8.0 ms.
• SGDM remains faster in these measurements: 7.9 ± 0.3 ms for ViT-S/16 and 40.0 ± 1.0 ms for ViT-H/14.
• The paper does not report total training-time improvements because cluster input/output and network effects complicate comparisons.

SGD-SaI improves more slowly early in training and approaches the baselines near the end. The curves show why faster optimizer steps alone do not determine time to a target accuracy.

SGD-SaI combines the lowest displayed final losses with a 9.31 ms update and 50.00% normalized optimizer-state memory. The approximately threefold speed comparison applies to Adam-mini’s optimizer update.

The table records optimizer-state memory and marks configurations that run out of memory under FP32 profiling with context length 1024 and batch size 1. Its Llama2-7B entries imply a 24.33 GB state-memory difference between AdamW and SGD-SaI, which differs from the abstract’s stated saving.

SGD-SaI has the highest reported grid-average accuracy and the lowest seed variation at its best setting. AdamW and Prodigy achieve slightly higher peak accuracies; the grid-average and best-setting statistics describe different sources of variation.

SGD-SaI approximately matches SGDM’s state-memory footprint and updates faster than the adaptive baselines. SGDM itself remains faster than SGD-SaI for both profiled ViT sizes.

5.3. Parameter Efficient Fine Tuning: PEFT

In 5.3.1. LLMS PARAMETER EFFICIENT FINE-TUNING, GPT-2 medium is fine-tuned on E2E using LoRA rank 4, learning rate 1e−3, and weight decay 1e−2. Table 4 gives SGD-SaI the highest reported scores across BLEU and NIST, measures based on matching word sequences; MET, the table’s label for METEOR, a text-generation evaluation metric; ROUGE-L, a sequence-overlap measure; and CIDEr, a consensus-based text-generation measure.

• SGD-SaI scores 69.9 BLEU, 8.81 NIST, 46.7 MET, 72.1 ROUGE-L, and 2.53 CIDEr.
• The strongest listed scaled-AdamW baseline scores 69.6, 8.77, 46.6, 71.8, and 2.52, respectively. The scaled baselines adjust learning rates for the two LoRA parameter matrices using the preconditioning method of Zhang & Pilanci (2024).
• These are small numerical improvements; Table 4 supplies no uncertainty estimates for assessing their repeatability.

In 5.3.2. DMS PARAMETER EFFICIENT FINE-TUNING, the paper applies LoRA to ChilloutMix, a diffusion image-generation model, using the Mix-of-Show concept-customization setup and compares SGD-SaI with scaled-AdamW at learning rate 0.1. Figure 9 shows generated examples for “a pencil sketch of ⟨Vpotter⟩,” with the authors identifying failures in the scaled-AdamW row.

• The assessment is qualitative and based on displayed samples.
• No numerical image-quality metric, concept-fidelity metric, or repeated evaluation is reported for this comparison.

SGD-SaI leads all five reported generation metrics, but its margins over scaled-AdamW are small. The table does not include uncertainty estimates.

The displayed samples illustrate the authors’ qualitative claim that SGD-SaI captures the target concept more consistently. Selected examples and marked baseline failures do not quantify overall generation quality or prompt fidelity.

5.4. Convolutional Neural Network(CNN)

For ResNet18 on CIFAR-10, SGD-SaI reaches 95.36% peak test accuracy versus 95.26% for SGDM, 94.74% for AdamW, and 94.23% for Adam. Additional experiments evaluate eleven architectures from NATS-Bench, a neural architecture benchmark, on CIFAR-10, CIFAR-100, and ImageNet16-120 using the same candidate hyperparameters across optimizers.

• Figure 5 shows sensitivity to learning rate and weight decay for ResNet18 and ViT-S/16.
• Figure 10 shows accuracy distributions that favor SGD-SaI across the selected architectures and hyperparameter settings.
• The NATS-Bench experiments use seed 777 and select ten architectures with top-10 validation accuracy plus one with bottom-1 validation accuracy. The plotted distributions therefore describe this selected set rather than an unrestricted architecture sample.

The histograms place more SGD-SaI trials toward higher accuracies across the three datasets; the smooth curves estimate the distributions from those trials. The results apply to the selected eleven architectures, candidate hyperparameters, and single reported NATS-Bench seed.

6. Conclusion

The conclusion presents fixed initialization-based block scaling as a way to obtain competitive final performance while retaining momentum SGD’s small optimizer-state footprint. The ViT-S/16 search results, GPT-2-Small losses, and isolated optimizer profiling support this claim within the tested settings; they do not establish a universal replacement for AdamW.

7. Limitation

The authors identify limited compute and slower early convergence as the main limitations. Complete large-model pretraining, including Llama-2-7B, remains untested, so memory feasibility and optimizer-step speed do not establish convergence quality or time to a target loss at that scale.

Appendix

• A. More Experiments Details describes the hyperparameter searches and training procedures for ViT and CNNs, including momentum SGD baselines with and without Nesterov, selected NATS-Bench architectures, and an extended ResNet18 search. The CNN comparisons control augmentation and use full training; the NATS-Bench comparison uses a single reported seed, while the extended ResNet18 search repeats across seeds. These choices delimit the reported robustness results.
• B. Optimizer Analysis contrasts momentum SGD’s stored momentum with Adam(W)’s two moment tensors, Adam-mini’s blockwise second moments with elementwise treatment for embeddings, and Prodigy’s additional parameter-sized states. The analysis concerns lower bounds on optimizer-state storage, rather than total training memory.
• C. Profiling Results on ”g-SNR Calculation” Stage measures the one-time initialization cost separately from optimizer-step time on different GPUs. Some calculation times exceed one reported optimizer step, despite the caption’s broader claim; the calculation occurs only once, and these measurements do not establish end-to-end training speed.
• D. Algorithm Overview explains that architecture-dependent block scales are computed before the first optimizer update and remain fixed throughout training. It distinguishes eliminating repeated local-gain calculations from the potential throughput benefit of using reduced memory to support larger batches.

Brief Thoughts

The method provides a simple baseline for testing how much of adaptive optimization’s benefit can be recovered by fixed blockwise scaling. The improvement from 63.80% to 72.92% over SGDM on ViT-S/16 is substantial within this experiment, while AdamW and Prodigy retain slightly higher peak accuracies.

A useful follow-up would measure time to a target validation loss in complete large-model training while varying the initialization batch and optimizer implementation. The current experiments leave those questions open: large models are profiled without full training, initial-scale stability is analyzed primarily on ViT-S/16, and optimizer-step speed is measured separately from slower early convergence.