ReLU's Revival: On the Entropic Overload in Normalization-Free Large Language Models

TL;DR

This study offers key insights for optimizing transformer architectures where LayerNorm introduces significant challenges, and highlights that smoother activations like GELU are ill-suited for LayerNorm-free architectures, whereas ReLU's geometrical properties lead to improved learning dynamics and better information retention in the absence of LayerNorm.

Abstract

LayerNorm is a critical component in modern large language models (LLMs) for stabilizing training and ensuring smooth optimization. However, it introduces significant challenges in mechanistic interpretability, outlier feature suppression, faithful signal propagation, and computational and communication complexity of private inference. This work explores desirable activation functions in normalization-free decoder-only LLMs. Contrary to the conventional preference for the GELU in transformer-based models, our empirical findings demonstrate an {\em opposite trend} -- ReLU significantly outperforms GELU in LayerNorm-free models, leading to an {\bf 8.2\%} perplexity improvement. We discover a key issue with GELU, where early layers experience entropic overload, leading to the under-utilization of the representational capacity of attention heads. This highlights that smoother activations like GELU are {\em ill-suited} for LayerNorm-free architectures, whereas ReLU's geometrical properties -- specialization in input space and intra-class selectivity -- lead to improved learning dynamics and better information retention in the absence of LayerNorm. This study offers key insights for optimizing transformer architectures where LayerNorm introduces significant challenges. The code and implementation are available at https://github.com/Nandan91/relu-revival-normfree

Figures (5)

• Figure 1: Entropy heatmaps of attention for baseline (a, b) and normalization-free (c, d) GPT-2 models with GELU and ReLU in the FFN. In the absence of LayerNorm, GELU in the FFN leads to significantly higher entropic overload (highlighted in yellow, c) compared to ReLU.
• Figure 2: Headwise entropy distribution and evaluation loss for baseline and normalization-free GPT-2 models, using GELU and ReLU activations, trained from scratch on CodeParrot dataset.
• Figure 3: Learnable negative slope for leaky ReLU in FFN of LN-free GPT-2 model. (a) Layerwise slopes showing initial variability and convergence towards zero. (b) Global slope trend towards zero over training steps, indicating a preference for zero negative slope in LN-free architectures.
• Figure 4: Evolution of Layerwise entropy when GPT-2 ($L$=12, $H$=12, $d$=768) models with various nonlinearity configurations are trained from scratch on CodeParrot dataset. Evolution of layer-wise entropy during training of GPT-2 models ($L$=12, $H$=12, $d$=768) with different nonlinearity configurations on the CodeParrot dataset. The near-identical entropy dynamics in Figures d, e, and f underscore a natural preference for a zero negative slope, similar to ReLU, in the FFN activation function of the normalization-free model.
• Figure 5: Training instability, indicated by NaNs, and corresponding entropy dynamics in Normalization-Free GPT-2 ($L$=12, $H$=12, $d$=768) models with fixed negative slopes in the leaky ReLU. The larger the negative slope, the earlier the training instability and entropy collapse occurred.