Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

Guía de lectura

Problema y motivación

Durante la optimización de redes profundas con SGD, los cambios en los parámetros de las capas anteriores alteran la distribución de las entradas de las capas posteriores. Los autores denominan a este fenómeno Internal Covariate Shift.

El internal covariate shift genera dos desafíos principales:

1. Las subredes deben adaptarse continuamente a distribuciones de activación cambiantes.
2. Las activaciones derivan fácilmente hacia los regímenes de saturación de las no linealidades (como las funciones sigmoides), provocando desvanecimiento de gradientes y exigiendo tasas de aprendizaje pequeñas y una inicialización meticulosa de los parámetros.

Aunque el blanqueamiento (decorrelación y estandarización) de las entradas de las capas en todo el conjunto de datos podría estabilizar las distribuciones, calcular matrices de covarianza completas y sus raíces cuadradas inversas en cada actualización de parámetros es computacionalmente prohibitivo y puede hacer que los gradientes exploten si la normalización no está integrada en la retropropagación.

Ideas clave

• Capa de normalización diferenciable: En lugar de estandarizar las activaciones fuera de la optimización, la normalización se integra directamente en la arquitectura del modelo, asegurando que la retropropagación considere las dependencias entre las activaciones de las capas y las estadísticas del mini-batch.
• Normalización por dimensión: En lugar de un blanqueamiento conjunto completo, cada dimensión escalar de características se normaliza de forma independiente a media cero y varianza unitaria.
• Restauración de la capacidad representacional: Para asegurar que la normalización no destruya la capacidad expresiva de una capa (como forzar a una sigmoide a comportarse de forma puramente lineal), se introducen parámetros aprendibles $\gamma$ (escala) y $\beta$ (desplazamiento) para permitir que la capa recupere la transformación identidad si es necesario.
• Inferencia determinista: Durante el entrenamiento, las estadísticas se calculan sobre cada mini-batch $\mathcal{B}$. En el momento de la prueba, estadísticas poblacionales fijas (calculadas mediante promedios móviles o sobre los lotes de entrenamiento) reemplazan a las estadísticas del mini-batch, garantizando una inferencia determinista de una sola muestra.
• Invarianza a la escala de parámetros: Las transformaciones de escala de los pesos de las capas satisfacen $\text{BN}((aW)\mathrm{u})=\text{BN}(W\mathrm{u})$, lo que significa que pesos más grandes conducen a gradientes efectivos más pequeños, previniendo la explosión de gradientes.

Método (con las ecuaciones clave)

Para un mini-batch $\mathcal{B}=\{x_{1\dots m}\}$ correspondiente a una sola dimensión de activación:

1. Calcular la media del mini-batch: $${\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}\sum _{i=1}^m{x}_i$$
2. Calcular la varianza del mini-batch: $${\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m}\sum _{i=1}^m(x_i-{\mu }_{\mathcal{B}}{)}^2$$
3. Normalizar las activaciones (con la constante de estabilidad $ϵ$): $${\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}$$
4. Escalar y desplazar usando los parámetros aprendibles $\gamma$ y $\beta$: $$y_i\leftarrow \gamma {\hat{x}}_i+\beta \equiv {\text{BN}}_{\gamma ,\beta }(x_i)$$

Aplicación a las capas

• Ubicación: Se aplica inmediatamente antes de las no linealidades: $z=g(\text{BN}(Wu))$. El sesgo original de la capa $b$ se omite porque la resta de la media lo cancela (su función es asumida por $\beta$).
• Capas convolucionales: Para preservar la propiedad de equivariancia a la traslación, la normalización se aplica por mapa de características en lugar de por activación espacial. Para un tamaño de lote $m$ y una resolución espacial $p\times q$, el tamaño efectivo del mini-batch pasa a ser $m^{\prime }=m\cdot p\cdot q$, con un único par de $(\gamma ,\beta )$ compartido en cada mapa de características.

Transformación en inferencia

En el momento de la prueba, la transformación se convierte en una operación lineal estática utilizando las estadísticas poblacionales $\mathrm{E}[x]$ y $\mathrm{Var}[x]$: $$y=\frac{\gamma }{\sqrt{\mathrm{Var}[x]+ϵ}}\cdot x+\left(\beta -\frac{\gamma \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\right)$$

Resultados (con las cifras más importantes)

• MNIST con sigmoides: Las redes equipadas con Batch Normalization conservaron distribuciones de activación estables a lo largo del entrenamiento, acelerando la convergencia y alcanzando una mayor precisión de prueba en comparación con las redes no normalizadas.
• Velocidad de entrenamiento en ImageNet (red individual):
  • Inception no normalizado alcanzó un 72,2% de precisión de validación tras $31,0\times 10^6$ pasos.
  • La línea base de Inception con Batch Normalization alcanzó el 72,2% en $13,3\times 10^6$ pasos.
  • Con una tasa de aprendizaje 5 veces mayor (BN-x5), alcanzó el 72,2% en $2,1\times 10^6$ pasos—una reducción de 14 veces en los pasos de entrenamiento requeridos, alcanzando finalmente un pico del 73,0%.
  • Con una tasa de aprendizaje 30 veces mayor (BN-x30), el modelo alcanzó una precisión máxima del 74,8% tras $6,0\times 10^6$ pasos.
• No linealidades saturadas: Mientras que Inception no normalizado con activaciones sigmoides falló en el entrenamiento (quedándose cerca del 0,1% de precisión aleatoria), BN-x5-Sigmoid logró un 69,8% de precisión top-1.
• SOTA Ensemble en ImageNet: Un ensemble de 6 redes normalizadas por lotes logró un error top-5 de validación del 4,9% y un error de prueba del 4,82% en el servidor de pruebas de ILSVRC, superando las líneas base de evaluadores humanos.

Limitaciones y preguntas abiertas

• Dependencia del tamaño de lote: Dado que la normalización se basa en las estadísticas muestrales del mini-batch, tamaños de mini-batch muy pequeños pueden degradar la calidad de las estimaciones de la varianza y la media.
• Redes recurrentes: El artículo evalúa modelos feedforward y convolucionales, señalando a las redes neuronales recurrentes (RNN) como una cuestión abierta no comprobada donde el internal covariate shift también podría ser crítico.
• Fundamentos teóricos: Los autores conjeturan que Batch Normalization hace que los jacobianos de las capas tengan valores singulares cercanos a 1, pero señalaron que la verificación teórica precisa del comportamiento del gradiente queda pendiente para análisis futuros.