Differentiable Programming from Scratch

Summary

Differentiable Programming from ScratchDifferentiable programming has become a hot research topic, and not only due to the popularity of machine learning frameworks like TensorFlow, PyTorch, and JAX. Many fields apart from machine learning are finding differentiable programming to be a useful tool for solving optimization problems. In computer graphics, differentiable rendering, differentiable physics, and neural representations are all gaining popularity.This article received an honorable mention in 3Blue1Brown’s Summer of Math Exposition 2!PrerequisitesDifferentiationIt all starts with the definition you learn in calculus class:\[ f^{\prime}(x) = \lim_{h\rightarrow 0} \frac{f(x + h) - f(x)}{h} \]In other words, the derivative computes how much $$f(x)$$ changes when $$x$$ is perturbed by an infinitesimal amount. If $$f$$ is a one-dimensional function from $$\mathbb{R} \mapsto \mathbb{R}$$, the derivative $$f^{\prime}(x)$$ returns the slope of the graph of $$f$$ at $$x$$.However, there’s another perspective that provides better intuition in higher dimensions. If we think of $$f$$ as a map from points in its domain to points in its range, we can think of $$f^{\prime}(x)$$ as a map from vectors based at $$x$$ to vectors based at $$f(x)$$.One-to-OneIn 1D, this distinction is a bit subtle, as a 1D “vector” is just a single number. Still, evaluating $$f^{\prime}(x)$$ shows us how a vector placed at $$x$$ is scaled when transformed by $$f$$. That’s just the slope of $$f$$ at $$x$$.Many-to-OneIf we consider a function $$g(x,y) : \mathbb{R}^2 \mapsto \mathbb{R}$$ (two inputs, one output), this perspective will become clearer. We can differentiate $$g$$ with respect to any particular input, known as a partial derivative:\[ g_x(x,y) = \lim_{h\rightarrow0} \frac{g(x+h,y) - g(x,y)}{h} \]The function $$g_x$$ computes the change in $$g$$ given a change in $$x$$. By combining $$g_x$$ with the corresponding $$g_y$$, we produce the gradient of $$g$$:\[ \nabla g(x,y) = \begin{bmatrix...