Closing the Curvature Gap: Full Transformer Hessians

arXiv:2510.16927 · cs.LG · Submitted 2026-08-21 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Closing the Curvature Gap: Full Transformer Hessians".

Tom: The provided text details complex mathematical derivations regarding the estimation of norms for LayerNorm derivatives and Hessians, specifically presenting Lemma 4.

Jane: First, who's behind it and why it matters.

Paper discussion segment 1: Tom: So, we’re talking about "Closing the Curvature Gap: Full Transformer Hessians." Essentially, the authors are filling a huge hole where we lacked theoretical results for LayerNorm and feedforward Hessians, which are core parts of every Transformer block. They derived explicit second-order expressions for these components to complete the Hessian characterization of full Transformer blocks.

Jane: That means they’ve moved beyond just looking at the self-attention part and now have a complete view of how the loss landscape behaves across an entire layer, including LayerNorm and FFNs. It’s like finally seeing every single bump and valley in a mountain, not just one ridge.

Lu: Exactly! They generalize prior self-attention analyses to give us estimations for the role of each sublayer in curvature propagation across the whole model structure. This is significant because it shows how curvature moves through the block, which informs scaling laws like those discussed by Chen et al. on compute-optimal training six seven.

Meng: I wonder if these explicit expressions are computationally expensive to derive, or if this is a purely theoretical exercise that gives us insights we can plug into practical optimization routines? We need something runnable for real training.

Lalam: I see it as providing the necessary blueprints for better learning dynamics. If we understand the curvature propagation, we can design AI systems whose training processes are inherently more robust and less prone to getting stuck in bad local minima early on.

Tom: That’s a great point, Lalam. It’s not just about theory; it’s about building better foundations for the actual AI training process. Jane, what do you make of their main contributions?

Jane: Their main contributions are deriving the first full Hessian expressions for Transformer blocks, explicitly including LayerNorm and FFNs, which fills a critical gap in prior analyses. Furthermore, they establish theoretical bounds on how the loss landscape evolves with dataset size, giving us a rigorous framework for understanding landscape stabilization.

Lu: That bound is key because it’s not just an observation; it’s a framework that allows us to predict convergence trajectories based on data size, which is what we need for robust AI design.

Tom: So, they aren't just describing the current state of optimization; they're building a predictive tool for how those landscapes change as we scale up. This paper sets the stage for understanding scaling laws more deeply. Where do you think this leads us next, Jane?

Jane: It really pushes us toward understanding critical batch size estimation and how data budgeting interacts with model complexity, which are huge practical hurdles in training large AI models twenty-one twenty-two.

Paper discussion segment 2: Tom: We’re diving into the summary of "Closing the Curvature Gap: Full Transformer Hessians" now. The authors summarize how they've assembled a complete blockwise Hessian for a Transformer layer by deriving expressions for m/FFN second derivatives and blockwise spectral-norm bounds, which closes a missing piece in second-order geometry for this architecture.

Jane: That means they’ve successfully aligned the high-level theory with the empirical curvature structure we observe in practice. They are taking what we’ve seen empirically and formalizing it mathematically using these derived Hessian structures.

Lu: What I find most exciting is that this assembly of components allows for a principled account of how Transformer curvature evolves with data and training, which is something previous studies just hinted at without a solid second-order treatment.

Meng: From an engineering standpoint, having these precise blockwise bounds helps us understand where the most sensitive parts of the architecture are during optimization, which could inform how we distribute our computational resources during training runs.

Lalam: I think this helps in designing systems with inherent stability. If we know exactly where the curvature is highest within a layer, we can potentially tune that specific part more carefully to ensure smooth learning.

Tom: So, it’s moving from descriptive analysis to prescriptive analysis; they are giving us the mathematical tools to prescribe how training should proceed based on the structure of the AI block itself. This is a big step for understanding generalization behavior.

Jane: And that moves us away from just hoping things converge nicely, toward having a principled way to analyze and control convergence trajectories based on the underlying geometry of the loss surface.

Paper discussion segment 3: Tom: Now we move into what the authors are actually suggesting as improvements. They propose a Taylor-expansion–based framework for analyzing loss differences, which is a new way to quantify convergence trajectories based on local geometry at the optimum w*.

Jane: That framework allows us to calculate exactly how much more data we need or how much stability we expect from a given training run by looking at the Hessian structure right around where the model settles. It turns abstract curvature analysis into something actionable for data budgeting.

Lu: This is incredibly powerful because it directly addresses the limitations of previous work that relied on just observing stabilization thresholds without a rigorous mathematical foundation thirty-four. They are providing a way to calculate those thresholds based on second-order information.

Meng: If we can use this to predict required sample sizes, it means we can optimize our compute budget much more intelligently, avoiding wasted training cycles that don't actually help the model learn. That’s a huge practical win for resource management.

