Geometric Interpretation of Layer Normalization and a Comparative Analysis with RMSNorm

TL;DR

This paper offers a geometric perspective on LayerNorm, revealing its connection to the uniform vector, and demonstrates that RMSNorm can replace LayerNorm for improved efficiency without loss of performance.

Contribution

It provides a new geometric interpretation of LayerNorm, compares it with RMSNorm, and shows that removing the uniform vector component is redundant, advocating for RMSNorm use.

Findings

LayerNorm is linked to the uniform vector in representation space.

Standardization in LayerNorm involves removing the uniform component, normalizing, and scaling.

LLMs operate orthogonal to the uniform vector during inference, making removal redundant.

Abstract

This paper presents a novel geometric interpretation of LayerNorm and explores how LayerNorm influences the norm and orientation of hidden vectors in the representation space. With these geometric insights, we prepare the foundation for comparing LayerNorm with RMSNorm. We show that the definition of LayerNorm is innately linked to the uniform vector, defined as $\boldsymbol{1} = [1, 1, 1, 1, \cdots, 1]^T \in \mathbb{R}^d$. We then show that the standardization step in LayerNorm can be understood in three simple steps: (i) remove the component of a vector along the uniform vector, (ii) normalize the remaining vector, and (iii) scale the resultant vector by $\sqrt{d}$, where $d$ is the dimensionality of the representation space. We also provide additional insights into how LayerNorm operates at inference time. Finally, we compare the hidden representations of LayerNorm-based LLMs with models trained using RMSNorm and show that all LLMs naturally operate orthogonal to the uniform vector at inference time, that is, on average they do not have a component along the uniform vector during inference. This presents the first mechanistic evidence that removing the component along the uniform vector in LayerNorm is a redundant step. These results advocate for using RMSNorm over LayerNorm which is also more computationally efficient.