Skip to content

Gradient Heterogeneity Complements Hessian Heterogeneity in Transformer Optimization

Jan 2025 · 6 citations · ⚡ 1 influential
Computer Science

TL;DR

This study provides a theoretical analysis showing that gradient heterogeneity, together with Hessian heterogeneity, degrades the convergence of gradient-based methods such as SGD, while sign-based methods are substantially less sensitive to this effect.

Abstract

Transformers are difficult to optimize with stochastic gradient descent (SGD) and largely rely on adaptive optimizers such as Adam. Despite extensive efforts, the mechanisms behind Adam's advantage over SGD in Transformer optimization are still not fully understood. In this study, we analyze the optimization of Transformer models in the fine-tuning setting through the lens of gradient heterogeneity, defined as the variation in gradient norms across parameter blocks. We provide a theoretical analysis showing that gradient heterogeneity, together with Hessian heterogeneity, degrades the convergence of gradient-based methods such as SGD, while sign-based methods are substantially less sensitive to this effect. Adam and SignSGD both perform coordinate-wise updates and are less sensitive to the scale of individual gradient coordinates than SGD. This motivates our use of SignSGD as an analytically tractable proxy for Adam-like behavior. Our analysis uses the fact that SGD and SignSGD follow steepest descent directions under different norms, and derives upper bounds on the iteration complexity with implications for learning-rate scaling for SignSGD. We further investigate the origin of gradient heterogeneity in Transformer architectures and show that it is strongly influenced by the placement of layer normalization, with Post-LN architectures exhibiting particularly pronounced heterogeneity. Experimental results from fine-tuning Transformers in both NLP and vision domains validate our theoretical analysis.

View source

Similar papers

Preprint Aug 2026

On the Principles Behind Neural Network Optimizers

Reliable optimization is central to neural network (NN) training, yet Adam, the default optimizer for modern LLMs, rests on a fragile foundation. This thesis develops a principled grounding for Adam and motivates new designs. First, we revisit Adam's divergence--convergence debate and show the existence of a problem-dependent phase transition: with properly chosen, batch-size-dependent hyperparameters, Adam converges, whereas under small-$\beta_2$ regimes it can diverge. Second, we investigate why Adam substantially outperforms SGD on Transformers through Hessian structure. We find that the Hessian evolves toward a near-block-diagonal form along training, accompanied by strong block heterogeneity. We prove that this structure makes Adam's diagonal preconditioner effective. We further show that this special Hessian structure originates from consecutive multiplications of large matrix variables, and we provide a rigorous analysis based on random matrix theory. Finally, these insights motivate Adam-mini, a new optimizer that reduces Adam's memory footprint by 50\% while preserving its performance. Our results also have broader implications beyond Adam: they reveal new local structures in matrix-based nonconvex problems, and also help understand and improve recent NN optimizers, such as Muon.

Yushun Zhang · 0 citations
Open access Aug 2026

Quartz Optimizer: Robust Gradient Shaping and Bounded Adaptive Steps for Stable Deep Learning Training

Optimization plays a critical role in training deep neural networks, directly impacting convergence speed, model generalization, and stability. While existing methods such as stochastic gradient descent (SGD) and adaptive optimizers like Adam and AdamW have achieved significant success, they exhibit limitations in handling extreme gradients and noisy updates, and maintaining stable convergence across diverse architectures and datasets. In this study, we propose Quartz, a novel optimizer that combines momentum accumulation with bounded adaptive scaling to improve convergence efficiency and robustness. Quartz introduces a gradient-saturation mechanism that prevents excessively large updates and enforces adaptive step size bounds, thereby addressing the key limitations observed with conventional optimizers. We evaluate Quartz across three benchmark datasets (MNIST, Fashion-MNIST, and Arabic Character Classification) using multiple convolutional neural network (CNN) architectures, including GoogLeNet, VGGNet, and ResNet-18. The experimental results demonstrate that Quartz achieves up to 99.54% test accuracy on MNIST, 91.6% test accuracy on Fashion-MNIST, and 98.27% test accuracy on Arabic Character Classification. It consistently outperformed or matched the results achieved using state-of-the-art adaptive optimizers under identical training conditions. In terms of efficiency, Quartz reduces training time in several settings while maintaining higher accuracy. Across all the experiments, Quartz also shows statistically significant improvements (p< 0.001 in most comparisons) and smoother convergence behavior, indicating improved optimization stability. These findings highlight Quartz’s potential as a reliable and efficient tool for use in optimization strategies for a broad range of deep learning tasks.

Ahmad Raza Khan, Sarab Almuhaideb · 0 citations
Open access Jul 2026

The Gradient Descent Based Adaptive Optimizer Algorithm in Deep Learning

