Entendendo diferenciação automática (Automatic Differentiation) em 30 linhas de Python

 (vmartin.fr)
3 pontos por GN⁺ 2023-08-27 | 1 comentários | Compartilhar no WhatsApp
  • Mostra como implementar diretamente a diferenciação automática, peça central do treinamento de redes neurais, com uma classe Tensor para escalares, demonstrando como o cálculo de valores e o cálculo de derivadas se conectam no mesmo grafo computacional
  • Com variáveis Python comuns, em z = x + y só resta o valor resultante e a relação desaparece; por isso, Tensor precisa armazenar tanto o valor quanto o histórico de operações
  • Com Children(a, b, op) e chamadas recursivas de forward(), cria um grafo computacional em árvore binária; ao redefinir adição e multiplicação, é possível recalcular a expressão mesmo inserindo valores depois
  • grad(deriv_to) define a derivada em relação a si mesmo como 1 e em relação a outro escalar como 0, aplicando recursivamente as regras de derivação das operações básicas para criar um novo grafo computacional
  • A implementação trata apenas escalares e pode ser lenta; operações com arrays, poda de ramos com multiplicação por 0, tratamento de nós constantes e um cache para reduzir cálculos repetidos ficam como pontos de melhoria

Com variáveis Python comuns, a relação desaparece

  • Ao calcular x = 3, y = 5, z = x + y, em z fica apenas o valor resultante 8
  • Mesmo que os valores de x ou y mudem depois, z não consegue rastrear a partir de quais variáveis foi criado
  • Como a relação entre as variáveis não permanece, é difícil calcular automaticamente a derivada em relação a uma variável específica

Preservando o histórico de operações com Tensor

  • O novo tipo Tensor armazena um valor (value) e redefine operadores para retornar um novo Tensor ao calcular entre objetos Tensor
  • A implementação inicial redefine apenas __add__, permitindo que Tensor(3) + Tensor(5) crie T:8
  • Nessa etapa, z ainda não consegue preservar o histórico de operações de que é resultado de x + y

Grafo computacional e forward()

  • Para preservar o histórico de operações, introduz Children = namedtuple('Children', ['a', 'b', 'op'])
    • a: tensor de entrada à esquerda
    • b: tensor de entrada à direita
    • op: a operação real, como np.add ou np.multiply
  • Cada Tensor pode ter não só um valor numérico, mas também children, formando assim um grafo computacional em forma de árvore binária
  • forward() visita recursivamente os nós filhos para calcular o valor real
    • Com x = Tensor(3), y = Tensor(5), z1 = x + y, z2 = z1 * y vira T:40
    • Mesmo criando primeiro o grafo com x = Tensor(None), y = Tensor(None) e depois definindo x.value = 3, y.value = 5, chamar z2.forward() calcula T:40

Transformando diferenciação automática em grafo computacional

  • A diferenciação automática é implementada adicionando regras de derivação para cada operação básica compatível com Tensor
  • grad(self, deriv_to) percorre recursivamente o grafo computacional e decompõe funções complexas em combinações de funções simples
  • As regras básicas são:
    • Derivar um tensor em relação a ele mesmo resulta em Tensor(1)
    • Derivar um escalar sem filhos em relação a outro tensor resulta em Tensor(0)
    • Adição: (a + b)' = a' + b'
    • Multiplicação: (ab)' = a'b + ab'
  • Ao derivar z2 = (x + y) * y em relação a y, o resultado g não é um valor simples, mas um novo grafo computacional que representa a derivada parcial
    • Em forma de expressão, g = ∂z2/∂y = x + 2*y
    • Quando x = 3 e y = 5, o valor de g é 13

Expansão para subtração, divisão e função exponencial

  • Para lidar com expressões mais complexas, são adicionadas a Tensor operações de subtração, divisão, exponencial e negação
  • grad() passa a incluir regras de derivação correspondentes a cada operação
    • Subtração: (a - b)' = a' - b'
    • Divisão: (a/b)' = (a'b - ab') / b²
    • Função exponencial: exp(a)' = a' * exp(a)
  • forward() também é alterado para tratar operações que precisam de apenas um termo
    • Ex.: exp(a) não precisa do segundo termo b
    • -x é tratado na forma 0 - x

Expressão de exemplo e verificação com Sympy

  • A expressão a seguir é escrita com Tensor, e as derivadas parciais em relação a x e y são calculadas
z = (12 - (x * e^y)) / (45 + x * y * e^-x)
  • No código, ela é expressa assim