Lalam: For me, this predictive capability is exciting because it suggests we can move toward AI systems that are inherently self-aware of their own learning needs and adjust their training dynamically based on how the landscape is evolving.

Tom: So, we’re talking about a framework where the geometry itself dictates the next steps in our training strategy, moving from reactive adjustments to proactive, curvature-aware strategies. This paper really deepens our understanding of convergence dynamics.

Conclusion: Jane: So, to wrap up this discussion on "Closing the Curvature Gap: Full Transformer Hessians," we’ve seen how the authors have provided explicit second-order expressions for LayerNorm and FFN Hessians, establishing a complete blockwise Hessian characterization. They also laid out a framework for analyzing loss differences based on local geometry at the optimum w*.

Tom: And they showed how this informs scaling laws by providing rigorous bounds on landscape evolution with dataset size, giving us actionable diagnostics for curvature-aware training and data budgeting. It’s a lot of heavy lifting in terms of theory, but the payoff is a much more principled way to approach optimization.

Lu: This work really solidifies the connection between architectural design and optimization geometry in Transformers, setting a new foundation for how we analyze these models going forward.

Meng: From an engineering perspective, this gives us tangible tools to manage training stability and resource allocation based on actual curvature metrics rather than just guessing.

Lalam: This research opens up avenues for designing AI systems that are more inherently stable because they can predict their own learning needs based on the geometry of the loss landscape.

Tom: Fantastic discussion, everyone. The work in "Closing the Curvature Gap: Full Transformer Hessians" is a major contribution to understanding how these models learn and scale. We’ll be keeping a close eye on how this theoretical foundation translates into real-world AI advancements. Thanks for tuning in!

Egor Petrov, Vladislav Meshkov, Nikita Kiselev, Andrey Grabovoy

Yandex · BRAIn Lab · Moscow State University

cs.LG

Submitted: 2026-08-21

Updated: 2026-08-25

Code: https://github.com/modernTalker/transformer_hessian

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 84/100

The gist: The provided text details complex mathematical derivations regarding the estimation of norms for LayerNorm derivatives and Hessians, specifically presenting Lemma 4.

Key concepts

Full Transformer Hessians
The paper derives explicit second-order expressions for LayerNorm and FFN Hessians within Transformer blocks. This provides a complete characterization of the Hessian for the entire block, moving beyond previous analyses that focused only on self-attention.
Curvature Propagation
This refers to how curvature—the shape of the loss landscape—moves through a Transformer model structure. The authors generalize prior self-attention analyses to estimate this role in curvature propagation across all sublayers of the model.
Scaling Laws
The paper establishes theoretical bounds on how the loss landscape evolves with dataset size. This framework allows researchers to predict convergence trajectories based on data size, which is crucial for understanding scaling laws in large AI models.

Terminology

Summary

The provided text details complex mathematical derivations regarding the estimation of norms for LayerNorm derivatives and Hessians, specifically presenting Lemma 4. The summary of this technical content is as follows:

Lemma 4 establishes specific estimations for the norm of the LayerNorm derivative, JLN(X), and its Hessian, HLN(X), where X in R m times n.

1. Bounds Statement:

The lemma asserts that the following estimation holds:

JLN(X) 2 sqrtX squared over 1 + X squared times sqrt 3 = X squared + sqrt 3

(Equation 28)

And for the Hessian norm:

HLN(X) 2 3 + sqrt 3 + sigma over n + sigma over m

(Equation 29)

Where sigma denotes M i 2, and the bounds are derived using properties established in the preliminaries and Theorems 2–3.

2. LayerNorm Jacobian Structure and Bound (JLN(X) 2):

The proof of the Jacobian bound proceeds by first determining the structure of JLN(X). By Theorem 2 (with L to m, dV to n), JLN(X) is structured as:

JLN(X) = (P I n) G + (I m M) H

where G = I mn - n 1 1 n times n, H = d X, and P = diag-1(sigma).

Using Properties 9, 10, and 8, the norm is bounded:

JLN(X) 2 P squared G squared + M squared H squared = P squared G squared + M squared H squared

The individual factors are bounded as follows:

  • ** G 2:** It is bounded by 1, since n 1 1 n times n is a projection, leading to I n - n 1 1 n times n 2 1.

  • ** P 2:** This factor is bounded by 1 over sigma, where D = diag(sigma).

  • ** M 2:** This factor is bounded by X squared, because M(X) = X (I n - n 1 1 n times n) and the right factor is a projector with norm 1.

Finally, the bound for JLN(X) 2 is collected to yield Equation 28:

JLN(X) 2X squared over 1 times 1 + X squared times sqrt 3

3. LayerNorm Hessian Structure and Bound (HLN(X) 2):

The Hessian HLN(X) is derived from Theorem 3 (with m, n) and has the structure:

