PyTorch nn.Embedding : dimensions, padding et gradients

Dans PyTorch, nn.Embedding associe des identifiants entiers aux lignes d’une matrice de poids. Avec num_embeddings=4 et embedding_dim=3, cette matrice contient quatre vecteurs de trois valeurs. La recherche d’un identifiant sélectionne son vecteur ; un lot d’identifiants produit un lot de vecteurs.

Ce guide vérifie le comportement de la couche : valeurs lues, dimensions de sortie, gradients des indices répétés, padding, poids gelés et max_norm. Nous avons exécuté 41 vérifications avec Python 3.12.14 et PyTorch 2.15.0+cpu. Les six exemples utilisent des poids fixes et présentent les sorties capturées. Ils ne mesurent ni les performances GPU ni la qualité des représentations apprises.

Pour comprendre la place des embeddings dans un modèle plus large, consultez notre guide d’architecture des réseaux de neurones, en anglais. Ici, l’objectif est de savoir quelles lignes la couche lit et quelles opérations peuvent les modifier.

Lire les dimensions des poids et de la sortie

Nous fournissons une petite matrice avec from_pretrained pour vérifier chaque résultat à la main. Le nom de ce constructeur ne signifie pas que ces nombres ont été appris : il s’agit de données de test originales.

import torch

weights = torch.tensor([[0., 0., 0.], [1., 2., 3.],
                        [4., 5., 6.], [7., 8., 9.]])
layer = torch.nn.Embedding.from_pretrained(weights.clone(), freeze=False)
ids = torch.tensor([[2, 1], [0, 2]], dtype=torch.long)
output = layer(ids)
print('Weight shape:', tuple(layer.weight.shape))
print('Output shape:', tuple(output.shape))
print('Output:', output.tolist())

Weight shape: (4, 3)
Output shape: (2, 2, 3)
Output: [[[4.0, 5.0, 6.0], [1.0, 2.0, 3.0]], [[0.0, 0.0, 0.0], [4.0, 5.0, 6.0]]]

L’entrée a pour dimensions (2, 2). La sortie conserve ces deux axes et ajoute une largeur de vecteur de trois, soit (2, 2, 3). L’identifiant 2 sélectionne la troisième ligne, [4, 5, 6], à chacune de ses occurrences. L’identifiant 0 sélectionne la première ligne.

Pour un lot de tokens de dimensions (batch, sequence_length), la même règle donne (batch, sequence_length, embedding_dim). La référence de l’API Embedding décrit ce contrat de dimensions. Nos vérifications couvrent aussi un identifiant scalaire, une liste à une dimension et un axe d’identifiants vide. Une multiplication indépendante par une matrice one-hot reproduit toutes les valeurs de cet exemple de lecture simple.

La correspondance entre identifiants et lignes doit rester cohérente. Si l’ordre du vocabulaire change, le même entier sélectionnera le vecteur d’un autre token. Conservez cette correspondance avec la configuration et le checkpoint du modèle. Inspectez les tenseurs d’identifiants produits par votre DataLoader PyTorch, présenté dans notre guide en anglais, avant leur arrivée dans la couche.

Vérifier le type et les bornes des indices

Les exemples utilisent torch.long. Les identifiants int64 et int32 fonctionnent dans nos tests CPU ; les identifiants à virgule flottante sont rejetés. Pour la matrice de quatre lignes, les valeurs valides vont de 0 à 3. Les indices -1 et 4 provoquent tous deux une erreur.

Un token inconnu doit recevoir explicitement un identifiant valide. La couche ne crée pas automatiquement une ligne supplémentaire pour une valeur inconnue et ne transforme pas les chaînes de caractères en identifiants. Le tenseur d’entiers n’est pas non plus le paramètre appris : les gradients portent sur la matrice de poids à virgule flottante.

Un padding_idx négatif a un sens distinct lors de la construction. Notre test avec padding_idx=-1 désigne la dernière ligne comme padding, mais une entrée contenant -1 échoue toujours. N’utilisez pas l’indexation négative de Python dans le tenseur d’entrée.

Observer les gradients des identifiants répétés

import torch

weights = torch.tensor([[0., 0.], [1., 2.], [3., 4.]])
ids = torch.tensor([1, 1, 2], dtype=torch.long)
plain = torch.nn.Embedding.from_pretrained(weights.clone(), freeze=False)
scaled = torch.nn.Embedding.from_pretrained(
    weights.clone(), freeze=False, scale_grad_by_freq=True)
plain(ids).sum().backward()
scaled(ids).sum().backward()
print('Default gradient:', plain.weight.grad.tolist())
print('Frequency-scaled gradient:', scaled.weight.grad.tolist())

