What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis
Authors: Weronika Ormaniec
Felix Dangel
Sidak Pal Singh
Venue: ICLR 2025 Spotlight (2025)
Keywords: Hessian
Transformers
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Abstract: The Transformer architecture has inarguably revolutionized deep learning, overtaking classical architectures like multi-layer perceptions (MLPs) and convolutional neural networks (CNNs). At its core, the attention block differs in form and functionality from most other architectural components in deep learning—to the extent that, in comparison to MLPs/CNNs, Transformers are more often accompanied by adaptive optimizers, layer normalization, learning rate warmup, etc. The root causes behind these outward manifestations and the precise mechanisms that govern them remain poorly understood. In this work, we bridge this gap by providing a fundamental understanding of what distinguishes the Transformer from the other architectures—grounded in a theoretical comparison of the (loss) Hessian. Concretely, for a single self-attention layer, (a) we first entirely derive the Transformer’s Hessian and express it in matrix derivatives; (b) we then characterize it in terms of data, weight, and attention moment dependencies; and (c) while doing so further highlight the important structural differences to the Hessian of classical networks. Our results suggest that various common architectural and optimization choices in Transformers can be traced back to their highly non-linear dependencies on the data and weight matrices, which vary heterogeneously across parameters. Ultimately, our findings provide a deeper understanding of the Transformer’s unique optimization landscape and the challenges it poses.

Published as a conference paper at ICLR 2025
WHAT DOES IT MEAN TO BE A TRANSFORMER? INSIGHTS FROM A THEORETICAL HESSIAN ANALYSIS
Weronika Ormaniec*
Felix Dangel
Sidak Pal Singh
ETH Zürich Zürich, Switzerland
Vector Institute
ETH Zürich
Toronto, Canada
Zürich, Switzerland
ABSTRACT
The Transformer architecture has inarguably revolutionized deep learning, over- taking classical architectures like multi-layer perceptrons (MLPs) and convolu- tional neural networks (CNNs). At its core, the attention block differs in form and functionality from most other architectural components in deep learning- to the extent that, in comparison to MLPs/CNNs, Transformers are more often accompanied by adaptive optimizers, layer normalization, learning rate warmup. etc. The root causes behind these outward manifestations and the precise mech- anisms that govern them remain poorly understood. In this work, we bridge this gap by providing fundamental understanding of what distinguishes the Trans- former from the other seshhitccturs-grruunce (loss) Hessian. Concretely, for single self-attention layer, (a) we first entirely derive the Transformer's s Hessian and express it in matrix derivatives; (b) we then characterize it in terms of data, weight, and attention moment dependencies; and (c) while doing so further highlight the important structural differences to the Hes- sian of classical networks. Our results suggest that various common architectural and optimization choices in Transformers can be traced back to their highly non- linear dependencies on the data and weight matrices, which vary heterogencously across parameters. Ultimately, our findings provide deeper understanding of the Transformer's unique optimization landscape and the challenges poses.
a theoretical comparison of the
INTRODUCTION AND RELATED WORK
The Transformer architecture (Vaswani et al., 2017) has shown remarkable success across natural language processing (Devlin et al., 2018; Radford et al., 2018) and vision (Dosovitskiy et al., 2020) tasks. In particular, its self-attention mechanism has become a mainstay of modern architectures, enabling parallelization while effectively capturing long-range and modality-agnose dependencies. Yet, despite their impressive performance, significant gap remains in our understanding of the properties of Transformer-base models relative to traditional architectures such as multi-layer perceptrons (MLPs, Rosenblatt, 1958) or convolutional neural networks (CNNs, LeCun et al., 1998).
Transformers' unique architectural built. Compared to the classical architectures, Transform- ers are unique in several ways. Firstly, the data (tokens) enters the Transformer, through the self- attention, multiple times. Secondly, the softmax nonlincarity inside self-attention differs from the piece-wise linear activations, like ReLU (Glorot et al., 2011). Thirdly, the query-key attention incor- porates two (instead of one) directly multiplied weight matrices within a single architectural block.
The side-factors making Transformers' unique architecture work. These architectural dif- ferences render practical approaches for training Transformer distinct compared to classical nets (Popel & Bojar, 2018; Bengio, 2012) F.g., Transformers are usually trained with adaptive op- timizers like Adam(W) (Kingma & Ba, 2015; Loshchilov & Hutter, 2019) and require architectural extensions such as skip connections (He et al., 2016) and layer norm (Xiong et al., 2020), learning rate warm-up (Goyal et al., 2017), and using different weight initializations (Huang et al., 2020).
Aim. Given these outward and indirect manifestations of the Transformer's presence, it is unclear how these are explicitly triggered due to particular loss landscape geometry endowed by Transform- ers. To address this, we theoretically investigate the Transformer's loss landscape by analyzing, in detail, the structure of the Hessian matrix as well as its data dependency and behavior across layers.
"Correspondence wormanice @ @ethz.ch, fdangel @ vectorinstituue...... contact sidakpal.com.

---

