La rétropropagation

Sixième article de la série Smallem. Précédent : Les mêmes opérations sur GPU.

On sait calculer, désormais il faut apprendre. Cet article construit le mécanisme qui rend l'apprentissage possible : le calcul automatique des gradients, ou autograd.

Ce qu'on cherche

Un réseau de neurones contient des millions de nombres réglables, les paramètres. Pour un lot de texte donné, le réseau produit une perte : un seul nombre, qui dit à quel point il s'est trompé.

Apprendre, c'est modifier les paramètres pour faire baisser la perte. Encore faut-il savoir dans quel sens modifier chacun d'eux. La réponse est le gradient :

• Gradient positif : augmenter le paramètre augmente la perte. Il faut le diminuer.
• Gradient négatif : il faut augmenter le paramètre.
• Gradient proche de zéro : le paramètre n'a presque pas d'influence.

Un pas d'apprentissage consiste donc à déplacer chaque paramètre un peu dans le sens opposé à son gradient. Tout le problème est de calculer ces millions de gradients suffisamment rapidement.

La règle de la chaîne, sans formule

Prenons un calcul minuscule :

c = a × b
perte = c + a

On veut savoir comment la perte réagit à a. Le paramètre a agit sur la perte par deux chemins : directement, par l'addition, et indirectement, à travers c.

Raisonnons à l'envers, en partant de la fin.

• La perte dépend d'elle-même avec un facteur 1. C'est le point de départ.
• perte = c + a : si c augmente de 1, la perte augmente de 1. Même chose pour a. Une addition transmet le gradient tel quel à ses deux opérandes.
• c = a × b : si a augmente de 1, c augmente de b. Une multiplication transmet le gradient à chaque opérande, multiplié par l'autre.

En remontant, a reçoit donc 1 par le chemin direct, et 1 × b par le chemin de c. Son gradient est 1 + b.

C'est toute la rétropropagation. Elle repose sur trois idées :

1. on part de la perte, avec un gradient de 1, et on remonte le calcul à l'envers
2. chaque opération sait comment transmettre le gradient de sa sortie à ses entrées : c'est sa règle locale
3. quand une valeur sert à plusieurs endroits, elle reçoit la somme des gradients venus de chacun.

Ce procédé a une propriété remarquable : Une seule remontée donne le gradient de tous les paramètres à la fois, pour un coût comparable à celui du calcul lui-même. C'est ce qui rend possible l'entraînement de réseaux à des millions de paramètres.

La tape

Pour remonter un calcul, il faut l'avoir noté. C'est le rôle de la tape : un enregistrement de chaque opération, au moment où elle est faite.

pub struct Tape<'b, B: Backend> {
    id: u64,
    backend: &'b B,
    nodes: RefCell<Vec<Node<B>>>,
}

struct Node<B: Backend> {
    value: B::Tensor,
    op: Op<B>,
    requires_grad: bool,
}

Chaque nœud garde la valeur calculée et l'opération qui l'a produite, avec ses opérandes :

enum Op<B: Backend> {
    Leaf,
    Constant,
    Add(Var, Var),
    Mul(Var, Var),
    MatMul(Var, Var),
    Transpose(Var, usize, usize),
    // … les autres primitives
    Operation {
        rule: Box<dyn Rule<B>>,
        inputs: Vec<Var>,
    },
}

L'énumération suit les deux étages de l'article 4. Les primitives y figurent une par une : il y en a peu, et la liste ne bouge pas. Toutes les opérations partagent la dernière variante, sur laquelle on revient plus bas.

Une Var "pointe" vers un noeud. Elle ne contient qu'un entier, le rang du noeud dans la liste :

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Var(usize);

Ce choix évite un problème bien connu en Rust. Un graphe de calcul est un enchevêtrement de références : chaque nœud pointe vers ses parents, et plusieurs noeuds partagent le même parent. Le représenter avec des références se heurte de front au borrow checker. Avec des indices dans un vecteur, il n'y a plus de références du tout : une Var se copie librement.

