Closing the Curvature Gap: Full Transformer Hessians

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.

Machine Learning (cs.LG)

arXiv:2510.16927 [cs.LG]

(or arXiv:2510.16927v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2510.16927 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.