What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis

Authors: Weronika Ormaniec, Felix Dangel, Sidak Pal Singh
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: hessian, transformers

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.

Decision: accept (spotlight)
Reviewer scores: 6, 8, 5, 8, 8

Reviews and discussion:
<thread>
creation date:
1730137832662

thread_type:
official_review

<thread_object>
creator: reviewer_gftd

<content>
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.

strengths:
- **Originality**: To my knowledge, this is the first paper deriving the full expression of the hessian for the self-attention operation.
- **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.   
- **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.
- **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.  
- **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.

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 ?
- 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 ?
- 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$ ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730445583358

thread_type:
official_review

<thread_object>
creator: reviewer_tfvz

<content>
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.

strengths:
$\bullet$ The derivation of the Hessian of Transformers provides a new framework of analyzing the dynamics of different components of self-attention.

$\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.

2. 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.

3. There is no empirical comparison between Transformers and MLPs/CNNs. Including such empirical comparisons would make the findings more compelling and straightforward to interpret.

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?

2. 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730667139219

thread_type:
official_review

<thread_object>
creator: reviewer_xwux

<content>
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.

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.
- 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.
- 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.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730703414971

thread_type:
official_review

<thread_object>
creator: reviewer_jto3

<content>
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.

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.

2. 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.

2. 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.

3. 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.

questions:
Please answer the questions mentioned in previous section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730789600402

thread_type:
official_review

<thread_object>
creator: reviewer_ikru

<content>
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.

strengths:
- This paper tackles an important theoretical question regarding the dynamics of Transformers by directly analyzing the Hessian. 
- A thorough theoretical derivation and analysis like this is novel and provides a valuable new perspective. 
- The categorization of Hessian dependencies offers a structured framework for understanding the complex interactions within the architecture.
- The derivations appear sound and are presented with sufficient detail. 
- The exploration of how different Transformer components impact the Hessian adds depth and rigor to the study.
- 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.
- 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)?
- 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.
- 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.

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?
- 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?
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732568625471

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to All Reviewers

comment:
We thank all reviewers for their thoughtful feedback and constructive comments. We are pleased that the reviewers find that our paper “tackles an important theoretical question” [ikrU] and that our discoveries are “useful for the community” [xwuX], “fundamental for future works” [TFvz] and “serv[ing] as foundation for better understanding the role of the self attention” [gfTd].

We reply to each of the reviewers’ questions and concerns individually below. To track the significant changes to the manuscript, we colored any additions to the PDF in blue if they go beyond minor edits.

In our global response, we address the multi-layer scenario that was requested by some of the reviewers.

**Multi-layer setting:** 
* Our goal was to derive and interpret in detail the closed-form expressions for the Hessian. Following other works on a single layer or a very shallow Transformer (for example [1, 2, 3]), we believe a single self-attention layer is an interesting object worth studying. Our work serves as an important building block for extending the theory to multiple layers. You have a point that we currently can only address certain scenarios with depth (see below). We highlighted this limitation of our theory in the revision.

* To facilitate the discussion of the multi-layer Transformer Hessian, we empirically check what scaling laws can be observed for its blocks:
    * New in Appendix G: Empirically, we find out that for a multi-layer GPT-2 Transformer and practically relevant ranges of $\sigma$ we can apply our block scaling laws **to every single layer in isolation**. We added the figures and the discussion of this experiment to Appendix G (Figure 6). 

* Our theoretical claims extend to some extent into the deep case. We made the following additions:
    * New in Appendix G: We empirically verified the discussion of depth in a Transformer with linear attention from L445-451 and also compared it with a deep linear MLP. As predicted by our theory, we observe the super-exponential (with depth) data dependency for the Transformer (Figure 7), in contrast to the constant (with depth) dependency from the MLP (Figure 8).
    * Theorems 3.1 and 3.2 directly apply to the last attention layer in a self-attention net, by replacing $\mathbf{X}$ with the second to last layer’s output. Similarly, parts of our analysis carry over to any self-attention layer in a deep Transformer by applying the chain rule between the layer and the remainder of the network. We have added these comments as a footnote in the main text.

References:

[1] Understanding Addition in Transformers

[2] Are Transformers with One Layer Self-Attention Using Low-Rank Weight Matrices Universal Approximators?

[3] Linear attention is (maybe) all you need (to understand Transformer optimization)
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper provides a theoretical analysis of the Hessian matrix of a (single) attention layer in transformer architectures. This analysis is highly relevant, as it opens up characterization of the relationship between data, weight and attention. In turn, new insights are obtained on why transformers may be particularly amenable to e.g. adaptive optimizers or layer normalization, but in more general why they have a unique optimization landscape different from other deep networks. Reviewers all agreed that this is a fundamental and important contribution to a field where transformers are prevalent in usage. The contributions of the paper thus lay out theoretical building blocks to further understanding and open up the way for more theory to be developed. Apart from some concerns on clarity and writing, that were mostly resolved throughout the discussion phase, a common reviewer concern was the focus of the analysis on a single layer. The authors have conducted tangible revision steps to ensure that their proposed growth rates are nevertheless valuable for the analysis of multiple layers and that they may empirically hold up. Overall, reviewers are happy to accept the paper to the conference. As a consequence of the significance of the contribution, the reviewers’ positive feedback, and the constructive revisions, the AC proposed to accept the paper. Because the paper provides a well-written analysis on a highly popular topic, the AC further suggests the paper for an oral presentation, as it is likely that a large part of the community will benefit from the obtained insights.