Différentiation automatique en mode inverse avec un graphe de calcul NetworkX

Ce notebook illustre comment :

Encoder des opérations primitives NumPy comme valeurs de noeuds dans un graphe orienté acyclique (DAG) NetworkX
Effectuer un tri topologique pour déterminer l’ordre d’évaluation (liste de Wengert)
Évaluer la fonction en parcourant le graphe dans cet ordre
Associer à chaque noeud une règle VJP (Vector-Jacobian Product) et créer des fermetures lors du parcours avant
Parcourir la liste fermée en ordre inverse pour effectuer la différentiation en mode inverse (rétropropagation)

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

Définition des primitives et de leurs règles VJP¶

Chaque primitive est définie par :

une fonction forward $f(x_1, \ldots, x_k) = z$
une règle VJP qui, étant donné le cotangent de la sortie $\bar{z}$ et les valeurs des entrées/sortie, retourne les cotangents des entrées $(\bar{x}_1, \ldots, \bar{x}_k)$

Rappel : si $z = f(x_1, \ldots, x_k)$ , la règle VJP donne

\bar{x}_i = \bar{z} \cdot \frac{\partial f}{\partial x_i}

(1)

Primitive	Forward	VJP
`add(x,y)`	$x + y$	$(\bar{z},\; \bar{z})$
`mul(x,y)`	$x \cdot y$	$(\bar{z} \cdot y,\; \bar{z} \cdot x)$
`sin(x)`	$\sin(x)$	$(\bar{z} \cdot \cos(x),)$
`exp(x)`	$e^x$	$(\bar{z} \cdot e^x,)$
`neg(x)`	$-x$	$(-\bar{z},)$

# Registre des opérations primitives
primitives = {}

def register(name, forward, vjp_rule):
    """Enregistre une primitive avec sa fonction forward et sa règle VJP.
    
    La signature de vjp_rule est:
        vjp_rule(z_bar, *inputs, z) -> tuple de cotangents des entrées
    où z_bar est le cotangent de la sortie, inputs sont les valeurs des
    entrées au moment de l'évaluation, et z la valeur de sortie.
    """
    primitives[name] = {"forward": forward, "vjp_rule": vjp_rule}

# z = x + y  =>  x̄ = z̄,  ȳ = z̄
register("add",
    forward=lambda x, y: x + y,
    vjp_rule=lambda z_bar, x, y, z: (z_bar, z_bar))

# z = x * y  =>  x̄ = z̄·y,  ȳ = z̄·x
register("mul",
    forward=lambda x, y: x * y,
    vjp_rule=lambda z_bar, x, y, z: (z_bar * y, z_bar * x))

# z = sin(x)  =>  x̄ = z̄·cos(x)
register("sin",
    forward=lambda x: np.sin(x),
    vjp_rule=lambda z_bar, x, z: (z_bar * np.cos(x),))

# z = exp(x)  =>  x̄ = z̄·exp(x) = z̄·z
register("exp",
    forward=lambda x: np.exp(x),
    vjp_rule=lambda z_bar, x, z: (z_bar * z,))

# z = -x  =>  x̄ = -z̄
register("neg",
    forward=lambda x: -x,
    vjp_rule=lambda z_bar, x, z: (-z_bar,))

Construction du graphe de calcul¶

On construit un graphe orienté acyclique (DAG) avec NetworkX. Chaque noeud possède :

op : le nom de la primitive (ou None pour les entrées)
parents : la liste ordonnée des noeuds parents (les opérandes)

Les arêtes vont des parents vers les enfants (dans le sens du flot de données).

Exemple : $f(x, y) = \sin(x \cdot y) + e^x$ ¶

La liste de Wengert correspondante est :

Variable	Opération
$v_0 = x$	entrée
$v_1 = y$	entrée
$v_2 = v_0 \cdot v_1$	`mul`
$v_3 = \sin(v_2)$	`sin`
$v_4 = e^{v_0}$	`exp`
$v_5 = v_3 + v_4$	`add`

G = nx.DiGraph()

# Noeuds d'entrée
G.add_node("x", op=None, parents=[])
G.add_node("y", op=None, parents=[])

# v2 = mul(x, y)
G.add_node("v2", op="mul", parents=["x", "y"])
G.add_edge("x", "v2")
G.add_edge("y", "v2")

# v3 = sin(v2)
G.add_node("v3", op="sin", parents=["v2"])
G.add_edge("v2", "v3")

# v4 = exp(x)
G.add_node("v4", op="exp", parents=["x"])
G.add_edge("x", "v4")

