Entendiendo la diferenciación automática en 30 líneas de Python

 (vmartin.fr)
3 puntos por GN⁺ 2023-08-27 | 1 comentarios | Compartir por WhatsApp
  • Muestra cómo implementar directamente la diferenciación automática, clave en el entrenamiento de redes neuronales, con una clase escalar Tensor, y cómo el cálculo de valores y derivadas se encadena sobre el mismo grafo computacional
  • Con variables normales de Python, en z = x + y solo queda el valor resultante y se pierde la relación, así que Tensor debe guardar tanto el valor como el historial de operaciones
  • Con Children(a, b, op) y llamadas recursivas a forward() se construye un grafo computacional en árbol binario; al redefinir suma y multiplicación, la expresión puede recalcularse aunque los valores se asignen después
  • grad(deriv_to) toma la derivada respecto de sí mismo como 1 y respecto de otro escalar como 0, y aplica recursivamente las reglas de derivación de las operaciones básicas para crear un nuevo grafo computacional
  • La implementación solo maneja escalares y puede ser lenta; quedan como mejoras pendientes las operaciones con arreglos, la poda de ramas multiplicadas por 0, el manejo de nodos constantes y un caché para reducir cálculos repetidos

Con variables normales de Python la relación se pierde

  • Si se calcula x = 3, y = 5, z = x + y, en z solo queda el valor resultante 8
  • Aunque luego cambien los valores de x o y, z ya no puede rastrear de qué variables fue creado
  • Como no se conserva la relación entre variables, resulta difícil calcular automáticamente la derivada respecto de una variable concreta

Preservar el historial de operaciones con Tensor

  • El nuevo tipo Tensor guarda un valor (value) y redefine operadores para que, al operar entre Tensor, devuelva un nuevo Tensor
  • La implementación inicial solo redefine __add__, así que Tensor(3) + Tensor(5) puede producir T:8
  • En esta etapa, z todavía no puede conservar el historial de operaciones de que es el resultado de x + y

Grafo computacional y forward()

  • Para conservar el historial de operaciones se introduce Children = namedtuple('Children', ['a', 'b', 'op'])
    • a: tensor de entrada izquierdo
    • b: tensor de entrada derecho
    • op: operación real como np.add o np.multiply
  • Cada Tensor puede tener no solo un valor numérico sino también children, y con eso construir un grafo computacional en forma de árbol binario
  • forward() visita recursivamente los nodos hijos para calcular el valor real
    • Con x = Tensor(3), y = Tensor(5), si z1 = x + y y z2 = z1 * y, el resultado es T:40
    • Incluso si primero se arma el grafo con x = Tensor(None), y = Tensor(None) y después se asigna x.value = 3, y.value = 5, al llamar z2.forward() también se obtiene T:40

Construir la diferenciación automática como un grafo computacional

  • La diferenciación automática se implementa agregando reglas de derivación para cada operación básica que soporte Tensor
  • grad(self, deriv_to) recorre recursivamente el grafo computacional y descompone funciones complejas en combinaciones de funciones simples
  • Las reglas básicas son estas
    • Si se deriva un tensor respecto de sí mismo, se obtiene Tensor(1)
    • Si se deriva un escalar sin hijos respecto de otro tensor, se obtiene Tensor(0)
    • Suma: (a + b)' = a' + b'
    • Multiplicación: (ab)' = a'b + ab'
  • Si se deriva z2 = (x + y) * y respecto de y, el resultado g no es un valor simple, sino un nuevo grafo computacional que representa la derivada parcial
    • Como fórmula, g = ∂z2/∂y = x + 2*y
    • Cuando x = 3 y y = 5, el valor de g es 13

Extensión a resta, división y función exponencial

  • Para manejar expresiones más complejas, se agregan a Tensor la resta, la división, la función exponencial y la negación
  • grad() incorpora las reglas de derivación correspondientes a cada operación
    • Resta: (a - b)' = a' - b'
    • División: (a/b)' = (a'b - ab') / b²
    • Función exponencial: exp(a)' = a' * exp(a)
  • forward() también cambia para manejar operaciones que solo requieren un operando
    • Por ejemplo, exp(a) no necesita un segundo operando b
    • -x se maneja como 0 - x

Fórmula de ejemplo y validación con Sympy

  • Se escribe la siguiente expresión con Tensor y se calculan las derivadas parciales respecto de x y y
z = (12 - (x * e^y)) / (45 + x * y * e^-x)
  • En el código se expresa así
