PyTorch LayerNorm : dimensions, epsilon et paramètres

Dans PyTorch, nn.LayerNorm normalise les dimensions finales indiquées par normalized_shape, puis applique éventuellement des paramètres appris d’échelle et de décalage. Pour une entrée de forme (lot, séquence, caractéristiques), LayerNorm(features) normalise chaque vecteur de caractéristiques indépendamment. La forme de sortie reste identique à celle de l’entrée.

Le choix déterminant porte sur les valeurs réunies dans un même groupe de normalisation. Deux configurations peuvent renvoyer la même forme tout en calculant des résultats différents. Ce guide utilise des tenseurs fixes en float64 sur CPU pour montrer cette différence, puis examine la variance, epsilon, les paramètres affines, le mode d’évaluation et les erreurs d’organisation des images.

Les six exemples s’exécutent indépendamment avec Python 3.12.14 et PyTorch 2.15.0+cpu. Nos 44 vérifications comprennent un calcul scalaire indépendant et une comparaison de gradients par différences finies. Elles ne mesurent ni la vitesse GPU ni la qualité d’un modèle. La référence LayerNorm documente les arguments utilisés.

Choisir normalized_shape à partir des axes finaux

Le premier tenseur a pour forme (2, 2, 3). Ses axes représentent le lot, la position dans la séquence et la caractéristique. L’entier 3 signifie que la dernière dimension doit avoir trois éléments. Le tuple (2, 3) sélectionne ensemble les deux dernières dimensions.

import torch

x = torch.tensor([
    [[1., 2., 4.], [10., 20., 40.]],
    [[3., 5., 9.], [2., 3., 8.]]
], dtype=torch.float64)
last = torch.nn.LayerNorm(3, elementwise_affine=False)(x)
last_two = torch.nn.LayerNorm((2, 3), elementwise_affine=False)(x)
print('Input and output shapes:', tuple(x.shape), tuple(last.shape), tuple(last_two.shape))
print('Last-axis means:', last.mean(dim=-1).round(decimals=6).tolist())
print('Two-axis means:', last_two.mean(dim=(-2, -1)).round(decimals=6).tolist())
print('Same values:', torch.allclose(last, last_two))
print('First row, last axis:', [round(v, 4) for v in last[0, 0].tolist()])
print('First row, two axes:', [round(v, 4) for v in last_two[0, 0].tolist()])

Input and output shapes: (2, 2, 3) (2, 2, 3) (2, 2, 3)
Last-axis means: [[-0.0, -0.0], [-0.0, 0.0]]
Two-axis means: [-0.0, 0.0]
Same values: False
First row, last axis: [-1.069, -0.2673, 1.3363]
First row, two axes: [-0.8612, -0.7884, -0.6429]

LayerNorm(3) effectue quatre normalisations distinctes de trois valeurs, une pour chaque paire élément du lot-position. LayerNorm((2, 3)) effectue deux normalisations de six valeurs, une par élément du lot. Aucune ne réunit les deux éléments du lot dans cet exemple.

La première ligne devient environ [-1.069, -0.2673, 1.3363] avec le dernier axe seul, mais [-0.8612, -0.7884, -0.6429] lorsqu’elle partage ses statistiques avec les valeurs plus élevées de la deuxième ligne. Les formes coïncident : une assertion sur les dimensions seules manquerait ce changement de sens.

Les valeurs -0.0 affichées sont des résultats flottants arrondis, proches de zéro. Nos tests les comparent avec une tolérance. Un autre test modifie le deuxième élément du lot et confirme que le résultat du premier, normalisé sur deux axes, reste inchangé.

Pour une séquence de forme (B, L, D), LayerNorm(D) laisse B et L varier. LayerNorm((L, D)) exige que ces deux dimensions finales correspondent à sa configuration. Décidez si les positions doivent partager leurs statistiques avant de choisir le tuple. Notre guide de l’architecture Transformer, en anglais, présente le contexte du modèle ; les exemples ici se concentrent sur l’opération de normalisation.

