Différentiation automatique

PyTorch Autograd Explained - In-depth Tutorial - YouTube 2.5. Automatic Differentiation — Dive into Deep Learning documentation

Différentiation automatique avec JAX

JAX est une bibliothèque de calcul numérique et scientifique haute performance qui exploite le matériel moderne, notamment les GPU et les TPU. Elle étend NumPy avec la différentiation automatique et calcule efficacement les gradients par des transformations de fonctions. Ces gradients jouent un rôle central dans l’optimisation et les algorithmes d’apprentissage automatique.

Différentiation automatique (autodiff)

La différentiation automatique évalue les dérivées en composant les règles des opérations élémentaires du programme par la règle de la chaîne. Contrairement au calcul symbolique, elle n'a pas besoin de construire une formule simplifiée ; contrairement aux différences finies, elle n'estime pas la dérivée à partir de valeurs voisines. Elle dérive le calcul implémenté, avec les erreurs d'arrondi et les règles disponibles. Une discontinuité, une décision portant sur des entiers, une opération non prise en charge ou une dérivée inexistante ne deviennent pas dérivables grâce à l'autodiff. Les dérivées supérieures exigent une différentiabilité suffisante et peuvent coûter cher. Revoyez les dérivées et les dimensions des applications linéaires avant les produits jacobiens.

Autodiff en mode avant et inverse

Pour $f:\mathbb R^n\to\mathbb R^m$, le jacobien $J$ a pour coefficients $J_{ij}=\partial f_i/\partial x_j$ et pour taille $m\times n$. Le mode avant propage une direction d'entrée $v\in\mathbb R^n$ vers $Jv\in\mathbb R^m$. Le mode inverse propage des poids de sortie $u\in\mathbb R^m$ vers $J^Tu\in\mathbb R^n$. Initialiser avec des vecteurs de base permet d’obtenir des colonnes ou des lignes individuelles ; aucun des deux modes n’impose une initialisation limitée à une seule coordonnée.

La transposée vient de la règle de chaîne multivariée, pas d’une convention propre à JAX. Soit $z=f(x)$ et une perte scalaire $L(z)$. Chaque entrée peut agir sur la perte par chaque sortie, d’où

$$\frac{\partial L}{\partial x_j}=\sum_i\frac{\partial L}{\partial z_i}\frac{\partial z_i}{\partial x_j}, \qquad \nabla_xL=J^T\nabla_zL.$$

Le mode inverse prend donc $u=\nabla_zL$. La somme explique aussi pourquoi les contributions de chemins distincts s’accumulent : la règle de chaîne scalaire s’applique à toutes les coordonnées intermédiaires.

Un jacobien complet demande donc environ $n$ parcours avant ou $m$ parcours inverses. Le mode avant convient à peu d'entrées et beaucoup de sorties ; le mode inverse à beaucoup d'entrées et peu de sorties, surtout une perte scalaire. Il fonde la rétropropagation, mais doit conserver ou recalculer les intermédiaires pour le parcours arrière. Ce sont des compromis entre travail et mémoire, pas un classement universel de vitesse.

Graphe de différentiation

Un graphe de différentiation est une représentation graphique du calcul d’une fonction. Il consigne toutes les opérations, de l’entrée à la sortie, ainsi que les variables intermédiaires et leurs dépendances.

La différentiation automatique utilise ce graphe pour appliquer la règle de la chaîne efficacement et précisément :

• En mode avant, les dérivées se propagent dans le sens du graphe.
• En mode inverse, les gradients s’accumulent en sens opposé.

Voir l’image en grand

Cet exemple PyTorch calcule Q = 3a³ − b² élément par élément. Suivez le graphe vers le bas pour le calcul, vers le haut pour la propagation des gradients. PowBackward0 désigne une fonction de rétropropagation enregistrée ; « (2) » indique les dimensions du tenseur d’entrée. Le code JAX ci-dessous traite une autre fonction avec le même principe de composition des dérivées.

PyTorch contributors, A Gentle Introduction to torch.autograd. · BSD 3-Clause

Comment JAX implémente l’autodiff

JAX prend en charge les modes avant et inverse pour calculer efficacement les gradients, les jacobiens et les hessiens. Ses principales interfaces sont :

• jax.grad :

  • Calcule le gradient d’une fonction.
  • Utilise le mode inverse.
  • Convient aux cas où la dimension d’entrée dépasse largement celle de sortie, comme la perte scalaire d’un réseau de neurones.

• jax.jvp (produit jacobien-vecteur) :

  • Calcule un produit jacobien-vecteur.
  • Constitue l’opération élémentaire du mode avant.
  • Calcule $J \cdot v$, où $J$ est le jacobien et $v$ un vecteur.

• jax.vjp (produit vecteur-jacobien) :

  • Calcule un produit vecteur-jacobien.
  • Constitue l’opération élémentaire du mode inverse.
  • Calcule $v^T \cdot J$, où $v$ est un vecteur et $J$ le jacobien.

Exemple d’utilisation

Voici comment calculer avec JAX le gradient d’une fonction :

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x) * jnp.cos(x)

# Obtenir la fonction gradient de f
grad_f = jax.grad(f)

# Calculer le gradient de f en x = 1.0
print(grad_f(1.0))

Ce code calcule la dérivée de $f(x) = \sin(x) \cos(x)$ au point $x = 1.0$.

Suivre une entrée partagée

Posons $a=xy$, $b=\sin x$ et $f=a+b$. En $(x,y)=(2,3)$, l’initialisation en mode avant $(\dot x,\dot y)=(1,0)$ donne $\dot a=y\dot x+x\dot y=3$ et $\dot b=\cos2$, donc $\dot f=3+\cos2\approx2.583853$. C'est $J(1,0)^T$.

En mode inverse, on initialise $\bar f=1$. L'addition transmet $\bar a=\bar b=1$. La multiplication apporte $y\bar a=3$ à $\bar x$ et $x\bar a=2$ à $\bar y$ ; le sinus apporte $\cos2\,\bar b$ à $\bar x$. Il faut additionner les contributions à une même entrée : $\nabla f=(3+\cos2,2)$. Les deux chemins passant par $x$ expliquent pourquoi écraser son gradient serait incorrect.

Dans l'exemple JAX précédent, la dérivée attendue est $\cos^2(1)-\sin^2(1)=\cos2\approx-0.416147$. jax.grad exige une sortie scalaire (forme (), pas un vecteur de longueur un) et dérive habituellement des entrées flottantes. Pour des sorties vectorielles, choisissez jax.jvp, jax.vjp, jax.jacfwd ou jax.jacrev selon le produit ou la matrice recherchés. En un point anguleux, comme zéro pour $|x|$, la règle choisie par une implémentation ne prouve pas l'existence d'une dérivée classique.

Références et liens utiles

• The Autodiff Cookbook — JAX documentation
• Automatic Differentiation Videolectures

Mes lectures · Les notes ouvertes sont mémorisées dans ce navigateur.