# v5 = add(v3, v4)  -- sortie finale
G.add_node("v5", op="add", parents=["v3", "v4"])
G.add_edge("v3", "v5")
G.add_edge("v4", "v5")

print("Noeuds :", list(G.nodes(data=True)))
print("Arêtes :", list(G.edges()))

Noeuds : [('x', {'op': None, 'parents': []}), ('y', {'op': None, 'parents': []}), ('v2', {'op': 'mul', 'parents': ['x', 'y']}), ('v3', {'op': 'sin', 'parents': ['v2']}), ('v4', {'op': 'exp', 'parents': ['x']}), ('v5', {'op': 'add', 'parents': ['v3', 'v4']})]
Arêtes : [('x', 'v2'), ('x', 'v4'), ('y', 'v2'), ('v2', 'v3'), ('v3', 'v5'), ('v4', 'v5')]

Visualisation du graphe¶

# Disposition en couches selon le tri topologique
pos = {
    "x":  (0, 1),
    "y":  (0, -1),
    "v2": (1, 0.3),
    "v3": (2, -1),
    "v4": (2, 1),
    "v5": (3, 0),
}

labels = {}
for n, d in G.nodes(data=True):
    if d["op"] is None:
        labels[n] = n
    else:
        args = ", ".join(d["parents"])
        labels[n] = f"{n} = {d['op']}({args})"

fig, ax = plt.subplots(figsize=(10, 4))
nx.draw(G, pos, ax=ax, with_labels=True, labels=labels,
        node_color="lightblue", node_size=3000,
        font_size=9, font_weight="bold",
        arrowsize=20, edge_color="gray")
ax.set_title(r"Graphe de calcul : $f(x,y) = \sin(x \cdot y) + e^x$", fontsize=13)
plt.tight_layout()
plt.show()

Passe avant : tri topologique, évaluation et création des fermetures VJP¶

On parcourt le graphe dans l’ordre topologique (la liste de Wengert). Pour chaque noeud intermédiaire :

On évalue la primitive en utilisant les valeurs des parents
On crée une fermeture (closure) qui capture :
- la règle VJP de la primitive
- les valeurs des entrées et de la sortie au moment de l’évaluation
- les noms des parents

Ces fermetures sont empilées dans une liste. Lors de la passe arrière, on les dépilera en ordre inverse pour propager les cotangents.

def forward(G, input_values):
    """Passe avant sur le graphe de calcul.
    
    Paramètres
    ----------
    G : nx.DiGraph
        Graphe de calcul.
    input_values : dict
        Valeurs des noeuds d'entrée, ex. {"x": 1.0, "y": 2.0}.
    
    Retourne
    --------
    values : dict
        Valeurs de tous les noeuds après évaluation.
    vjp_tape : list of (str, closure)
        Liste ordonnée de fermetures VJP pour la passe arrière.
        Chaque fermeture a la signature : closure(z_bar) -> (cotangents, parents)
    """
    order = list(nx.topological_sort(G))
    print(f"Ordre topologique (liste de Wengert) : {order}")
    
    values = dict(input_values)
    vjp_tape = []
    
    for node in order:
        data = G.nodes[node]
        op_name = data["op"]
        
        if op_name is None:
            # Noeud d'entrée, pas d'opération à évaluer
            continue
        
        prim = primitives[op_name]
        parents = data["parents"]
        parent_vals = [values[p] for p in parents]
        
        # 1. Évaluation forward
        z = prim["forward"](*parent_vals)
        values[node] = z
        print(f"  {node} = {op_name}({', '.join(parents)}) = {z:.6f}")
        
        # 2. Création de la fermeture VJP
        #    On capture les valeurs par défaut pour éviter les problèmes
        #    de portée dans les boucles Python.
        def make_vjp_closure(vjp_rule, captured_inputs, captured_output, captured_parents):
            def closure(z_bar):
                cotangents = vjp_rule(z_bar, *captured_inputs, captured_output)
                return cotangents, captured_parents
            return closure
        
        closure = make_vjp_closure(
            prim["vjp_rule"],
            list(parent_vals),
            z,
            list(parents)
        )
        vjp_tape.append((node, closure))
    
    return values, vjp_tape

# Évaluation pour x=0.5, y=3.0
x_val, y_val = 0.5, 3.0
values, vjp_tape = forward(G, {"x": x_val, "y": y_val})

output_node = "v5"
print(f"\nRésultat : f({x_val}, {y_val}) = {values[output_node]:.6f}")
print(f"Vérification NumPy : {np.sin(x_val * y_val) + np.exp(x_val):.6f}")