Reproduire le calcul avec la bonne variance

Dans chaque groupe, soustrayez la moyenne et divisez par la racine carrée de la variance augmentée de eps. Avant l’étape affine facultative, l’expression est (x - mean) / sqrt(variance + eps). Le calcul suivant utilise la moyenne des écarts au carré, en divisant par le nombre de valeurs plutôt que par ce nombre moins un :

import torch

x = torch.tensor([[1., 2., 4.], [8., 8., 8.]], dtype=torch.float64)
eps = 1e-5
mean = x.mean(dim=-1, keepdim=True)
variance = ((x - mean) ** 2).mean(dim=-1, keepdim=True)
manual = (x - mean) / torch.sqrt(variance + eps)
layer = torch.nn.LayerNorm(3, eps=eps, elementwise_affine=False)
y = layer(x)
print('Means:', mean.flatten().tolist())
print('Variances:', variance.flatten().tolist())
print('Matches manual:', torch.allclose(y, manual, rtol=1e-12, atol=1e-12))
print('Constant row:', y[1].tolist())
print('Output variances:', y.var(dim=-1, correction=0).tolist())

Means: [2.3333333333333335, 8.0]
Variances: [1.5555555555555554, 0.0]
Matches manual: True
Constant row: [0.0, 0.0, 0.0]
Output variances: [0.9999935714698979, 0.0]

Pour [1, 2, 4], la moyenne vaut 7/3 et la variance 14/9. Notre calcul Python indépendant correspond à la sortie du module. La variance d’échantillon avec correction=1 ne convient pas ici ; utilisez correction=0 lorsque vous vérifiez cette opération avec torch.var.

Le groupe constant [8, 8, 8] a une variance nulle et produit des zéros avec eps positif, sans transformation affine. Sa variance de sortie ne vaut pas un. Même celle du groupe non constant est légèrement inférieure à un : avant l’échelle et le décalage, elle vaut v/(v + eps), où v désigne la variance du groupe d’entrée.

Cette distinction compte dans les assertions. Vérifiez l’expression réelle avec une tolérance numérique plutôt que d’imposer une variance exactement égale à un à chaque groupe. Les paramètres appris peuvent encore modifier les statistiques de sortie.

Comprendre l’échelle et le décalage appris

Avec elementwise_affine=True, LayerNorm conserve un poids et un biais de forme normalized_shape. Le poids initial vaut un et le biais zéro. Cet exemple les remplace par des valeurs connues pour en vérifier l’effet :

import torch

x = torch.tensor([[1., 2., 4.], [2., 4., 8.]], dtype=torch.float64)
layer = torch.nn.LayerNorm(3, dtype=torch.float64)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([2., 0.5, -1.], dtype=torch.float64))
    layer.bias.copy_(torch.tensor([1., -2., 3.], dtype=torch.float64))
z = torch.nn.LayerNorm(3, elementwise_affine=False)(x)
y = layer(x)
expected = z * layer.weight + layer.bias
scale_only = torch.nn.LayerNorm(3, bias=False, dtype=torch.float64)
no_affine = torch.nn.LayerNorm(3, elementwise_affine=False)
print('Matches scale and shift:', torch.allclose(y, expected, rtol=1e-12, atol=1e-12))
print('Output row means:', [round(v, 4) for v in y.mean(dim=-1).tolist()])
print('Parameter counts:', *(sum(p.numel() for p in m.parameters())
                              for m in (layer, scale_only, no_affine)))
print('Scale-only bias:', scale_only.bias)

Matches scale and shift: True
Output row means: [-0.536, -0.536]
Parameter counts: 6 3 0
Scale-only bias: None

La sortie correspond aux valeurs normalisées multipliées par [2, 0.5, -1], puis décalées de [1, -2, 3]. Ces trois paramètres de caractéristiques sont partagés entre les lignes. Les moyennes des lignes deviennent environ -0.536 : la sortie affine apprise n’a donc pas nécessairement une moyenne nulle.