Deux sortes de noeuds n'ont pas d'opérandes. Une feuille (Leaf) est une valeur dont on veut le gradient : un paramètre. Une constante est une valeur fixe, qui n'en reçoit pas.

L'ordre est gratuit

Pour remonter, il faut traiter chaque noeud après tous ceux qui l'utilisent. Dans un graphe quelconque, il faudrait le trier.

Ici, c'est inutile. Les noeuds sont ajoutés dans l'ordre où ils sont calculés, donc un noeud vient toujours après ses opérandes. La liste est déjà dans le bon ordre : il suffit de la parcourir de la fin vers le début.

Utiliser la tape

Chaque opération de la tape fait deux choses : elle calcule la valeur, en appelant le backend, et elle note ce qu'elle a fait.

pub fn mul(&self, a: Var, b: Var) -> Var {
    self.binary(a, b, Op::Mul(a, b), |backend, x, y| backend.mul(x, y))
}

Le calcul de tout à l'heure s'écrit donc :

let tape = Tape::new(&backend);
let a = tape.leaf(valeur_de_a);
let b = tape.leaf(valeur_de_b);
let perte = tape.add(tape.mul(a, b), a);
let grads = tape.backward(perte);
let gradient_de_a = grads.get(a);

Les méthodes prennent &self et non &mut self, grâce au RefCell. Sans cela, on ne pourrait pas imbriquer les appels comme dans tape.add(tape.mul(a, b), a).

La tape ne connaît les tenseurs qu'à travers le trait Backend. Elle est donc écrite une seule fois, et calcule les gradients sur CPU comme sur GPU.

Remonter

Le coeur de backward est une boucle à l'envers, avec un grand match : une règle locale par opération.

for index in (0..=loss.0).rev() {
    let Some(grad) = grads[index].take() else {
        continue;
    };
    let node = &nodes[index];
    let out = &node.value;

    match &node.op {
        // c = a · b  →  da = dc · b, db = dc · a
        Op::Mul(a, b) => {
            let da = backend.mul(&grad, value(*b));
            let db = backend.mul(&grad, value(*a));
            accumulate(&mut grads, *a, unbroadcast(backend, da, shape(*a)));
            accumulate(&mut grads, *b, unbroadcast(backend, db, shape(*b)));
        }
        // c = exp(a)  →  da = dc · c
        Op::Exp(a) => accumulate(&mut grads, *a, backend.mul(&grad, out)),
        // …
    }
}

grad est le gradient de la sortie du noeud, déjà connu puisqu'on vient de la fin. La règle en déduit le gradient de chaque opérande.

Voici les règles principales. Chacune se retrouve en se demandant : si l'entrée bouge un peu, de combien bouge la sortie ?

Plusieurs règles réutilisent la sortie c : la dérivée de l'exponentielle est l'exponentielle elle-même, déjà calculée. C'est pour cela que la tape garde les valeurs.

Deux pièges guettent quiconque écrit cette boucle.

Piège 1 : accumuler, jamais remplacer

Dans l'exemple du début, a recevait deux contributions. La fonction accumulate les additionne :

grads[var.0] = Some(match grads[var.0].take() {
    Some(existing) => backend.add(&existing, &grad),
    None => grad,
});

S'il y a déjà un gradient, on lui ajoute la nouvelle contribution. Écrire simplement grads[var.0] = Some(grad), qui remplacerait l'ancien gradient par le nouveau, donnerait un programme qui tourne, qui produit des nombres, mais dont certains gradients sont faux : seule la dernière contribution resterait.

Piège 2 : défaire le broadcasting

Quand on ajoute un vecteur de biais [3] à une matrice [4, 3], le biais est étiré : ses trois valeurs servent chacune quatre fois. Le gradient qui redescend a la forme [4, 3], mais celui du biais doit avoir la forme [3].

Étirer une valeur, c'est l'utiliser plusieurs fois. D'après la troisième idée, son gradient est donc la somme des gradients de toutes ses copies. La fonction unbroadcast fait cette somme le long des dimensions étirées :

