Chain Rules for Matrix-Valued Compositions

Why Composition Changes the Derivative

Many objectives are not written directly in terms of the variable being optimized. Instead, a variable is transformed, and a scalar loss is then evaluated on the transformed value. A residual may be formed before it is squared; a matrix may be multiplied on both sides before a norm is taken. The derivative must account for every stage of that construction. The chain rule provides this accounting while keeping the dimensions of each derivative explicit.

The central idea is to differentiate the inner map first, then feed its change into the differential of the outer function. In matrix calculus, this viewpoint is especially useful because it avoids guessing where transposes belong: the gradient is determined by how the differential acts on an arbitrary perturbation.

Definition: Let \(g:\mathbb{R}^p\to\mathbb{R}^q\) be differentiable at \(x\), and let \(f:\mathbb{R}^q\to\mathbb{R}\) be differentiable at \(u=g(x)\). The composite function is \(F=f\circ g\), so \(F(x)=f(g(x))\). Its differential is the linear map \(DF(x)[h]=Df(g(x))[Dg(x)[h]]\). If \(J_g(x)\) is the \(q\times p\) Jacobian of \(g\), and gradients are column vectors, then \(\nabla F(x)=J_g(x)^T\nabla f(g(x))\).

For matrix variables, use the Frobenius inner product \(\langle A,B\rangle_F=\operatorname{tr}(A^TB)\). The gradient of a scalar function \(\phi\) at a matrix \(Y\) is defined by \(D\phi(Y)[H]=\langle\nabla_Y\phi,H\rangle_F\) for every perturbation \(H\) of the same dimensions as \(Y\). Thus a matrix gradient has the dimensions of the matrix being differentiated.

Definition: For a differentiable matrix-valued map \(G\), its differential at \(X\) is a linear map \(DG(X)\) taking a perturbation \(\Delta X\) to the first-order change \(DG(X)[\Delta X]\). Its adjoint \(DG(X)^*\) is defined by \(\langle Z,DG(X)[\Delta X]\rangle_F=\langle DG(X)^*[Z],\Delta X\rangle_F\) for all conformable \(Z\) and \(\Delta X\). If \(\Phi\) is scalar-valued, then \(\nabla_X(\Phi\circ G)=DG(X)^*[\nabla_Y\Phi]\), where \(Y=G(X)\).

The Vector Chain Rule, Derived

The Jacobian convention fixes the orientation in the vector formula: \(J_g(x)\) has one row per component of \(g\) and one column per component of \(x\). The outer gradient is a \(q\)-vector, so multiplication by \(J_g(x)^T\) produces the required \(p\)-vector. The formula follows directly from the first-order definition of differentiability.

Theorem: Suppose \(g:\mathbb{R}^p\to\mathbb{R}^q\) is differentiable at \(x\), and \(f:\mathbb{R}^q\to\mathbb{R}\) is differentiable at \(g(x)\). Then \(F=f\circ g\) is differentiable at \(x\), and \(\nabla F(x)=J_g(x)^T\nabla f(g(x))\).

Proof. Write \(u=g(x)\). Differentiability of \(g\) means that for a perturbation \(h\), \(g(x+h)=u+J_g(x)h+r(h)\), where \(\|r(h)\|/\|h\|\to0\) as \(h\to0\). Set \(k=g(x+h)-u\). Since \(J_g(x)\) is a fixed linear map and \(r(h)=o(\|h\|)\), we have \(k=O(\|h\|)\). Differentiability of \(f\) gives \(f(u+k)=f(u)+\nabla f(u)^Tk+s(k)\), where \(s(k)=o(\|k\|)\).

Substituting the expansion for \(k\) yields \(F(x+h)=F(x)+\nabla f(u)^TJ_g(x)h+\nabla f(u)^Tr(h)+s(k)\). The term \(\nabla f(u)^Tr(h)\) is \(o(\|h\|)\), because \(\nabla f(u)\) is fixed and \(r(h)=o(\|h\|)\). Also \(s(k)=o(\|h\|)\): when \(k\ne0\), \(s(k)/\|h\|=[s(k)/\|k\|][\|k\|/\|h\|]\), whose first factor tends to zero and whose second factor is bounded; when \(k=0\), the remainder is zero. Therefore the composite has derivative \(h\mapsto\nabla f(u)^TJ_g(x)h\). Taking the gradient gives the stated formula. \(\square\)

Worked Example: A Weighted Residual Quadratic

Let \(r(x)=Ax-b\), where \(A=\begin{pmatrix}1&2\\-1&1\end{pmatrix}\), \(b=\begin{pmatrix}1\\0\end{pmatrix}\), and let \(W=\begin{pmatrix}2&0\\0&1\end{pmatrix}\). Define \(F(x)=\tfrac12 r(x)^TWr(x)\). Find the gradient at \(x=\begin{pmatrix}1\\2\end{pmatrix}\).

