Modo forward en diferenciación automática

El modo forward en la diferenciación automática, como su nombre lo dice, se enfoca en encontrar la derivada de una función representada por un grafo computacional haciendo cómputos en avance. En este sentido, el grafo computacional alamcena una traza tangente, donde mientras se computan los valores de la función, se obtienen también los valores de las diferentes derivadas. Posteriormente, las derivadas de nodos superiores en el grafo se computan utilizando la regla de la cadena.

De manera general, el modo forward computa una derivada en un nodo $f_i$ a partir de un nodo precedente $f_{i-1}$ por medio de la regla de la cadena de la siguiente forma:

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

De esta forma, la derivada se propaga hacia adelante, en avance, hasta obtener la derivada de las salidas de la función. En un solo paso, se pueden obtener las derivadas parciales para todas las salidas sobre una de las entradas.

El modo forward se inicializa asignando un 1 a una de las variables de entrada. O de forma más general, se inicializa asignando un vector de entrada (de la misma dimensión que la entrada de la función), que generalmente es un vector base $e_j$, de tal forma que la derivada parcial se estime sobre la $j$-ésima entrada $\frac{\partial f}{\partial x_j}$. Cuando se inicializa con un vector que no es base, el modo forward se interpreta como el producto entre la matriz Jacobiana $J_f$ y el vector de entrada.

Implementación del modo forward

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

    def forward(self):
        # 1. Calcula el gradiente del nodo padre
        # 2. Con el gradiente del nodo padre se obtiene el gradiente del nodo actual

        return # Gradiente del nodo actual

El gradiente se obtendrá al final en cada una de las salidas de la función como gráfica computacional, en f.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 forward(self):
        self.parent.forward()
        node_grad = -self.n*np.sin(self.n*self.parent.value)
        self.grad = node_grad*self.parent.grad
        
        return 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 forward(self):
        self.parent.forward()
        node_grad = self.pow*(self.parent.value**(self.pow-1))
        self.grad = node_grad*self.parent.grad
        
        return 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 forward(self):
        self.parentA.forward()
        self.parentB.forward()
        self.grad = self.parentA.grad + self.parentB.grad
        
        return self.grad

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

    def __str__(self):
        return str(self.value)
        
    def forward(self):
        self.grad = self.grad

A partir de los nodos definidos, podemos definir funciones y obtener sus derivadas. Por ejemplo, hacemos una función sencilla, con un sólo nodo, que computa la función:

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

#Nodo de entrada
x = Node(3, grad=1)
#Nodo de operación coseno
cos = Cos(n=2)

#Construcción del nodo
result = cos(x)
print(result)

0.960170286650366

Podeos aplicar, entonces, la traza derivativa para obtener el gradiente aplicando el método forward. Como se puede observar, el gradiente final está en el nodo de salida.

result.forward()
print('Gradiente {}'.format(result.grad))

Gradiente 0.5588309963978517

Podemos construir una función más compleja para ver cómo trabaja el modo forward:

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

Para esto, definimos un nodo con para la variable $y$, así como los nodos para la potencia al cuadrado y la suma. Nótese que debemos definir dos nodos de potencia, pues se aplican con resultados diferentes en la gráfica dinámica.

#Nodo de entrada y
y = Node(2, grad=1)
#Potencia nodo 1
pow = Pow(p=2)
#Potencia nodo 2
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

Finalmente, podemos obtener el gradiente de la función. En este caso, el modo forward obtiene la suma de las derivadas parciales $\frac{\partial f(x,y)}{\partial x} + \frac{\partial f(x,y)}{\partial y}$. Para obtener la derivada parcial de una sola entrada, el gradiente (parámetro grad) de la otra entrada debe ser 0.

result.forward()
print('Gradiente {}'.format(result.grad))

Gradiente 5.07314583600087