micrograd: el pase hacia adelante y el grafo de cómputo

Esta lección abre un ciclo sobre micrograd, una pequeña biblioteca de diferenciación automática, alrededor de la cual se construye casi todo lo que haremos más adelante en el curso. Hoy analizaremos solo la mitad de la mecánica: cómo un número se envuelve en un objeto Value y cómo las operaciones sobre dichos objetos construyen un grafo de cómputo sin que nos demos cuenta. La segunda mitad, el pase hacia atrás y la regla de la cadena, la dejaremos para la próxima lección. Después de esta sesión, podrás leer el pase hacia adelante de cualquier fórmula en términos de micrograd y recorrer de forma independiente el grafo de cómputo que generó.

Curso no oficial de AI University sobre código abierto (MIT). La lección incluye código de karpathy/micrograd © Andrej Karpathy, licencia MIT; los comentarios han sido traducidos al español, las explicaciones han sido escritas por nuestro equipo. El curso no está relacionado con el autor del código ni ha sido aprobado por él.

¿Por qué necesitamos un grafo de cómputo?

Para entrenar una red neuronal, necesitamos saber cómo cambia la función de pérdida con un pequeño cambio en cada parámetro, es decir, necesitamos las derivadas parciales ∂L/∂w para cada peso w. Hay tres formas de obtenerlas.

La diferenciación numérica calcula la derivada por definición: (f(x+h) - f(x)) / h para un h pequeño. Es fácil de implementar, pero el resultado es aproximado, depende de la elección de h, y para cada parámetro se necesita un pase separado a través de toda la función. Si hay millones de parámetros, tendremos que recalcular la función completa millones de veces, lo cual es inaceptablemente lento.

La diferenciación simbólica toma una fórmula y deriva analíticamente la fórmula de la derivada, como en una clase de cálculo. El resultado es exacto, pero la fórmula de la derivada para una función compleja con millones de operaciones puede crecer a un tamaño inutilizable, y este enfoque no funciona bien con construcciones de control como las ramificaciones dentro del cálculo.

La diferenciación automática en modo inverso, que es lo que implementa micrograd, toma lo mejor de ambos mundos. No almacena la fórmula simbólica de la derivada por completo, sino que durante el pase hacia adelante, recuerda de qué operaciones y qué entradas se obtuvo cada valor intermedio. Después de que se calcula el valor final (por ejemplo, la función de pérdida), recorremos este registro en orden inverso y en cada paso aplicamos la regla de la cadena, utilizando solo la derivada local de la operación específica. Como resultado, obtenemos derivadas exactas para todos los parámetros a la vez en un solo pase hacia atrás, y no una por parámetro, como en la diferenciación numérica.

Esta "grabación de lo que resultó de qué" es precisamente el grafo de cómputo: los nodos son valores, los bordes apuntan del resultado de la operación a sus entradas.

Clase Value: qué almacena

Abramos la definición de la clase y veamos su constructor.

# Fragmento: micrograd/micrograd/engine.py

class Value:
    """ guarda un valor escalar y su gradiente """

    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0
        # variables internas, necesarias para construir el grafo de diferenciación automática
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op # operación que generó este nodo, para graphviz / depuración, etc.

Analicemos los campos individualmente.

data es un número de punto flotante ordinario, el mismo valor que participa en los cálculos. Es lo que se imprime cuando, por ejemplo, miras la salida de la red.

grad es el acumulador de gradiente. Inicialmente es cero: antes del pase hacia atrás, aún no sabemos cómo este valor afecta el resultado final. En la próxima lección veremos cómo se escribe ∂L/∂(este valor) aquí.

_prev es un conjunto de nodos padre, es decir, aquellos Value de los que se obtuvo el valor actual. Estos son los bordes del grafo, solo que dirigidos "del resultado a las entradas". ¿Por qué un conjunto y no una lista o tupla? Primero, el orden de los padres no es importante para construir el grafo, lo importante es el hecho de la conexión. Segundo, el conjunto elimina automáticamente los duplicados: si el mismo nodo participa en una operación dos veces (como en la expresión x * x, donde ambos operandos son el mismo objeto x), solo aparecerá en _prev una vez, y al recorrer el grafo en el pase hacia atrás, no querremos procesarlo dos veces.

_op es una cadena con el nombre de la operación que generó el nodo: '+', '*', '**2', etc. No afecta los cálculos, es solo una etiqueta para depuración y visualización del grafo.