x = Tensor(3)
y = Tensor(5)
z = (Tensor(12) - (x * y.exp())) / (Tensor(45) + x * y * (-x).exp())
  • Los valores calculados para las derivadas parciales son los siguientes
    • z.grad(x) → T:-3.34729777301069
    • z.grad(y) → T:-9.70176956641438
  • El resultado de calcular la misma expresión con Sympy usando diff() y evalf() también coincide
    • Con xs = 3, ys = 5, el valor de la derivada respecto de x es -3.34729777301069
    • El valor de la derivada respecto de y es -9.70176956641438

Limitaciones de esta implementación simple y puntos de optimización

  • Esta implementación se parece a un sistema de diferenciación automática de la forma más simple posible y, al mismo tiempo, puede ser muy lenta
  • La clase actual solo maneja escalares
    • Para convertirse en una biblioteca más útil, tendría que agregar operaciones con arreglos de tamaño arbitrario
  • Al observar el grafo computacional, se ven varias optimizaciones posibles
    • En un nodo de multiplicación, si uno de los hijos es 0, no hace falta seguir explorando más profundo
    • Si un nodo y sus hijos no dependen del tensor x respecto del cual se deriva, ese nodo puede tratarse como una constante y detener el recorrido
    • Cuando se repite la misma operación, se puede usar un caché para evitar hacer el mismo cálculo varias veces

Lecturas relacionadas

