On the Expressivity Role of L ayer N orm in Transformers’ Attention (2023)

Abstract: Layer Normalization (LayerNorm) is an inherent component in all Transformer-based models. In this paper, we show that LayerNorm is crucial to the expressivity of the multi-head attention layer that follows it. This is in contrast to the common belief that LayerNorm’s only role is to normalize the activations during the forward pass, and their gradients during the backward pass. We consider a geometric interpretation of LayerNorm and show that it consists of two components: (a) projection of the input vectors to a d-1 space that is orthogonal to the [1,1,...,1] vector, and(b) scaling of all vectors to the same norm of d. We show that each of these components is important for the attention layer that follows it in Transformers:(a) projection allows the attention mechanism to create an attention query that attends to all keys equally, offloading the need to learn this operation in the attention; and(b) scaling allows each key to potentially receive the highest attention, and prevents keys from being “un-select-able”.We show empirically that Transformers do indeed benefit from these properties of LayeNorm in general language modeling and even in computing simple functions such as “majority”. Our code is available at https://github.com/tech-srl/layer_norm_expressivity_role .

## introduction layernorm @xcite is the most commonly used normalization technique in modern neural networks such as transformers @xcite . originally, @xcite motivated layer-norm as an efficient way of normalizing the activations during the forward pass or providing distribution stability as in batch normalization @xcite . later, @xcite and @xcite argued that more importantly than normalizing forward activations, layernorm stabilizes the gradients during the backward pass. however, in this work, we show that layernorm, which was originally proposed for rnns, has an additional crucial role in the theoretical and practical expressivity of the multi-head attention layer that follows it in transformers. 1 that is, layer-norm makes it easier for the transformer to learn certain functions during training. layernorm can be seen as two independent components: projection and scaling, that were merged into a single operator. first, layernorm projects its inputs onto a particular d -1 space that is orthogonal to the "ones" ⃗ 1 = [1, 1, ..., 1] vector. this allows the attention layer that follows the layer-norm to create queries that are close to ⃗ 1, and thus attend to all keys equally, when needed, regardless of the identity of the keys. in section 3 we show that this projection helps, for example, computing the "majority" among token types in a sequence. figure 1a shows how without layernorm, the keys and queries in a transformer's attention have no apparent geometric structure. in contrast, figure 1b shows that layernorm has projected all keys to the hyperplane that is orthogonal to the ⃗ 1 vector. further, the attention mechanism has learned queries that are close to ⃗ 1, making them attend equally to any possible key, when trained to compute "majority". we analyze and prove this in section 4.1. the second component of layernorm is scaling: we show that layernorm scales the projected input to have an ℓ 2 norm of exactly √ d. in section 3, we show that scaling the input vectors prevents the problem of "unselectable" keys @xcite , where some key vectors are contained in the convex hull formed by the other keys, and thus can never get the highest attention score. figures 1c and 1d show the average fraction w/o layernorm figure 1 : figures 1a and 1b show the effect of projection in layernorm, which makes all keys lie on the hyperplane that is orthogonal to the ⃗ 1 vector. figures 1c and 1d show the effect of scaling, where n is the number of vectors, d is the dimension, and the color represents the average fraction of "unselectable" key vectors. ## decomposing layernorm given an input @xmath0, layernorm is defined as the following quotient: foot_1 we start with the numerator xµ and show that it corresponds to the projection of x onto the hyperplane h defined by the normal vector that is, xµ is always orthogonal to the ⃗ 1 vector. next, we show that the denominator scales the projected vector to have a norm of exactly √ d: then, σ in the denominator of layernorm scales the projected vector to have a norm of exactly √ d. layernorm can thus be seen as two independent components: (a) projection of the input vectors onto the hyperplane orthogonal to ⃗ 1, and (b) scaling of the projected vectors to have a norm of √ d. ## expressivity role in attention each of the components of layernorm supports the transformer's attention in a different way. projection helps creating a query that attends to all keys equally, when needed, while scaling helps the model to avoid the problem of "unselectable" keys. recall that in transformers, given vectors q and k, the attention scoring function is defined as: ## experimental results in this section, we empirically show the effects of the layernorm components -projection and scaling -on the attention mechanism in transformers. we first show how a transformer learns to use the projection of layernorm to compute "majority"; then, we show that scaling allows the model to avoid the problem of "unselectable" keys, allowing the attention mechanism to focus on any key. ## conclusion in this paper, we show that the commonly used layernorm component is crucial not only for the optimization process but also for the expressivity of attention in transformers. we decompose lay-ernorm into two geometric operations: projecting the input vectors onto a subspace that is orthogonal to the ⃗ 1 vector, and scaling the projected vectors to have the same norm √ d. we show that projection helps to compute even simple tasks such as "majority" by performing an exact average of the keys and that scaling helps to avoid the problem of "unselectable" keys. these findings are important for the community's understanding of attention and expressivity in transformers. further, these results raise a variety of follow-up questions, such as: why should the keys be orthogonal to the ⃗ 1 vector, instead of any, learnable, different vector for every layer? and what would happen if we force each layer's keys to be orthogonal to multiple normal vectors? to this end, we make our code publicly available at https://github.com/tech-srl/ layer_norm_expressivity_role .