The inner map has Jacobian \(J_r(x)=A\). For the outer function \(\phi(r)=\tfrac12r^TWr\), the matrix \(W\) is symmetric, so its gradient is \(\nabla_r\phi=Wr\). The chain rule therefore gives \(\nabla_xF=A^TWr\).

At the specified \(x\), the residual is \(r=Ax-b=\begin{pmatrix}1+4\\-1+2\end{pmatrix}-\begin{pmatrix}1\\0\end{pmatrix} =\begin{pmatrix}4\\1\end{pmatrix}\). Then \(Wr=\begin{pmatrix}8\\1\end{pmatrix}\), and \(A^TWr=\begin{pmatrix}1&-1\\2&1\end{pmatrix}\begin{pmatrix}8\\1\end{pmatrix} =\begin{pmatrix}7\\17\end{pmatrix}\). Thus \(\nabla_xF=\begin{pmatrix}7\\17\end{pmatrix}\). As a check, the differential is \(dF=(Wr)^TA\,dx\); its coefficient row is \((8,1)A=(7,17)\), agreeing with the column gradient.

Matrix Chain Rules Through Adjoints

For matrices, the Jacobian can be represented using vectorization, but that may obscure the structure of a transformation. The adjoint formulation expresses the same rule directly in terms of Frobenius inner products. It also tells us how to obtain the gradient: take the outer gradient, pass it backward through the adjoint of the inner differential, and check that the result has the dimensions of the original variable.

Theorem: Let \(G\) map matrices \(X\) to matrices \(Y=G(X)\), and suppose \(G\) is differentiable at \(X\). Let \(\Phi\) be scalar-valued and differentiable at \(Y\). Then \(\nabla_X\Phi(G(X))=DG(X)^*[\nabla_Y\Phi(Y)]\).

Proof. For any perturbation \(\Delta X\), the differential chain rule gives \(D(\Phi\circ G)(X)[\Delta X]=D\Phi(Y)[DG(X)[\Delta X]]\). By the definition of the outer gradient, the right-hand side is \(\langle\nabla_Y\Phi,DG(X)[\Delta X]\rangle_F\). By the definition of the adjoint, this equals \(\langle DG(X)^*[\nabla_Y\Phi],\Delta X\rangle_F\). The matrix multiplying \(\Delta X\) in this inner product is therefore the gradient with respect to \(X\), proving the result. \(\square\)

Worked Example: A Quadratic Loss After Two Matrix Multiplications

Let \(Y=AXB\), with \(A=\begin{pmatrix}2&0\\0&1\end{pmatrix}\), \(B=\begin{pmatrix}1&0\\0&3\end{pmatrix}\), and \(X=\begin{pmatrix}1&2\\0&1\end{pmatrix}\). Define \(\Phi(Y)=\tfrac12\|Y\|_F^2\). Compute \(\nabla_X\Phi(AXB)\).

First, \(Y=AXB=\begin{pmatrix}2&12\\0&3\end{pmatrix}\), so the outer gradient is \(\nabla_Y\Phi=Y\): indeed, \(d\Phi=\langle Y,dY\rangle_F\). The inner differential is \(dY=A(dX)B\). For any \(Z\), \(\langle Z,A(dX)B\rangle_F=\operatorname{tr}(Z^T A(dX)B)=\operatorname{tr}(B Z^T A(dX))=\langle A^TZB^T,dX\rangle_F\). Thus the adjoint sends \(Z\) to \(A^TZB^T\), and the chain rule gives \(\nabla_X\Phi=A^TYB^T\).

Here \(A\) and \(B\) are symmetric, so \(A^TYB^T=\begin{pmatrix}2&0\\0&1\end{pmatrix} \begin{pmatrix}2&12\\0&3\end{pmatrix} \begin{pmatrix}1&0\\0&3\end{pmatrix} =\begin{pmatrix}4&72\\0&9\end{pmatrix}\). This is a \(2\times2\) gradient, matching \(X\). The scalar value is \(\Phi(Y)=\tfrac12(2^2+12^2+0^2+3^2)=78.5\).

Nested Quadratics: Apply the Rule at Each Layer

A nested objective can contain more than one scalar transformation. It is useful to name each intermediate quantity and differentiate one link at a time. For instance, if \(q(x)=x^TQx-c\) and \(F(x)=\tfrac12q(x)^2\), then the outer square contributes a factor \(q(x)\), while the derivative of the inner quadratic contributes its own gradient. The result is their product, not either derivative alone.

Worked Example: Squaring a Quadratic Function

Let \(Q=\begin{pmatrix}1&2\\0&1\end{pmatrix}\), \(c=2\), and \(F(x)=\tfrac12(x^TQx-c)^2\). Find \(\nabla F(x)\) at \(x=\begin{pmatrix}1\\1\end{pmatrix}\).

