🤖 machine learning

Closing the Curvature Gap: Full Transformer Hessians

This paper closes a theoretical gap in understanding Transformer optimization by deriving the exact, closed-form Hessian for the complete Transformer block, explicitly accounting for interactions between Layer Normalization, Feed-Forward Networks, and residual connections, while validating these formulas with empirical speedups and spectral norm bounds.

Original authors: Egor Petrov, Nikita Kiselev, Vladislav Meshkov, Andrey Grabovoy

Published 2026-08-25

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Egor Petrov, Nikita Kiselev, Vladislav Meshkov, Andrey Grabovoy

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the vast landscape of modern artificial intelligence, a specific architecture known as the Transformer has become the dominant force, powering everything from language translation to image recognition. These systems are built from layers of mathematical operations that process information in parallel, but their inner workings remain somewhat mysterious to the scientists who design them. While we know these models work, the precise mathematical reasons for their stability and the specific ways they learn are not fully mapped. A central tool for understanding this behavior is the Hessian, a complex mathematical object that describes the curvature of the learning process. Imagine the training of a model as a journey across a hilly terrain; the Hessian tells us exactly how steep the slopes are and how the ground curves beneath our feet. This information is crucial because it determines whether an algorithm will slide smoothly to a solution or get stuck in a difficult spot. Until now, researchers had only managed to calculate this curvature for isolated parts of the system, leaving the full picture of how the entire machine behaves incomplete.

A team of researchers has now closed this gap by deriving the exact mathematical description of the curvature for a complete Transformer block. This block is the fundamental unit of the architecture, composed of several interacting parts: a mechanism that weighs the importance of different pieces of information, a normalization step that keeps values in check, a feed-forward network that adds non-linear processing, and residual connections that allow information to flow around these layers. The team did not rely on approximations or simulations; instead, they used rigorous mathematical techniques to write down the precise formulas for how the entire block curves in response to changes in its weights. They treated the system as a whole, accounting for how the normalization layer and the feed-forward network interact with the attention mechanism, rather than studying them in isolation.

The analysis reveals that the curvature of the learning landscape is not uniform but is shaped by distinct contributions from each architectural component. The attention mechanism, which allows the model to focus on specific parts of the input, introduces a type of curvature that is highly sensitive to the size of the input data. The normalization layer, which stabilizes the training process, adds its own curvature based on the variance of the data, making the landscape sensitive to how spread out the values are. The feed-forward network, which uses a simple rule to decide which signals to pass through, contributes curvature primarily through the interactions between its different weight matrices, while the residual connections act as a bridge that controls how these curvatures propagate through the system. The researchers found that these different sources of curvature combine in specific ways, creating a complex landscape where some directions are very steep and others are relatively flat.

To ensure their formulas were correct, the team compared their theoretical results against the standard methods used by computer software to calculate derivatives. The match was exact, down to the limits of computer precision, confirming that their closed-form equations accurately describe the system. Beyond verifying the math, the study demonstrated that using these explicit formulas is significantly faster than the standard computational methods for certain calculations. In tests, the new formulas provided a speedup of over eighty times for some components and hundreds of times for others. This efficiency suggests that understanding the exact curvature of these models is not just a theoretical exercise but a practical tool that could help engineers analyze and improve the training of large artificial intelligence systems more effectively. The work provides a clear, unified view of how the different parts of a Transformer work together to shape the learning process, moving the field from partial understanding to a complete mathematical characterization.

Technical Summary: Closing the Curvature Gap: Full Transformer Hessians

Problem Statement

Despite the dominance of Transformer architectures in modern deep learning, the mathematical understanding of their optimization landscape remains incomplete. While recent theoretical work has successfully derived curvature properties for isolated self-attention mechanisms, a comprehensive characterization of the full Transformer block is lacking. Specifically, existing analyses often omit the complex interactions between self-attention, residual connections, Layer Normalization (LayerNorm), and Feed-Forward Networks (FFNs). These components are widely believed to be critical for training stability, yet their exact second-order contributions to the Hessian have not been unified in a single analytical framework. This gap hinders a rigorous understanding of how architectural choices influence conditioning, curvature, and the stability of gradient-based optimization.

Methodology

The authors address this gap by deriving the exact, closed-form Hessian for a complete post-normalization Transformer block under arbitrary twice-differentiable loss functions. The methodology relies on rigorous matrix calculus and a specific notation system to handle the non-linearities and structural complexities of the model:

