Closing the Curvature Gap: Full Transformer Hessians and Their Implications for Scaling Laws

TL;DR

This work closes a key gap in the second-order understanding of Transformers by deriving explicit Hessian expressions for full Transformer blocks, including LayerNorm and FFN components. By combining Gauss-Newton decompositions with Kronecker-structured representations, it provides per-block spectral-norm bounds that tie curvature to input statistics, LayerNorm scales, and architectural choices, and it offers a Taylor-series-based framework to analyze loss differences as dataset size grows. The contributions include the first full block-level Hessian formulas, theoretical bounds on curvature evolution with data, and empirical validation on Vision Transformers, linking optimization geometry to practical scaling laws. The results support curvature-aware training and provide mechanistic insight for data- and model-scale decisions in large-scale Transformers and related architectures.

Abstract

The lack of theoretical results for Layer Normalization and feedforward Hessians has left a gap in the study of Transformer optimization landscapes. We address this by deriving explicit second-order expressions for these components, thereby completing the Hessian characterization of full Transformer blocks. Our results generalize prior self-attention analyses and yield estimations for the role of each sublayer in curvature propagation. We demonstrate how these Hessian structures inform both convergence dynamics and the empirical scaling laws governing large-model performance. Further, we propose a Taylor-expansion-based framework for analyzing loss differences to quantify convergence trajectories. By extending Hessian theory to the full Transformer architecture, this work establishes a new foundation for theoretical and empirical investigations of optimization in large-scale deep learning.

Figures (12)

• Figure 1: Overview of our observations. Part (a) shows the loss function landscape, which is a surface in the parameters space, and how it changes as the dataset size increases. Part (b) shows the schematic view of a proposed method --- carry out an analysis of a Transformer's Hessian, which greatly impacts on a loss landscape convergence, leading to a sample size determination framework.
• Figure 2: Transformer architecture we use in our experiments
• Figure 3: Hessian entries visualization for an initialized model with one Transformer Block. We see the entire magnitudes' heterogeneity, while the Values corresponding blocks have larger values.
• Figure 4: Hessian entries visualization for a model trained for a number of epochs with one Transformer Block. We see the entire magnitudes' heterogeneity, while the Values-Values corresponding block has the largest values.
• Figure 5: Parameters' blocks norms and their Hessians' norms, calculated exactly on one batch containing 128 examples from the MNIST training dataset.

Theorems & Definitions (46)

• Theorem 1
• Theorem 2: Jacobian of LayerNorm
• Theorem 3: Hessian of LayerNorm
• Lemma 1: ReLU derivative and Hessian
• Theorem 4: Transformer block derivative
• Theorem 5: Hessian of the Transformer block \ref{['eq:transformer']}
• Theorem 6: Spectral-norm estimate of the Transformer Hessian
• Theorem 7: Convergence of Self-Attention and Transformer Blocks
• Definition 1: Matrix Norms
• Definition 2: Vectorization and Element-wise Operations