Default gradient: [[0.0, 0.0], [2.0, 2.0], [1.0, 1.0]]
Frequency-scaled gradient: [[0.0, 0.0], [1.0, 1.0], [1.0, 1.0]]

La perte est la somme de toutes les valeurs lues : chaque occurrence contribue donc un gradient de un à chaque composante de sa ligne. L’identifiant 1 apparaît deux fois et donne [2, 2] ; l’identifiant 2 apparaît une fois et donne [1, 1]. La ligne 0, inutilisée, reçoit des zéros.

Avec scale_grad_by_freq=True, le gradient de la ligne répétée est divisé par sa fréquence dans cette entrée. Les deux lignes utilisées ont alors [1, 1]. Nos vérifications confirment que les valeurs de la passe avant restent identiques. Cette option modifie le calcul du gradient, pas le vecteur sélectionné.

Chaque exemple crée de nouvelles couches. Dans une boucle d’entraînement, effacez les gradients accumulés à l’étape prévue. Les nombres ci-dessus décrivent un seul appel à backward avec cette somme comme perte ; ils ne s’appliquent pas à toutes les fonctions de perte ni à tous les schémas d’accumulation.

Tester le padding avec un optimiseur explicite

import torch

layer = torch.nn.Embedding(4, 2, padding_idx=0)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([[0., 0.], [1., 2.],
                                     [3., 4.], [5., 6.]]))
ids = torch.tensor([0, 1, 1, 2], dtype=torch.long)
optimizer = torch.optim.SGD(layer.parameters(), lr=0.1,
                            momentum=0, weight_decay=0)
optimizer.zero_grad(set_to_none=True)
output = layer(ids)
output.sum().backward()
gradient = layer.weight.grad.clone()
before = layer.weight.detach().clone()
optimizer.step()
print('Padding output:', output[0].tolist())
print('Gradient:', gradient.tolist())
print('Padding row unchanged:', torch.equal(before[0], layer.weight[0]))
print('Unused row unchanged:', torch.equal(before[3], layer.weight[3]))
print('Used row changed:', not torch.equal(before[1], layer.weight[1]))

Padding output: [0.0, 0.0]
Gradient: [[0.0, 0.0], [2.0, 2.0], [1.0, 1.0], [0.0, 0.0]]
Padding row unchanged: True
Unused row unchanged: True
Used row changed: True

La ligne de padding renvoie des zéros et ne reçoit aucun gradient de lecture. La ligne 1 reçoit deux contributions, la ligne 2 une seule et la ligne 3 inutilisée aucune. Avec cette configuration de SGD, les lignes de padding et inutilisée restent inchangées, tandis qu’une ligne utilisée est modifiée. Nous avons aussi vérifié la mise à jour complète par rapport à before - 0.1 * gradient.

Un vecteur de padding nul ne retire pas automatiquement les positions de padding des réductions suivantes. Pour la sortie ci-dessus, la moyenne ordinaire sur les quatre positions vaut [1.25, 2.0]. En ne sélectionnant que les positions dont l’identifiant est non nul avant la moyenne, le résultat est environ [1.667, 2.667]. Le premier calcul divise par quatre, le second par trois.

Choisissez un masque ou un traitement des séquences adapté à la couche suivante et à la perte. La règle de padding de l’embedding ne garantit pas, à elle seule, que tout le modèle ignore ces positions.

Distinguer un gradient nul d’une ligne immuable

Cet exemple fournit volontairement un vecteur de padding non nul et active la pénalisation weight_decay de SGD :

import torch

weights = torch.tensor([[5., 5.], [1., 2.]])
layer = torch.nn.Embedding.from_pretrained(
    weights.clone(), freeze=False, padding_idx=0)
optimizer = torch.optim.SGD(layer.parameters(), lr=0.1,
                            momentum=0, weight_decay=0.1)
optimizer.zero_grad(set_to_none=True)
layer(torch.tensor([0, 1], dtype=torch.long)).sum().backward()
padding_gradient = layer.weight.grad[0].clone()
optimizer.step()
print('Padding gradient:', padding_gradient.tolist())
print('Padding weights after decay:',
      [round(v, 3) for v in layer.weight[0].tolist()])

Padding gradient: [0.0, 0.0]
Padding weights after decay: [4.95, 4.95]

Le gradient de lecture de la ligne 0 est nul, mais ses valeurs enregistrées passent de 5 à 4.95. Avec ces réglages, chaque composante diminue de 0.1 * 0.1 * 5. Cette modification provient de l’optimiseur, indépendamment du gradient apporté par la lecture.