_backward es una función sin argumentos que puede propagar el gradiente de out a sus padres. Por defecto es lambda: None, un marcador de posición que no hace nada, necesario para nodos que no tienen padres, es decir, para las "hojas" del grafo, creadas directamente a partir de un número. Cómo se llena esta función para cada operación, lo analizaremos en detalle en la próxima lección; hoy solo nos importa que exista y se adjunte a cada nuevo nodo.

Cómo las operaciones construyen el grafo: add y mul

Veamos la suma.

# Fragmento: micrograd/micrograd/engine.py

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')

        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward

        return out

La primera línea resuelve un problema práctico importante: ¿qué pasa si escribimos a + 5, donde a es Value y 5 es un número ordinario? Entonces other resultaría ser el número int, y no Value, y no tendría ni .data ni ._prev. La línea other = other if isinstance(other, Value) else Value(other) envuelve el número puro en Value sobre la marcha, creando un nuevo nodo para él sin padres (un nodo hoja del grafo). Después de esto, ambos operandos están garantizados de ser Value, y el código posterior no necesita preocuparse por tipos mixtos.

Luego se crea out, un nuevo nodo Value, cuya data es la suma ordinaria self.data + other.data, y los padres se especifican como la tupla (self, other), que dentro del constructor se convierte en el conjunto _prev. La etiqueta de la operación es '+'. Con esto, el pase hacia adelante para la suma ha terminado: obtuvimos un número-resultado y un nodo que recuerda que surgió como la suma de dos nodos específicos.

La función _backward dentro de __add__ no la necesitamos en detalle hoy, basta con notar que se define por separado para cada operación y se adjunta a out._backward. Cómo exactamente usa out.grad para distribuir el gradiente a self y other, lo analizaremos en la próxima lección.

La multiplicación está organizada de forma completamente simétrica:

# Fragmento: micrograd/micrograd/engine.py

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')

        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward

        return out

El mismo envoltorio del número en Value, la misma creación de out con datos-producto, padres (self, other) y etiqueta '*'. La única diferencia está en la fórmula para data y en qué función _backward se coloca allí, pero eso lo veremos la próxima vez.

La exponenciación funciona de manera ligeramente diferente, porque el exponente no es Value, sino un número ordinario (micrograd solo admite la exponenciación a una potencia constante, no Value a la potencia Value):

# Fragmento: micrograd/micrograd/engine.py

    def __pow__(self, other):
        assert isinstance(other, (int, float)), "only supporting int/float powers for now"
        out = Value(self.data**other, (self,), f'**{other}')

        def _backward():
            self.grad += (other * self.data**(other-1)) * out.grad
        out._backward = _backward

        return out

Aquí out tiene un solo padre, self, ya que el exponente no es un nodo del grafo, es simplemente un número-parámetro de la operación. La etiqueta de la operación se forma dinámicamente, por ejemplo '**2' o '**-1', lo cual es conveniente para la depuración.

La función relu sigue el mismo patrón, pero contiene una ramificación en el pase hacia adelante:

# Fragmento: micrograd/micrograd/engine.py

    def relu(self):
        out = Value(0 if self.data < 0 else self.data, (self,), 'ReLU')

        def _backward():
            self.grad += (out.data > 0) * out.grad
        out._backward = _backward

        return out

El pase hacia adelante de relu es simple: si el número es negativo, el resultado es cero; de lo contrario, el resultado es el número mismo. El nodo out nuevamente recibe un único padre self y la etiqueta 'ReLU'.

Observa el patrón general que se repite en cada operación: toma los datos de entrada, calcula la data del resultado usando la fórmula matemática habitual, crea un nuevo Value con esos datos y la lista de padres, adjúntale la función _backward, devuélvelo. Este patrón es la esencia misma de la construcción del grafo de cómputo: el grafo se construye solo, junto con los cálculos habituales, no necesitas describirlo por separado.

Operaciones derivadas: menos, resta, división

Para no escribir _backward de nuevo para cada operación imaginable, micrograd expresa parte de las operaciones a través de las ya definidas:

# Fragmento: micrograd/micrograd/engine.py

    def __neg__(self,): # cambio de signo: -self
        return self * -1

    def __radd__(self, other): # suma, cuando Value a la derecha: other + self
        return self + other

    def __sub__(self, other): # resta: self - other
        return self + (-other)

    def __rsub__(self, other): # resta, cuando Value a la derecha: other - self
        return other + (-self)

    def __rmul__(self, other): # multiplicación, cuando Value a la derecha: other * self
        return self * other

    def __truediv__(self, other): # división: self / other
        return self * other**-1

    def __rtruediv__(self, other): # división, cuando Value a la derecha: other / self
        return other * self**-1

