Module 5 — Connexions résiduelles et normalisation par couche

Un Transformer profond, disons douze couches, ne s'entraîne pas si l'on retire les résidus et la normalisation. Ce n'est pas une question d'accessoire : ce sont deux mécanismes structurels qui rendent le gradient utilisable jusqu'aux premières couches et stabilisent les activations d'un bloc à l'autre. Ce module explique pourquoi, et pose enfin le squelette du bloc Transformer que nous assemblerons au module 10.

Le raccourci qui ne traverse aucune couche

Une connexion résiduelle ajoute l'entrée d'un bloc à sa sortie :

$$y = x + \mathrm{Bloc}(x)$$

L'idée, popularisée par ResNet en 2015 pour la vision, est simple à énoncer mais profonde dans ses effets. Le bloc n'apprend plus la transformation complète $y = f(x)$, mais seulement la correction $y - x = f(x) - x$. Si la meilleure transformation est proche de l'identité — et cela arrive souvent dans les couches profondes — le bloc peut se contenter d'apprendre une petite modification autour de zéro, ce qui est bien plus facile.

L'effet sur le gradient est encore plus important. Sans résidu, la dérivée d'une composition de $N$ blocs $f_N \circ \cdots \circ f_1$ enchaîne $N$ produits jacobiens : si chacun est inférieur à 1, le gradient s'évanouit ; s'il est supérieur, il explose. Avec un résidu, la dérivée par rapport à $x$ contient un terme identité — le raccourci — qui laisse toujours passer une part du signal, quelle que soit la profondeur. C'est ce qui rend les Transformeurs à 12, 24 ou 96 couches entraînables.

LayerNorm plutôt que BatchNorm

Le second ingrédient est la normalisation par couche (LayerNorm), qui centre et met à l'échelle les activations. Contrairement à la normalisation par lots (BatchNorm), elle calcule ses statistiques sur les traits d'un même exemple, pas sur les exemples d'un lot :

$$\mathrm{LayerNorm}(x) = \gamma \cdot \frac{x - \mu(x)}{\sqrt{\sigma^2(x) + \epsilon}} + \beta$$

où $\mu$ et $\sigma^2$ sont calculés sur la dernière dimension du vecteur d'un jeton, et $\gamma, \beta$ sont deux vecteurs appris de taille $d$.

Ce choix a trois conséquences pratiques :

• il fonctionne avec un lot de taille 1 en inférence, ce qui est indispensable pour un modèle servi qui répond une requête à la fois ;
• il est indépendant de la longueur de séquence : chaque jeton se normalise seul ;
• il conserve les mêmes statistiques à l'entraînement et à l'inférence, alors que la normalisation par lots doit maintenir des moyennes courantes qui divergent parfois, cf. module 8 du cours 08.

Aspect	LayerNorm	BatchNorm
Axe de normalisation	Traits d'un même jeton	Lot pour un même trait
Robuste à un lot de taille 1	Oui	Non
Utile pour les séquences de longueurs variables	Oui	Non
Adopté dans les Transformeurs	Systématiquement	Presque jamais

Pré-norm ou post-norm

Il existe deux façons d'agencer résidu et normalisation dans un bloc. La version post-norm de l'article original applique la norme après le résidu :

$$y = \mathrm{LayerNorm}(x + \mathrm{Bloc}(x))$$

La version pré-norm, adoptée par GPT-2 et devenue standard depuis 2019, place la norme avant le bloc, à l'intérieur du résidu :

$$y = x + \mathrm{Bloc}(\mathrm{LayerNorm}(x))$$

La pré-norm produit un chemin de gradient qui traverse la pile de bout en bout sans passer par une fonction non linéaire, ce qui la rend beaucoup plus stable en début d'entraînement. Elle permet de se passer du réchauffage du taux d'apprentissage (warm-up), obligatoire avec la post-norm sur les gros modèles. C'est pourquoi la quasi-totalité des Transformeurs modernes utilisent la pré-norm.

Post-norm et gros modèles ne font pas bon ménage

