Autograd 自动微分的动机和实现

Omitted

0.1 Motivation

• 如何在计算机中进行微分?

  • Manual Differentiation: 人工推导导数公式, 然后写进代码.

  • Symbolic Differentiation: 计算机符号化地推导导数公式, 然后生成代码 (Mathematica 就是这么干的)

  • Numerical Differentiation: 用数值方法近似地计算导数.

  • Automatic Differentiation: 通过链式法则和计算图, 自动地计算导数.

    

    Figure 1: 计算机中微分的 4 种方式 [1].

• 放在 ML 的场景下, 一个神经网络可能有几十亿的参数, 每个参数都要求梯度, 手动推导不可能完成! 符号化推导也会导致表达式爆炸, 数值方法精度又不够好. 但是神经网络的函数不是随意的函数, 它是高度算子化和分层的. 我们希望导数的信息能自动沿着计算的路径 (Computational Graph) 反向传播!
  • 用「传播」这个词是因为神经网络是高度分层的结构, 数据之间有明显的依赖关系 (partially-ordered).
  • 用「自动」这个词是因为我们希望一句 loss.backward() 就能像遍历一个树一样完成所有的梯度计算.
  • 而且我们希望在程序看来每个梯度计算都是「局部」的, 运行的函数并不知道自己在传播一个大计算图.
    • 这个理论基础是 Chain Rule 链式法则! 我们这样解读链式法则: 要求 \(y\) 关于变量 \(v\) 的导数, 只要知道下面三类信息 (在 Figure 2 标为蓝色):
      • \(v\) 参与了哪些算子的运算;
      • 这些算子的输出 (当然每个算子只有一个) 分别对 \(y\) 的导数;
      • 这些算子的其它输入的值.

      

      Figure 2: Chain Rule 的局部视角, 求 \(\partial_v y\) 只需要知道蓝色的信息, 对于 \(a,b, \ldots, c\) 与 \(y\) 的关系不需要知道!

EXAMPLE: A scalar computational graph

考虑式子 [2]: \[ y := \ln x_1 + x_1 x_2 - e^{x_2} \tag{1}\]

可以由以下图表示:

Figure 3: Computational graph demo.

我们希望计算 \(y\) 关于 \(x_1, x_2\) 在 \((x_1, x_2) = (2, 5)\) 处的梯度.

• Forward propagation: 首先进行前向传播, 目的是计算出中间变量的数值 (为什么要算呢, 因为链式法则需要知道中间变量的值!):

  

  Figure 4: Forward propagation.

• Backward propagation: 如 Figure 5, 蓝色数值代表 \(y\) 对该节点的导数.

  

  Figure 5: Backward propagation.

可知:

\[ \begin{aligned} \partial_{x_1} y &= 5.5 \\ \partial_{x_2} y &= -146.4 \end{aligned} \]

0.2 Minimal Implementation

我们将 Figure 4 的棕色方块建模为 Tensor 类, 青色圆形建模为算子比如 Add. 这是 autograd 重点关注的两个结构. 现在我们实现 Equation 1 的 autograd 库 (只实现了必要的算子, adapted from [3]). 主要关注 Tensor 类的 backward() 方法和各个算子的 backward() 方法的相互调用!

mytorch/tensor.py

from __future__ import annotations
from typing import Optional
import numpy as np

class Tensor:
    def __init__(self, data: np.ndarray, requires_grad=False, operation=None):
        # Payload
        self.data = data
        # Metadata
        self.requires_grad = requires_grad
        if self.requires_grad:
            self.grad = np.zeros_like(data, dtype=np.float32)
        self.operation = operation  # What operation cls created this tensor
        self.children = []          # What other tensors are created from this tensor

    # This method is called by the operation backward()
    def backward(self, 
                 grad_: Optional[np.ndarray] = None,    # The downstream operator pass this gradient for you
                 z: Optional[Tensor] = None):           # Which child tensor is passing the gradient
        if not self.requires_grad:
            return "Cannot backpropagate on a tensor that does not require gradients."
        
        if grad_ is None:           # Called only the first time from loss.backward()
            grad_ = np.ones_like(self.data, dtype=np.float32)   # Set a tiny nudge of ones
            
        self.grad += grad_          # Aggregate gradients from that children (but possibly not all yet)
        
        if z is not None:           # NOT called only the first time
            self.children.remove(z) # I heard the gradient from you, no need to wait for you anymore
        
        if self.operation:
            if not self.children:   # Received grad_ from all children, ready to pass grad upstream
                self.operation.backward(self.grad, self)

    def zero_grad(self):
        ''' Sets the gradient of this tensor to zero. '''
        if self.requires_grad:
            self.grad = np.zeros_like(self.data)
            
    # Some basic operators, more operators must be called via their classes below
    def __add__(self, other: Tensor) -> Tensor:
        op = Add()
        return op.forward(self, other)

    def __neg__(self) -> Tensor:
        op = Neg()
        return op.forward(self)

    def __sub__(self, other: Tensor) -> Tensor:
        return self + (-other)
    
    def __mul__(self, other: Tensor) -> Tensor:
        op = Mul()
        return op.forward(self, other)
    