HLN(X) = (I mn G) d squared P over d + (I m M) d squared P over d X squared + (I mn H) d squared P over d X squared

The proof bounds the three terms separately using Properties 10 and 9:

  • (i) First Term: The bound is found to be G 2 I m vecr(I n) 2 = 1 times n times sqrt 3 = 3.

  • (ii) Second Term: This term involves bounding d squared P over d X squared. Following the chain of differentiation, it is established that:

d squared P over d X squared 1 over 3 X squared

This leads to a bound for the second term:

(I m M) I mn d squared P over d X squared 3 X squared + 5 n sigma squared

  • (iii) Third Term: The bound is found to be H 2 vecr(I m) I mn 2 = sqrt 3 times m times 1 = sqrt 3.

Summing these three components using Property 8 yields the final bound for the Hessian norm, Equation 29:

HLN(X) 2 3 + sqrt 3 + 5 over 3 + sigma over n + sigma over m

Improvements for AI systems

The findings in this paper provide a complete, rigorous characterization of the optimization landscape for the full Transformer block, bridging the gap between theoretical Hessian analysis and practical scaling laws. This allows for several highly specific improvements to AI systems that go far beyond current state-of-the-art optimization techniques.

Here are the specific improvements and capabilities of an improved AI system based on this research:

Improvement: Implementation of a dynamic, block-wise preconditioning strategy for large Transformer models, moving beyond standard adaptive optimizers (like Adam or AdaGrad).

Mechanism: The paper establishes that the Hessian structure is highly heterogeneous (H ij), with specific blocks—such as the Key and Value attention matrices—exhibiting significantly higher curvatures than others.

The improved system uses the derived spectral norm estimates (M and M tr) to calculate a dynamic precondition P(w) that is tailored to the local curvature of each sub-block. Instead of applying one global precondition, the system applies:

Update = P-1(w) times grad L(w)

This preconditioning is calculated using the specific Hessian bounds derived in Theorem 6 (H ij squared) and allows for a tailored learning rate scaling factor gamma ij for each parameter set W i, ensuring that optimization progresses at an optimal rate without premature stalling or divergence, even when dealing with extreme curvature variance.

Improvement: A predictive framework for determining the optimal balance between dataset size (k) and model architecture complexity, minimizing computational waste.

Mechanism: The system utilizes Theorem 7, which quantifies the loss difference L k+1(w) - L k(w) based on a Taylor expansion of the local geometry at the optimum w*. This allows us to calculate:

Required k about Max Curvature M

For a given target performance level, the system can calculate the minimum number of samples (k) required to ensure that further data collection does not yield significant improvements in local loss stability. Conversely, if computational resources are constrained, it can determine the minimum model size required to achieve a stable loss landscape under a fixed data budget. This replaces heuristic scaling with a rigorous, curvature-based determination of resource requirements.

Improvement: Integrating Hessian information into Neural Architecture Search (NAS) algorithms to guide the design of robust and efficiently trained Transformer blocks.

Mechanism: The system uses the detailed Jacobian and Hessian components (H ij from Theorem 5) to identify structural weaknesses in a proposed architecture. If a specific block's curvature falls below a critical threshold (i.e., sigma squared is too small, indicating a flat loss landscape), the the NAS algorithm is penalized or forced to iterate on that component. This ensures that the resulting architecture has curved enough to prevent vanishing gradients and guarantee stable convergence, leading to inherently more robust models than those designed without second-order analysis.

Improvement: A real-time monitoring system for detecting potential training instability or premature convergence using local curvature metrics.

Mechanism: The system continuously calculates the spectral norm of the empirical Hessian (H(w) squared) during training. By comparing this real-time value to the theoretical bounds derived in Theorem 1, the system can trigger alerts when:

a) The gradient magnitude is excessively high relative to local curvature (suggestive of instability).

b) The Hessian norm stabilizes too quickly, indicating that the model may have reached a local minimum before achieving global optimality, prompting an adaptive restart or architectural modification.

Abstract

The optimization landscape of Transformer models remains poorly understood despite their widespread adoption. While recent studies have derived curvature properties for isolated self-attention mechanisms, a comprehensive theoretical characterization of the full Transformer block, accounting for the interactions between Layer Normalization, Feed-Forward Networks (FFNs), and residual connections, is missing. In this work, we close this gap by deriving the exact, closed-form Hessian for the complete Transformer block under arbitrary twice-differentiable loss functions. We utilize rigorous matrix calculus to handle the non-linearities of LayerNorm and row-wise activations, establishing explicit spectral norm bounds for the resulting Hessian blocks. Our analysis reveals how different architectural components contribute distinct curvature mechanisms, identifying the specific curvature contributions of particular sub-layers. Furthermore, empirical validation against automatic differentiation confirms the exactness of the derived formulas up to numerical precision and shows substantial computational speedups for the closed-form Jacobian evaluations.

Sources

Related papers