1 comentarios

 
GN⁺ 2023-08-27
Opiniones en Hacker News

  • Me gustan estas demos de código pequeñas y elegantes, porque te permiten entender el concepto ensuciándote las manos
    Los puzzles de GPU y de tensores de Sasha Rush son ejemplos parecidos
    https://github.com/srush/GPU-Puzzles
    https://github.com/srush/Tensor-Puzzles

  • Si crees que con esto ya entendiste por completo la diferenciación automática, te estás engañando
    Cuando el grafo es un árbol, todo es muy simple, como en este artículo. Pero si el grafo es un grafo acíclico dirigido más general, por ejemplo x = 5; y = 2x; z = xy, la implementación sigue siendo muy simple, pero entender por qué esa implementación es correcta no lo es. Si piensas que es “solo la regla de la cadena normal”, también te estás engañando
    Una de las primeras explicaciones fue de Paul Werbos, quien llamó a la regla necesaria regla de la cadena para derivadas ordenadas, y la demostró por inducción a partir de la regla de la cadena normal. Aun así, no se desprende de forma inmediata y evidente de la regla de la cadena normal. Si alguien cree lo contrario, me gustaría que lo demostrara; me alegraría mucho

    • Entonces, ¿dónde conviene leer más? Quienes crearon frameworks como autograd, PyTorch y mxnet seguramente lo aprendieron con detalle en algún lado, y me interesa saber cuál fue esa fuente. Tengo entendido que mxnet salió del ámbito académico, quizá de CMU
    • Honestamente, no tengo muy claro qué busca la gente en este tipo de discusiones, y quizá sea porque la abstracción implícita de derivadas ordenadas no es ideal
      Si aplicas la regla de la cadena normal a lo largo de las aristas del grafo computacional, es decir, de un grafo acíclico dirigido, en cada paso obtienes el valor correcto. La regla adicional necesaria es algo como: “si una variable se usa varias veces en el cálculo, es decir, si salen varias aristas del mismo nodo o, en sentido inverso, entran varias aristas, hay que sumar los gradientes calculados para cada una”; pero eso también me parece bastante básico e intuitivo
      Por ejemplo, si pasas z tanto como x como y a f(x, y), entonces d/dz f(z, z) = f_x(z, z) + f_y(z, z), donde el subíndice indica una derivada parcial. Para mí, esta forma es más simple matemáticamente que mezclar ambas cosas y hacer que parezca “algo más allá de la regla de la cadena”; además, parece más cercana a la implementación real, en especial a lo que hace PyTorch, que es con lo que estoy más familiarizado
    • La regla de la cadena está definida para derivadas parciales, así que técnicamente aún puede verse simplemente como la regla de la cadena

  • La diferenciación automática se siente como magia
    Muchos científicos de la computación se han fascinado con esto y han escrito artículos que presentan la técnica desde una perspectiva más amplia. Mi artículo es uno de ellos e incluye una “variante para pobres” que usa números complejos sin sobrecarga de operadores
    https://pizzaseminar.speicherleck.de/automatic-differentiati...

    • Cuando hacía machine learning en 1994-1995, no conocía la diferenciación automática, y el profesor que creó la función objetivo también calculó a mano las derivadas analíticas. Me enteré de esto recién hace unos años, y me sorprendió al pensar en todo el tiempo que pasé a fines de los 90 aprendiendo suficiente Mathematica para generar derivadas analíticas por mi cuenta
    • Esto parece remontarse a la aproximación de derivadas por paso complejo de J. Martins, P. Sturdza y J. Alonso, de 2003. Ese paper vale la pena
      [0]: https://doi.org/10.1145/838250.838251
    • Realmente se siente como magia. Me gustaría conocer algún material introductorio sobre backpropagation escrito de manera similar

  • Tengo una implementación de diferenciación automática en Python en 26 líneas: https://gist.github.com/sradc/d9d66e3898ffe3a02e0b6b266629b0...

    • Que sea breve está bien, pero creo que mi cerebro funciona mucho mejor cuando hay una cantidad razonable de espacios en blanco. Tendré que practicar un poco estas otras formas de escribir código

  • Es muy parecido a una técnica usada en sistemas de ingeniería basados en conocimiento, donde la llaman seguimiento de dependencias. Usada junto con caché de nodos o tensores, puede reducir la cantidad de cálculo, algo especialmente útil en modelos 3D paramétricos grandes
    Al obtener un valor, se llama recursivamente al árbol binario/de dependencias para verificar qué variables cambiaron y se recalcula solo lo necesario. Si usas objetos y atributos personalizados de Python con métodos __set__ y __get__, puedes hacer que se sienta como una funcionalidad incorporada de un modelo orientado a objetos
    x = Tensor(3)
    y = Tensor(5)
    z = x + y
    print(x, y) # 3, 5
    print(z) # 8
    x.value = 4 # al establecer el valor no se recalcula nada
    print(z) # 9, porque la dependencia modificada se recalcula en el momento de obtener el valor

  • Andrej Karpathy tiene un video interesante donde construye un motor de autograd, y es bastante revelador
    https://youtu.be/VMj-3S1tku0?si=wuKhELwOwoYbzpt7
    Repositorio:
    https://github.com/karpathy/micrograd

  • La variante de diferenciación automática que conozco no construye un grafo de operaciones. En cambio, calcula el valor al vuelo.

    • Probablemente estés pensando en la diferenciación automática en modo directo. Es más útil cuando la dimensión de salida de la función es relativamente grande, y es distinta de la diferenciación automática en modo inverso, que es más útil cuando la dimensión de salida es relativamente pequeña.
      Ambas funcionan, pero una u otra es más eficiente según el caso. En cosas como el “entrenamiento de redes neuronales”, a menudo se optimiza una única salida de pérdida sobre muchos objetivos, así que normalmente se usa el modo inverso.

  • Ojalá a la diferenciación automática simplemente se le llamara regla de la cadena numérica, o al menos se explicara así. Literalmente eso es todo, con algunos trucos añadidos para no tener que calcular explícitamente la matriz jacobiana en ciertas operaciones, así que sería mucho más claro.

    • El “autodiff” que se explica aquí y que se usa con más frecuencia en implementaciones de retropropagación es la diferenciación automática en modo inverso, pero también existe el modo directo y estrategias entre ambos extremos. Al final todo se reduce a la regla de la cadena, pero elegir el enfoque a nivel de algoritmo no es para nada trivial.
      De hecho, si le pides a alguien que use la regla de la cadena para propagar gradientes a través de un grafo computacional, creo que la mayoría pensaría intuitivamente en el modo directo como opción predeterminada. Yo también.
      https://en.wikipedia.org/wiki/Automatic_differentiation#Beyo...
      Visto así, parece útil usar el término para referirse a un método específico de acumular gradientes al recorrer las expresiones que proporciona la regla de la cadena.
    • Técnicamente es incorrecto. La regla de la cadena numérica usa diferencias finitas, y los errores se acumulan a medida que avanza el cálculo.
      Mira la sección “Comparación con otros métodos”: https://en.m.wikipedia.org/wiki/Automatic_differentiation
      Como dicen comentarios cercanos, la clave es que la implementación realmente importa y vale la pena estudiarla. Está bien decir que la diferenciación automática es un conjunto de métodos para implementar la regla de la cadena, pero decir que es “simplemente” la regla de la cadena numérica es incorrecto.
    • Podría ser más preciso, pero no diría que es más claro.

  • ¿Cuál es el problema? La diferenciación automática no es más que una lente cartesiana de la matriz jacobiana y la diferencial total en la categoría de funciones suaves: https://www.youtube.com/watch?v=ne99laPUxN4

  • Me da curiosidad por qué la clase se llama Tensor. ¿Hay alguna forma de pensar una expresión o su derivada como un tensor? ¿O es porque un escalar también es un tensor, y esto se puede extender para admitir otros tipos de tensores?

    • Podría estar equivocado, pero matemáticamente creo que a los objetos 2D se les llama matrices y a los objetos de 3D o más se les llama tensores.
      Como el algoritmo de diferenciación automática descrito funciona con objetos arbitrarios de alta dimensión, parece tener sentido llamar tensores a estos objetos.