El menos unario __neg__ es simplemente una multiplicación por -1: no se necesita un nuevo _backward para esta operación, todo el trabajo ya está hecho dentro de __mul__. La resta __sub__ se expresa como una suma con negación: self - other se convierte en self + (-other). La división __truediv__ se expresa mediante la exponenciación a la potencia -1: self / other es self * other**-1, es decir, una multiplicación por el recíproco. Este es un truco elegante: en lugar de pensar por separado en la regla de diferenciación para la resta y la división, basta con reducirlas a la suma, la multiplicación y la potencia, para las cuales la regla ya está escrita.

Los métodos con el prefijo r merecen una atención especial: __radd__, __rsub__, __rmul__, __rtruediv__. Son necesarios porque Python procesa las expresiones del tipo a + b de manera diferente según el lado en que se encuentre el objeto del tipo deseado. Cuando escribes a + b, donde a es Value, Python primero intenta llamar a a.__add__(b), y esto funciona. Pero si escribes 2 + a, Python primero intenta llamar a (2).__add__(a), y un número ordinario int no tiene un método que sepa qué hacer con Value. Entonces Python, como opción de respaldo, intenta llamar a a.__radd__(2), que es precisamente lo que define micrograd. Sin __radd__, la expresión 2 + a lanzaría una excepción TypeError.

La misma historia con __rmul__, es especialmente importante en la práctica: cuando sumemos una lista de valores a través de sum(values), la función incorporada sum por defecto comienza con cero y calcula 0 + values[0] + values[1] + .... La primera suma es 0 + values[0], es decir, int.__add__ no funcionará, y entrará en juego values[0].__radd__(0). Sin este método, la suma de la lista Value a través de sum() no funcionaría.

Ejemplo: grafo para tu propia fórmula

Tomemos un problema más sencillo que las redes neuronales, digamos, una estimación del costo de envío de un pedido. Sea weight el peso del paquete en kilogramos, distance la distancia en kilómetros, y el costo total se compone de una tarifa base, una tarifa por peso y una tarifa por distancia, pasada por una relu para eliminar un resultado negativo:

# Nuestro ejemplo
from micrograd.engine import Value

base_fee = Value(50.0)          # tarifa base por la entrega
weight = Value(3.0)             # peso del paquete, kg
distance = Value(12.0)          # distancia, km
discount = Value(-8.0)          # descuento de temporada: reduce el costo final

weight_fee = weight * 15.0      # tarifa por peso: 15 por kilogramo
distance_fee = distance * 2.0   # tarifa por kilómetro

raw_cost = base_fee + weight_fee + distance_fee + discount
cost = raw_cost.relu()          # el costo no puede ser negativo

print(cost.data)

Cada línea con aritmética en realidad no solo calcula un número, sino que crea un nuevo nodo Value y lo vincula con sus padres a través de _prev. Para el momento en que imprimimos cost.data (que es 111.0, calcula tú mismo: 50 + 45 + 24 - 8), ya hay un grafo completo en la memoria: cost apunta a raw_cost, este apunta a cuatro sumandos, y weight_fee y distance_fee a su vez apuntan a weight, distance y los números 15.0 y 2.0 envueltos en Value.

Para ver este grafo con tus propios ojos, sin recurrir a graphviz, escribamos un simple recorrido recursivo:

# Nuestro ejemplo
def print_graph(value, depth=0, visited=None):
    """recursivamente imprime el grafo de cómputo, comenzando desde el nodo dado"""
    if visited is None:
        visited = set()
    indent = '  ' * depth
    op_label = value._op if value._op else 'leaf'  # las hojas no tienen operación
    print(f"{indent}Value(data={value.data:.2f}) op={op_label}")
    if value in visited:
        return  # el nodo ya fue expandido arriba, no repetimos sus hijos
    visited.add(value)
    for child in value._prev:
        print_graph(child, depth + 1, visited)

print_graph(cost)

La función imprime el nodo actual con una sangría que corresponde a la profundidad en el grafo, y luego desciende recursivamente a cada uno de sus _prev. El conjunto visited protege contra la expansión repetida de los hijos del mismo nodo si aparece varias veces en el grafo (por ejemplo, si hubiéramos usado weight en otro lugar). Al ejecutarlo en cost, verás un árbol de operaciones: ReLU en la parte superior, debajo una cadena +, y en la parte inferior las hojas, es decir, los números originales sin operación.

