Discussion Regarding Prior Work
We would like to thank the reviewers for their time and effort on evaluating our paper, and want to take this opportunity to provide clarification on how the proposed interpretation differs from prior work cited in the paper.
As mentioned by the reviewers, Brody et al. were the first to introduce the geometric interpretation of LayerNorm and its relationship to vector projection and scaling. Nonetheless, their main focus was on how this interpretation allows for the computability of certain functions within a single attention block (specifically the _majority_ function), a goal the authors explicitly refer to in the title of their paper as the "Expressivity Role of LayerNorm".
We extend this interpretation to the entire transformer architecture by analyzing its consequences on the internal representation space of every layer:
* First, we slightly modify the decomposition proposed by Brody et al. to derive a formulation of LayerNorm using only the geometric primitives (vector projection and scaling), such that orthogonality to the $\overrightarrow{\mathbf{1}}$ vector does not arise as a numerical consequence of subtracting the mean (as discussed in Brody et al.), but is instead proposed as an intrinsic property of the normalization step. To illustrate, note that under our formulation (Equations 3 & 4) is easy to see how this projection can be generalized to vector/hyperplane combinations different from $\frac{1}{\sqrt{d}}\overrightarrow{\mathbf{1}}$, while it is not straightforward to do so under Brody et al.'s interpretation.
Equiped with a strong geometric intuition on what LayerNorm does, we revisit the ideas of Xiong et al., Geva et al. and Millidge & Black:
* [Xiong et al.] We discuss at the end of Section 2.1 how the arrangement of transformer layers proposed by Xiong et al. (commonly known as Pre-LN Transformers) affects the model's internal representations, such that **all attention and feed forward layers within a transformer share the same input representation space**, which is different from the representation space of the residual stream, although it is related to it via projection onto the $d-1$ hypersphere (defined by $\frac{1}{\sqrt{d}}\overrightarrow{\mathbf{1}}$). This is a notable result, as common intuition would suggest that the representation space of each layer after normalization would be unique, given that the mean and standard deviation are data-dependent variables. Under the proposed geometric interpretation of LayerNorm is easy to see this is not the case, as all layers apply exactly the same projection to the representation space of the residual stream.
* [Geva et al. and Millidge & Black] In Section 2.2.2 we discuss how the SVD method proposed by Millidge & Black can be interpreted in relation to the key-value stores introduced by Geva et al. In both of these works, linear probing has been limited to the value "side" of either the VO matrix or the feed forward layer, that is, the parameter groups that directly interact with the residual stream. From Millidge & Black: _"Our method is extremely simple. Take a weight matrix M of the network. Take the SVD of this matrix to obtain left and right singular vectors M=USV. Take whichever matrix has the same dimensionality as the residual stream (typically the right singular matrix V)"_. Similarly, Geva et al. only probe the weights of the MLP value matrices, and rely on trigger examples (instead of probing) to interpret the information encoded in the MLP key matrices. We complement their work by **introducing a probing methodology to understand both sides of these key-value matrices** in Section 3.3.1 and show the results in Appendix E.