Le nombre de paramètres distingue trois configurations : la couche ordinaire de longueur trois possède six scalaires ; bias=False conserve trois paramètres d’échelle ; elementwise_affine=False supprime l’échelle et le décalage. Pour normalized_shape=(2, 3), notre test supplémentaire trouve un poids et un biais de forme (2, 3), soit douze scalaires.

L’API fonctionnelle layer_norm reçoit explicitement weight, bias et eps. Notre appel fonctionnel correspond au module configuré lorsque ces valeurs sont identiques. Utilisez le module lorsque vous voulez qu’il possède et enregistre les paramètres appris dans votre modèle.

Les exemples définissent des paramètres float64 lorsqu’ils sont présents, comme les tenseurs d’entrée. Dans un modèle réel, gardez les entrées et les paramètres sur des périphériques et dans des types compatibles ; changer normalized_shape ne corrige pas une incompatibilité de périphérique ou de type numérique.

Observer l’effet de eps sur un groupe presque constant

Epsilon est ajouté à l’intérieur de la racine carrée. Son effet apparaît clairement lorsque la variance d’entrée est faible :

import torch

x = torch.tensor([[1., 1.001, 0.999], [7., 7., 7.]], dtype=torch.float64)
small = torch.nn.LayerNorm(3, eps=1e-5, elementwise_affine=False)(x)
large = torch.nn.LayerNorm(3, eps=0.1, elementwise_affine=False)(x)
print('Small eps:', [[round(v, 6) for v in row] for row in small.tolist()])
print('Large eps:', [[round(v, 6) for v in row] for row in large.tolist()])
print('Larger eps reduces magnitudes:', bool(large[0].abs().max() < small[0].abs().max()))
print('Both finite:', bool(torch.isfinite(small).all() and torch.isfinite(large).all()))

Small eps: [[0.0, 0.306186, -0.306186], [0.0, 0.0, 0.0]]
Large eps: [[0.0, 0.003162, -0.003162], [0.0, 0.0, 0.0]]
Larger eps reduces magnitudes: True
Both finite: True

Les valeurs du premier groupe s’écartent de leur moyenne d’environ 0.001 seulement. Avec eps=1e-5, les sorties non nulles ont une amplitude d’environ 0.306186. Augmenter eps à 0.1 la réduit à environ 0.003162. Les deux résultats correspondent à notre calcul indépendant. Le groupe constant reste nul et toutes les sorties restent finies.

La grande valeur sert volontairement à rendre l’effet visible ; ce n’est pas une recommandation de réglage par défaut. Modifier eps change le calcul. Consignez ce paramètre avec la configuration du modèle, surtout lors d’une comparaison d’implémentations ou du chargement d’un checkpoint dont l’architecture est définie ailleurs.

Distinguer mode d’évaluation et suivi des gradients

Cette couche utilise les statistiques des groupes de l’entrée courante en entraînement comme en évaluation. L’exemple compare les modes, puis propage les gradients à travers la sortie d’évaluation :

import torch

x = torch.tensor([[1., 2., 4.], [3., 7., 8.]], dtype=torch.float64, requires_grad=True)
layer = torch.nn.LayerNorm(3, dtype=torch.float64)
layer.train()
training = layer(x)
layer.eval()
evaluation = layer(x)
weights = torch.tensor([1., 2., 4.], dtype=torch.float64)
(evaluation * weights).sum().backward()
print('Train/eval equal:', torch.equal(training, evaluation))
print('Running-stat buffers:', list(layer.named_buffers()))
print('Eval output requires gradients:', evaluation.requires_grad)
print('Input gradient finite:', bool(torch.isfinite(x.grad).all()))
print('Input gradient nonzero:', bool(x.grad.abs().max() > 0))

