RMSNorm/es
| Article | |
|---|---|
| Topic area | neural-network-components |
| Prerequisites | Layer Normalization, Transformer |
Visión general
La normalización por capas con raíz cuadrática media (RMSNorm) es una capa de normalización utilizada en redes neuronales profundas, introducida por Zhang y Sennrich en 2019 como una simplificación de la Layer Normalization. RMSNorm reescala cada vector de entrada por su raíz cuadrática media (RMS) y aplica una ganancia aprendida por característica, pero a diferencia de LayerNorm omite el paso de centrado de la media (re-centrado) y el bias aditivo. El resultado es una capa con aproximadamente la mitad del coste aritmético de LayerNorm, conservando la mayor parte de sus beneficios de estabilidad durante el entrenamiento.
RMSNorm se ha convertido en un componente estándar de las arquitecturas modernas de Transformer, en particular de los grandes Language Models. Es la capa de normalización utilizada en T5, LLaMA, LLaMA 2, LLaMA 3, Mistral, Gemma, Qwen y muchos otros modelos de pesos abiertos. Su popularidad stems de una combinación de simplicidad, un cómputo ligeramente más rápido, una estabilidad de entrenamiento marginalmente mejor con la disposición Pre-LayerNorm, y la observación empírica de que el re-centrado de las activations rara vez modifica la calidad en tareas posteriores.
Motivación
La Layer Normalization estandariza un vector $ x \in \mathbb{R}^d $ restando su media y dividiendo por su desviación estándar, para luego aplicar una escala aprendida $ \gamma $ y un desplazamiento $ \beta $. El paso de re-centrado (la sustracción de la media) se motivó originalmente por analogía con la Batch Normalization, que centra cada característica a lo largo del lote. Zhang y Sennrich observaron que los beneficios prácticos de LayerNorm en redes profundas provienen casi por completo del re-escalado y no del re-centrado: el papel de la capa es mantener acotadas las magnitudes de las activation para que los gradientes no se desvanezcan ni exploten. El paso de sustracción de la media implica una pasada adicional sobre el vector, un parámetro adicional ($ \beta $) y aporta poco a la optimización.
RMSNorm es la respuesta a la pregunta: "¿qué ocurre si conservamos solo la parte de re-escalado?". Es invariante al re-escalado por vector de la entrada pero, a diferencia de LayerNorm, no es invariante a un desplazamiento constante. Empíricamente, esta invariancia faltante no ha tenido un coste medible en el entrenamiento moderno de Transformer.
Formulación
Dado un vector de entrada $ x \in \mathbb{R}^d $, RMSNorm calcula
$ {\displaystyle \operatorname{RMS}(x) = \sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2}} $
y produce como salida
$ {\displaystyle \operatorname{RMSNorm}(x)_i = \frac{x_i}{\operatorname{RMS}(x) + \varepsilon}\, g_i} $
donde $ g \in \mathbb{R}^d $ es un vector de ganancia aprendida (inicializado a unos) y $ \varepsilon $ es una constante pequeña (típicamente $ 10^{-6} $ o $ 10^{-5} $) añadida por estabilidad numérica. No hay bias aprendido ni sustracción de la media.
Para una secuencia de vectores de tokens, RMSNorm se aplica de forma independiente por token, exactamente como se hace con LayerNorm en los Transformers. En forma matrixicial, si $ X \in \mathbb{R}^{n \times d} $ apila $ n $ vectores de tokens como filas, entonces
$ {\displaystyle \operatorname{RMSNorm}(X) = \operatorname{diag}(g)\, X \oslash \sqrt{\frac{1}{d}(X \odot X)\mathbf{1} + \varepsilon}} $
donde $ \odot $ y $ \oslash $ son la multiplicación y la división elemento a elemento, y $ \mathbf{1} $ es el vector de unos.
Relación con la normalización por capas
LayerNorm calcula la media $ \mu = \frac{1}{d}\sum_i x_i $ y la desviación estándar $ \sigma = \sqrt{\frac{1}{d}\sum_i (x_i - \mu)^2} $, y produce como salida $ \gamma \odot (x - \mu)/(\sigma + \varepsilon) + \beta $. RMSNorm es exactamente LayerNorm con $ \mu $ forzada a cero y $ \beta $ eliminada. De forma equivalente, RMSNorm es LayerNorm restringida al caso en que la distribución de entrada ya está centrada, lo cual se vuelve aproximadamente cierto tras unas pocas capas de entrenamiento en Transformers profundos.
El coste aritmético de LayerNorm por vector de dimensión $ d $ es aproximadamente dos reducciones (media y variance) más una transformación afín elemento a elemento; RMSNorm tiene una sola reducción (suma de cuadrados) más una operación de escalado elemento a elemento. En las GPU modernas, el ahorro en tiempo de reloj es pequeño de forma aislada, pero acumulado a lo largo de cientos de capas y billones de tokens resulta no despreciable. Más importante aún, el kernel más simple facilita escribir implementaciones fusionadas y mantenerlas numéricamente estables en precision mixta.
Ubicación en los Transformers
RMSNorm se utiliza casi siempre en la configuración residual Pre-LayerNorm (pre-norma): la capa se aplica a la entrada de cada subbloque, y la residual connection suma la entrada sin normalizar a la salida del subbloque. Esquemáticamente, para el subbloque de Self-attention,
$ {\displaystyle y = x + \operatorname{Attention}(\operatorname{RMSNorm}(x))} $
y de forma análoga para el subbloque Feedforward Network. Esto contrasta con la disposición Post-LayerNorm (post-norma) del Transformer original, donde la normalización sigue a la suma residual. La pre-norma con RMSNorm entrena de forma estable sin trucos de warmup de tasa de aprendizaje a escalas en las que la post-norma típicamente diverge, y es el valor por defecto en prácticamente todos los LLM modernos de pesos abiertos.
Convencionalmente se aplica una RMSNorm final a la salida del último bloque, antes de la proyección de des-incrustación.
Variantes y extensiones
Varias variantes de RMSNorm aparecen en la literatura y en modelos de producción.
Partial RMSNorm calcula la RMS utilizando únicamente las primeras $ k < d $ coordenadas de la entrada. La intuición es que la RMS de un vector de alta dimensión se concentra fuertemente alrededor de su media, por lo que una suma parcial es casi igual de precisa. Zhang y Sennrich reportaron una pérdida de calidad insignificante con $ k = d/8 $ y un entrenamiento más rápido. Partial RMSNorm no ha visto una adopción amplia en la práctica porque la aceleración absoluta es pequeña y los kernels modernos de atención dominan el tiempo de ejecución.
Gated RMSNorm multiplica la salida de RMSNorm por una función de compuerta de la entrada o de otro tensor. Se utiliza en algunas arquitecturas de State Space Model, en particular Mamba 2, donde se aplica como compuerta a la salida del SSM antes de la proyección.
Group RMSNorm divide la dimensión de características en $ G $ grupos y normaliza cada grupo de forma independiente con su propio vector de ganancia, en analogía con la Group Normalization para convolutional networks. Group RMSNorm se utiliza dentro de cabezas de atención agrupadas en algunas arquitecturas recientes.
QK Normalization aplica RMSNorm a los vectores de consulta y clave dentro de Self-attention antes del dot product. Esto estabiliza el entrenamiento a escalas muy grandes al impedir que los logits previos al softmax se desplacen hacia magnitudes extremas, y se utiliza en Gemma 2, en varios modelos de frontera y en el Vision Transformer ViT-22B.
Consideraciones de implementación
Una implementación de referencia en PyTorch es una de las capas más simples del aprendizaje profundo:
<syntaxhighlight lang="python"> class RMSNorm(nn.Module):
def __init__(self, dim, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(dim))
self.eps = eps
def forward(self, x):
# x: (..., dim)
rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
return self.weight * (x * rms)
</syntaxhighlight>
Las implementaciones de producción fusionan toda la capa en un único kernel CUDA para minimizar el tráfico de memoria, y calculan la suma de cuadrados en fp32 aun cuando las entradas estén en bf16 o fp16, para evitar el desbordamiento cuando $ d $ es grande (p. ej., 8192 en LLaMA 70B). El vector de ganancia se almacena típicamente en la misma precision que las activations y se convierte a fp32 únicamente dentro del kernel fusionado. La $ \varepsilon $ se añade dentro de la raíz cuadrada, no después, para mantener un gradiente bien comportado cuando la entrada está cerca de cero.
Apex, FlashAttention, and the major training frameworks all ship optimized RMSNorm kernels. The forward and backward passes are bandwidth-bound rather than compute-bound on accelerators.
Empirical Performance
Across machine translation, language modeling, and downstream evaluation tasks, RMSNorm matches or slightly exceeds LayerNorm in final task quality while training 5-10% faster end-to-end on common Transformer sizes. The original paper reported wins of 0.1-0.3 BLEU on WMT translation benchmarks at no quality cost, with ablations on RNN-based models showing similar trends. Subsequent large-scale studies, including the empirical work behind the T5 and LLaMA model families, found no scenario in which LayerNorm provided a clear quality advantage. Combined with its simpler kernel and good interaction with pre-norm residual connections, this evidence has made RMSNorm the default normalization choice in essentially every Transformer-based Language Model released since 2022.
Limitations
RMSNorm is not invariant to a constant additive shift of its input, so models that rely on encoding information in the mean of a representation could in principle behave differently under RMSNorm than under LayerNorm. In practice, Self-attention and Feedforward Network sub-blocks do not exploit the mean of their input, and the activations in pre-norm Transformers are approximately mean-zero by the second or third block, so the missing invariance is not measurable in real models. RMSNorm shares with LayerNorm the property that all coordinates of a vector influence the normalization of every other coordinate, which couples activations across the feature dimension and complicates certain forms of model surgery (such as pruning or per-feature quantization) more than would a per-feature normalizer like Batch Normalization.