x = Tensor(3)
y = Tensor(5)
z = (Tensor(12) - (x * y.exp())) / (Tensor(45) + x * y * (-x).exp())
  • Os valores calculados das derivadas parciais são:
    • z.grad(x) → T:-3.34729777301069
    • z.grad(y) → T:-9.70176956641438
  • O resultado do cálculo da mesma expressão com diff() e evalf() do Sympy também é idêntico
    • Com xs = 3, ys = 5, o valor da derivada em relação a x é -3.34729777301069
    • O valor da derivada em relação a y é -9.70176956641438

Limitações da implementação simples e pontos de otimização

  • Esta implementação se aproxima de um sistema de diferenciação automática mais simples possível e, ao mesmo tempo, pode ser muito lenta
  • A classe atual trata apenas escalares
    • Para se tornar uma biblioteca mais útil, seria necessário adicionar operações com arrays de tamanho arbitrário
  • Observando o grafo computacional, algumas otimizações são possíveis
    • Se um dos filhos em um nó de multiplicação for 0, não é necessário explorar mais profundamente
    • Se um nó e seus filhos não dependem do tensor x que é o alvo da derivação, esse nó pode ser considerado uma constante e a travessia pode parar
    • Quando a mesma operação se repete, pode-se usar um cache para evitar executar o mesmo cálculo várias vezes

Leituras relacionadas