fn unbroadcast<B: Backend>(backend: &B, grad: B::Tensor, target: &[usize]) -> B::Tensor {
    let mut grad = grad;
    // Dimensions ajoutées à gauche.
    while backend.shape(&grad).len() > target.len() {
        grad = backend.sum(&grad, 0, false);
    }
    // Dimensions de taille 1 qui ont été étirées.
    for (dim, &size) in target.iter().enumerate() {
        if size == 1 && backend.shape(&grad)[dim] != 1 {
            grad = backend.sum(&grad, dim, true);
        }
    }
    grad
}

La somme sur un axe et l'étirement sont donc l'inverse l'un de l'autre, dans les deux sens : le gradient d'un étirement est une somme, et le gradient d'une somme est un étirement.

Vérifier : les différences finies

Chaque règle locale est une petite démonstration, et il est très facile de s'y tromper d'un signe ou d'une transposée. Un gradient faux ne provoque aucune erreur : le modèle apprend simplement mal, ou pas du tout. Il faut donc un moyen indépendant de contrôler.

Ce moyen revient à la définition : le gradient dit de combien la perte change quand un paramètre bouge un peu. Alors bougeons-le, et regardons :

gradient ≈ (perte(x + ε) - perte(x - ε)) / 2ε

On calcule la perte deux fois, avec le paramètre légèrement augmenté puis légèrement diminué, et on divise l'écart par la taille du déplacement. C'est la méthode des différences finies. Elle est très lente : deux calculs complets par paramètre, mais elle n'utilise aucune des règles qu'on cherche à vérifier.

let numeric = (perturbed(EPSILON) - perturbed(-EPSILON)) / (2.0 * EPSILON);
let analytic = analytic.data()[element];
let scale = analytic.abs().max(numeric.abs()).max(1.0);
assert!((analytic - numeric).abs() <= TOLERANCE * scale);

C'est ici que le f64 devient indispensable. Avec ε = 1e-5, la différence entre les deux pertes est minuscule. En f32, qui ne garde que sept chiffres, elle se noierait dans les erreurs d'arrondi. En f64, on peut exiger une erreur relative inférieure à 1e-6. Voilà pourquoi le backend CPU est générique sur le type flottant : c'était pour ce test.

Un raffinement rend le test plus sévère. Plutôt que de sommer simplement la sortie pour obtenir un scalaire, on la pondère par des poids aléatoires tous différents. Chaque élément de la sortie reçoit alors un gradient différent, et une règle qui confondrait deux éléments serait repérée.

Chaque opération de la tape a son test, et le projet compte une cinquantaine de tests de gradient.

Le gradient d'une opération

Le match ci-dessus couvre les primitives. Les opérations de l'article 4, elles, portent leur propre règle. C'est la seconde méthode du trait Run, laissée de côté jusqu'ici :

fn run_backward(&self, op: &O, inputs: &[&Self::Tensor],
                output: &Self::Tensor, grad: &Self::Tensor) -> Vec<Option<Self::Tensor>>;

À partir du gradient de la sortie, elle rend le gradient de chaque entrée, ou rien pour une entrée qui n'en reçoit pas. Pour une opération définie par sa référence, la règle s'écrit elle aussi avec les primitives, dans le même fichier, juste sous la définition :

// dx = s × (g - moyenne(g) - x̂ × moyenne(g × x̂))
let mean_grad = backend.mean(grad, last, true);
let mean_projection = backend.mean(&backend.mul(grad, output), last, true);

Sur la tape, toutes les opérations passent par une seule fonction :

pub fn apply<O: Operation>(&self, op: O, inputs: &[Var]) -> Var
where
    B: Run<O>,
{
    // … calcule `backend.run(&op, valeurs)`, puis range l'opération elle-même
    let op = Op::Operation { rule: Box::new(op), inputs: inputs.to_vec() };
    self.push(value, op, requires_grad)
}

Et en remontant, une seule branche du match les traite toutes :