Name the inner scalar \(q(x)=x^TQx-c\). At the specified point, \(Qx=\begin{pmatrix}3\\1\end{pmatrix}\), so \(x^TQx=4\) and \(q(x)=4-2=2\). Its differential is \(dq=d(x^TQx)=dx^TQx+x^TQ\,dx=((Q+Q^T)x)^Tdx\). Consequently \(\nabla q=(Q+Q^T)x=\begin{pmatrix}2&2\\2&2\end{pmatrix}\begin{pmatrix}1\\1\end{pmatrix} =\begin{pmatrix}4\\4\end{pmatrix}\).

The outer function is \(s(q)=\tfrac12q^2\), whose derivative is \(q\). Applying the scalar chain rule gives \(\nabla F=q\nabla q\), hence \(\nabla F(x)=2\begin{pmatrix}4\\4\end{pmatrix}=\begin{pmatrix}8\\8\end{pmatrix}\). To check by differentials, \(dF=q\,dq=2(4\,dx_1+4\,dx_2)=8\,dx_1+8\,dx_2\), which has the same gradient.

Worked Example: A Nested Residual Quadratic

Let \(r(x)=Ax-b\), \(q(x)=\tfrac12r(x)^TWr(x)\), and \(F(x)=\tfrac12(q(x)-\tau)^2\), where \(A\), \(b\), \(W\), and \(x\) are as in the first worked example, and \(\tau=2\). Find \(\nabla F(x)\) at \(x=\begin{pmatrix}1\\2\end{pmatrix}\).

The first example’s calculations give \(r=\begin{pmatrix}4\\1\end{pmatrix}\) and \(\nabla q=A^TWr=\begin{pmatrix}7\\17\end{pmatrix}\). The value of the inner quadratic is \(q=\tfrac12(4\cdot8+1\cdot1)=\tfrac12(33)=16.5\), so \(q-\tau=14.5\). The outer scalar function \(s(q)=\tfrac12(q-\tau)^2\) has derivative \(s'(q)=q-\tau\). Therefore \(\nabla F=(q-\tau)\nabla q=14.5\begin{pmatrix}7\\17\end{pmatrix} =\begin{pmatrix}101.5\\246.5\end{pmatrix}\).

The differential confirms the factor: \(dF=(q-\tau)dq=14.5(7\,dx_1+17\,dx_2)\). Omitting the factor \(q-\tau\) would give the gradient of \(q\), not the gradient of the nested objective \(F\).

Common Mistakes and Why the Chain Rule Matters

Composition is not merely a notation issue. A missing factor or incorrectly placed transpose changes the direction and scale of a gradient, which can lead to an incorrect optimization step. The differential gives a reliable check: pair the proposed gradient with an arbitrary perturbation and verify that the resulting scalar matches the differential of the objective.

Using the Jacobian in the wrong orientation. With \(J_g\) defined as \(q\times p\), the vector gradient is \(J_g^T\nabla f\), not \(J_g\nabla f\). The correct expression has \(p\) entries, as a gradient with respect to \(x\) must.
Stopping at an intermediate derivative. For \(F(x)=\tfrac12(q(x)-\tau)^2\), writing \(\nabla F=\nabla q\) omits the outer derivative \(q-\tau\). The correction is \(\nabla F=(q-\tau)\nabla q\).
Guessing matrix transposes from memory. For \(Y=AXB\), test with the Frobenius inner product: \(dY=A(dX)B\), so the adjoint is \(Z\mapsto A^TZB^T\). The order reverses when passing through the adjoint.
Returning a gradient with the wrong dimensions. A gradient with respect to \(X\) must have the dimensions of \(X\). If the expression has the dimensions of \(Y\), the inner map’s differential has not yet been passed through its adjoint.

Key takeaway: Differentiate a composition by tracking perturbations through its layers. For vectors, this yields \(J_g^T\nabla f\); for matrices, it yields the adjoint of the inner differential applied to the outer gradient. In either case, the differential provides a direct check of the final gradient.

Check Your Understanding

For each question, identify the inner and outer maps before writing the gradient.

If \(g:\mathbb{R}^p\to\mathbb{R}^q\) has Jacobian \(J_g\) and \(f:\mathbb{R}^q\to\mathbb{R}\), what are the dimensions of \(J_g^T\nabla f\), and why?
For \(F(x)=\tfrac12(x^TQx-c)^2\), what two factors appear in the gradient, and where does each come from?
If \(Y=AXB\) and \(dY=A(dX)B\), what map is the adjoint of the differential under the Frobenius inner product?
Why does differentiating \(\tfrac12(q(x)-\tau)^2\) produce a factor \(q(x)-\tau\)?
What dimension check can reveal that a proposed matrix chain-rule expression is a gradient with respect to the wrong variable?