Ordre topologique (liste de Wengert) : ['x', 'y', 'v4', 'v2', 'v3', 'v5']
  v4 = exp(x) = 1.648721
  v2 = mul(x, y) = 1.500000
  v3 = sin(v2) = 0.997495
  v5 = add(v3, v4) = 2.646216

Résultat : f(0.5, 3.0) = 2.646216
Vérification NumPy : 2.646216

Passe arrière : différentiation en mode inverse¶

On parcourt la bande VJP (vjp_tape) en ordre inverse. Pour chaque fermeture :

On lit le cotangent $\bar{z}$ du noeud courant
On appelle la fermeture VJP pour obtenir les cotangents des parents
On accumule ces cotangents dans le dictionnaire (un parent peut apparaître dans plusieurs opérations)

On initialise le cotangent du noeud de sortie à $\bar{v}_5 = 1$ (car $\frac{\partial f}{\partial f} = 1$ ).

Déroulement attendu¶

Pour $f(x,y) = \sin(xy) + e^x$ :

\frac{\partial f}{\partial x} = y \cos(xy) + e^x, \qquad \frac{\partial f}{\partial y} = x \cos(xy)

(2)

def backward(G, vjp_tape, output_node):
    """Passe arrière (mode inverse) en utilisant les fermetures VJP.
    
    Paramètres
    ----------
    G : nx.DiGraph
        Graphe de calcul.
    vjp_tape : list of (str, closure)
        Fermetures VJP créées lors de la passe avant.
    output_node : str
        Nom du noeud de sortie.
    
    Retourne
    --------
    cotangents : dict
        Cotangent (gradient) de chaque noeud.
    """
    # Initialisation : tous les cotangents à zéro
    cotangents = {n: 0.0 for n in G.nodes}
    # Le cotangent de la sortie est 1 (on dérive f par rapport à f)
    cotangents[output_node] = 1.0
    
    print("Passe arrière (ordre inverse de la liste de Wengert) :")
    
    # Parcours inverse de la bande
    for node, closure in reversed(vjp_tape):
        z_bar = cotangents[node]
        input_bars, parents = closure(z_bar)
        
        print(f"  {node} : z̄={z_bar:.6f} → ", end="")
        parts = []
        for parent, x_bar in zip(parents, input_bars):
            cotangents[parent] += x_bar
            parts.append(f"{parent}̄ += {x_bar:.6f}")
        print(", ".join(parts))
    
    return cotangents

cotangents = backward(G, vjp_tape, output_node)

print(f"\n--- Gradients calculés par mode inverse ---")
print(f"∂f/∂x = {cotangents['x']:.6f}")
print(f"∂f/∂y = {cotangents['y']:.6f}")

# Gradients analytiques
df_dx_exact = y_val * np.cos(x_val * y_val) + np.exp(x_val)
df_dy_exact = x_val * np.cos(x_val * y_val)
print(f"\n--- Gradients analytiques ---")
print(f"∂f/∂x = y·cos(xy) + exp(x) = {df_dx_exact:.6f}")
print(f"∂f/∂y = x·cos(xy) = {df_dy_exact:.6f}")

Passe arrière (ordre inverse de la liste de Wengert) :
  v5 : z̄=1.000000 → v3̄ += 1.000000, v4̄ += 1.000000
  v3 : z̄=1.000000 → v2̄ += 0.070737
  v2 : z̄=0.070737 → x̄ += 0.212212, ȳ += 0.035369
  v4 : z̄=1.000000 → x̄ += 1.648721

--- Gradients calculés par mode inverse ---
∂f/∂x = 1.860933
∂f/∂y = 0.035369

--- Gradients analytiques ---
∂f/∂x = y·cos(xy) + exp(x) = 1.860933
∂f/∂y = x·cos(xy) = 0.035369

Vérification par différences finies¶

Pour confirmer nos résultats, on compare avec une approximation numérique par différences finies centrées :

\frac{\partial f}{\partial x} \approx \frac{f(x+\epsilon, y) - f(x-\epsilon, y)}{2\epsilon}

(3)

def f_numpy(x, y):
    """La même fonction implémentée directement en NumPy."""
    return np.sin(x * y) + np.exp(x)

eps = 1e-7
df_dx_num = (f_numpy(x_val + eps, y_val) - f_numpy(x_val - eps, y_val)) / (2 * eps)
df_dy_num = (f_numpy(x_val, y_val + eps) - f_numpy(x_val, y_val - eps)) / (2 * eps)

