Geometric Interpretation of Layer Normalization and a Comparative Analysis with RMSNorm

2024-09-19

AI Summary

LayerNorm’s “mean subtraction” step is widely assumed necessary, yet its functional necessity—particularly in large language model (LLM) inference—is poorly understood. Method: We analyze LayerNorm geometrically as a three-step operation: orthogonal projection onto the uniform vector direction (i.e., mean removal), normalization of the residual, and affine scaling. Leveraging vector space decomposition and empirical analysis of hidden states across mainstream LLMs, we examine the alignment of activations with the uniform vector during inference. Contribution/Results: We find that LLM hidden states are naturally orthogonal to the uniform vector at inference time, rendering the projection (mean subtraction) step redundant. This provides the first mechanistic explanation and empirical validation for RMSNorm’s efficacy: by omitting projection while preserving normalization and scaling, RMSNorm achieves comparable performance with reduced computational overhead. Our work unifies the representational understanding of LayerNorm and RMSNorm, revealing an intrinsic structural simplification in normalization for LLM inference.

Technology Category

• Natural Language Processing: (Large) Language Models
• Machine Learning: Large Multimodal Models (LMMs)
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions

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 $oldsymbol{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.

Problem

• LayerNorm
• RMSNorm
• Normalization Techniques

Innovation

• LayerNorm Simplification
• RMSNorm Efficiency
• Intrinsic Vector Relation