# Paramaters are exact the same as Tensors except they always require gradients to be updated.
class Parameter(Tensor):
    ''' Subclass of Tensor which always tracks gradients. '''
    def __init__(self, data, requires_grad = True, operation = None) -> None:
        super().__init__(data, requires_grad=requires_grad, operation=operation)


# These are just operators
class Add:
    def forward(self, a: Tensor, b: Tensor) -> Tensor:
        # Record which two tensors were added
        self.parents = (a, b)
        
        # Create result tensor z
        requires_grad = a.requires_grad or b.requires_grad
        data = a.data + b.data
        z = Tensor(data, requires_grad=requires_grad, operation=self) 
        
        # Side effects to a, b
        a.children.append(z)
        b.children.append(z)

        return z
        
    def backward(self, dz: np.ndarray, z: Tensor):
        a, b = self.parents
        
        if a.requires_grad:
            da_ = dz
            a.backward(da_, z)
        if b.requires_grad:
            db_ = dz
            b.backward(db_, z)

class Neg:
    def forward(self, a: Tensor) -> Tensor:
        # Record which tensor was negated
        self.parent = a
        
        # Create result tensor z
        requires_grad = a.requires_grad
        data = -a.data
        z = Tensor(data, requires_grad=requires_grad, operation=self) 
        
        # Side effects to a
        a.children.append(z)

        return z
        
    def backward(self, dz: np.ndarray, z: Tensor):
        a = self.parent
        
        if a.requires_grad:
            da_ = -dz
            a.backward(da_, z)

class Mul:
    def forward(self, a: Tensor, b: Tensor) -> Tensor:
        # Record which two tensors were multiplied
        self.parents = (a, b)
        
        # Create result tensor z
        requires_grad = a.requires_grad or b.requires_grad
        data = a.data * b.data
        z = Tensor(data, requires_grad=requires_grad, operation=self) 
        
        # Side effects to a, b
        a.children.append(z)
        b.children.append(z)

        return z
        
    def backward(self, dz: np.ndarray, z: Tensor):
        a, b = self.parents
        
        if a.requires_grad:
            da_ = dz * b.data
            a.backward(da_, z)
        if b.requires_grad:
            db_ = dz * a.data
            b.backward(db_, z)
            
class Exp:
    def forward(self, a: Tensor) -> Tensor:
        # Record which tensor was exponentiated
        self.parent = a
        
        # Create result tensor z
        requires_grad = a.requires_grad
        data = np.exp(a.data)
        z = Tensor(data, requires_grad=requires_grad, operation=self) 
        
        # Side effects to a
        a.children.append(z)

        return z
        
    def backward(self, dz: np.ndarray, z: Tensor):
        a = self.parent
        
        if a.requires_grad:
            da_ = dz * z.data   # since d(exp(a))/da = exp(a)
            a.backward(da_, z)
            
class Log:
    def forward(self, a: Tensor) -> Tensor:
        # Record which tensor was logged
        self.parent = a
        
        # Create result tensor z
        requires_grad = a.requires_grad
        data = np.log(a.data)
        z = Tensor(data, requires_grad=requires_grad, operation=self) 
        
        # Side effects to a
        a.children.append(z)

        return z
        
    def backward(self, dz: np.ndarray, z: Tensor):
        a = self.parent
        
        if a.requires_grad:
            da_ = dz / a.data   # since d(log(a))/da = 1/a
            a.backward(da_, z)

• Tensor 类中包含的信息:
  • data: 该 tensor 的数值.
  • grad: 该 tensor 的梯度 (y 对它的导数).
  • operation: 创建它的 operator. 用来在反向传播时
  • children:

  

  Figure 6: 每个 tensor 包含 data, grad (y 对它的导数), 创建它的 operator 和依赖它的 tensors (红色的 tensor 含有红色区域的信息, etc.)

写一个测试程序:

tensor-backward.py

import mytorch as torch
import numpy as np

x1 = torch.Tensor(np.array([2]), requires_grad=True)
x2 = torch.Tensor(np.array([5]), requires_grad=True)

v1 = torch.Log().forward(x1)
v2 = x1 * x2
v3 = torch.Exp().forward(x2)

v4 = v1 + v2

loss = v4 - v3
print("loss =", loss.data)

loss.backward()

print("x1.grad =", x1.grad)
print("x2.grad =", x2.grad)

运行结果:

loss = [-137.72001192]
x1.grad = [5.5]
x2.grad = [-146.41316]

Figure 7: 每个 operator 包含它所有的 operands (红色的算子含有红色区域的信息, etc.)

References

1.

Baydin AG, Pearlmutter BA, Radul AA, Siskind JM (2018) Automatic differentiation in machine learning: A survey

2.

deep_thoughts (2021) Derive autograd forward and reverse propagation from scratch

3.

Leitão E (2024) Documented and unit tested educational deep learning framework with autograd from scratch.