Train/eval equal: True
Running-stat buffers: []
Eval output requires gradients: True
Input gradient finite: True
Input gradient nonzero: True

La couche LayerNorm isolée renvoie des sorties exactement identiques pour cette entrée fixe et ces paramètres inchangés. Elle ne stocke aucun tampon de statistiques courantes dans la configuration testée. L’appel eval() ne détache pas le résultat : l’entrée reçoit des gradients finis et non nuls.

Nous pondérons la sortie par [1, 2, 4] pour éviter une simple somme de valeurs centrées, peu informative comme exemple de gradient. La suite de vérification compare chaque dérivée d’entrée à une différence finie calculée indépendamment et contrôle séparément les gradients du poids et du biais.

D’autres couches d’un modèle complet peuvent encore changer de comportement entre train et eval. Choisissez séparément l’enregistrement des gradients selon la tâche ; notre guide PyTorch sur eval, no_grad et inference_mode, en anglais, décrit ces contrôles. Des sorties LayerNorm identiques ne rendent pas eval() inutile pour tout le modèle.

Ne pas confondre canaux d’image et dernier axe

Pour un tenseur organisé en (N, C, H, W), LayerNorm(C) vérifie toujours le dernier axe, ici W. Il n’interprète pas C comme le nom d’un axe de canaux. Notre entrée a deux canaux et une largeur de trois : cet appel échoue.

import torch

x = torch.arange(12, dtype=torch.float64).reshape(1, 2, 2, 3)
try:
    torch.nn.LayerNorm(2, elementwise_affine=False)(x)
except RuntimeError as error:
    print('NCHW channel-count attempt:', type(error).__name__)
channels_last = x.permute(0, 2, 3, 1)
normalized = torch.nn.LayerNorm(2, elementwise_affine=False)(channels_last)
back = normalized.permute(0, 3, 1, 2)
print('Restored shape:', tuple(back.shape))
print('Channel means:', back.mean(dim=1).round(decimals=6).tolist())
print('First pixel:', [round(v, 6) for v in back[0, :, 0, 0].tolist()])

NCHW channel-count attempt: RuntimeError
Restored shape: (1, 2, 2, 3)
Channel means: [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]
First pixel: [-0.999999, 0.999999]

La permutation place les axes dans l’ordre logique (N, H, W, C). LayerNorm(2) normalise alors la paire de valeurs de canaux à chaque pixel. La permutation inverse rétablit (1, 2, 2, 3). Au premier pixel, les valeurs 0 et 6 ont une moyenne de 3 et une variance de 9, donnant environ [-0.999999, 0.999999].

Notre appel supplémentaire à LayerNorm(3) sur le tenseur NCHW initial réussit, mais normalise la largeur et produit des valeurs différentes. Si le nombre de canaux est égal à la largeur, l’appel incorrect peut aussi réussir le contrôle des dimensions. Nommez les axes au lieu d’attendre qu’une exception révèle toutes les erreurs.

La référence permute décrit le réordonnancement des dimensions. Un reshape ne remplace pas le déplacement explicite de l’axe des canaux. Distinguez aussi la normalisation des canaux à chaque pixel de celle de tous les canaux et positions spatiales avec une forme (C, H, W) : les groupes sont différents.

Vérifier une petite entrée avant l’intégration au modèle

• Notez le sens des axes et la forme finale exacte des valeurs qui doivent partager leurs statistiques.
• Comparez un groupe à un calcul manuel avec correction=0 et eps à l’intérieur de la racine carrée.
• Inspectez les formes des paramètres et distinguez bias=False de elementwise_affine=False.
• Testez séparément le mode d’évaluation et les gradients du modèle complet.

Les exemples et le code de vérification sont disponibles dans l’archive de revue. Nos tests utilisent des tenseurs denses float64 sur CPU ; CUDA, la précision mixte, les gains d’apprentissage et les performances nécessitent leur propre évaluation. LayerNorm ne remplace pas universellement toutes les autres méthodes de normalisation.