Ten en cuenta que _prev es un conjunto, por lo que el orden en que se imprimen los padres durante el recorrido no está garantizado y puede variar de una ejecución a otra. Para la depuración, esto no es un problema: lo importante es el hecho de las conexiones, no el orden de su enumeración.

Limitaciones de micrograd y por qué sigue siendo necesario

micrograd tiene limitaciones conscientes. Cada Value almacena un solo número, no un arreglo de números, por lo que para entrenar una red real con miles de parámetros, tendrías que crear miles de objetos Value separados y realizar operaciones una por una en un ciclo de Python puro. Esto es órdenes de magnitud más lento que las operaciones vectorizadas sobre tensores en PyTorch o NumPy, donde una instrucción del procesador o la tarjeta de video procesa un arreglo completo de números a la vez.

Pero precisamente esta simplicidad es el valor de micrograd como herramienta de aprendizaje. PyTorch está estructurado bajo el mismo principio: cada torch.Tensor con requires_grad=True almacena un grafo de cómputo, cuyos nodos son tensores, no escalares, y las operaciones entre tensores también crean nuevos nodos con reglas adjuntas para el pase hacia atrás. Al comprender con ejemplos escalares simples cómo funciona la combinación de datos, gradiente, padres y operación, leerás las entrañas de cualquier framework de aprendizaje profundo con mucha más confianza.

Intenta tú mismo

1. Agrega a la clase Value el método exp(self), que en el pase hacia adelante calcula out = Value(math.exp(self.data), (self,), 'exp'), es decir, la exponencial del valor actual, con un único padre self y la etiqueta 'exp'. No olvides agregar import math al principio del archivo engine.py, de lo contrario math.exp no estará disponible. El pase hacia atrás (_backward) en este paso puede dejarse como un marcador de posición o, con un asterisco, intentar escribirlo tú mismo, recordando que la derivada de la exponencial es la exponencial misma. Criterio de éxito: Value(1.0).exp().data está cerca de 2.71828.

2. Usando las operaciones ya existentes y el método exp de la tarea 1, expresa la tangente hiperbólica a través de la fórmula tanh(x) = (e^(2x) - 1) / (e^(2x) + 1), sin tocar el interior de la clase Value, es decir, escribe una función ordinaria my_tanh(x) que toma Value y devuelve Value, usando solo +, -, *, / y exp(). Criterio de éxito: para x = Value(0.0) el resultado está cerca de 0.0, y para x grandes en valor absoluto el resultado está cerca de 1.0 o -1.0.

3. Inventa tu propia fórmula con cinco o seis operaciones (suma, multiplicación, potencia, relu), construye un grafo de cómputo para ella y usa la función print_graph de esta lección para imprimirlo. Cuenta manualmente cuántos nodos Value se obtuvieron en total, incluyendo las hojas, y compáralo con lo que realmente imprime el recorrido (pista: es conveniente crear otra función que simplemente recoja todos los nodos únicos en un conjunto y devuelva su longitud).

Resultados

• La diferenciación automática en modo inverso memoriza el historial de cálculos durante el pase hacia adelante y luego calcula todas las derivadas en un solo pase hacia atrás, a diferencia de la lenta diferenciación numérica y la engorrosa diferenciación simbólica.
• La clase Value almacena el número data, el acumulador de gradiente grad, el conjunto de padres _prev, la etiqueta de la operación _op y la función _backward, que por ahora sigue siendo un marcador de posición hasta la próxima lección.
• Cada operación (__add__, __mul__, __pow__, relu) envuelve los números en Value, crea un nuevo nodo con los padres y la etiqueta correctos, ocultando al mismo tiempo en _backward la regla para el futuro pase hacia atrás.
• La resta, la división y el menos unario no requieren una implementación separada de _backward: se expresan a través de la suma, la multiplicación y la potencia.
• Los métodos __radd__ y __rmul__ son necesarios para que las expresiones del tipo 2 * a y la suma a través de sum() funcionen cuando Value está a la derecha del operador.
• El grafo de cómputo es un DAG de objetos Value, que se puede recorrer fácilmente con una recursión normal sobre _prev, sin recurrir a graphviz, y el mismo principio subyace al autodiff de tensores en PyTorch.

Fuentes originales

• Fuente original: video de Andrej Karpathy (en inglés): The spelled-out intro to neural networks and backpropagation: building micrograd
• Código de la lección en GitHub: karpathy/micrograd/micrograd/engine.py

Los videos se proporcionan solo como enlaces a la fuente original; el texto de la lección no es una traducción ni un resumen de ellos. Textos completos de las licencias: en la lección "Qué sigue, licencias y fuentes".