La référence de SGD décrit son terme de weight decay. Ce test utilise des gradients denses, aucun momentum et une seule mise à jour. Il ne démontre pas le comportement de tous les optimiseurs ni de tous leurs états. Si une valeur exacte de padding est nécessaire, vérifiez-la après les mises à jour réellement effectuées par votre application.

Dans cet exemple, from_pretrained conserve la ligne de padding fournie. Une nouvelle couche nn.Embedding initialise sa ligne de padding à zéro, ce que nous avons vérifié séparément. La référence de la fonction embedding distingue également la lecture de l’initialisation.

Choisir explicitement de geler les poids

import torch

weights = torch.tensor([[1., 2.], [3., 4.]])
frozen = torch.nn.Embedding.from_pretrained(weights.clone())
trainable = torch.nn.Embedding.from_pretrained(weights.clone(), freeze=False)
ids = torch.tensor([1], dtype=torch.long)
frozen_output = frozen(ids)
trainable_output = trainable(ids)
trainable_output.sum().backward()
print('Frozen weight requires grad:', frozen.weight.requires_grad)
print('Frozen output requires grad:', frozen_output.requires_grad)
print('Trainable weight requires grad:', trainable.weight.requires_grad)
print('Trainable gradient:', trainable.weight.grad.tolist())

Frozen weight requires grad: False
Frozen output requires grad: False
Trainable weight requires grad: True
Trainable gradient: [[0.0, 0.0], [1.0, 1.0]]

from_pretrained gèle la matrice de poids par défaut. Dans cette lecture isolée, les poids gelés et la sortie ne nécessitent pas de gradients. Le réglage freeze=False les active, et seule la ligne sélectionnée en reçoit dans l’exemple.

Nous passons weights.clone() à chaque couche pour séparer leur stockage de celui des données de test. Le gel contrôle le suivi des gradients pour ce paramètre ; il ne transforme pas le tenseur en objet accessible uniquement en lecture. L’exemple suivant montre pourquoi cette distinction compte.

Lors de l’enregistrement d’un modèle, conservez la correspondance du vocabulaire et la configuration de l’embedding avec les poids. Notre guide de sauvegarde et de chargement PyTorch, en anglais, traite de l’état du modèle et des checkpoints d’entraînement. Même correctement restaurée, la matrice doit conserver le même sens pour chaque ligne.

Tenir compte des modifications provoquées par max_norm

import torch

weights = torch.tensor([[3., 4.], [0., 2.], [6., 8.]])
layer = torch.nn.Embedding.from_pretrained(
    weights.clone(), freeze=True, max_norm=1.0)
before = layer.weight.detach().clone()
output = layer(torch.tensor([0, 1], dtype=torch.long))
rounded = lambda rows: [[round(v, 3) for v in row] for row in rows]
print('Output:', rounded(output.tolist()))
print('Stored weights:', rounded(layer.weight.tolist()))
print('Frozen:', not layer.weight.requires_grad)
print('Stored weights changed:', not torch.equal(before, layer.weight))

Output: [[0.6, 0.8], [0.0, 1.0]]
Stored weights: [[0.6, 0.8], [0.0, 1.0], [6.0, 8.0]]
Frozen: True
Stored weights changed: True

Les deux vecteurs demandés ont initialement des normes de cinq et deux. La passe avant ramène leurs normes à environ un et enregistre les vecteurs ajustés dans la matrice. La ligne 2, dont le vecteur a une norme de dix, n’a pas été demandée et reste [6, 8].

Cela se produit avec freeze=True, sans appel à backward ni étape d’optimiseur. Si les vecteurs fournis doivent rester inchangés pendant la lecture, n’activez pas ce comportement sans tenir compte de cette modification.

La référence de l’API avertit que max_norm entraîne une modification sur place, avec des conséquences pour les opérations différentiables qui utilisent les poids avant la lecture. Suivez ces indications si le modèle partage ou transforme ces poids. Notre test vérifie les valeurs enregistrées dans une couche gelée ; il n’examine pas tous les graphes autograd à poids partagés.

Utiliser de petits exemples avant un entraînement complet

Vérifiez un identifiant connu, un identifiant répété, une position de padding et un indice invalide. Confirmez les dimensions de sortie et examinez les gradients d’une ligne utilisée, d’une ligne inutilisée et de la ligne de padding. Contrôlez ensuite les poids enregistrés après votre étape d’optimiseur et toute lecture avec max_norm.

Les exemples ont réussi sur CPU, y compris les tests d’erreur de type et de dépassement des bornes. Nous n’avons évalué ni les optimiseurs à gradients creux, ni l’exécution distribuée, ni les performances GPU, ni la similarité sémantique. Ces questions sont distinctes de la vérification des lignes lues et modifiées par la couche.