Row-Wise Vectorization: To align with the token-wise processing nature of Transformers, the authors adopt row-wise vectorization ( $\text{vec}_r$ ) rather than standard column-wise vectorization. This allows derivatives to be expressed as explicit linear operators with consistent dimensions.
Gauss-Newton Decomposition: The Hessian of the composite loss function $L(F(w))$ $L (F (w))$ is decomposed into two distinct terms:
- The Outer Term: Induced by the curvature of the loss function projected onto the parameter space via the model's Jacobians.
- The Functional Term: Induced by the curvature of the model map $F$ itself, interacting with the loss gradient. This term is particularly significant in Transformers due to non-polynomial operations like Softmax and LayerNorm.
Component-wise Derivation: The authors systematically derive Jacobians and Hessians for each sub-layer:
- Self-Attention: Extending prior work on isolated attention to include interactions with other blocks.
- LayerNorm: Deriving exact formulas for the Jacobian and Hessian of row-wise normalization, treating affine parameters as constants to focus on weight-induced curvature.
- Feed-Forward Networks (FFN): Modeling the standard non-linear formulation (e.g., ReLU) and deriving Jacobians and interaction terms, noting that ReLU has a zero second derivative almost everywhere.
Spectral Norm Analysis: Using the derived closed-form expressions, the authors establish upper bounds on the spectral norms of the Hessian blocks to analyze scaling properties and curvature heterogeneity.

Key Contributions

Exact Second-Order Characterization: The paper provides the first closed-form Jacobians and Hessian blocks for a full Transformer block, encompassing self-attention, residual connections, LayerNorm, and nonlinear FFNs. This extends previous analytical studies limited to isolated self-attention layers.
LayerNorm and FFN Second-Order Calculus: The authors derive explicit matrix-calculus formulas for the Jacobian and Hessian of row-wise LayerNorm and the interaction terms of the FFN. These formulas explicitly show how inverse variance, activation gates, and residual streams enter the full Hessian.
Spectral-Norm Bounds and Curvature Heterogeneity: The study establishes upper bounds on the spectral norms of Hessian blocks. These bounds reveal distinct curvature mechanisms:
- Self-Attention: Exhibits high-order input dependence (scaling as $O(\|X\|_2^6)$ ), indicating extreme sensitivity to input norms.
- LayerNorm: Introduces curvature sensitive to row-wise variance, making the Hessian susceptible to variance collapse.
- FFN: Contributes primarily through activation-gated Jacobians and cross-layer interactions, as the ReLU second derivative vanishes almost everywhere.
Empirical Validation and Computational Efficiency: The derived formulas are validated against PyTorch automatic differentiation (autograd), showing numerical agreement up to floating-point precision. Furthermore, the closed-form expressions yield substantial computational speedups (e.g., 82 $\times$ for LayerNorm Jacobians, 444 $\times$ for FFN Jacobians) compared to autograd for componentwise evaluations.

Results

Validation: Experiments confirm that the theoretical formulas reproduce autograd results with relative errors below $10^{-7}$ for complex components like LayerNorm Hessians and Transformer blocks, confirming the correctness of the chain-rule assembly without approximation.
Curvature Scaling: The analysis confirms that the self-attention Hessian is dominated by terms scaling with the sixth power of the input norm ( $O(\|X\|_2^6)$ ). This suggests that pre-LayerNorm placement (controlling input norms before attention) is crucial for stabilizing curvature, aligning with empirical observations.
Sequence Length: Several curvature terms scale inversely with sequence length (e.g., $1/L^3$ ), implying that the intrinsic curvature of the attention mechanism may naturally dampen for very long sequences.
Non-Convexity: The decomposition reveals that even if the loss function $L$ is convex, the full optimization landscape can become non-convex due to the "Functional Term" arising from the architecture (specifically LayerNorm and Softmax).

Significance and Claims

The paper claims to "close the gap" in the theoretical understanding of Transformer optimization by providing a unified, exact analytical framework for the full block. The significance of this work lies in:

Mechanistic Insight: It explicitly resolves the "cryptic" functional term in the Gauss-Newton decomposition, showing how specific architectural components (residuals, normalization, activations) modulate the curvature landscape.
Training Stability: The derived spectral bounds provide a principled way to reason about Transformer conditioning and training stability, highlighting that full-block curvature is not merely inherited from self-attention but is modulated by normalization and residual pathways.
Practical Utility: The closed-form expressions serve as practical tools for curvature analysis, offering significant computational advantages over automatic differentiation for specific derivative evaluations, which could facilitate more efficient second-order optimization methods or curvature-based diagnostics in large-scale models.

The authors maintain a modest stance, noting that their primary contribution is theoretical derivation and validation, with the computational speedups serving as a practical benefit of the derived formulas rather than a proposal for a new training algorithm.