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

Guide de lecture

Problème et motivation

Lors de l'optimisation des réseaux profonds par SGD, les modifications des paramètres des premières couches modifient la distribution des entrées des couches suivantes. Les auteurs nomment ce phénomène l'Internal Covariate Shift.

L'internal covariate shift engendre deux difficultés majeures :

1. Les sous-réseaux doivent s'adapter en permanence à des distributions d'activations changeantes.
2. Les activations dérivent facilement vers les régimes de saturation des fonctions non linéaires (telles que les sigmoïdes), ce qui provoque la disparition du gradient (vanishing gradient) et requiert des taux d'apprentissage faibles ainsi qu'une initialisation méticuleuse des paramètres.

Bien que le blanchiment (décorrélation et standardisation) des entrées de chaque couche sur l'ensemble du jeu de données puisse stabiliser les distributions, le calcul des matrices de covariance complètes et de l'inverse de leurs racines carrées à chaque mise à jour est prohibitif sur le plan computationnel et peut provoquer une explosion des gradients si la normalisation n'est pas intégrée dans la rétropropagation.

Idées clés

• Couche de normalisation différentiable : Au lieu de standardiser les activations en dehors de l'optimisation, la normalisation est directement intégrée dans l'architecture du modèle, garantissant que la rétropropagation prenne en compte les dépendances entre les activations de la couche et les statistiques du mini-lot.
• Normalisation par dimension : Plutôt qu'un blanchiment conjoint complet, chaque dimension scalaire des caractéristiques est normalisée indépendamment pour avoir une moyenne nulle et une variance unitaire.
• Restauration de la capacité de représentation : Afin de s'assurer que la normalisation ne détruise pas la capacité expressive d'une couche (comme forcer une sigmoïde à fonctionner uniquement dans sa zone linéaire), des paramètres apprenables $\gamma$ (échelle) et $\beta$ (décalage) sont introduits pour permettre à la couche de retrouver la transformation identité si nécessaire.
• Inférence déterministe : Durant l'entraînement, les statistiques sont calculées sur chaque mini-lot $\mathcal{B}$. Lors du test, des statistiques globales fixes sur la population (calculées par moyennes glissantes ou sur l'ensemble des lots d'entraînement) remplacent les statistiques de mini-lot, garantissant une inférence déterministe échantillon par échantillon.
• Invariance à l'échelle des paramètres : Les changements d'échelle des poids d'une couche vérifient $\text{BN}((aW)\mathrm{u})=\text{BN}(W\mathrm{u})$, ce qui signifie que des poids plus élevés conduisent à des gradients effectifs plus faibles, évitant ainsi l'explosion du gradient.

Méthode (avec les équations clés)

Pour un mini-lot $\mathcal{B}=\{x_{1\dots m}\}$ correspondant à une seule dimension d'activation :

1. Calculer la moyenne du mini-lot : $${\mu }_{\mathcal{B}}\leftarrow \frac{1}{m}\sum _{i=1}^m{x}_i$$
2. Calculer la variance du mini-lot : $${\sigma }_{\mathcal{B}}^2\leftarrow \frac{1}{m}\sum _{i=1}^m(x_i-{\mu }_{\mathcal{B}}{)}^2$$
3. Normaliser les activations (avec une constante de stabilité $ϵ$) : $${\hat{x}}_i\leftarrow \frac{x_i-{\mu }_{\mathcal{B}}}{\sqrt{{\sigma }_{\mathcal{B}}^2+ϵ}}$$
4. Appliquer le changement d'échelle et le décalage via les paramètres apprenables $\gamma$ et $\beta$ : $$y_i\leftarrow \gamma {\hat{x}}_i+\beta \equiv {\text{BN}}_{\gamma ,\beta }(x_i)$$

Application aux couches

• Positionnement : Appliqué immédiatement avant les non-linéarités : $z=g(\text{BN}(Wu))$. Le biais original $b$ de la couche est omis car la soustraction de la moyenne l'annule (son rôle est repris par $\beta$).
• Couches convolutionnelles : Pour préserver la propriété d'équivariance par translation, la normalisation est appliquée par carte de caractéristiques (feature map) plutôt que par position spatiale. Pour une taille de lot $m$ et une résolution spatiale $p\times q$, la taille effective du mini-lot devient $m^{\prime }=m\cdot p\cdot q$, avec une seule paire $(\gamma ,\beta )$ partagée par carte de caractéristiques.

Transformation lors de l'inférence

Au moment du test, la transformation devient une opération linéaire statique utilisant les statistiques de population $\mathrm{E}[x]$ et $\mathrm{Var}[x]$ : $$y=\frac{\gamma }{\sqrt{\mathrm{Var}[x]+ϵ}}\cdot x+\left(\beta -\frac{\gamma \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\right)$$

Résultats (avec les chiffres les plus importants)

• MNIST avec sigmoïdes : Les réseaux dotés de Batch Normalization ont conservé des distributions d'activation stables tout au long de l'entraînement, accélérant la convergence et atteignant une meilleure précision de test par rapport aux réseaux non normalisés.
• Vitesse d'entraînement sur ImageNet (modèle unique) :
  • Inception sans normalisation a atteint 72,2 % de précision de validation après $31,0\times 10^6$ étapes.
  • Inception avec Batch Normalization de base a atteint 72,2 % en $13,3\times 10^6$ étapes.
  • Avec un taux d'apprentissage 5 fois plus élevé (BN-x5), il a atteint 72,2 % en $2,1\times 10^6$ étapes, soit une réduction d'un facteur 14 du nombre d'étapes d'entraînement requises, pour plafonner finalement à 73,0 %.
  • Avec un taux d'apprentissage 30 fois plus élevé (BN-x30), le modèle a atteint une précision maximale de 74,8 % après $6,0\times 10^6$ étapes.
• Non-linéarités saturantes : Alors qu'Inception non normalisé avec des activations sigmoïdes a échoué à converger (restant proche du niveau du hasard à 0,1 %), BN-x5-Sigmoid a obtenu une précision top-1 de 69,8 %.
• Ensemble SOTA sur ImageNet : Un ensemble de 6 réseaux normalisés par lots a obtenu une erreur top-5 de validation de 4,9 % et une erreur de test de 4,82 % sur le serveur d'évaluation ILSVRC, surpassant les performances d'évaluateurs humains.

Limites et questions ouvertes

• Dépendance à la taille du lot : La normalisation reposant sur les statistiques empiriques du mini-lot, des tailles de lot très réduites peuvent dégrader la précision de l'estimation de la moyenne et de la variance.
• Réseaux récurrents : L'article teste des modèles feedforward et convolutionnels, désignant les réseaux de neurones récurrents (RNN) comme une piste non explorée où l'internal covariate shift pourrait également s'avérer critique.
• Fondements théoriques : Les auteurs formulent l'hypothèse que Batch Normalization amène les jacobiens des couches à avoir des valeurs singulières proches de 1, mais soulignent qu'une démonstration théorique rigoureuse du comportement des gradients reste un sujet de recherche futur.