Published as a conference paper at ICLR 2025
H(Wy, Wv)
H(Wq. Wal
Tranfornce: Black
Self-Attection
Self Ant Eibh.
2500
M
Logalt lank Aboohite Edttiios
Reqpasty
WV
5000
MLP
10-7
10-3
Absolute Entrios
(a) Histogram of (absolute) en- tries corresponding the value and query diagonal blocks.
(b) On the left, the Hessian matrix for minimal Transformer. and, the right, the zoomed-in block query, key, and value parameters.
Figure Disparity in the Hessian blocks of Transformer seen quantitatively and qualita- tively. We used a single-block GPT-2 Transformer initialization applied to the next token predic- tion task (for details see appendix F). We observe block heterogeneity in the magnitudes of Hessian antries-those of the query block are significantly smaller than those of the value block.
Prior Hessian-related work. The Hessian matrix is fundamental object for optimization (Martens, 2010; Schaul et al., 2013; Cohen et al., 2021), generalization (Keskar et al., 2016; Jiang et al., 2019), and more (LeCun al., 1989; Singh & Alistarh, 2020). (a) Empirical work: For traditional architectures, the Hessian has received a lot of attention through empirical studies on its bulk-outlier spectrum and eigenvalue density (Sagun al., 2016; 2017; Ghorbani et al., 2019; Papyan, 2020; Yao et al., 2020) and in turn how it is affected by architectural components and how it changes during training. (b) Theoretical studies classical architectures: More recently, several theoretical works have also analyzed, in detail, the Hessian structure and rank (Singh et al., 2021; 2023), beyond Random- Mandom-Matrix-TTine Matrix-Theory based approximations (Pennington Bahri, 2017). However, this latter line of work has been restricted to folly-connected and convolutional architectures.
reettee Hessian studies. Park & Kim (2022) provide empirical evidence that the multi-head self-attention mechanism in Transformers leads to a more non-conver but smoother loss landscape than that of CNNs, the Hessian has
negative eigenvalues, but of smaller magnitude. Zhang et al. (2024) studied the Hessian spectra of Transformers to explain the need for adaptive optimizers for successful training. They empirically showed that although the full Hessian spectra of Transformers are quite similar to those of CNNs, the Hessian diagonal blocks in Transformers are much more heterogeneous-a possible cause for the need for adaptive optimizers. Barring these few empirical studies, to the best of our knowledge, theoretical treatment of the Hessian in Transformers remains lacking.
Goals and contributions. By theoretically analyzing the Hessian of Transformers, we aim iden- tify and explicitly state the differences to classical architectures. We believe this provides the foun- dation for a deeper understanding of the unique loss landscape features and challenges the Trans- former poses. Our detailed contributions are:
(i) We derive the Hessian of single self-attention layer and exhibit the structure contained within its well-known positive-definit and indefinite parts (theorems 3.1 and 3.2).
(ii) We categorize its dependencies into data, weights, and attention moments, find that the Hessian highly non-linear and heterogeneous for different parameter groups within a self-attention block, and show how various Transforrmer-elltted tricks can be understood through these characteristics.
(iii) We explicitly establish how individual components of self-attention, like the softmax and query-key- guera-kee-parrmmtteiiiaiii parameterization, result in more non-linear and heterogeneous structure within the Hessian (see fig. 1). then contrast the Transformer's Hessian with that of traditional architectures and uncover pronounced differences.
2 SETUP AND BACKGROUND
Single self-attention layer. We consider single self-attention layer with slightly generalized definition, which will later enable us to investigate the impact of specific components. The layer maps token embeddings X Rlxdy of sequence of length I. with embedding dimension dy
2

---

Published as a conference paper at ICLR 2025
OF
i vecr aw;
OF
vec,
OF
F²F
F
Vec,-(F)
OW;
i vec, W;
vec,
aw;
aW WWWW
b vvec W
Figure 2: Construction of the second derivative matrix of matrix-valued network F. Taking the second derivative of F using row-wise vectorization and numerator layout is equivalent to computing the second derivatives of each entry separately and stacking them into column block matrix.
into a new sequence F(X) Rxxdy by computing the values X XWv using value weights Wy Rd xd and re-weighting them with the self-attention matrix (or map) A(X) E R LLx f.e.
F(X) = X where
= a
(I)
Here, we decompose the attention map into two components, the query-key similarity transfor- mation T(X) E R 1.x which may introduce additional trainable parameters, and the activation function a : R.x. + RExI. The classical self-attention from Vaswani et al. (2017) uses T(X) = XWQ W Vdxx, where a softmax (row-wise), with learnable query and key weight matrices Wa, WK E Rdvxdx, Further,
single self-attention layer from eq. (1) feeds into a loss function f R fxxdv Exdy + R that measures the discrepancy between the predicted sequence F ((X) and the sequence labels Y. We will assume square loss (where ||r denotes the Frohenius norm of a matrix), ((F(X), Y) [F(()) YII (T.dy)
Hessian, block structure, and Gauss-Newton decomposition. Our goal is to compute the full Hessian of the loss w.r.t. the learnable, matrix-shaped, parameters {W; E R&xxP.) J of the attention layer (for instance E {K, Q, V for canonical self-attention). To break down this task, first notice that this matrix consists of blocks O² (e o F)/OW, WWW,, j.
Further, the Hessian decomposes into two components as a consequence of the chain rule applied to the functional split 8o F.
S (E F) aw aw () aw,
OF ())
02 OF
OF ow
OE
8 F
F
())
(F(-)
(2)
OF
pcgs
ow,ow aaw,
E-outer-product Hessian H, (W;.W.)
F-functional Hessian H, (W.,W,)
Equation (2) is known as Gauss-Newton decomposition and such splits can be widely seen in the literature (Sagun et al., 2017; Kunstner et al., 2019; Martens, 2020; Dangel et al., 2020; Singh al., 2021). We adopt the terminology provided in Singh et (2021; 2023) by referring to the term containing second-order derivatives of é as outer-product Hessian (often called generalized Gauss-Newton (Schraudolph, 2002) matrix), and to the term capturing second-order derivatives of the network as functional Hessian. To emphasize the split which the decomposition through the chain rule was applied, we added the prefix F above. This will later become useful to identify expressions in the Hessian stemming from different splits (see section 3 and appendix D). But until then and unless stated otherwise, we will consider the split E F and thus omit specifying this hereafter. For classical self-attention, this yields the Hessian decomposition H H. H with
H. (Wo, Wo)
H. (Wo, Wk) H. (WK, WK)
H. (Wo, Wv)
H,- H. (WK, Wo)
H. (WK, Wv) H. (Wv, Wv)
(3)
H. (Wy, WQ)
H, (Wy, Wk)
for E fo, and where H, (W,. (W,,W,)- = (H, (W,, W.)) due to symmetry.
Matrix calculus. So far, we have slightly abused notation and did not define formally what means to take derivatives of a matrix-valued object w.r.t. matrix-shaped argument. We will follow the recipe of matrix calculus (Magnus & Neudecker, 2019) which reduces those derivatives to the vector case by introducing a flattening convention (we will use row-stacking, indicated by the operation vec,). Consider two vectors a E R and b(a) E R. The Jacobian Ôb/da E R Bx collects the first-order derivatives such that [āb/da),, 0b,/day. The Hessian 0 b/baaa E R BAX
column block matrix that concatenates the Hessians for each entry of b w.n.t. a; equivalent to flattening the Jacobian, then differentiating it second time. To generalize these definitions
3

---

Published as a conference paper at ICLR 2025
to matrix-valued objects, simply flatten all arguments before differentiation (see fig. 2 for visualization). Applied to the expressions in eq. (2), we have,
OF
OFF veer oW.
aF Ô vec, F ow; = vve, W;
E RLd xp:g and aWOWw = vec, W
E R BLdypiq::ps (4)
For the square loss discussed above, the Jacobian 0f/0F R and Hessian F R Ldy w.r.t. the network's prediction simplify vec, (F(.) Y)" (())))d) and 21 ( (i.e., scaled identity matrix) respectively.
Hessian of MLPs and CNNs. Since one of our main interests stating explicit differences between the Transformer Hessian and that of traditional architectures, we briefly recap the Hessian structure from Singh et al. (2021; 2023), for the case of linear MLPs, F(x) = w. Wix. tion, we will only consider the outer-product Hessian and look at its diagonal block corresponding to the i-th layer parameter matrix, W;, which takes the form
xx. As an illustra-
H. H.(W,, W. W.) = W WT,,. W & Wi-1...W, W WT
with Σxx E [x J. the uncentered data covariance over the dataset. Observe that the above expression depends quadratically on the data through Σxx. Likewise, the functional Hessian for MLPs, has quadratic dependence the data, but through a different matrix which carries the covariance of the residuals dx.y F(x) with the input x. Importantly, the Hessian diagonal blocks are entirely comprised of the contribution from the outer-product Hessian, as the functional Hessian has hooloww structure (zero diagonal blocks), which also extends to any piecewise nonlinearity in the sense of almost everywhere.
A notable difference from our setup is that these works aggregate the Hessian over the dataset, while, for brevity, we focus on a single sample setting. As side-benefit, this allows us to cast their (non-sequential) setup one-to-one here by considering their whole dataset in the form of a single sequence, where the MLP applied separately to each element of the sequence. This point of view helps facilitate a precise comparison between the two settings.
3 EXACT STRUCTURE OF THE SELF-ATTENTION HESSIAN
We start by studying the Hessian of standard self-attention with square loss. We present our main re- sults regarding the exact Hessian in the form of two theorems, each of them targeting one term in the Gauss-Newton decomposition from eq. (2), further broken down into parameter blocks from eq. (3). Later, our goal will be to analyze and interpret them further through simplifications. For brevity, we omit the blocks w.r.t. the key weight matrix as they are essentially symmetric to the ones w.rt. the query weight matrix (modulo differences in arrangement) and defer the proofs to the appendix.
Theorem 3.1. Outer-produce Hessian Ho For single self-attention layer, eq. (1), with classical self-attention that feeds into the square loss, the blocks of H, are
2
H. (Wv, Wv) = M Mi @ Lav:
Ldy
2
H. (Wo: Wo) = Tdydk
(Ia, 8 Wx) ZJ (I. @ W.) Z, (Iay & WK),
2
H. (Wv, WQ) = Ldv
(MT 8 WT) Z (Lav & Wk),
2
H. (WQ, Wk) Idydk (Ier & wx) 2 (I & W. Z (Wo I Ky
with the first attention moment matrix M AX E Rixdv (see section 3.2) and where Z = (I @ X (0A//OT) & X) E R Idy xdy contains first derivatives of the softmax.
To arrive at these expressions we start from eq. (2), insert 8e/OFF. and use the self-attention Ja- cobians derived by Noci et al. (2022). See appendix B for details. We note that the value diagonal
Here we discuss the Hessian in single data point setting to improve exposition, but our insights translate directly into batch setting-one just needs average the Hessian formulas over the data points.
4

---

Published as a conference paper at ICLR 2025
block is structurally identical to MLP's outer- outer-product Hessian, while the query diagonal and the mixed query-key blocks display similar Kronecker structures, but with varying weight matrices.
Theorem 3.2. Functional Hessian HІ. For the setup of theorem 3.1, the functional Hessian w.r.t. the value weight matrix H₃ (Wv, Wv) is zero and the remaining blocks are given by
2
Hr (Wa, Wo) Ldydk
Rayax (Ir @ W @ Lay @ wx) Ze (Lay @ Wk),
Hr (Wv, Wo) Ldv
R (Lz S) Z (Lay @ Wk),
2
H₂ (. Wk) Ldydk Ravdk (I & W. & LL w Ze (Wa @ Iay dy
Idvvdk Rdv (I W lar) (Z) lay ( Ldxt
with the duplicated residual R., := vec, (F(X) Y) @ Im E mmmmfd shuffling trix = (Iay & ddv.dy) (voc, Iay @ Idy. E Rdvxdv where Kay is commutation matrix (see e.g. lemma B.2). Z1 defined as theorem 3.1, and Ze := (I @ X X & x (02AAATTT)) A/OT)) & X) E IR Lag xd containing second-order softmax derivatives.
For theorem 3.2, we differentiate the self-attention Jacobian, and multiply it with the square loss gradient 0E/0F. See appendix B for details. Similarly to the MLP functional Hessian, the diagonal value block of the self-attention functional Hessian disappears. Again, structurally similar expres- sions appear in the query-key mixed term, and the query diagonal block, with the query-key mixed term being the most involved and consisting two summands.
Categorizing the constituent terms. To simplify the analysis of theorems 3.1 and 3.2, we divide the terms into four categories. First, the least important are layout matrices that copy or shuffle entries via Kronecker products with identity permutation matrix. These serve correctly arrange the entries in the Hessian from our matrix layout in eq. (4) and therefore we will often ignore them. Second, we have the matrices Z₁ Zz, which contain first- and second-order derivatives of the softmax and have explicit dependencies on the data X (section 3.1). Third, we identify first-order moments M2 of the attention matrix, well higher moments in the attention derivatives in Z1,2 (section 3.2). Fourth, the Hessian depends on the weight matrices Wv, WQ. WK (section 3.3).
As one of our main interests is the scaling behavior of blocks w.r.t. different dependencies, we will simplify expressions using the Landau notation from Martens (2020, Section 14), where for two matrices A, B, denoting f E ((AB) means all entries of f are ( Bx.1) for any j,k,k k. Due to their boundedness to [0, 1). we can ignore the softmax expressions in the attention matrix.
3.1 DATA DEPENDENCE VARIES ACROSS HESSIAN BLOCKS
One characteristic of self-attention is that the data X enters multiple times: keys, queries, and values. This leads to highly non-linear data dependencies in the Hessian. For brevity omitting the dependence of all the functional Hessian blocks the residual ōxy vec, (F(X) Y) find that for the expressions in theorems 3.1 and 3.2
K
K
V
Q
T
O(X9
o(x")
Q
5
+ o(x²)
O(x))
H. E
K
O(X))
o(x)
H E
K
o(3)
(5)
V
(x
V
((1)
Data heterogeneity. Throughout different Hessian blocks, we observe varying depen- dence on the embedding matrix X. The most data-dependent blocks are the key and query ones. For the outer-product Hessian Ho each key and query weight contributes cubic de- pendence on X while value weight matrix Wy brings a linear dependence on X. This obser- vation confirms the finding of Noci al. (2022) who show |@F/aWvll O(IX while
5

---

Published as a conference paper at ICLR 2025
H
H1
109
10P
H(W., Wy)
101
H(Wq, Wo)
10"
100
Block [| lle
10
10-2
10-2
H
10
10
10
10-1
10-2
10
10
10
10"
10
10--
10
102
10
10
10
10
10
10
10
102
10
(a) No LN. We omit Hr (Wv, Wv) as it equals zero.
(b) Pre-LN.
Figure 3: (Plotted in log-log scale.) Empirical verification with a CE loss confirms derived growth rates w.rf. magnitude a of X from eq. (5). We show the growth rates through the Frobe- nius norm || ||e of value and query diagonal blocks. The dashed lines correspond to the trend (a) predicted by theory as in eq. (5), (b) estimated from the Frobenius norm measurements on the log- log scale by the linear regression slope. For details on the experimental setting, see appendix F. a (LHS of the vertical line) corresponds to practical values of o.
|ØF) /ow K le. WF//WWar WIlle O(I The block -eterogettoo dependence is also visible in the functional Hessian Hr, which includes an additional dependence on X through the residual dxy- If we ignore oXy, then the outer-product Hessian dominates the Hessian for large X.
Empirical validation. We test our theoretical derivations empirically using a single GPT-2 Trans- former block (Radford et al., 2019) on a digit addition task adapted from Quirke & Barez (2024). The problem is set up as a next-token prediction task, with cross-entropy (CE) loss (see appendix F)-a realistic Transformer application. We initialize the embedding X from a distribu-
setting closer tion with varying standard deviation a and measure the Frobenius norm of the Hessian's 's diagonal blocks (fig. 3a). The growth rates from eq. (5) transfer to the dependence of the block Frobenius norm on a. We omit the value functional Hessian block H HW(Wv, Wv) from the plot because it equals zero in both theory and our empirical evaluation. For the value block, the full Hessian follows the o² trend from the outer-produe Hessian from eq. (5), since the functional block is zero. The query block follows the functional Hessian's 's σ trend, and not the σ trend from the outer-product Hessian, which is likely suppressed by other non-data dependencies. Overall, this confirms that al- though our theoretical results were derived for MSE loss and only a single self- self-attention layer, they extend faithfully to CE loss and full Transformer block.
3.2 ATTENTION MOMENTS AFFECT THE HESSIAN
Attention scores induce distributions. Each row of the self-attention matrix A (XX is normalized
distribution over the input tokens, and we already saw the first moment M, (i.e. multiplying the sequence by the attention matrix) of those distributions emerge in theorems 3.1 and 3.2. To further simplify the Zı, matrices, we now introduce the natural generalization to higher-order centered attention moments. The i-th centered attention moment matrix M, should be thought of as an L dy x... x dy tensor (where dy appears times) that stacks the i-th centered moments for each distribution, which is then flattened into Layy" x dy matrix. In matrix notation, have:
Definition 3.1. Attention moment matrices. Let A E R be attention matrix and X E 4xdy be the token embeddings. We define the first attention moment matrix as
Mj = AX [A]]
ISKL E R Lxdv
Similarly, the second and the third central moment matrices are
Mg
Σ Aig (X,,: : [M.).)) (X,,: [M.](.) 1sish.
E RLddyxd xdv
M3
E/=n (X,, [M.).) (X,,: [Mi)..) 100..M[M -(M.).))
E RRili xdv.

---

Published as a conference paper at ICLR 2025
Transformer Hessian dependency on attention moments. Definition 3.1 allows to further simplity the matrices Z from theorems 3.1 and 3.2. In addition to the dependency on the first attention moment matrix, obtain dependencies on the second and third moment matrices.
Remark 3.1. The data terms emerging in the self-attention Hessian can be expressed as func- tions of the self-attention central moment matrices (proof in appendix C),
1 = X M2; and Z2 (I @ K dv H lav (x x MM3),
where * is the Khatri-Rao product (Khatri & Rao, 1968), see definition A.2.
Regarding attention moments, the value outer outer-product Hessian depends on the first moment, while the query-key outer-product Hessians are influenced by the second central moment. The query-key functional Hessian even exhibits a dependency on the third central moment.
Influence of the attention moment matrices on the Hessian. If the attention scores were data- independent (similar to the uniform attention assumption in Noci et al. (2022)), which happens almost surely at initialization for large dx. sequences with more similar words will result in lower contribution of the query-key block. Considering an orthogonal setting with a fixed input sequence and studying the influence of varying attention scores, the query-key outer-product Hessian will dominate the attention scores are highly dispersed across tokens. If instead, the attention matrix sparse and some tokens attend to only token (attention row is a one-hot vector), the contribu- tion of the query-key part of the Hessian diminishes because the second and third central moment matrices of such one-hot self-attention distributions equal zero. To the best of our knowledge, work is the first one to notice the relationship between the self-attention Hessian and self-attention moments. Knowing this dependence can contribute better understanding and interpretability of the Hessian especially its evolution during training.
3.3 DEPENDENCE ON THE WEIGHT MATRICES
The blocks in the self-attention Hessian vary significantly when it comes to the dependence on the weight matrices. Among the terms on the diagonal of the outer-product Hessian, only the query and key components explicitly (and quadratically) depend on the weights beyond their influence through the attention scores. This dependence is on both the value weight matrix and one of the key and query matrices. The mixed terms involving queries and keys depend linearly on the keys and queries weights and quadratically on the value weights. Finally, the mixed terms involving values depend linearly on the selected weight matrices.
Also, the blocks of the functional Hessian depend on the weight matrix in varied manner. The first observation that the diagonal value functional Hessian block is always zero. Furthermore, the diagonal query block depends quadratically the key weight matrix and linearly on the value weight matrix. The mixed value-query block depends only linearly on the key weight matrix.
@ We find highly non-linear dependencies and block heterogeneity in terms of data, de- grees of attention moments, and weights. We expect that identified sources of heterogeneity influence different algebraic properties, such as traces, norms, and eigenvalues of blocks- possibly explaining the varying block spectra, previously observed by Zhang et al. (2024).
IMPACT OF TRANSFORMER DESIGN COMPONENTS ON THE HESSIAN
The Transformer architecture has several key components that distinguish it from classical models like MLPs and CNNs. In the previous section, we analyzed how data dependence is one big distin- guishing factor between the Hessian for Transformers and classical architectures. Now, we proceed to the other two significant departures within Transformer, namely, the effects of having (i) a soft- max nonlinearity within a layer block and (ii) a quadratic weight interaction through the query-key parameterization, by disabling them one at time.
7

---

Published as a conference paper at ICLR 2025
Figure 4: Softmax results in heterogeneity in magnitudes of Hessian block entries. Histogram of (absolute) entries corresponding to the value and query diagonal Hessian blocks of a single block Transformer. High and low sat- uration correspond to linear and classical self-attention respectively.
H(Wv, WyJ
H(WQ. We
linear
classical
linear
Fequanry
6000-
classicul
2500
10"2
10
10
10
10"
10"a
Aleolute Eairies
4.1 SOFTMAX ACTIVATION
Apart from the Transformer architecture, softmax is not a common choice for an internal activation function. It is usually used as the final layer activation in classification tasks to model a distribution. It is not an easy activation to work with, it is prone to numerical instabilities and vanishing gradient problems Goodfellow et al. (2016); Noci et al. (2022); Wang et al. (2021).
Hessian of linear self-attention. To study the influence of softmax, we compare self- aff-attention to lin- ear self- attention, i.e. a = id and T(X) = XWWWRR /vdx in eq. (1). In theorems 3.1 and 3.2, that the first and second derivatives of A w.r.t. T are Læ and 0, respectively:
Remark 4.1. Assume linear self-attention F with id and T(X) X WWW /Vdk followed by square loss. Then the outer outer-product Hessian blocks are
21
H. (Wv, Wv) = dydk
Exx Wx Exx @ Lav.
2L"
H. (Wo, Wo) dydk
Exx 8 Wr WK,
2L²
H. (Wv, WQ) dydk
(Exx & WX
2L
H. (Wq: Wk) = dvdk
(Exx & W Exx Wy W.
Ka dv:
where Exx = - is the empirical (uncentered) intra-sequence covariance. Moreover, only the off- off-diagonal functional Hessian blocks are non-zero and equal to
Hr (Wv, Wo) dyvdk
(oxx @ LLte) (Ir. & S) (X @ Exx Wk)
H (W WK) dy dv
(xxx (X W & Iay)
Lix
where 8XY == vec, (F(X) Y) and defined as in theorem 3.2.
Since linear self-attention is simple matrix multiplication, we observe a Kronecker product structure akin to that emerging from the outer product of matrix multiplication derivatives (eq. (15)). Moreover, we note that the functional Hessian has block hollow structure, similar to the MLP functional Hessian as discussed in section 2.
Softmax makes the Hessian block-heterogennus. Removing the softmax from the self-attention definition makes the dependence of the Hessian blocks on the data more uniform across blocks. All blocks of the outer-product Hessian depend cubically on Σxx, while the non-zero functional Hessian blocks depend linearly X Σxx. Similarly, the number of weight matrices entering the expressions for the blocks for the Hessian without softmax is the same for the outer- outer-product Hessian and across non-zero blocks of the functional Hessian.
Figure empirically demonstrates that removing softmax from self-attention makes the magnitudes of the Hessian entries more homogeneous across blocks. For the classical attention (see also fig. 1a), we see that the distribution of Hessian block entries varies between query and value blocks the entries of the query block are two orders of magnitude smaller than the entries of the value block.
8

---

Published as a conference paper at ICLR 2025
We could explain this by X being initialized to values 1 and the query Hessian block having higher order dependence on X than the value Hessian block (see section 3.1). After removing the that both query and value histograms are largely similar from remark 4.1. Figure also informs us that softmax has a particularly strong influence on the magnitude of the query Hessian block and not on the value Hessian block.
softmax (fig. 4),
would expect
Softmax makes the Hessian diagonal blocks more indefinite. When we remove the softmax from the self-attention definition, the functional Hessian has a block-hollow structure with zero blocks on the diagonal. This means that the Hessian blocks are fully defined by the outer-product Hessian, which ensures that they are rrsitive-semidefniie Park & Kim (2022) empirically noticed the more indefinite nature of Transformer Hessian compared to CNNs our our observation suggest that the reason for that could lie in the ubiquitous use of softmax in Transformers.
Linear self-attention vs MLP. The loss Hessian of an MLP and CNN has been previously theoreti- cally studied in the setting without nonlinearities, so analyzing the linear self-attentin Hessian lets us compare the two in similar setting. The Transformer Hessian much more data-dependent than that MLPs and CNNs, and we summarize this in section 4.1 using the matrix big-O notation.
MLPICNN
Transformer
Table 1: Dependence of the Hessian of linear self-attention layer on the intra-sequence covariance Exx and the input-residual covariance f2yy (F(X)- in big-O notation (remark 4.1).
H. ( O( Exx)
O
H
One crucial difference is the asymptotic growth rate w.r.t. depth. we stack D linear self-attention layers, the block-diagonal matrices of the Hessian will contain 3 input intra-sequence covariances Σxx. This result follows directly from the fact that a D-layer linear self-attention network, as well its Jacobian w.r.t. any weight matrix, is matrix chain involving 3D input sequence matrices X. due to the recurrence that the output of a layer is used three times by the consecutive layer. This is in stark contrast to deep linear MLP, whose Jacobian is only linear in x, and therefore Hessian diagonal block contains only 2 data instances, independent of depth.
4.2 SIERY-KEY PARAMETERIZATION
In practice, self-attention parameterizes T using two matrices Wo and WK. This ensures having clearly defined, interpretable token embeddings on the self-attention level and for small dx results in fewer parameters. From a function class perspective, we could equivalently replace the product Wo W with single matrix W WQK. However, this changes the landscape of loss.
Query-key parameterization induces additional Hessian decomposition. From theorem 3.2 we know that we can further decompose the query-key blocks of the functional Hessian of classical self- attention. We present this in detail in appendix D. This decomposition and specifically T f-functional Hessian from remark D.I are by-products of parameterizing the self-attention matrix with two ma- trices Wq and WK. To see that, let consider self-attention parameterized with a single matrix as another control model. the definition self-attention from eq. (1) we assume a = softmax applied row-wise and T(X) XWQKX In lemma 4.1 present the Hessian of this model w.r.t. Wok.
Lemma 4.1. Hessian of self-attention with single matrix attention parameterization. As- sume the self-attention definition from eq. (1), where softmax is applied row-wise, and T(X) = XW W T followed by an MSE loss function. The loss Hessian w.r.t W Qr is
H(Wo,, Wok) = z (I & WT Z + (1xx (L (W) Ze
where Z,, Zz, 8xy are defined as in theorems 3.1 and 3.2 and remark D.I.
The expression for the Hessian H( QK. WΩk) is part of the T-outer- T-outer-product Hessian the clas- sical self-attention in remark D.I, while there is counterpart expression to the T-functional Hes- sian. Double matrix parameterization of self-atte attentio implies also that the query-key Hessian part explicitly depends the query and key weight matrices Wo Q and WK.

---

Published as a conference paper at ICLR 2025
4.3 OTHER COMMON DESIGN CHOICES
The influence of temperature on the Hessian terms varies across blocks. Assume that in the definition of self-attention from eq. (1) we employ the classical self-attention but with an additional temperature scaling, namely T T(X)= X WW W W X ((tdd)). This impacts the scaling factors in the T-Gauss-Newton decomposition from remark D.I. The T-outer-product Hessian H will be scaled by 1/t², and the T-functional Hessian H by 1/t. To see that, note that is just an extra multiplier in front of vd and the formulas for H, and H. depend on 1/vdx linearly and quadratically, respectively. This implies that as L grows, H. dominates the query-key part of the Hessian, and when i + 0, H becomes the more prominent part.
Layer norm reduce inter-block Hessian data heterogeneity. Layer norm is used either in the original Post-LN version, where it is applied between residual blocks, or in the Pre-LN version (see Baevski & Auli (2018); Wang et al. (2019); Xiong al. (2020)), where it is placed inside the residual connection and additionally after the final layer. The heavy dependence of the Hessian on the input matrix X, growing super-exponentially with network depth (see section 4.1), means that, unless the data is standardized, we can observe exploding or vanishing phenomena, which will be more pronounced than in MLPs. This highlights the importance of layer norm in the Transformer architecture. Moreover, the proper placement of layer norm, as in the Pre-LN setting, addresses the data heterogeneity across Hessian blocks. In fig. 3b we plot the Frobenius norm of the Hessian blocks for the Transformer block including the Pre-LN, We verify that Pre-LN indeed addresses the block- hook-hheerogeneity w.r.t data growth rates the difference between the trend exponent in the two compared blocks is much smaller than in fig. 3.
CONCLUSION
Summary. In this work, we theoretically derived the entire structure of the self-attention Hessian and discussed, in detail, how it behaves. In particular, we explicitly characterized that the self- attention Hessian blocks have heterogeneous dependence on the dala, weight, and degree of the attention moment matrices. Further, we identified that Transformer-specific design decisions, such as query-key parameterizaaton and softmax, result in a more non-linear and heterogeneous Hessian structure. Thus, thcoretical works should be mindful of not discarding these design choices for the sake of mathematical convenience, as they can significantly alter the Hessian structure across blocks.
Discussion. Our results contribute to the theoretical and practical discussion:
(i) Understanding Transformer optimization. The research community has been exploring why op- timizing Transformers is more challenging than other architectures (Zhang et al., 2025; Xiong et al., 2020: Liu et al., 2019; Pan & Li, 2022). Ahn et al. (2024) proposed linear self-attention model with a single weight matrix parameterization to examine the Transformer's loss landscape. Using this simplified model, they reproduced key optimization phenomena, such as the performance gap between Adam and SGD. Our work provides a Hessian- perspective on this model, providing new hypotheses on which components may drive Transformer optimization challenges.
(ii) Transformer-speciffc optimizers. We believe that analyzing the biock-diagonal structure can be beneficial for developing optimization algorithms specifically tailored to Transformer models. For example, there evidence (Zhang et al., 2024; 2025) suggesting that the memory consumption of adaptive optimizers can be significantly reduced by assigning single adaptive learning rate to parameters corresponding to a homogeneous block on the Hessian diagonal. Knowing the exact block structure in self-attention layers, we could now use it to adapt the learning rate.
Limitations. While our theoretical setting limited to a simple, single-layer' model, we saw that it displays rich algebraic structure that is contained within the Hessian. Nevertheless, extending it to multi-layer networks resembling fully-fledged Transformers. at least along some specific axes, would be interesting. Moreover, since usually the embedding matrix X is also learned, it would be worth studying its Hessian blocks.
Notwithstanding, our work lays crucial foundation for the theoretical analysis of the Transformer Hessian-to to our knowledge, we are the first to study its exact expressions and their dependencies.
2 Theoormms 3.1 and 3.2 directly generalize to the last layer of any deep self-attention network if we replace the data matrix X with the output from the penultimate layer. Moreover, thanks to the chain rule, the second derivatives of self-attention derive useful for the analysis of Transformers of any depth.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
We attach proof' all theorems, lemmas, and more involved remarks presented this manuscript in the appendix. Specifically, Appendix discusses prerequisites, and Appendices B and outline the proofs. Appendix F contains details on the experimental setup. The code used to generate numerical results is available at: https//gghhbb.ooomdaaaaaaattteee
ACKNOWLEDGMENTS
We would like thank Thomas Hofmann, Antonio Orvieto, and the members of the DALab for their insightful comments. We also thank Michael Vollenweider for proofreading the manuscript. Felix Dangel would like to recognize that resources used preparing this research were provided, part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute. Sidak Pal Singh would like to acknowledge the financial support from Max Planck ETH Center for Learning Systems.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Lin- ear attention is (maybe) all you need (to understand transformer optimization). In International Conference Learning Representations, 2024. 10
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference on Learning Representations, 2018. 10
Yoshua Bengio. Practical recommendations for gradient-based training of deep architectures. In Neural networks: Tricks of the trade: Second edition, pp. 437-478. Springer, 2012.
Jacob Bridgeman and Christopher Chubb. Hand-waving and interpretive dance: an introductory course tensor networks. Journal of physics A: Mathematical and theoretical, 50(22):223001, 2017. 19
Jeremy Cohen, Simran Kaur, Yuanzhi Li, J Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In International Conference on Learning Representations, 2021. 2
Felix Dangel, Stefan Harmeling, and Philipp Hennig. Modular block-diagonal curvature approxi- mations for feedforward architectures. International Conference on Artificial Intelligence and Statistics, 2020. 3
Felix Dangel, Runa Eschenhagen, Weronika Ormaniec, Andres Fernandez, Lukas Tatzel, and Agustinus Kristiadi. Position: Curvature matrices should be democratized via linear operators. arXiv preprini axX1v:2501.19183, 2025. 32
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, arXiv preprint arXiv:18/0.048805 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 1999,, 2020.
Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. International Conference on Machine Learning, 2019.
Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep sparse rectifier neural networks. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pp- 315-323, JMLR Workshop and Conference Proceedings, 2011.
Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. http:
//wwww dwwpweeeeeeeee.. earningoonnnooooo org. 8

---

Published as a conference paper at ICLR 2025
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: training imagenet in 1 hour. arXiv preprint arXiv:1706.0222
2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. Proceedings of the IEEE Conference on Computer Vision and Patter Recognition, 2016. 1
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimiza- tion through better initialization. In International Conference on Machine Learning, 2020.
Yiding Jiang. Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantas- tic generalization measures and where find them. In International Conference Learning Representations, 2019. 2
Nitish Shirish Keskar, Dheevatsa Mudigere, Jorge Nocedal, Mikhail Smelyanskiy, and Ping Tak Pe- ter Tang. On large-batch training for deep learning: Generalization gap and sharp minima. In International Conference on Learning Representations, 2016. 2
CG Khatri and Radhakrishna Rao. Solutions to some functional equations and their applications to characterization of probability distributions. Sankhyā: the Indian journal of statistics, series A. pp. 167-180, 1968. 7, 26
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approxima- tion for natural gradient descent. In Advances in Neural Information Processing Systems, 2019.
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. Advances in Neural Information Processing Systems, 1989. 2
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 1998.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. In International Conference on Learning Representations, 2019. 10
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the diffi- culty of training transformers. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5747-5763, 2020. 22
Shuangzhe Liu. Matrix results on the khatri-rao and tracy-singh products. Linear Algebra and its Applications, 1999. 15, 25
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Confer- ence on Learning Representations, 2019.
Jan R Magnus and Heinz. Neudecker. Matrix differential calculus with applications in statistics and econometrics. John Wiley & Sons, 2019. 3, 16, 17, 19, 30
James Martens. Deep learning via hessian-free optimization. Proceedings of the 27th Interna- fional Conference on International Conference Machine Learning, 2010.
James Martens. New insights and perspectives on the natural gradient method. Journal of Machine
Learning Research, 2020. 3,5
Neel Nanda and Joseph Bloom. Transformerlens. github
Trans formerLeeesss 2022. 32
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 2022. 4, 7, 8, 17, 18, 22
12

---

Published as a conference paper at ICLR 2025
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for trans- formers. In OPT 2022: Optimization for Machine Learning (NeurIPS 2022 Workshop), 2022.
10
Vardan Papyan. Traces of class/cross-class structure pervade deep learning spectra. Journal of Machine Learning Research, 2020. 2
Namuk Park and Songkuk Kim. How do vision transformers work? In International Conference on Learning Representations, 2022. 2.9
Jeffrey Pennington and Yasaman Bahri. Geometry of neural network loss surfaces via random matrix theory. In International Conference on Machine Learning, 2017. 2
Martin Popel and Ondřej Bojar. Training tips for the transformer model. The Prague Bulletin of Mathematical Linguistics, 2018.
Philip Quirke and Fazl Barez. Understanding addition in transformers. In International Conference on Learning Representations, 2024. 6,32
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019, 6, 32
Frank Rosenblatt. The perceptron: probabilistic model for information storage and organization in the brain. Psychological Review, 1958.
Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues the hessian in deep learning: Singu- larity and beyond. arXiv preprint arXiv: 2016. 2
Levent Sagun, Utku Evci, V Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over--arametein paraperraetrize neural networks. arXiv preprint
7444004444, 2017. 2,
Tom Schaul, Sixin Zhang, and Yann LeCun. No more pesky learning rates. In International Con- ference Machine Learning, 2013.
Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural Computation, 2002. 3
Sidak Pal Singh and Dan Alistarh. Woodfisher: Efficient second-order approximation for neural
network compression. Advances in Neural Information Processing Systems, 2020.
Sidak Pal Singh, Gregor Bachmann, and Thomas Hofmann. Analytic insights into structure and rank of neural network hessian maps. Advances in Neural Information Processing Systems, 2021.
2, 3. 4, 16, 30, 36
Sidak Pal Singh, Thomas Hofmann, and Bernhard Schölkopf. The hessian perspective into the nature of convolutional neural networks. In International Conference on Machine Learning. 2023. 2, 3,
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- fion Processing Systems, 2017. 1, 3, 18, 30, 31
Qiang Wang, Bei Li, Tong Xiao, Jingho Zhu, Changliang Li, Derek F Wong, and Lidia S Chao. Learning deep transformer models for machine translation. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Shulun Wang, Feng Liu, and Bin Liu. Escaping the gradient vanishing: Periodic alternatives of softmax in attention mechanism. IEEE Access, 9:168749-168755 2021. 8
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing. Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning, 2020.
10
13

---

Published as a conference paper at ICLR 2025
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE International Conference on Big Data, 2020. 2
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why trans- formers need adam: A hessian perspective. Advances in Neural Information Processing Systems, 2024. 2, 7, 10
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Diederik P Kingma, Yinyu Ye, Zhi- Zhi-Quan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. In The Thirteenth International Conference on Learning Representations, 2025. 10
14

---

Published as a conference paper at ICLR 2025
CONTENTS
A Known Definitions and Properties
15
Jacobian and Hessian Expressions
17
B.1 Classical Self-Attention
17
B.2
Sclf-Attention with Single-Matris Parameterization
22
Self-Attention Moment Matrices
23
C.1 The Second Derivative of the Row-Wise Softmax
23
C.2 Self-Attention Moment Matrices
26
D More on the Query-Key Hessian Block
29
Multi-Head Self-Attention Hessian
30
E.1 Influence of the Low-Rank Nature of Wv the Self-Attention Hessian
31
Experimental Setup
32
Additional Experimental Results
32
A KNOWN DEFINITIONS AND PROPERTIES
This manuscript extensively uses the Kronecker product, as introduced in definition A.1. Some of our results rely on the generalization of the Kronecker product named inition A.2. This section lists matrix calculus and linear algebra properties used in this manuscript.
product from def-
Definition A.1. Kronecker product. Let A E m and B E RPx matrix. The Kronecker product A & Rmppxn is block matrix given by:
Q12 B
O1n B
021
B
Q22B B
02n B
ABB=
0m1
B
Om
B
Omn B
Definition A.2. Khatri-Rao product Liu (1999), Let and B be block matrices with n blocks. The Khatri-Rao product A B is a block matrix with blocks defined as:
A1,1 & B
A 1 @ B
A + B (Ais & Bajls =
A 2 B
An.m @ Ban
where A, and Bı are the blocks from 1th block row and 1tth block column of A and B, respectively.
15

---

Published as a conference paper at ICLR 2025
Basic Kronecker product properties. Here list some useful properties of the Kronecker prod- uct. discussion on the Kronecker product and proofs of the listed properties can found in Magnus & Neudecker (2019). For all the listed properties, we assume that the matrices are of the appropriate dimensions to perform the operations.
A 2 C)
B C and (B + - -
(6)
(kA) & @ (kB) ((A & B). kER
(7)
(A & B C A ( &B C)
(8)
(A (C D) (AC) 8 (BD)
(9)
(A B T A &B
(10)
For b being column vectors
a
0 b= ba = bba
(11)
Vector-matrix Kronecker product. Assume that and can be multiplied. Based on the mixed product property and the fact that taking Kronecker product of matrix and scalar is equivalent to only scaling matrix entries we can write
A & B = (1 &18 A) @ B (v 0 AB).
(12)
Vectorization & Kronecker product.
Assume that E Rmxx then
vec, A (A In) vec, In.
(13)
Proof. For a column-wise vectorization, theorem 2.2 from Magnus & Neudecker (2019) tells us that for matrix B E R²*0 it holds that vec.
(I, B) vec, Lg. Hence, we can write
vec, A Km, Vec. A Km, (I, @ A) vec In
= (A & I,) Kn,n vec. In = (A @ In) vec, I,
Product of block-diagonal matrix and Kronecker product. Let for any i E {1,..., m} A E Raxr, B E Rm and E R Then, thanks to the vector-matrix Kronecker product property (eq. (12)) the following expression holds
A
0
B
blockdiag (A.)
(B 8 C)
8
A
mm
A1
0
B i,
C
A (B 2 C)
(14)
Am
B m,
8 C
Am
B
8
C)
B
A
B
A
C
In the above equation B,, is the i-th row of B in column vector format.
Derivative of a matrix product. If = AXB, then
VxF A
(15)
The proof of the above formula can be found in Singh et al. (2021).
16

---

Published as a conference paper at ICLR 2025
Derivative of Kronecker product. RPx", Then
Let's take F(X) X Y, where X R and Y VxF (In X p.g I,) (Ing & vec, Y),
(16)
and analogously
Fy (I, 00 K p.g 8 I,) (vec, X 2 Ipr).
(17)
Proof. The proof makes use of the same formula but for the column-wise definition of derivative that be found in Magnus & Neudecker (2019) and states that the column-wise derivative of F w.r.t. X is given by (Ig X K,, @ Ip) (Ing @ vec. Y). We will use the above formula with X Y instead of X, Y respectively together with the first identification theorem (theorem 5.6 from Magnus & Neudecker (2019)).
d(vec, (X ② Y)) d(vec, (X 2 Y) d(vec, (X & YT
(I, & p. I,) (Ing & vec. YT) dl( d(vece X (In @ Kp I,) (Ing @ vec,, Y) d(vec, X),
so again by the first identification theorem, we obtain eq. (16).
To prove eq. (17) we follow similar steps as above and use the formula for column-wis derivative of F w.r.t. Y from Magnus & Neudecker (2019), namely (I, % K, @ Ip) (vec, X @ Ipr). As for the previous formula, we again use the first identification theorem.
d(vec, (X & Y)) = d(vece (X & Y) = d(vec, (X & ()
(I, X K p.g & I,) (vec, X & Ipr) d(vec. Y (In & K p.g & I,) (vec, X @ Ipr) d(vee, Y),
Derivative of transposition.
Let F(X) X where X E R"xq, Then
d vec, F(X) = vec, X K d vec, X, so from the first identification theorem (thcorem 5.6 from Magnus Neudecker (2019))
Vx F K
(18)
B JACOBIAN AND HESSIAN EXPRESSIONS
In this section, derive the expression for the self-attention Hessian. Firstly, foeus the clas- sical self-attention which we discuss in section 3. Next, we move on to self-attention parameterized by single query-key matrix, as discussed in section 4.2, For clarity, this section, will refer to the functional and outer anter-product Hessians as F-functional and F-outer- r-outer-product Hessians.
B.1 CLASSICAL SELF-ATTENTION
Knowing the Jacobian of self-attention is crucial step in deriving the Hessian (the loss gradient's Jacobian). We start by recalling the formulas for self-attention Jacobians w.r.t. weight matrices derived in Noci et al. (2022).
17

---

Published as a conference paper at ICLR 2025
Lemma B.1. Jacobians of self-attention Noci et al. (2022). The Jacobians of the classical self-attention layer (Vaswani et al., 2017) have the following form:
OF
(XWoWk)
softmax
X Idv
dWv
vdk
OF @Wq = (I @ w(x
0A
X X
OT
vdk
where the Jacobian of the row-wise softmax w.r.t. its inputs is as follows:
aA = blockdiag
aAi
(19)
öT
UT
and where
0A diag(A,,) (A.,) A
aT;,:
with A,, being the i-th row of A in column vector format.
With the Jacobians of self-attention w.r.t the weight matrices established, introduce the Hessian. We start with the F-outer-product Hessian and then proceed to the F-functional Hessian.
Theorem 3.1. Outer-product Hessian Ho- For single self-attention layer, eq. (1), with classical self-attention that feeds into the square loss, the blocks of H. are
2
H. (Wv, Wv) = MM M, @ Idv,
Ldy
2
H. (WQ. WQ) = Ldydk
(Iav & w Z (I & W. Z (Iay @ WK),
2
H. (Wv, WQ) LdvVdx
(MT & w Z (Lav & Wk),
2
H. (Wo. WR) Ldydk (Lay 8 wx) z (I, 8 WvWT) Zi (Wo @ Lay) Käg
with the first attention moment matrix M == AX E R4xdv (see section 3.2) and where Z = (I, @ x (0AA/88)))) & X) E Xdy contains first derivatives of the softmax.
Proof. Hessian of the mean- mean-square error loss w.r.t. network prediction is an identity matrix scaled 2 by Ldv Hence, computing the F-outer-product Hessian block is just taking the outer products of the appropriate self-attention Jacobians from lemma B. and scaling them by Ldy The expressions can be further simplified using the mixed productt property of the Kronecker product (eq. (9)).
Let us prove the exact formula for the mixed query-value block. The remaining formulas follow the same steps.
2 OF OF
H. (Wv, WT) Ldy OW aWa
2
(x T & Iav) (I 8 WTXT)
0A
X&XWA
Ldv
ØT
vdk
2
X A W.X.)
0A
X & XW
Ld
Vdk
AA
T & WT) (I, @ X
X ② X)
(Iay & Wk)
Ldy 2
ÖT
(T WT Z, (La, @ Wk) E Ldy
Lddy
18

---

Published as a conference paper at ICLR 2025
Before derive the F-functional Hessian expressions. we prove two helper lemmas which allow us to unify the structuring expressions that emerge in the F F-functional Hessian. The proof lemma B.2 is translation of tensor networks (Bridgeman & Chubb, 2017) to index notation. For more transparent presentation, in this lemma, we denote the matrix sizes using capital letters.
First, we briefly discuss commutation matrices. For more information computation matrices, Magnus & Neudecker (2019). Let us consider matrix A R RM and its row-flattens version vec, A RMN The commutation matrix KN.M provides way to transpose the indices of A using the flattened convention.
Precisely, if we apply K N.M the flattened version of A, get the flattened version of A
N. vec, A = vec, (A
The commutation matrix can also be used to permute two indices in a super-index, i.c. suppose E IRMN is vector indexed by super-index (m,n) where 1 S N and 1 S s M. Then [K, (n.m) = (m..)) i.e. the commutation matrix swaps the index order inside the super-index. To see this, note that for any R.M. R.M. we can find A E IR MxN such that a vec, A, and
K KN,Mal)
K N.M vec, (
[vec, (A T 1m.m))
[vec, A (m.))
|al (m).))
Finally, let us note that in index notation, the entries KN.M are given by
K KN,M] (x.M)( = mm.m'
Lemma B.2. Let K NN,M& E I MN N denote the commutation matrix, then
(Ins @ K KN,M) (vec, IM @ IN) (KM.N @ IM) (IN & vec, IM)
(20)
Specifically, for M N we obtain
(Ia & (v,ar)(vee, Im IM) (KM. Iar) vec, IM).
Proof. Consider the left-hand side of eq. (20) in index notation and simplify, which yields
[IM @ KN.M) (vec, IM @ IN (my ma)
I & KN,M] (my ma) (ma [vec, Ias @ Inl(mq m4. m. Σ
mi. mi.mb.ng
Σ mi .m...n
mm m Sn mn dmg mL mg 8nk
8 mi 8ng mi.
mi 11..m..0
Simplifying the right-hand side of eq. (20) in index notation yields the same expression and thereby establishes the equality:
[KM @ IM) (IN @ vec, IM) (m) (m))mmm22).22
Σ
KMNN & IM] [M]mm me.. IN & vec,
my mmgon
Σ
mm mi, ony,ng mmg m Eng, my On
mi mim..
Σ
mni Sng ma
mi m/....n
19

---

Published as a conference paper at ICLR 2025
Lemma B.3. For any matrix E R it holds that
(Im @ Km,n (vec, A & A) = (A
X In) (In @ Kn,x) (vec, In 2 In).
Proof. The equality follows from the chain of transformations
(Im & Km,n (vec, A & A)
(Lm @ Km.n) ((A X I,) vec, In & A)
(Im Km,n) (A I @ A) (vec, In I.)
A @ Km, (In @ A)) (vec, In @ In)
(A A I.) K K (vec, I, @ In)
(A & A X I, (I, @ Kn.x) (vec, I, IL))
were the first equality comes from eq. (13) and the remaining ones consequences of the mixed product property (eq. (9)) and fundamental properties of the commutation matrix.
Theorem 3.2. Functional Hessian HІ. For the setup of theorem 3.1, the functional Hessian w.r.t. the value weight matrix H4 (Wv, Wv) is zero and the remaining blocks given by
2
Hr (Wa, Wo) Ldvdk Ravdæ (I 2 W & Lav @ wx) Ze (Lav @ Wk),
H₂ (Wv, WQ) Ldv
Ra (I (S Z (Iay @ Wk),
2
H₂ (WQ, Wk) Ldydk Rdydk (I & W L. @ wx) Ze (Wa @ Iay Kdx.dv
Ruy (Ir & WT 8 lay) (Z) 8 lay ] 2 Lix:
LdvVdx
with the duplicated residual Rm := vec, (F(X) Y) @ Im E mmmfddy shuffling ma- trix = (Iay & Kdv,dv) (vec,, @ Iav) E xdv where Kav is a commutation matrix (see e.g. lemma B.2), Z₁ defined as theorem 3.1, and Zy == (I @ X @ X &
x A82)) @ X) E IR La? xdb containing second-order softmax derivatives.
Proof. We derive the F-functional Hessian block by block. For each block, derive the second derivative matrix of F, as in fig. 2. Each expression can be obtained by multiplying the second derivative matrix by the gradient of the loss according to the definition of F-functional Hessian.
F-functional Hessian w.r.t Wv & WQ. Let us find the formula for the second derivative of the Hessian w.r.t. value and key weight matrices Wy & Wo. To do that take the Jacobian of OF y w.r.t the query weight matrix Wo. To simplify, consider the following assignments
OF
F
F2 Idys
aW
XWOWK TXT Vdk
F softmax
X.
By the chain rule, we know that
F
OF,
OF, OF, OF, dWa
aW.OW,
aW
From eq. (16) obtain
OF
(Iz 18 Kiy dy 3 Iay) (ILdy @ vec,
ØF
I X Kiv & Iay (Ldv & vec,
where S is in the theorem statement.
20

---

Published as a conference paper at ICLR 2025
Moreover, by observing that F.2 differs from A only by the presence of matrix Wv, from lemma B.1 we get the derivative
2 (I, x OT dWQ
OA
X [
vdk
Plugging the formulas together yields the formula for the second derivative. The formula for second derivative w.r.t value and key weight matrices can be derived in the same way, with the only exception being that we take a derivative of F2 w.r.t. Wk instead of w.r.t WQ.
F-functional Hessian w.r.t Wo. To obtain the query-query block, we will start by differentiating the Jacobian to obtain the formula for the second derivative of self-attention w.r.t. query weight matrix W Let
OF
X & x
F
WL G
dWq
vdk
0A
TOA,: OT,,
G
blockdiag
OT
T XWo(XWk)
dd
as in lemma B.1. By the chain rule, we obtain
O
OF
OF JG UT
aw @W aWo
aG
UT dWa
Thanks to eq. (15) and basic properties of the Kronecker product (eqs. (8) and (10)) that
write
OF
W1.
[8XWX
&
aG
Vdlk 11 LOWV
X 8 X WK
Vdk
OT
X X&XW
aWa Q
Vdx
Additionally, we note that 8G/0T is the second derivative of row-wise softmax
aG
0²A
8T STAT
which we describe in detail in appendix C.1. Plugging in the above formulas the chain rule yields the formula for the second derivative were looking for.
F-functional Hessian w.r.t. Wo WK. proceed by differentiating the Jacobian of w.r.t W Let us again start with defining some helper terms
Fr
I,
Wy Wv TxT).
A
OA ÖT.
G OT
blockdiag
X XW
F2 =
Vdx
Note that only and F2 depend on Wk so by applying the chain rule to the product of functions and using derivative of matrix product formula (eq. (15)) we get
OF
OF GF
aG
OF,
FF F
OW + (F, & Taydx)
(21)
OWoDWK
OWk
OWK
21

---

Published as a conference paper at ICLR 2025
Now, thanks to the chain rule, derivative of matrix product (eq. (15)) and the derivative of transposition (eq. (18))
0G OA
XWQ X
Kdg.a
OWk
OTOT
vdx
Additionally, due to the chain rule and the formula for the Kronecker product derivative (eq. (17))
OF 2
X & Iak vdk
(1z & Kr,d @ Lax) (vec, X & ILdx)
aWk
(I KI @ Idx (vec, X X ② Idx)
vdk
(I @ KI.dv (vec, X & X) & Lax
vdk
(X 69 Iav) (lay & Kay (vec, Lav @ lav) @ Lax
Vdk
2 X Iav) (Kay,dy @ Iay (Tay so ver, (Iav)) &lag
vdk
where the first two transformations follow from the mixed product property (eqs. (9) and (12)) and the remaining two from lemmas B.2 and B.3. Plugging in all the terms obtain the Hessian block.
B.I.I THE JACOBIAN
We highlight that our observations regarding the Hessian made throughout the paper align well with those concerning the loss gradient w.r.t. self-attention parameters and the self-attention Jacobian. Liu et al. (2020) observed that the gradient w.r.t. the self-attention parameters is unbalanced, with the gradients w.r.t. the key and query parameters being of smaller norm than those w.r.t. the value weight matrix. From the self-attention Jacobian formulas derived by Noci et al. (2022) (see lemma B.1), observe heterogeneous data dependence in the self-attention Jacobians. These formulas suggest that the Frobenius norm of the Jacobian w.r.t. the query parameter should scale cubically with the magnitude of X, whereas the Jacobian w.r.t. the value parameter only linearly. Moreover, the value Jacobian depends on the first moment of attention, while the query and key Jacobian depend on the second moment matrix. Finally, the Jacobians exhibit heterogeneous de- pendence on the scale of the weight matrices, as previously noted by Noci et al. (2022) in theorem 3.1.
B.2 SELF-ATTENTE WITH SINGLE-MATRIL PARAMETERIZATION
Lemma 4.1. Hessian of self-attention with single matrix attention parameterization. As- sume the self-attention definition from eq. (1), where softmax is applied row-wise, and
T(X) = XW QooX" X is followed by an MSE loss function. The loss Hessian w.r.t. W WQK is
2 H(WQK Wok) = z (I, & W. Z (I @ Zo
where Z1, Zz, 6xy are defined as in theorems 3.1 and 3.2 and remark D.I.
Proof. The Jacobian of self-attention for W Wok matrix is given by
OF
(I Wv X
0A
xxx),
(22)
dW QK
ÖT
where 0A/0T is defined as in lemma B.I. The proof follows from applying the chain rule and the derivative of matrix multiplication in eq. (15).
22

---

Published as a conference paper at ICLR 2025
Now, to obtain the F-functional Hessian formula, let
OF
O Wr x G(XX)),
F1 = OW QK
G = 0A
(OA UT₁
blockdiag
OT
T = XW XWO.X.
From the chain rule, we obtain
OFF
OF 8G OT
OWkkoWok QK
8G θT dWok
0AA
I, 2 Wv x 8 8X
OTAT (x 8X),
where to get the second line we use the formula for the derivative a matrix product (eq. (15)) twice. The formula from the lemma statement follows directly from plugging in the above components to eq. (2).
C SELF-ATTENTION MOMENT MATRICES
In this section, we prove remark 3.1. To notice the dependence of the Hessian on the second moment matrix, need to derive and simplify the second derivative of the row-wise softmax- we focus on that in appendix C.I. In appendix C.2 we move on to interpreting matrices Z₁ and Z₂ through the lens of the self-attention moment matrices.
C.I THE SECOND DERIVATIVE OF THE ROW-WISS SOFTMAX
To gain some intuition the structure of the second derivative of the row-wise softmax, begin with the simplest possible case, assuming sequence length of 2. Later on, we will generalize the expressions sequences of any length.
23

---

Published as a conference paper at ICLR 2025
Remark C.1. Assume that the sequences are of length 2, meaning that the attention matrix is of the form
A
01,1
01,2
42,1
02.2
aA
By lemma B.1 the Jacobian of the row-wise softmax is given by matrix
of the form
OT
11,1
of
1
0
01,101,2
0122
a
0
0
022
22
(/2,1
02,2
0
-2,
0222
22,2
By vectorizing the above matrix row-wise, computing derivatives w.r.t. corresponding softmax inputs, and defining bi.j == 20,,j we obtain that the second derivative matrix is of the form
(a1,1 o. 1b.,
01,101,b111
11,1011,1111
01,10112b1.
0
O
0
O
0
0
O
21, a1 01,101.21,2
0
(a1
a
161
0
0
0
OFF
оTAT
0
(02,1
a 102,1
02,1
0
0
02,1022,22,,
0221 02,1022222,2
O
0
0
0
0
0
-2,1102.2b2.1 022 02,1002,222,
0221
(02,2 - 0z 2222
To obtain each entry of the matrix, we simply apply the chain rule. This means that each entry is sum of products, where the first factor is the derivative of the entries of 0A/0T w.r.t. di,i and the second factor is the derivative of aj w.r.t the entries of T. These derivatives can be found directly in the 0A/0T matrix. The blocks of zeros occur because the softmax function is applied
independently to each row, so the second derivative for mixed-row entries is always zero.
Lemma C. generalizes the above observations any sequence length T.
24

---

Published as a conference paper at ICLR 2025
Lemma C.1. Second derivative of row-wise softmax. Assume self-attention matrix allsssn. Then the second derivative of the row-wise softmax has a block-diagonal structure, namely
F
blockdiag D; R
OTAT
where D; R L'x is of the form
D
8 Arj
D;
with D.j e; & OT, T
D
and e; being unit vectors in the standard basis. Moreover, the single-element second derivatives can be expressed as
Aij OT, TT;
Ard AA + E diag(A,,) - T.:.:-Amp A...e) R
where E = e E RRmx a matrix filled with zeros except for entry in jth row and jth column which is 1.
Additionally, if we structure the single entry second derivatives into a block column matrix
0² A
Ai OT
@AiL
rewrite D, concisely
OT. OT.
OT, oT.
D. = (1L.1 (LLI) 0c)
A
(23)
where represents Khatri-Rao product Liu (1999) (see definition A.2) with blocks defined by the standard basis vectors in the LHS matrix and by the second derivatives of a single row entry in the RHS matrix.
Proof. Let start with discussing the general structure of the second derivative matrix. Later will focus on specific matrix entries.
Block-column structure. Recall that function that takes a matrix T as an argument. In the layout and vectorization scheme assumed in this manuscript, the Hessians of every entry of A in the row-wise order are placed consecutively into a block-column matrix (see fig. 2).
computing the second derivative of matrix-valuee
Block-diagonal structure. Let consider single block from the block-column structure discussed in the previous paragraph. Since the softmax acts on every row of T separately, and the Hessian of Ai, is computed w.r.t. all L entries of T. the Hessian will have potentially non-zero values only in one sub-block on the diagonal whose rows end columns have indices in the range from (i - 1)L iL inclusive, This translates into the derivative of the row- row-wise softmax pfr/aTaT also having the block- alock-diagonal structure. Specifically, let us enumerate the L L single-entry Hessians placed into this block-column matrix with an index For any integer such that S S L when (i T)T. +1 1 < J < iT. we have possible non-zero values only in columns iT. to (i + 1)T. - inclusive. Now let us group these Hessians into T. larger blocks corresponding to single row of Ahehe whole such block corresponding A,, has possible non-zero entries only in columns from (i-1)L + to iL inclusive. Hence the structure the second derivative of self-attention block-diagonal if consider only the possibly non-zero sub-blocks. We refer the sub-block corresponding to row A,, w.r.t T,, by D,.
Block-column structure of D.. Finally, D, consists of blocks, cut out of the element-wiss Hessians we have just discussed, which explains the Khatri-Rao product in eq. (23).
Non-zero elements. The non-zero entries of D, correspond
8
A
ÖT, .T.
25

---

Published as a conference paper at ICLR 2025
We now find its exact formula.
Recall from lemma B.1 that
0A diag(Ar,) A,, OT,
which means that
0A.,) (e, A,,) E RIXL
OT.
To get the second derivative it is enough to differentiate the transpose of the above function using the Leibniz product rule, namely
8 A o(e, Ai.)) Ar, + (e,- -A,))
aA
8T ST.
OT,,
OT, i,:
A.
Aij + (ej Ai,)) 0A.s
ST:
OT;,:
Since we already know both derivatives present in this expression from lemma B.1 it is enough to substitute them to obtain
0 Ais (diag(A : - A., Ais + (e, Ar,)(e; - A Ais ØT, OT.
Ais (2A t - - Ai,,e)).
C.2 ATTENTINNN MOMENT MATRICES
Remark 3.1. The data terms emerging in the self-atention Hessian can be expressed as func- tions of the self-attention central moment matrices (proof in appendix C),
Z X M2, and Z (I Kdy dy Iav (x X Ms).
where is the Khatri-Rao -ao product (Khatri & Rao, 1968), see definition A.2.
Proof. Before starting the derivations, let us specify that for the Khatri-Rao product in the theo- rem statement, X split row-wise and M; is split block-row-wise into central moment matrices corresponding to attention rows as in definition 3.1. The proof follows straight from transforming matrices Z₁ and Zz.
Dependence on the second central moment matrix M,. It holds that
A
Z, (I,
&
X
(X &
X)
OT
dA
blockdiag
X
(XxX)
OT,
XT. & X AA
x
JT..
X M
I & X
OTT
where the first line follows from the two first matrices in the product being block diagonal and the second from eq. (14).
26

---

Published as a conference paper at ICLR 2025
Dependence the third central moment matrix Mg. Let's recall from theorem 3.2 that
8A OTOT (X x)
Z [I
X X x 8x (I 8 X
R
blockdiag (D:) (X
The first two matrices in the above matrix product are block-diagonal with L equal-sized blocks, their product is also block-diagonal matrix with L equal-sized blocks of form [X & X @ D Hence, by (14) we know that Z₂ has block structure Z2 = (Z2slisist where
Zz = (x7 3 X (XT)D, (xI, 0x),
with X,, Rdv being the 7th row of X as column vector.
Let us recall lemma C.1 that gives us the D, formula transform the expression for Zz:
Ze: = [ X X x D. X ox)
AA OT, UT.
= =(x (x
8 x (1L,1 [e,) e,)
X ox)
2 A OT.
- xXX &X
(xT,
8 AiLL
e; X
8T .T...
0 A e; 2 OT, .T.
Σ Xx, x
(x.. 8x)
Σ (x,,0XT) e; OXXT
PAi
(XT,
UT.
8 PAi X ÜT, ,T,
Xx,001] ex X
PAij X, JT, ÜT,
X 8 X, X 8 T
=1
where the last three equalities follow from the mixed product property of the Kro- necker product. More precisely, the last equality follows from (X,, & xT)e exx =
(x), 8XT) (1 (10eXTT.) X, TT:X =x,, @ X
27

---

Published as a conference paper at ICLR 2025
Using known properties of the commutation matrix and the mixed product property (eq. (9)), further simplify the expression
8
X
X
X
TT
f=1
0 PA OT.
Σ dy (x @ x
X
Σ dK @ Lov) (X.:X. 8x 8 T 0
L.
X
J.. OT.
X j @ XT 0 A x
dv @ Idy)
X 2
ØT, Ti::
Finally, after plugging in 0² Kronecker product we simplify the summation expression N
from lemma C.1 and extracting As, in front of the
N
X 2 AT (2A., + E diag (A..) -., C,A,T AI..0)) x
==1
Σ A,, @ (2|M) + +X,, - T (A IX- M MM.l.....) j=1
This can be further simplified. We firstly note that
MM)a..MM, X,.X,, - x [M, (M,). - x..) (Mi)i, - -x,.)) T.
Moreover, from the fact that Kronecker product distributes over addition (eq. (6))
Σ A, @ (lMi)a.[M.).. -X diag((Ar:)))
Σ AisX @ [ - T diag(Ar,))
Mike: & MMI]. - -X diag (A).
These two equations give us
N Σ A x,, (MM), M.]. X diag(Ar.))))
==1
Σ A, (MMI. + HX,,JJ. MM (MI).
j=1
Mi @ (M..., - X diag(A,,
Σ A,yX,, (Mi).- - X,.)) (Mil. Xz.))
We now at the finish line of obtaining the desired formula. Let us note that similar expression appear in both summands of the equation above, specifically
( (iag(A,.)) - MM M = Σ A (Mil,, - X,,)) (Mi).s -X,.)
==1
28

---

Published as a conference paper at ICLR 2025
Hence,
N Mila; @ [MMI,,,MMI., - -x diag(As,))X)
Σ J, 2 -M,)):-X,..). (Mili: -x;,)
/=1
A Miiki & (Mik, X3.)) (Mi]), -x..)
A (Mil. - x,,)) (Mili: (M.)) - xz.:)
j=1
Aij (Xj,: - Mila.)) 8 (Mil: - x,.)) (M1)k - -x,.)
==1
M
After plugging in the above into the Zz formula, we obtain
ZZ (Kdy.dv Idv) (X.:X., @ Msle.).
Finally, note that
X * M3 XT @ MMla,sss xdy
and hence
Z₂ = (Ig X Kdv,dv @ Lav) lov) ( Ma))
D MORE ON THE QUERY-KEY HESSIAN BLOCK
Nested structure of the query-key Hessian. The query-key Hessian blocks exhibit nested struc- ture which prompts us to define a Gauss-Newton-like decomposition of the query-key Hessian blocks. The mixed query-key block of the Hessian consists of three summands one coming from the F-outer-product Hessian and two from the F-functional Hessian. The F-outer-product term as well one of the F-functional terms are structurally similar the query blocks can find on the diagonal. Together they can be expressed as an outer product. This observation results in the decomposition from remark D.I.
Remark D.1. The query-key part of the Hessian from theorems 3.1 and 3.2 can be equivalently decomposed into sum of -ouuter mattur-product and T-functional Hessians, respectively given by
H (WoK, Wok)
=
1 V
z (I (I & W.) Z (X & WT) & L
Zz V.
dx
0
B & T Idx =
0 B
IFWQo. Wok) =
Idx-
vdx
B
Lax
Vdk
W
In the above formulas W Wok
OXY vec, (F(X)- Y) and.
WK
V = [(Q Iay) ddx Iay & WK], B B: (I, @ w @ Iay) (Z, @ Iay)
This decomposition can be thought of as Gauss-Newton decomposition when we split the function composition (o + AXW.) v) oao T at the level of T(X) = XWo W /vdk.
29

---

Published as a conference paper at ICLR 2025
Query-key Hessian eigenspectrum. The structure of the query-key Hessian block implies spe- cific configuration the eigenvalues of its summands. The T-functional Hessian has a character- istic block-hollow structure with blocks of zeros on the diagonal, which makes it responsible for the bulk eigenvalues of the query-key Hessian. This specific structure allows us to reason about the eigenspectrum of the query-key Hessian. Eigenvalues of TT T-functional Hessian come in pairs +λ, for 1 i S dv. To see that it is enough to note that if A is an eigenvalue with an eigenvector [v. V2 JT then also [-v] v2 is an eigenvector with corresponding eigenvalue -
Moreover, by theorem 2.1 from Magnus & Neudecker (2019) eigenvalues of the Kronecker product are products of the factor matrices' eigenvalues. Since all dx eigenvalues of Idx are ones, eigenval- ues of f-functional Hessian are the same as cigenvalues of
0
B 0
vdk
B
each with multiplicity dx. This is exactly like the eigenvalue structure of the functional Hessian of a two-layer MLP, where the eigenvalues are known to come in positive-negative pairs Singh et al. (2021), each with multiplicity dx-
Furthermore, the T-outer-prodie dy dx and most di non-zero eigenvalues when dv S dg. This because of the rank bound from lemma D.2 and the fact that rank is equal to the number of non-zero eigenvalues for symmetric matrices.
Hessian has at most 2dxdv d non-zero eigenvalues when
Lemma D.1. Singh et al. (2021) Let A Rmxx and E IRxx
Then
I B
A 2 l
rk
rrk (A) mrk (B) - (A))
Lemma D.2. Under the same assumptions as remark D.I, and additionally assuming that WK and Wo are full rank, the rank of the T-outer-produet Hessian can by bounded by
1dxxy dx d2
if if < dy
rk (H. QK. Wok)) <
if if 2 dv.
Proof.
rk (H. (Wo,, Wok)) = ( S min (rk (U)) rk (VT)) s rk (VT).
dx
To get the rank of T lemma D.1.
Wo
2 Iay
rk (VT) < rk
Lay 2
WK
Lay X Wk
rk
(W Q
Lav
dy rk (Wk) + dy rk (Wo) rk (Wk) (WQ)
[2dy - dx when dx < dv, when dx 2 dv.
- Lac
E MULTI-HEAD
HESSIAN
Multi-head self-attention, mechanism that allows the model jointly attend to information from different representation subspaces (Vaswani et al., 2017), enforces an interesting structure in self- attention Jacobian and the Hessian of the self-attention loss. multi-head self-attention layer can
30

---

Published as a conference paper at ICLR 2025
be defined as
H Σ
H
F(X)
Fh (X)
Σ (X)XXX where A² (T*
where T do not share weights, For example, in the classical definition of multi-head self-attention, have T X W x /vdx with different W and W for every head.
Multi-head classical self-attention Jacobian blocks depend only on single-head parameters. Note that is sum of completely independent functions of X, as every weight matrix enters F' for exactly one h. This implies that block of the Jacobian corresponding to weights parameterizing Fh depends only on W W and W Recalling the formulas for the gradient from lemma B.1, arrive the following remark E.1.
Remark E.1. The Jacobians of the multi-head classical self-attention layer (Vaswani et al., 2017) have the following form:
OF
OF
XW WARTA XT
softmax
Q vdx
X &
Iay
kW
aw
OF
OF
a h TTh
X X
(I @ x
aw
aw
Vdk
where the Jacobian of the row-wise softmax w.r.t. its inputs is defined lemma B.1.
With an increasing number of heads, the outer-product Hessian dominates the Hessian. Note that the mixed inter-head terms the Hessian are fully defined by the outer-product Hessian part. This is because the inter-head functional Hessian blocks are always zero, as the second derivative of F w.r.t. arbitrary inter-head weight matrices Wh. E Ree and wh: R**xx* for h; * hy is zero. This results from the Jacobian expressions in remark E.I, which always depend on weight matrices parameterizing just single self-attention head. Hence,
He Wh: w!:)
Oe OF (F(-) & Ipg-
OF ow hu o 0.
=0
=0, when Is liggt
This implies that, for a fixed dy and dx = dv H, the functional Hessian becomes increasingly sparse as the number of heads grows. This leads to most entries of the Hessian being fully defined by the outer-product Hessian. Notably, this observation holds irrespective of the loss function used.
E.1 INFLUENCE OF THE LOw-RANE NATURE OF Wy ON THE SELF-ATTENTION HESSIAN
practical implementation of the multi-head self-attention assumes that Wh is low-rank, namely, it is parameterized as w = W W for some W E dd xdx and W E Rdx xdy In this setting, we compute the Hessian w.r.t. W and W instead of W and the only self-ateention Hessian blocks that get affected are the ones w.r.t W or W
The (Wo, W Hessian block resembles the Hessian a linear two-layer MLP. Note that Fht (X) = A (X))WW w = M, W w
(24)
resembles two-layer linear MIP if we treated M, single input matrix. Hence, the discussion from section 2 on the Hessian of MLPs applies to the value matrix block of the self-attention Hessian. For instance, its diagonal consists entirely of the outer-product Hessian entries, since its functional Hessian counterpart has a block-hollow structure.
The main difference compared full rank parameterization by single matrix Wh is that now the Hessian also includes mixed blocks w.r.t. W and Wh. For the MSE loss, the functional part
31

---

Published as a conference paper at ICLR 2025
of these blocks has quadratic dependence the first moment matrix M1. as both the derivative of the loss w.r.t the model output /e DE/OF, and the mixed second derivative matrix of F w.t.. W and W (see eq. (2)) depend on it linearly. Additionally, the outer outer-pooduct Hessian block w.nt. the value matrices W and W has an extra quadratic dependence on the selected value matrices themselves and on the first moment matrix M
F EXPERIMENTAL SETUP
For the numerical experiments, we adapt the setting from Quirke & Barez (2024). They frame number addition as the next token prediction task and generate a custom dataset, which we also use in our experiments. For 5-digit number addition that we use, the sequence length is equal to L 17. The model we consider is (unless stated otherwise) a single-block GPT-2 Transformer (Radford et al., 2019) from TransformerLLens (Nanda & Bloom, 2022), in most experiments without layer normalization unless noted otherwise. The tokens and their positions are embedded and passed into single-head self-attention layer, followed by two-layer MLP. In all figures except for figs. 7 and we use cross- coss-py (CE) as loss function.
To obtain figs. and we initialize the weights of the model the same way as GPT-2, so biases are initialized to 0 and weights are initialized by sampling from A00 0.64/dv). In fig. we use the classical definition of self-attention, while in fig. 4 we additionally compare to self-attention without softmax. The size of the model is dy = dx 16, and the latent dimension of the MLP is 64. The Hessian computed using 64 data samples.
To obtain figs. 3 and 5 we vary the standard deviation of the initialization of the embedding layer while leaving the initialization scheme of the rest of the model unchanged. We compare (fig. 3a) the model without layer normalization with (fig. 3b) model with Pre-LN (meaning that layer normal- ization is applied before the self-attention layer, before the MLP layer, and after the last Transformer block as in GPT-2 Radford et al. (2019)). The size of the Transformer block is dv dx = 128 and the Hessian is computed using 64 sequences. We estimate the block Frobenius norm using the Hutchinson trace estimator of the block outer product, implemented in the CurvLinOps library (Dangel et al., 2025). Every configuration is repeated 20 times and we report the mean and standard error of the mean. For the Pre-LN we select the slope of the dashed trend lines by fitting linear regression model into pairs of points (log a, log (()))) where f((r) is the mean Frobenius norm of the block, estimated at
The setting of fig. 6 is the same as in fig. 3a with the only difference that we consider multi-layer GPT-2 Transformers without layer normalization.
To obtain fig. 7 we use multi-layer linear self-attention networks, meaning that we simply chain self-attention layers w/o softmax. The loss we use in this experiment is mean squared error (MSE). Similar to the experiment in fig. 3, we vary the standard deviation a used to initialize the embedding matrix X.
To obtain fig. 8 we use linear MLPs with the hidden dimension 128. The loss used in this experiment also MSE.
G ADDITIONAL EXPERIMENTAL RESULTS
Hessian blocks different layers multi-layer Transformer follow the growth rafes pre- dicted for single-layer self-attention model. With the experiment in fig. 6 we try to answer question: How well do our results for single-layer network generalize to an isolated layer inside a deep network? To do that empirically demonstrate how the self-attention (with softmax) Hessian blocks of multi-layer Transformers scale with the input matrix X scale σ.
We observe that for realistic range of E (0, 1), which we refer to the practical range, the query outer-product and functional blocks as well as the value outer-prodeen block scale exactly as the prediction for single layer. Moreover, the lines corresponding to earlier layers are lower on the log-log scale for e (0, 1). which means that their Hessian blocks have smaller multipliers as part of their expressions.
32

---

Published as a conference paper at ICLR 2025
H
H
Hy W H(Wq. Wel
400
HW Wy)
H(Wq. Wo)
200
Block -hee
2-
Block || 11l
H
H.
7.
70000
10000
30000
2.-
001
0.50
1.00
0.01
0.50
1.00
10
10
(a) Practical range, i... or E (0,
(b) Mon-Practical & Practical range, i... a E (0, 10)
Figure 5: (Plotted in linear scale.) Empirical verification with CE loss confirms derived growth rates w.rff magnitude of X from eq. (5). We demonstrate the growth rates through the Frobenius norm || || of value and query diagonal blocks for (a) practical range a (0, 1) and (b) bigger o values (0, 10). The dashed lines correspond to the trend predicted by theory in eq. (5). For details on the experimental setting, see appendix F. This figure presents the same data as in fig. 3a but using linear scale both axes instead of log-log scale.
For and deeper layers we start observing some higher-order dependencies on a. This suggests that similarly to the deep linear attention networks (see section 4.1), the Hessian of multi-layer Transformer with softmax attention exhibits a dependence on that grows with depth. These higher order dependencies are not visible for < 1, because they converge to zero quicker than the lower order dependencies. We also note, that the Frobenius norm of the query Hessian block and the value outer-prodee Hessian block corresponding to the top layer quite closely follow our prediction for a single self-attention layer, also for deeper networks and 1.
For layers other than the top one, the value functional Hessian block does not follow our theoretical prediction, because, for deeper layers, this Hessian block no longer equals zero. Note that the top layer (largest layer identifier in the plot) is not present in the figure because it equals zero.
Empirical comparison between self-attention networks and MLPs. In figs. 7 and we compare the diagonal blocks of the linear self-attention and MLP Hessians. In fig. 7a we observe the growth rates predicted by remark 4.1 for a single layer of linear self-attention. As expected, the diagonal blocks are fully driven by the outer-product Hessian, as the functional Hessian blocks equal zero For growing network depth D. we observe that the growth rates of the outer-product Hessian H. blocks become super-exponential. The complete Hessian H blocks also follow this trend tthe last layer for any and the remaining ones for bigger This is inline with the discussion in section 4.1.
In contrast, the diagonal blocks of linear MLP Hessian fig. grow quadratically with a, for any of the tested network depths.
33

---

Published as a conference paper at ICLR 2025
H(Wy, Wv)
H(W We)
H.
H
H.
H
10P
10
10P
10
10-1
10
Ir
100
10-7
10
Book
10P
10
10F
100
10F-
101-5
10-3
10
10
10
110
10
10
10
10
20
211
11
111
10
100
100
(a) Depth 2
(b) Depth
H
H
H.
H.
10
10
101
1.0
101-
10
II- lIr
101
10-
10
Block
100
2011
10P
100
10
10-5
10
10-1
10
10
10P
10
10
10
10
11P
101
10
100
10
(c) Depth 4
(d) Depth
Figure 6: (Plotted in log-log scale.) Value and query Hessian diagonal blocks at different lay- follow the predicted theoretical growth rates for practical ranges of the input. Frobenius norm || || of the self-attention Hessian blocks for multi-layer GPT-2 Transformers without layer normalization the token prediction task, split by Transformer block (1 corresponds to the input Transformer block). We indicate the growth rates predicted by theorems 3.1 and 3.2 with the gray dashed lines and the annotation in the bottom right corners. As for the single layer, the com- plete Hessian H value and query blocks follow the trend of the outer- anter-product and functional Hessian blocks respectively.
34

---

Published as a conference paper at ICLR 2025
H
H
H
Block ||| l|
10-30
10a
10-17
10
10y
109
0.1
0.2
0.5
1.0
0.1
0.2
07.5
1.0
0.1
0.2
0.5
1.0
(a) Depth
There is no prediction for the functional Hessian because (as predicted by remark 4.1) the diagonal blocks of the functional Hessian equal zero.
H
H,
107
10
10-7
10
10-17
Block
Lagsn
10a
10
10
10-10
0.2
0.5
0.1
0.5
0.1
0.5
1.0
(b) Depth
H
H:
H
10
10a
100
||
10 17
10 37
10-11
Block
10a
10P
10P
10
10
10
10-17 17
10-17
10-17
0.1
0.2
0.5
0.1
0.2
0.5
1.0
0.1
0.2
0.5
1.0
(c) Depth
Figure 7: (Plotted in log-log scale.) In linear self-attention networks value and query Hessian di- agonal blocks at different layers partially follow the super-exponential growth rate with depth. Frobenius norm || |[ of the self-attention Hessian blocks for multi-layer self-attention network on the next token prediction task, split by layer (1 corresponds to the input layer). We limit the range of sigma and the network depth, due to numerical problems caused by the super-exponential growth rate for larger and deeper networks. The dashed lines indicate the trend o 2-3 where D is the network depth.
35

---

Published as a conference paper at ICLR 2025
H
H
H
Book lle
10*
10"
109
10
0.01
(a) Depth =
(b) Depth =
(c) Depth 3
(d) Depth
Figure 8: (Plotted in log-log scale.) Diagonal blocks of linear MLP grow the
with
irrespective of network depth. Frobenius norm | || of diagonal Hessian blocks for a linear MLP on the next token prediction task, split by layer MLP the diagonal blocks of a functional Hessian are always zero (Singh et al., 2021), so we simply plot the complete Hessian diagonal blocks, without splitting them into outer-product and functional Hessians. We plot the block Frobenius norm separately for every layer, but they perfectly overlap.
corresponds to the input layer). For linear
36


Reviews:
[{"review_id": "hxsdUJCAAs", "reviewer": "Reviewer_ikrU", "summary": "This paper derives and analyzes the Hessian matrix of a single Transformer self-attention layer. It examines how the Hessian depends on the data, the weights, and the attention mechanism's internal moments. The Hessian is found to be highly non-linear and varies significantly across different parts of the self-attention layer. This variation is caused by the way data enters the attention mechanism as keys, queries, and values. It is also due to the softmax function and how the attention mechanism's query and key components are parameterized. These factors create complex relationships between the data, weights, and the Hessian. The authors believe this analysis helps explain why Transformers have a unique optimization landscape. They also suggest it explains why certain architectural choices, such as using adaptive optimizers and layer normalization, are beneficial for training Transformers.", "questions": "- In Figure 3, several plots show a mismatch between the predicted and observed Hessian scaling. The top right plot in Figure 3a doesn't display a prediction at all. Could the authors elaborate on these discrepancies?\n- Some analysis is presented more like a log book without explaining why is it important. For example, what are the key takeaways from Figure 4? More broadly, could the authors clarify the overarching message and how the different analyses contribute to it?\n- The paper claims to provide a Hessian-based perspective on the performance gap between Adam and SGD, referencing Ahn et al. (2023). However, this explanation isn't explicitly provided in the paper. Could the authors elaborate on this point and clarify how their analysis explains this performance gap?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- This paper tackles an important theoretical question regarding the dynamics of Transformers by directly analyzing the Hessian. \n- A thorough theoretical derivation and analysis like this is novel and provides a valuable new perspective. \n- The categorization of Hessian dependencies offers a structured framework for understanding the complex interactions within the architecture.\n- The derivations appear sound and are presented with sufficient detail. \n- The exploration of how different Transformer components impact the Hessian adds depth and rigor to the study.\n- The paper is well written and is generally a pleasure to read. The authors incorporate the existing literature nicely. While the Hessian structure is inherently complex, the authors have made a good effort to explain the key takeaways in an accessible way.", "weaknesses": "- The paper analyses only single layer, without saying much about multi-layer.\n- A lot of important aspects are not addressed, e.g. multi-layer, role of residual connection in the Hessian, multi-head attention. Additionally, can you comment on the implications of (W_V) often being a low-rank matrix with rank (d_k)?\n- The paper doesn't have a solid narrative and rather presents a reader with a bag of tricks. See some of the examples in the Question section below. It also makes claims that are not justified, e.g. that it can help  explaining the performance gap between Adam and SGD in lines 516-519.\n- To strengthen the paper's narrative, the author should have started with the analysis of the gradient before delving into the Hessian, since it is much simpler. Comparing and contrasting the properties of the gradient and Hessian could provide a more comprehensive understanding.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "fl660ECVfW", "reviewer": "Reviewer_jto3", "summary": "The paper compares the self-attention Hessian to classical networks such as CNN to better understand the unique optimization landscape  of self-attention based transformer architectures. The paper provides a understanding self-attention from hessian perspective, which is an interesting line to understand the inner workings of transformers. The empirical experiments on digit addition task validates the theoretical observations by considering CE loss.", "questions": "Please answer the questions mentioned in previous section.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "1. The paper makes an attempt to understand self-attention based models using hessian analysis. This allows authors to compare transformers with architectures such as CNN.\n\n2. The empirical evidence on digit addition task framed as next token prediction task validates the theoretical observations.", "weaknesses": "1.  The paper is not well written and is difficult to follow.\n\n2. Authors should clearly state how their observations leads to better understanding of self-attention. It will also be beneficial for the readers if author mentions the consequences of their observations, such does it lead to better interpretability, or sparse attention or stable training.\n\n3. In section 4.2 author discuss alternative to standard query-key parameterization and discusses change in loss landscape when single matrix W_{QK} is used instead of W_{Q}W_{K}^{\\top}. Authors should discuss it briefly about how this change effects the overall performance in transformers, does it even make any difference in terms of overall performance for specific task or does it have any effect on interpretability of self-attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hVPWGwgrfG", "reviewer": "Reviewer_xwuX", "summary": "This paper derived the expression of the Hessian of one self-attention layer and discussed how the special structure of Hessian makes transformer special.", "questions": "See Weaknesses.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- This paper provides a detailed expression of the Hessian of self-attention, which might be useful for the community for the theoretical understanding of Transformers.\n- The presentation is good. I especially appreciate that the authors write different symbols in different colors.", "weaknesses": "- Except the expression of the Hessian, I don't see any deep and detailed analysis in this paper. For example, the authors claim that understanding the structure of Hessian can help understand the optimization of Transformers, such as why Transformers have to be trained by Adam(W). However, I don't see any detailed discussion on this point in the paper. I would like to see a deeper discussion showing that how the stucture of Hessian derived in this paper connects to real Transformer behaviours.\n- This whole analysis is based on a single-layer self-attention. it is unclear how this analysis (or the conclusions drawn from this one-layer model) can possibly extend to deeper models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jSevsTh68A", "reviewer": "Reviewer_TFvz", "summary": "This work derives the Hessian of Transformers to analyze the data dependence of different components in attention and to compare Transformers with other classical models.", "questions": "1. How do you justify the omission of $\\delta_{XY}$ in Equation (5)? If the elements of $\\delta_{XY}$ are significantly larger than those of $X$, wouldn't the dependency on $X$ in Equation (5) become trivial?\n\n2. Could you clarify the experimental settings used for Figure 4? You mentioned that Softmax significantly reduces the magnitude of the query Hessian block entries, but this effect isn't very apparent in Figure 4.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "$\\bullet$ The derivation of the Hessian of Transformers provides a new framework of analyzing the dynamics of different components of self-attention.\n\n$\\bullet$ The discovery of the data dependencies among key, query, and value matrices is fundamental for future works, both in theoretical understanding and practical applications.", "weaknesses": "1. The omission of the $\\textbf{F}$–functional Hessian blocks ($\\delta_{XY}$) weakens the overall results, as the influence of $\\delta_{XY}$ on the Hessian remains unclear, and there is no detailed discussion about its role.\n\n2. The analysis is focused on a single self-attention layer and does not extend directly to more complex and practical Transformer architectures. The results are insightful but could benefit from further extensions to deeper, more realistic Transformer models.\n\n3. There is no empirical comparison between Transformers and MLPs/CNNs. Including such empirical comparisons would make the findings more compelling and straightforward to interpret.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NdpABRQrn5", "reviewer": "Reviewer_gfTd", "summary": "The paper is interested in deriving the full expression of the Hessian for a single self-attention layer, wrt the learned matrix parameters of query, key and values. The hessian is decomposed into two terms, the outer product and functional hessians, and their expressions are respectively given in Theorems 3.1 and 3.2. Then, the paper analyzes the dependence on the data and how different components of the architecture affect the hessian, such as the softmax activation or the position of the layer normalization.", "questions": "- I would be interested in having more details about the settings of the experiments leading to the figures shown in the paper, more precisely for Figure 1 and Figure 4, and the dashed lines in Figure 3. What is exactly plotted ? What kind of data were used to obtain these ? How does it confirm the insights derived from the theoretical derivations ?\n- In Figure 3b, what does \"the dashed lines correspond to the trend estimated from the data points by the linear regression coefficient\" mean ? Can the authors describe the setting behind this experiment and how the dashed lines are obtained ?\n- In Figure 3, all the trends in dashed lines are linear, even though the order of the dependence is changing. This makes me think that the range of values considered for $\\sigma$ is too small to clearly evaluate whether the empirical dependence are following the theoretical ones. Can the authors discuss that, and if possible, show results with a bigger range of values for $\\sigma$ ?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- **Originality**: To my knowledge, this is the first paper deriving the full expression of the hessian for the self-attention operation.\n- **Significance**: As mentioned in the conclusion of the paper, this work can serve as foundation for better understanding the role of the self attention operation in Transformers. As discussed and shown throughout the paper, the self attention layer has a singular behavior compared to better-understood convolutional or feed-forward layers in neural networks.   \n- **Quality**: Although I did not check all the proofs in details, a lot of work has been put to derive Theorems 3.1 and 3.2. The experiments presented in Figure 3 also validates to some extent the theoretical results obtained, in terms of dependence to the training data of two of the diagonal terms.\n- **Clarity**: I appreciated the color-coding of the terms within equations throughout the paper. It makes the reading and understanding of the results easier.", "weaknesses": "- **Clarity**: The links between the empirical results shown in the various figures and the insights derived from the expressions of the Hessian are not always clear. For instance, the experiments and what is plotted in Figure 1 are never described.  \n- **Quality**: It is difficult to evaluate the validity of all theoretical insights derived from the hessian since the settings of the experiments are not always described. More specifically, settings behind experiments to obtain Figure 1, Figure 3 and Figure 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)