Adam-SGD Gap in Language Modeling

• The Adam-SGD gap in language modeling refers to the phenomenon where adaptive optimizers like Adam achieve faster training but yield models with inferior generalization compared to SGD, particularly in large neural networks.
• This performance difference stems from Adam's component-wise updates causing directional misspecification and challenges with parameterization ill-conditioning, especially concerning weight magnitudes versus directions.
• Solutions like Normalized Direction-preserving Adam (ND-Adam), which controls update direction and normalizes weight vectors, alongside logit regularization, demonstrate improved generalization by better aligning optimization with the functional invariance of neural network parameters.

The Adam-SGD gap in LLMing refers to the persistent and, in many cases, substantial difference in optimization efficacy and generalization performance exhibited by Adam and its adaptive variants versus standard stochastic gradient descent (SGD). This gap, observed in empirical studies and substantiated by theoretical analysis, emerges most conspicuously in the training of deep neural networks for LLMing, especially on large-scale architectures such as Transformers. Below, the phenomenon is examined through foundational principles, causes, rigorous formulations, and implications for neural network training, with focus given to the findings and methods of "Normalized Direction-preserving Adam" (ND-Adam) and its context.

1. Foundations of the Adam-SGD Generalization Gap

Adam and similar adaptive optimizers have demonstrated superior optimization speed and practical robustness compared to SGD, particularly in the early stages of training deep neural networks. Paradoxically, these optimizers tend to yield models with inferior generalization: although Adam often achieves faster and higher initial training accuracy, final validation/test performance can lag behind that of SGD. This Adam-SGD generalization gap is well-documented in image classification and, as inferred, extends to LLMing tasks where parameterization and update geometry share similar pathologies and properties.

Critical to understanding this gap are two observations made in the ND-Adam work:

• Directional Misspecification: Adam adapts learning rates component-wise, causing the effective update direction to deviate from the direction given by the sum of current and historical gradients. As a result, optimization paths can diverge significantly from those taken by SGD, which maintains updates in the span of empirical gradients.
• Parameterization Ill-Conditioning: In architectures typical of LLMing (incorporating batch normalization, rectifiers, and large softmax layers), the magnitude of a weight vector often has little functional effect—only its direction matters. Adam’s lack of control over the coupling between weight magnitude and effective learning rate exacerbates instability and can guide optimization towards sharper minima, which historically have been associated with poorer generalization.

2. The ND-Adam Algorithm and Directional Control

ND-Adam is proposed as a solution that addresses both directionality and parameterization issues by restructuring the optimization update at the level of weight vectors rather than individual parameters. Key steps are as follows:

1. Tangent Space Gradient Projection: For each weight vector $w_i$, the raw gradient $\bar{g}_t(w_i)$ is projected onto the tangent space of the unit sphere:

$$g_t(w_i) = \bar{g}_t(w_i) - (\bar{g}_t(w_i) \cdot w_{i,t-1}) w_{i,t-1}$$

ensuring that updates remain orthogonal to $w_i$, preventing changes to weight norm during the update.

1. Per-Vector Adam Moments: Adaptivity is maintained at the vector level:

$$m_t(w_i) = \beta_1 m_{t-1}(w_i) + (1-\beta_1)g_t(w_i)$$

$$v_t(w_i) = \beta_2 v_{t-1}(w_i) + (1-\beta_2) \|g_t(w_i)\|_2^2$$

1. Vector-Normalized Update: The update is performed and then renormalized:

$$\bar{w}_{i, t} = w_{i, t-1} - \frac{\alpha_t^v}{\sqrt{\hat{v}_t(w_i)}+\epsilon} \hat{m}_t(w_i)$$

$$w_{i, t} = \frac{ \bar{w}_{i, t} }{ \| \bar{w}_{i, t} \|_2 }$$

thereby enforcing $\| w_{i, t} \|_2 = 1$ after each step. This shifts the locus of learning from weight magnitude to weight direction, aligning with the functional invariance of the network.

Other parameters not structured as hidden unit weight vectors are updated with standard Adam.

3. Empirical Results and Practical Efficacy

In benchmarking on Wide ResNet models for CIFAR-10 and CIFAR-100, the following outcomes were established:

• Adam attains high accuracy early but stagnates at higher test error rates compared to SGD.
• ND-Adam matches SGD's asymptotic generalization and typically outperforms Adam. For example, test error rates (CIFAR-10/CIFAR-100): | Method | CIFAR-10 (%) | CIFAR-100 (%) | |----------|-------------|---------------| | SGD | 4.61 | 20.60 | | Adam | 6.14 | 25.51 | | ND-Adam | 4.53 | 21.45 |
• The gap in performance is even more pronounced with specialized regularization at the softmax output.

The improvement is attributed to ND-Adam’s ability to match SGD’s effective learning rate for every weight vector, while retaining Adam’s robustness and convergence speed.

4. The Role of Softmax Logit Regularization

Another critical source of generalization degradation is the variability in the absolute scale of softmax logits. In cross-entropy classification, scaling logits does not affect predictions but drastically alters the scale and structure of backpropagated gradients:

• Small logits distribute gradients almost equally among non-target classes, diminishing "confusable" class learning;
• Large logits isolate learning to fewer classes, neglecting the structure among remaining classes.

Solution strategies:

• BatchNorm Softmax: Batch normalize logits with a fixed, non-trainable scale.
• Logit $L_2$ Regularization: Penalize the $L_2$ norm of logits in the loss.

Both improve generalization, with Adam particularly benefiting due to its internal normalization. ND-Adam also profits, achieving further reductions in test error: | Method | CIFAR-10 (%) | CIFAR-100 (%) | |----------|-------------|---------------| | SGD | 4.49 | 20.18 | | Adam | 5.43 | 22.48 | | ND-Adam | 4.14 | 19.90 |

5. Implications for LLMing and Broader Context

Although the empirical demonstration in "Normalized Direction-preserving Adam" centers on image classification, the identified mechanisms extend to LLMing:

• Weight directionality and normalization issues pervade transformer architectures, word embeddings, and massive-vocabulary softmax layers.
• Proper logit regularization is crucial in LLMs given the large vocabulary and prevalence of rare words.
• The core principle—that aligning update directions with the span of gradients, and decoupling effective learning rate from weight magnitudes, leads to improved generalization—holds when applied to deep NLP models.

Moreover, these findings serve as a foundation for improved optimization recipes and regularization strategies in LLMing, especially in scenarios where traditional Adam exhibits stagnating perplexity and poorer generalization relative to carefully tuned SGD.

6. Summary Table: Comparison of Approaches

Method	Generalization	Direction Control	Magnitude Normalization	Softmax Regulation
SGD	Good	Span-preserving	Weight decay/L2 approx	Optional
Adam	Often poor	Misspecified	No	Needed
ND-Adam	Good/Best	Explicit	Yes	Synergistic

ND-Adam combines the convergence advantages of Adam with the generalization and update geometry advantages of SGD, especially when combined with logit regularization.

7. Conclusion

The Adam-SGD gap in LLMing arises from fundamental issues in update directionality and the conditioning of parameter spaces inherent to deep architectures. Adaptive optimizers such as Adam, by deviating from data-aligned updates and failing to control for weight scale, are susceptible to sharp minima and poor generalization. ND-Adam provides a principled solution by restoring update directionality and normalizing weight vectors, thereby bridging the Adam-SGD gap. Logit regularization further supports robust optimization and generalization, with practical significance in the context of LLMing—where softmax structure and scaling are central to model behavior and performance.