print("--- Comparaison des gradients ---")
print(f"{'Méthode':<25} {'∂f/∂x':>12} {'∂f/∂y':>12}")
print(f"{'-'*25} {'-'*12} {'-'*12}")
print(f"{'Mode inverse (graphe)':<25} {cotangents['x']:>12.6f} {cotangents['y']:>12.6f}")
print(f"{'Analytique':<25} {df_dx_exact:>12.6f} {df_dy_exact:>12.6f}")
print(f"{'Différences finies':<25} {df_dx_num:>12.6f} {df_dy_num:>12.6f}")

# Vérification
assert np.allclose(cotangents["x"], df_dx_exact), "Erreur sur ∂f/∂x"
assert np.allclose(cotangents["y"], df_dy_exact), "Erreur sur ∂f/∂y"
print("\nTous les gradients correspondent.")

--- Comparaison des gradients ---
Méthode                          ∂f/∂x        ∂f/∂y
------------------------- ------------ ------------
Mode inverse (graphe)         1.860933     0.035369
Analytique                    1.860933     0.035369
Différences finies            1.860933     0.035369

Tous les gradients correspondent.

Deuxième exemple : $g(x, y) = -(\sin(x) \cdot \sin(y))$ ¶

Pour illustrer la généralité de l’approche, construisons un second graphe.

\frac{\partial g}{\partial x} = -\cos(x) \cdot \sin(y), \qquad \frac{\partial g}{\partial y} = -\sin(x) \cdot \cos(y)

(4)

G2 = nx.DiGraph()

G2.add_node("x", op=None, parents=[])
G2.add_node("y", op=None, parents=[])

# w1 = sin(x)
G2.add_node("w1", op="sin", parents=["x"])
G2.add_edge("x", "w1")

# w2 = sin(y)
G2.add_node("w2", op="sin", parents=["y"])
G2.add_edge("y", "w2")

# w3 = mul(w1, w2)
G2.add_node("w3", op="mul", parents=["w1", "w2"])
G2.add_edge("w1", "w3")
G2.add_edge("w2", "w3")

# w4 = neg(w3)
G2.add_node("w4", op="neg", parents=["w3"])
G2.add_edge("w3", "w4")

# Passe avant
x2, y2 = 1.0, 2.0
vals2, tape2 = forward(G2, {"x": x2, "y": y2})
print(f"\ng({x2}, {y2}) = {vals2['w4']:.6f}")

# Passe arrière
print()
cot2 = backward(G2, tape2, "w4")

# Vérification
dg_dx_exact = -np.cos(x2) * np.sin(y2)
dg_dy_exact = -np.sin(x2) * np.cos(y2)

print(f"\n--- Gradients ---")
print(f"{'Méthode':<25} {'∂g/∂x':>12} {'∂g/∂y':>12}")
print(f"{'-'*25} {'-'*12} {'-'*12}")
print(f"{'Mode inverse':<25} {cot2['x']:>12.6f} {cot2['y']:>12.6f}")
print(f"{'Analytique':<25} {dg_dx_exact:>12.6f} {dg_dy_exact:>12.6f}")

assert np.allclose(cot2["x"], dg_dx_exact)
assert np.allclose(cot2["y"], dg_dy_exact)
print("\nTous les gradients correspondent.")

Ordre topologique (liste de Wengert) : ['x', 'y', 'w1', 'w2', 'w3', 'w4']
  w1 = sin(x) = 0.841471
  w2 = sin(y) = 0.909297
  w3 = mul(w1, w2) = 0.765147
  w4 = neg(w3) = -0.765147

g(1.0, 2.0) = -0.765147

Passe arrière (ordre inverse de la liste de Wengert) :
  w4 : z̄=1.000000 → w3̄ += -1.000000
  w3 : z̄=-1.000000 → w1̄ += -0.909297, w2̄ += -0.841471
  w2 : z̄=-0.841471 → ȳ += 0.350175
  w1 : z̄=-0.909297 → x̄ += -0.491295

--- Gradients ---
Méthode                          ∂g/∂x        ∂g/∂y
------------------------- ------------ ------------
Mode inverse                 -0.491295     0.350175
Analytique                   -0.491295     0.350175

Tous les gradients correspondent.

Résumé¶

Ce notebook a montré les étapes clés de la différentiation automatique en mode inverse :

Graphe de calcul : chaque noeud encode une opération primitive et ses opérandes
Tri topologique : détermine l’ordre d’évaluation (liste de Wengert)
Passe avant : évalue le graphe et crée une fermeture VJP par noeud, capturant les valeurs intermédiaires
Passe arrière : parcourt les fermetures en ordre inverse, propage les cotangents par accumulation

Le coût de la passe arrière est proportionnel au coût de la passe avant (à un facteur constant près), quel que soit le nombre de variables d’entrée. C’est l’avantage fondamental du mode inverse.