Beaucoup d'implémentations pédagogiques copient encore la post-norm de l'article de 2017. Cela fonctionne sur un petit modèle comme le nôtre, mais devient instable au-delà de six ou sept couches sans warm-up soigneux. Pour ce cours, nous adopterons la pré-norm dès le module 10.

Le bloc feed-forward complète chaque couche

Chaque couche d'un Transformer alterne deux sous-couches : une attention multi-têtes et un petit réseau à propagation avant appliqué position par position. Ce feed-forward est un simple sandwich :

$$\mathrm{FFN}(x) = W_2 \, \phi(W_1 x + b_1) + b_2$$

avec $W_1 \in \mathbb{R}^{d \times d_{\mathrm{ff}}}$, $W_2 \in \mathbb{R}^{d_{\mathrm{ff}} \times d}$, et une fonction d'activation $\phi$ — ReLU dans l'article original, GELU dans BERT, SwiGLU dans les modèles récents. La dimension intermédiaire $d_{\mathrm{ff}}$ vaut classiquement $4 d$.

C'est ce bloc qui détient la majorité des paramètres de la couche. Pour $d = 512, d_{\mathrm{ff}} = 2048$, le feed-forward pèse $2 \times 512 \times 2048 \approx 2$ millions de paramètres, soit environ le double de l'attention multi-têtes ($4 d^2 \approx 1$ million, cf. module 3). Autrement dit, l'attention mélange l'information entre positions, le feed-forward la transforme à chaque position isolément, et c'est ce dernier qui domine le budget.

Le bloc complet, en code

Voici le squelette d'une couche d'encodeur en pré-norm, avec ses deux résidus et ses deux normes :

import torch
import torch.nn as nn

class CoucheEncodeur(nn.Module):
    def __init__(self, d: int, h: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.norme_1 = nn.LayerNorm(d)
        self.attention = AttentionMultiTetes(d, h, dropout)   # module 3
        self.dropout_1 = nn.Dropout(dropout)

        self.norme_2 = nn.LayerNorm(d)
        self.feed_forward = nn.Sequential(
            nn.Linear(d, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d),
        )
        self.dropout_2 = nn.Dropout(dropout)

    def forward(self, x, masque=None):
        # Pre-norm : la norme est appliquee AVANT le sous-bloc.
        y, _ = self.attention(self.norme_1(x), self.norme_1(x), self.norme_1(x), masque)
        x = x + self.dropout_1(y)

        y = self.feed_forward(self.norme_2(x))
        x = x + self.dropout_2(y)
        return x

couche = CoucheEncodeur(d=64, h=8, d_ff=256)
x = torch.randn(2, 12, 64)
print(couche(x).shape)   # (2, 12, 64)

On y retrouve tout le vocabulaire des modules précédents : l'attention multi-têtes, la pré-norm, le résidu, un dropout après chaque sous-bloc, et un feed-forward avec une activation GELU. La couche préserve la forme $(B, L, d)$, ce qui la rend empilable : dix, vingt ou cent couches identiques s'enchaînent sans transformation supplémentaire.

Un dropout par sous-couche, jamais dans le résidu direct

Le dropout doit être appliqué avant l'addition résiduelle, sur la sortie du sous-bloc uniquement. Le placer sur $x$ après le résidu casse la propriété d'identité et déstabilise la pile.

En résumé

• Les connexions résiduelles transforment un bloc en apprenant une correction plutôt qu'une transformation, et fournissent un chemin de gradient direct qui rend la pile profonde entraînable.
• LayerNorm normalise les traits d'un même jeton, ce qui la rend robuste aux lots de taille 1 et aux séquences de longueurs variables, contrairement à BatchNorm.
• La pré-norm — norme avant le sous-bloc — est le standard depuis 2019 ; elle supprime le besoin de warm-up sur les gros modèles.
• Chaque couche alterne attention multi-têtes et feed-forward de dimension $d_{\mathrm{ff}} = 4 d$ ; c'est ce dernier qui porte la majorité des paramètres.

Module suivant : empiler ces couches pour former un encodeur — la moitié « compréhension » du Transformer, incarnée par BERT et ses variantes.