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 de un ejemplo en modo forward

Para ejemplificar el modo forward definiremos un grafo computacional que compute la función $b:\mathbb{R}^3 \to \mathbb{R}$ dada como:

$$b = x(y-z)^2$$

En el grafo, definiremos tres nodos que nos servirán para computar la función final. Estos tres nodos son en orden los siguientes:

$$a = y-z$$$$c = a^2$$$$b = xc$$

Cada uno de estos nodos contiene una función más simple, cuyas derivadas son fáciles de estimar; asociaremos a cada nodo una variable con la respectiva derivada de la siguiente forma:

$$c' = x$$$$a' = 2ac'$$$$x' = c$$$$y' = a'$$$$z' = -a'$$

Dado que sólo tenemos una salida, la variable para este nodo de salida será siempre 1. Al final, con el modo reverse podemos obtener el vector gradiente correspondiente a la función en un solo paso.

class ReverseMode(object):
    def forward(self, x, y, z):
        # Define el avance en el grafo
        a = y-z
        c = a**2
        b = x*c
        
        print(b)
        return a,b,c
    
    def backward(self, x,y,z):
        #Calcula los valores d ela función
        a,b,c = self.forward(x,y,z)
        #Incializa las salidas
        db = 1
        #Derivadas
        dc = x
        da = 2*a*dc
        dx = c
        dy = 1*da
        dz = -1*da
        
        return dx, dy, dz

En el caso del modo reverse, antes de aplicar propiamente la diferenciación, debemos correr el grafo hacia adelante en el paso que llamamos 'forward'. Esto nos ayudará a obtener los valores que posteriormente podemos usar para los cálculos de las derivadas. En este caso, vemos que al correr la función hacia adelante obtenemos el resultado esperado bajo los valores de entrada determiandos. Asimismo, regresa los valores obtenidos en cada uno de los nodos, pues estos son los que nos interesarán para obtener la derivada.

#Genera el grafo computacional
f = ReverseMode()
#Corre la función hacia adeltante
f.forward(2,5,3)

8

(2, 8, 4)

Ahora podemos ejecutar el paso backward, que es donde el modo reverse obtiene el gradiente de la función. Como podemos ver, sólo basta correr una vez el modo reverse para obtener todo el vector gradiente, y por tanto, para obtener la derivada de toda la función en este caso.

#Aplica paso backward
print('Gradiente de la función: [{}]'.format(f.backward(2,5,3)))

8
Gradiente de la función: [(4, 8, -8)]

El modo reverse es efectivo cuando se cuenta con una función $f: \mathbb{R}^n \to \mathbb{R}^m$ tal que $m < n$, pues en este caso, realizará solo $m$ pasos para obtener la derivada de la función. El ejemplo anterior muestra este caso, puesto que al ser un funcional en que sólo se obtiene un valor de salida, basta correr sólo una vez el modo reverse para obtener el vector gradiente esperado.