A note on the multivariable chain rule

Teaching machine learning, I have found that many students are unprepared for the level of vector calculus required, particularly when it comes to doing backprop calculations, which require the chain rule. Here I attempt to review the chain rule for computing gradients, and related concepts such as derivatives and Jacobians, in a cohesive way.

Review: single-variable chain rule

In single-variable calculus, the chain rule is often written

where $z$$z$ is some function of $y$$y$, which is itself a function of $x$$x$. Students can remember this easily by "canceling terms" (although they are reminded that this is not technically correct).

We would like to extend this type of relationship to vector-valued functions of several variables. But first, we need to reinterpret the derivative.

The derivative as a linear map

The classical derivative is a number, denoted $f^{\prime }(x)$$f'(x)$, which corresponds to the instantaneous rate of change of $f$$f$ at $x$$x$ if you zoom in infinitely close to $x$$x$. This is conceptually equivalent to drawing a tangent line to $f$$f$ at $x$$x$; the slope of the tangent line is the derivative $f^{\prime }(x)$$f'(x)$. The tangent line can be written as the graph of the function

This leads to a more general notion of a derivative as the best local linear approximation of the change of a function with respect to its input. This is typically called the total derivative or differential of a function, or sometimes just derivative. For a function $f:X\to Y$$f : X \to Y$ where $X$$X$ and $Y$$Y$ are normed vector spaces, we say that $f^{\prime }(x):X\to Y$$f'(x) : X \to Y$ is the total derivative of $f$$f$ at $x$$x$ if it is linear and satisfies

assuming such a map exists. In the special case of $f:\mathbb{R}\to \mathbb{R}$$f : \R \to \R$, we abuse notation by writing $f^{\prime }(x)$$f'(x)$ for both the scalar and for the linear map, but it holds that

so we recover $f^{\prime }(x)(\mathrm{\Delta })=f^{\prime }(x)\mathrm{\Delta }$$f'(x)(\Delta) = f'(x) \Delta$.

Jacobian and gradient

For vector-valued functions $f:{\mathbb{R}}^m\to {\mathbb{R}}^n$$f : \R^m \to \R^n$, the numerical implementation of the total derivative turns out to be the Jacobian matrix $J_f(x)\in {\mathbb{R}}^{n\times m}$$J_f(x) \in \R^{n \times m}$, which contains all the partial derivatives:

In the special case of scalar-valued functions $f:{\mathbb{R}}^m\to \mathbb{R}$$f : \R^m \to \R$, we can define the gradient $\mathrm{\nabla }f(x)\in {\mathbb{R}}^m$$\nabla f(x) \in \R^m$ by

Note that if we interpret $\nabla f(x)$ as a column vector, i.e. an $m\times 1$$m \times 1$ matrix, then

Or, more generally, for $f:{\mathbb{R}}^m\to {\mathbb{R}}^n$$f : \R^m \to \R^n$,

This may seem like a coincidence if your only knowledge of the transpose operator is that it flips a matrix across the diagonal. The conceptual interpretation of transposition (operating on a vector) is that it takes a vector and produces a linear functional (also called linear form), i.e. a linear function that acts on other vectors to produce a scalar. Namely, for any vector $v$$v$ we can define the linear functional ${\ell }_v(w)=v^{\mathrm{\top }}w=\sum _i{v}_i{w}_i$$\ell_v(w) = v\T w = \sum_{i} v_iw_i$. There is essentially no difference between ${\ell }_v$$\ell_v$ and $v^{\mathrm{\top }}$$v\T$; while they are not literally the same "type" of object, they act the same way, and there is a one-to-one correspondence between them.

From the equations above, we see that the gradient $\nabla f(x)$ is the vector corresponding to the directional derivative functional

which quantifies how $f$$f$ changes along the direction $v$$v$ when starting from $x$$x$. Note that the directional derivative – considered as a function of the direction – coincides with the total derivative of $f$$f$ when $f$$f$ is scalar-valued.

Generalization: gradients in Hilbert spaces

In fact, the gradient can be defined more generally in a Hilbert space $\mathcal{H}$$\H$, which is a vector space equipped with an inner product $⟨\cdot ,\cdot {⟩}_{\mathcal{H}}$$\inner{\cdot}{\cdot}$ that generalizes the dot product from ${\mathbb{R}}^n$$\R^n$. Any Hilbert space is also a normed vector space with norm $\parallel v{\parallel }_{\mathcal{H}}=\sqrt{⟨v,v{⟩}_{\mathcal{H}}}$$\|v\|_\H = \sqrt{\inner{v}{v}}$.

The Riesz representation theorem guarantees that, for any linear functional $\ell :\mathcal{H}\to \mathbb{R}$$\ell : \H \to \R$, there exists a unique vector $v_{\ell }\in \mathcal{H}$$v_\ell \in \H$ such that $\ell (x)=⟨x,v_{\ell }{⟩}_{\mathcal{H}}$$\ell(x) = \inner{x}{v_\ell}$. Thus, the gradient of $f:\mathcal{H}\to \mathbb{R}$$f : \H \to \R$ can be defined as the unique vector $\mathrm{\nabla }f(x)\in \mathcal{H}$$\nabla f(x) \in \H$ satisfying

The multivariable chain rule

The chain rule generalizes nicely to the total derivative. Let $f:X\to Y$$f : X \to Y$ and $g:Y\to Z$$g : Y \to Z$, then their composition $g\circ f:X\to Z$$g \circ f : X \to Z$ has derivative

Observe that the order of composition of the derivatives matches the order of composition of the original functions. That is, $g\circ f$$g \circ f$ first applies $f$$f$ then $g$$g$, and similarly $(g\circ f{)}^{\prime }$$(g \circ f)'$ first applies $f^{\prime }$$f'$ then $g^{\prime }$$g'$.

This can also be expressed in terms of Jacobians:

Note that if $y=f(x)$$y = f(x)$ and $z=g(y)$$z = g(y)$, and we abuse notation by writing $\frac{\mathrm{d}\cdot }{\mathrm{d}\cdot }$$\dd{\,\cdot}{\,\cdot}$ for the Jacobian, the above equation becomes

just as before.

Example: ordinary least squares

A prototypical example in ML classes is ordinary least squares:

As we know, the solution can be found by identifying critical points of $L$, i.e. $w$$w$ where $\mathrm{\nabla }L(w)=0$$\nabla L(w) = 0$. The gradient can be found by expanding the square and applying a couple of common matrix calculus identities:

so we arrive at the normal equations $X^{\mathrm{\top }}Xw=X^{\mathrm{\top }}y$$X\T X w = X\T y$.

However, one can also use the chain rule to obtain $\mathrm{\nabla }L$$\nabla L$. Let's denote $z=f(w)=Xw-y$$z = f(w) = Xw-y$ and $L=g(z)=\parallel z{\parallel }_2^2$$L = g(z) = \|z\|_2^2$. We know $\mathrm{\nabla }g(z)=2z$$\nabla g(z) = 2z$, so $\frac{\mathrm{d}L}{\mathrm{d}z}=J_g(z)=2z^{\mathrm{\top }}$$\dd{L}{z} = J_g(z) = 2z\T$, and it should be apparent that $\frac{\mathrm{d}z}{\mathrm{d}w}=J_f=X$$\dd{z}{w} = J_f = X$, as this is the best linear approximation, so

Then

which is exactly what we got before.