Gradient descent-based optimization algorithms are widely used in deep learning, where adaptive optimizers dynamically adjust parameter updates to improve training efficiency. However, most existing adaptive optimizers mainly rely on historical gradients while overlooking short-term gradient fluctuations. Although methods such as diffGrad and AngularGrad attempt to address this issue through gradient-difference or angle-based correction mechanisms, their performance improvements are often task-dependent, and some approaches introduce additional computational overhead due to complex vector operations. To address these limitations, this paper proposes a lightweight short-term gradient-aware mechanism based on gradient sign consistency, which improves optimization stability while maintaining low computational complexity. Based on this idea, two Adam-based optimizers, sigSignGrad and tanhSignGrad, are developed by incorporating adaptive friction coefficients constructed through Sigmoid and Tanh mappings, respectively. In addition, the proposed friction coefficient can be flexibly integrated into the Adam optimizer family as a lightweight plug-in module. Extensive experiments are conducted across multiple computer vision and natural language processing tasks, including image classification, object detection, semantic segmentation, machine translation, language modeling, and language understanding benchmarks. To evaluate the robustness and generalization capability of the proposed methods, multiple datasets and model architectures are considered within the same task settings, covering both training-from-scratch and pre-trained fine-tuning scenarios. Experimental results show that sigSignGrad achieves the best overall performance on multiple visual tasks, including all evaluated Vision Transformer benchmarks, Pascal VOC object detection, and semantic segmentation datasets. In natural language processing tasks, tanhSignGrad achieves the best BLEU scores on the IWSLT14 De-En translation task and consistently outperforms existing optimizers in language modeling experiments under both pre-trained and training-from-scratch settings. Furthermore, the proposed methods maintain computational efficiency comparable to diffGrad while providing more stable optimization behavior across diverse architectures and datasets.

Rongjian Chen, Ruyan Xiong, Xiaodong Yang · 0 citations
Preprint Jul 2026

Leveraging Extragradient for Effective Sharpness-Aware Minimization in Deep Learning

Generalization remains a pivotal challenge in deep learning, where traditional optimizers like Stochastic Gradient Descent (SGD) often converge to sharp minima, leading to overfitting and reduced performance on unseen data. Building on Sharpness-Aware Minimization (SAM), for seeking flat minima associated with improved generalization, we propose the Extragradient-Inspired Sharpness-Aware Minimization (EISAM), a novel optimizer that enhances generalization via the extragradient technique. EISAM uses a two-step update process: a prediction step investigating the geometry of the loss landscape and a perturbation step that refines updates with a base optimizer. This approach achieves better generalization performance than SAM. Crucially, EISAM reduces sensitivity to the perturbation radius, enhancing robustness, and simplifying the tuning across diverse settings. Extensive experiments on benchmark datasets demonstrate that EISAM consistently outperforms SGD, Adaptive Moment Estimation (Adam), and SAM in test accuracy and training efficiency across various architectures. Theoretical analysis further confirms that EISAM tightens the generalization bound by steering parameters toward flatter minima with reduced curvature. Accompanied by a thorough hyperparameter analysis, EISAM offers practical tuning guidance, establishing it as a robust, scalable, and broadly applicable optimization solution that advances both the theory and practice in deep learning.

Yao Fu, Chunxia Zhang, Junmin Liu et al. · 0 citations
Preprint Aug 2026

Strong error analysis for the stochastic momentum optimizer

Stochastic gradient descent (SGD) optimization schemes are the methods of choice for the optimization of deep neural networks (DNNs) in artificial intelligence (AI) systems. Often not the standard SGD method is used but instead suitable accelerated, adaptive, and/or normalized variants of standard SGD such as Adam, AdamW, and MUON are employed to train large scale AI systems in practically relevant settings. The acceleration (higher order convergence speed) in all these popular optimizers relies on the momentum SGD optimizer. In this work we provide a rigorous error analysis for the momentum SGD optimizer. In particular, we establish convergence rates for the momentum optimizer in terms of the size of the learning rate (step size), the size of the mini-batch, and the size of the one-point convexity constant.

Davide Gallon, Arnulf Jentzen · 0 citations
Open access Jul 2026

Hybrid Switching Optimization Strategy for Efficient Training of Deep Neural Networks on the MNIST Dataset

Deep neural networks often encounter non-convex optimization challenges during training due to the presence of local minima, saddle points, and complex loss surfaces. Existing optimization algorithms such as Adam and Stochastic Gradient Descent (SGD) offer complementary advantages—Adam provides faster convergence, while SGD tends to achieve better generalization. However, neither optimizer alone effectively balances both properties in non-convex settings. To address this limitation, this paper proposes a phase-switch hybrid optimization strategy that combines the strengths of Adam and SGD. The proposed method employs Adam during the initial phase of training to enable rapid convergence and efficient exploration of the loss landscape, and then switches to momentum-based SGD in the later phase to improve generalization and ensure stable convergence. The effectiveness of the proposed approach is evaluated on one benchmark dataset, MNIST dataset, under different learning rate settings. Experimental results demonstrate that the proposed method achieves performance that is superior or comparable to existing optimizers in terms of accuracy and loss minimization. These results indicate that the proposed hybrid optimization strategy provides a simple and effective solution for handling non-convex optimization problems in deep learning.

Harish Kunder, Manjunath Kotari · 0 citations

Related blog posts