Op::Operation { rule, inputs } => {
    let values: Vec<&B::Tensor> = inputs.iter().map(|var| value(*var)).collect();
    let input_grads = rule.backward(backend, &values, out, &grad);
    for (var, input_grad) in inputs.iter().zip(input_grads) {
        if let Some(input_grad) = input_grad {
            accumulate(&mut grads, *var, input_grad);
        }
    }
}

La tape ne sait donc rien des opérations. Elle sait seulement qu'un noeud a été produit par « quelque chose qui connaît sa règle ». Ajouter une opération au projet ne change pas une ligne de l'autograd.

Deux cas qui demandent un kernel de plus

La rétropropagation des primitives n'utilise que d'autres primitives. Elle fonctionne donc sur GPU presque sans rien ajouter, à deux exception près :

Le scatter-add

Un embedding est une table qui associe un vecteur à chaque token. L'opération directe, gather_rows, lit des lignes de la table. Si le token 42 apparaît trois fois dans le lot, la ligne 42 est lue trois fois.

En remontant, la ligne 42 doit donc recevoir la somme de trois gradients. L'opération inverse, le scatter-add, additionne chaque ligne du gradient dans la ligne de la table désignée par son indice.

Le schéma « un thread par élément de sortie » n'est pas respecté. Trois threads veulent écrire dans la même case. Une addition ordinaire se fait en trois temps, lire, ajouter, écrire, et deux threads peuvent s'entrecroiser :

thread A lit 5,0
thread B lit 5,0
thread A écrit 5,0 + 1,0 = 6,0
thread B écrit 5,0 + 2,0 = 7,0     la contribution de A est perdue

Le DisjointSlice ne convient plus, puisqu'il garantit justement qu'aucune case n'est partagée. Il faut une addition atomique, qui fait les trois temps d'un bloc :

let cell = unsafe { DeviceAtomicF32::from_ptr(out.add(target)) };
cell.fetch_add(src[flat], const { AtomicOrdering::Relaxed });

Le test correspondant utilise des indices répétés exprès, sans quoi il ne prouverait rien.

Le const { … } autour de l'ordre est obligatoire. Le compilateur de kernels exige que ce soit une constante au point d'appel, et sans ce bloc la compilation échoue en mode non optimisé, celui des tests.

Softmax et entropie croisée, fusionnés

La perte d'un modèle de langage est un softmax suivi d'une entropie croisée. On pourrait les enchaîner, et laisser la tape remonter les deux règles l'une après l'autre.

Il vaut mieux les fusionner en une seule opération, CrossEntropy, pour deux raisons. La première est numérique : le logarithme d'une probabilité minuscule est instable. La seconde est que le gradient combiné des deux est simple :

dlogits = (softmax(logits) - cible) / N

où cible vaut 1 pour le bon token et 0 ailleurs. Le gradient est l'écart entre ce que le modèle a prédit et ce qu'il aurait dû prédire. Si le modèle donnait une probabilité de 0,7 au bon token, le gradient de ce logit est -0,3 : il faut l'augmenter un peu. S'il donnait 0,1 à un mauvais token, le gradient vaut 0,1 : il faut le baisser.

Vérifier sur GPU

Les différences finies valident les règles, sur CPU. Il reste à vérifier que le GPU les applique correctement. On compare donc les gradients des deux backends sur le même calcul, en f32. Cela valide les kernels utilisés en remontant, et en particulier le scatter-add.

Les deux vérifications se complètent : la première prouve que les formules sont justes, la seconde que les kernels les exécutent fidèlement.

Ce qu'il faut retenir

• Le gradient d'un paramètre dit dans quel sens le déplacer pour faire baisser la perte.
• La rétropropagation remonte le calcul à l'envers, applique à chaque opération une règle locale, et additionne les contributions.
• La tape enregistre le calcul dans une liste ; les variables sont de simples indices, et la liste est déjà dans l'ordre.
• Deux pièges : oublier d'additionner, oublier de défaire le broadcasting.
• Les différences finies, en f64, vérifient chaque règle sans en utiliser aucune.

On a maintenant de quoi calculer et de quoi apprendre. L'article suivant assemble les deux pour entraîner un premier réseau : Un premier réseau qui apprend.