1 comentários

 
GN⁺ 2023-08-27
Opiniões no Hacker News

  • Gosto dessas demonstrações de código pequenas e elegantes. Elas ajudam a entender o conceito colocando a mão na massa
    Os puzzles de GPU e de tensores do Sasha Rush são exemplos semelhantes
    https://github.com/srush/GPU-Puzzles
    https://github.com/srush/Tensor-Puzzles

  • Se você acredita que entendeu completamente diferenciação automática só com isso, está enganando a si mesmo
    Quando o grafo é uma árvore, tudo é bem simples, como neste texto. Mas, se o grafo for um grafo acíclico direcionado mais geral, por exemplo x = 5; y = 2x; z = xy, a implementação continua muito simples, mas entender por que ela está correta não é tão simples. Se você pensa que é “apenas a regra da cadeia comum”, também está se enganando
    Uma das primeiras explicações foi dada por Paul Werbos; ele chamou a regra necessária de regra da cadeia para derivadas ordenadas e a provou por indução a partir da regra da cadeia comum. Ainda assim, ela não decorre de forma imediatamente óbvia da regra da cadeia comum. Se alguém acredita no contrário, eu gostaria que provasse que estou errado; ficaria muito feliz com isso

    • Então onde dá para ler mais sobre isso? As pessoas que criaram frameworks como autograd, PyTorch e mxnet devem ter aprendido isso em detalhes em algum lugar, e tenho curiosidade sobre essa fonte. Pelo que sei, o mxnet veio do meio acadêmico, talvez da CMU
    • Sinceramente, não sei bem o que as pessoas querem nesse tipo de discussão, e talvez seja porque a abstração implícita, as derivadas ordenadas, não seja ideal
      Se você aplicar a regra da cadeia comum ao longo das arestas do grafo computacional, isto é, de um grafo acíclico direcionado, obterá o valor correto em cada etapa. A regra adicional necessária é algo como: “se uma variável for usada várias vezes no cálculo, ou seja, se houver várias arestas saindo do mesmo nó ou, no sentido reverso, várias arestas entrando, some os gradientes calculados para cada uma delas”. Isso também me parece bastante básico e intuitivo
      Por exemplo, se você passar z tanto para x quanto para y em f(x, y), então d/dz f(z, z) = f_x(z, z) + f_y(z, z), em que os subscritos indicam derivadas parciais. Para mim, essa abordagem é matematicamente mais simples do que misturar as duas coisas e fazer parecer que é “algo além da regra da cadeia”, além de parecer mais próxima da implementação real, especialmente do que o PyTorch, com o qual tenho mais familiaridade, faz
    • A regra da cadeia é definida para derivadas parciais, então, tecnicamente, ainda dá para dizer que é apenas a regra da cadeia

  • Diferenciação automática parece magia
    Muitos cientistas da computação ficaram fascinados por isso e escreveram textos apresentando a técnica de uma perspectiva mais ampla. O meu texto é um deles, e também inclui uma “variante de pobre” que usa números complexos sem sobrecarga de operadores
    https://pizzaseminar.speicherleck.de/automatic-differentiati...

    • Quando eu trabalhava com machine learning em 1994~1995, eu não conhecia diferenciação automática, e o professor que criou a função objetivo também calculava as derivadas analíticas manualmente. Só fiquei sabendo disso alguns anos atrás, e foi surpreendente pensar no tempo que passei, no fim dos anos 90, aprendendo Mathematica o suficiente para gerar derivadas analíticas por conta própria
    • Isso parece remontar à aproximação de derivadas por passo complexo, de J. Martins, P. Sturdza e J. Alonso, em 2003. Vale a pena ler esse artigo
      [0]: https://doi.org/10.1145/838250.838251
    • Realmente parece magia. Eu gostaria de conhecer algum material introdutório sobre backpropagation escrito de uma forma parecida

  • Tenho uma implementação de diferenciação automática em Python em 26 linhas: https://gist.github.com/sradc/d9d66e3898ffe3a02e0b6b266629b0...

    • Código curto é bom, mas acho que minha cabeça funciona muito melhor quando há espaçamento suficiente. Preciso praticar um pouco mais esses outros estilos

  • É muito parecido com uma técnica usada em sistemas de engenharia baseada em conhecimento, onde é chamada de rastreamento de dependências. Usada junto com cache de nós ou tensores, ela pode reduzir o volume de cálculo, sendo especialmente útil para grandes modelos 3D paramétricos
    Ao obter um valor, chama-se recursivamente a árvore binária/de dependências para verificar quais variáveis mudaram e recalcular apenas o necessário. Usando objetos e atributos Python customizados com métodos __set__ e __get__, dá para fazer isso parecer um recurso embutido de um modelo orientado a objetos
    x = Tensor(3)
    y = Tensor(5)
    z = x + y
    print(x, y) # 3, 5
    print(z) # 8
    x.value = 4 # ao definir o valor, nada é recalculado
    print(z) # 9, porque a dependência alterada é recalculada no momento em que o valor é obtido

  • Andrej Karpathy tem um vídeo interessante em que constrói um motor de autograd, e ele é bastante esclarecedor
    https://youtu.be/VMj-3S1tku0?si=wuKhELwOwoYbzpt7
    Repositório:
    https://github.com/karpathy/micrograd

  • A variação de diferenciação automática que eu conheço não cria um grafo de operações. Em vez disso, calcula o valor na hora

    • Provavelmente está pensando em diferenciação automática em modo direto. Ela é mais útil quando a dimensão da saída da função é relativamente grande, e é diferente da diferenciação automática em modo reverso, que é mais útil quando a dimensão da saída é relativamente pequena
      As duas funcionam, mas, dependendo da situação, uma é mais eficiente que a outra. Em casos como “treinamento de redes neurais”, muitas vezes se otimiza uma única saída de perda para muitos alvos, então normalmente se usa o modo reverso

  • Eu gostaria que diferenciação automática fosse chamada simplesmente de regra da cadeia numérica, ou pelo menos explicada assim. Literalmente é só isso, com alguns truques para evitar calcular explicitamente a matriz jacobiana em determinadas operações, o que deixaria tudo muito mais claro

    • O “autodiff” explicado aqui e mais usado em implementações de backpropagation é a diferenciação automática em modo reverso, mas também existe o modo direto e estratégias entre esses dois extremos. No fim, tudo se reduz à regra da cadeia, mas escolher a abordagem no nível do algoritmo está longe de ser trivial
      Na prática, se você pedir para alguém usar a regra da cadeia para propagar gradientes por um grafo computacional, acho que a maioria pensaria intuitivamente no modo direto como padrão. Eu também pensaria
      https://en.wikipedia.org/wiki/Automatic_differentiation#Beyo...
      Vendo por esse lado, parece útil usar o termo para se referir a um método específico de acumular gradientes ao percorrer as expressões fornecidas pela regra da cadeia
    • Tecnicamente, isso está errado. Regra da cadeia numérica usa diferenças finitas, e os erros se acumulam ao longo do cálculo
      Veja a seção “Diferenças em relação a outros métodos”: https://en.m.wikipedia.org/wiki/Automatic_differentiation
      Como em um comentário próximo, o ponto central é que a implementação realmente importa e vale a pena estudá-la. Dizer que diferenciação automática é um conjunto de métodos para implementar a regra da cadeia está bem, mas dizer que é “apenas” a regra da cadeia numérica está errado
    • Pode até ser mais preciso, mas eu não chamaria de mais claro

  • Qual é o problema se diferenciação automática, na categoria das funções suaves, é apenas uma lente cartesiana da matriz jacobiana e da diferencial total? https://www.youtube.com/watch?v=ne99laPUxN4

  • Fico curioso por que a classe é chamada de Tensor. Existe alguma forma de pensar na expressão ou na sua derivada como um tensor? Ou é porque escalares também são tensores, e isso pode ser estendido para dar suporte a outros tipos de tensores?

    • Posso estar errado, mas matematicamente acho que objetos 2D são chamados de matrizes, e objetos de 3D ou mais são chamados de tensores
      Como o algoritmo de diferenciação automática descrito funciona para objetos de alta dimensão arbitrária, parece fazer sentido chamar esses objetos de tensores