Modo reverse en diferenciación automática

El modo reverse, en contraste con el modo forward, obtiene las derivadas de una función compleja a partir de recorrer un grafo computacional en retroceso; i.e., desde las salidas de la función hacia las entradas. De esta forma, se calcula la derivada de las funciones de salida y se usan esas derivadas para las funciones en los nodos anteriores en el grafo usando la regla de la cadena para obtener las nuevas derivadas.

Así, si $f_i$ es la función en un nodo del cual queremos obtener su derivada y si $f_{i+1}$ es la función de un nodo superiro, entonces la derivada en el modo reverse se obtiene como:

$$\frac{\partial f_{i+1}}{\partial x} = \frac{\partial f_{i+1}}{\partial f_i} \frac{\partial f_i}{x}$$

De esta forma, la derivación se propaga hacia atrás, generando una traza derivativa que para obtener las derivadas parciales sobre las entradas de la función. En este caso, el modo reverse obtiene el vector gradiente $\nabla_x f_i$ para una de las salidas de la función; es decir, cada vez que corremos el modo reverse para obtener la diferenciación, obtenemos las derivadas parciales sobre todas las variables de entrada, pero sólo para una de las salidas.

El modo reverse almacena las derivadas de las funciones de los nodos del grafo computacional a partir de variables específicas. En este caso, la variable de salida se le asigna un valor de 1, si es esa la salida sobre la que queremos derivar; de otra forma, le asignamos el valor 0 a las salidas que no nos interesen en ese momento. De igual forma, se puede iniciar con un vector distinto a los vectores base; en este caso, interpretamos el resultado del modo reverse como el producto de la matriz Jacobiana $J_f$ transpuesta por el vector.

Implementación del modo reverse

class Operation:
    def __call__(self, x):
        ...
        return #operación sobre x

    def backward(self, consumer_grad):
        # 1. Calcula el gradiente del nodo actual utilizando el gradiente del consumer
        # 2. Envía el gradiente actual hacia atrás al nodo de entrada

El gradiente se obtendrá en las entradas, al inicio de la gráfica computacional en los valores x.grad.

Ejemplo

Para ejemplificar el uso del modo forward, generamos nodos para las funciones de coseno, potencia y suma. En cada caso, definimos la derivada en el nodo con base en el nodo padre o nodo de entrada.

import numpy as np

class Cos:
    def __init__(self, n=2):
        self.n = n
        self.grad = None

    def __call__(self, x):
        self.value = np.cos(self.n*x.value)
        self.parent = x
        return self

    def __str__(self):
        return str(self.value)

    def backward(self, grad=1):
        self.grad = -self.n*np.sin(self.n*self.parent.value)*grad
        self.parent.backward(grad=self.grad)

class Pow:
    def __init__(self, p=2):
        self.pow = p
        self.grad = None

    def __call__(self, x):
        self.value = x.value**self.pow
        self.parent = x
        return self

    def __str__(self):
        return str(self.value)

    def backward(self, grad=1):
        self.grad = self.pow*(self.parent.value**(self.pow-1)) * grad
        self.parent.backward(grad=self.grad)

class Sum:
    def __init__(self):
        self.grad = None

    def __call__(self, x, y):
        self.value = x.value + y.value
        self.parentA = x
        self.parentB = y
        return self

    def __str__(self):
        return str(self.value)

    def backward(self, grad=1):
        self.grad = grad
        self.parentA.backward(grad=self.grad)
        self.parentB.backward(grad=self.grad)

class Node:
    def __init__(self, value):
        self.value = value
        self.grad = None

    def __str__(self):
        return str(self.value)
        
    def backward(self, grad=1):
        self.grad = grad

Para ejemplificar el modo reverse podemos usar un sólo nodo para la operación coseno:

$$f(x) = cos(2x)$$

#Nodo de entrada x
x = Node(3)
#Nodo coseno
cos = Cos(n=2)

#Construcción de la gráfica
result = cos(x)
print(result)

0.960170286650366

Ahor podemos obtener el gradiente de la función con el método de reverse aplicando el backward. Como se puede observar, el gradiente se guarda en el nodo de entrada $x$, no el de salida como en el modo forward.

result.backward(grad=1)
print('Gradiente {}'.format(x.grad))

Gradiente 0.5588309963978517

Ahora podemos explorar una función más complejada dada por la fórmula:

$$f(x,y) = \cos^2(x) + y^2$$

#Nodo entrada y
y = Node(2)
#Nodo de potencia
pow = Pow(p=2)
#Nodo de potencia
pow2 = Pow(p=2)
#Nodo de suma
sum = Sum()

#Construcción de la gráfica dinámica
result = sum(pow(cos(x)), pow2(y))
print(result)

4.921926979366246

Ahora podemos aplicar el método backward para obtener el gradiente con modo reverse. En este caso, el flujo del gradientes es de la salida a la entrada. Como se puede observar, a diferencia del modo forward, el gradiente final se almacena en los nodos de entrada $x$ e $y$. Es decir, aquí obtenemos el vector gradiente de la función. La suma de las derivadas parciales se tiene que realizar aparte.

result.backward(grad=1)
print('Gradiente [{} {}]'.format(x.grad, y.grad))
print('Suma de derivadas: {}'.format(x.grad + y.grad))

Gradiente [1.0731458360008699 4]
Suma de derivadas: 5.07314583600087