Why Linear Trace Gradients Need Care
A scalar expression involving a matrix can look deceptively simple. For example, \(\operatorname{tr}(AX)\) is a scalar when the product is square, but its gradient must have the same dimensions as the variable \(X\), not necessarily the dimensions of \(A\). The position of a transpose matters, and so does the order of the fixed matrices.
The key is to define the gradient by how the function changes under a small matrix perturbation. Once that definition is in place, trace identities turn each differential into a Frobenius inner product. This tutorial derives the gradients of three basic linear trace functions and checks them in concrete rectangular examples.
Here \(dX\) denotes an arbitrary first-order perturbation, with the same dimensions as \(X\). Equivalently, the directional derivative in a direction \(H\in\mathbb{R}^{m\times n}\) is \(Df(X)[H]=\operatorname{tr}((\nabla_X f)^T H)\). The gradient is determined uniquely: if two matrices represent the same differential for every \(H\), their difference has zero Frobenius inner product with every \(H\), including itself, and therefore is zero.
Two Basic Linear Trace Functions
First suppose \(X,A\in\mathbb{R}^{m\times n}\) are the same size and \(A\) is fixed. The scalar \(\operatorname{tr}(A^T X)\) is the Frobenius inner product of \(A\) and \(X\). Its differential directly identifies the gradient.
This formula also follows entry by entry: \(\operatorname{tr}(A^T X)=\sum_{i=1}^{m}\sum_{j=1}^{n}A_{ij}X_{ij}\), so the partial derivative with respect to \(X_{ij}\) is \(A_{ij}\). The differential argument is more useful when expressions are written as products and traces.
Now let \(X\in\mathbb{R}^{m\times n}\) and \(A\in\mathbb{R}^{n\times m}\). Then \(AX\) is \(n\times n\), so its trace is defined. The gradient has dimensions \(m\times n\), while \(A\) has dimensions \(n\times m\); the transpose is essential.
To prove the formula, perturb \(X\) by \(dX\). Linearity gives \(df=\operatorname{tr}(A\,dX)\). The product \(A\,dX\) is \(n\times n\). By the definition of the Frobenius inner product, \(\operatorname{tr}(A\,dX)=\operatorname{tr}((A^T)^T dX)\). Thus the gradient is \(A^T\), which is \(m\times n\), as required. Since this equality holds for every perturbation, it proves the claimed gradient. \(\square\)
Worked Example: Gradient of tr(ATX)
Let \(X,A\in\mathbb{R}^{2\times 3}\), with \(A=\begin{pmatrix}1&-2&0\\3&4&5\end{pmatrix}\) and \(X=\begin{pmatrix}2&1&-1\\0&3&2\end{pmatrix}\). Find the value of \(f(X)=\operatorname{tr}(A^T X)\), its gradient, and its directional derivative in the direction \(H=\begin{pmatrix}1&0&2\\-1&2&0\end{pmatrix}\).
Both \(A^T\) and \(X\) are \(3\times2\) and \(2\times3\), respectively, so \(A^T X\) is \(3\times3\) and its trace exists. Using the entrywise inner product, \(f(X)=1(2)+(-2)(1)+0(-1)+3(0)+4(3)+5(2)=2-2+0+0+12+10=22\).
The matrices \(A\) and \(X\) have the same dimensions, so the first formula gives \(\nabla_X f=A\), a \(2\times3\) matrix. The directional derivative is \(Df(X)[H]=\operatorname{tr}(A^T H)=1(1)+(-2)(0)+0(2)+3(-1)+4(2)+5(0)=1+0+0-3+8+0=6\).
As a direct check, the inner product of \(A\) and \(H\) is also 6. Thus the value of the function is 22, while its rate of change in the specified direction is 6; these are different quantities.
A Product Inside the Trace
For a product with fixed matrices on both sides of \(X\), dimensions determine the valid expression and the order in the gradient. Let \(X\in\mathbb{R}^{m\times n}\), \(A\in\mathbb{R}^{p\times m}\), and \(B\in\mathbb{R}^{n\times p}\). Then \(AXB\in\mathbb{R}^{p\times p}\), and \(\operatorname{tr}(AXB)\) is scalar.
Proof. With \(A\) and \(B\) fixed, the differential is \(df=\operatorname{tr}(A\,dX\,B)\). The factors have sizes \(p\times m\), \(m\times n\), and \(n\times p\), so this trace is defined. The trace identity \(\operatorname{tr}(UVW)=\operatorname{tr}(WUV)\), when the products are dimensionally compatible, gives \(df=\operatorname{tr}(BA\,dX)\). Here \(BA\) is \(n\times m\), making \(BA\,dX\) an \(n\times n\) matrix. Finally, \(\operatorname{tr}(BA\,dX)=\operatorname{tr}(((BA)^T)^T dX)\). Matching this with the gradient definition gives \(\nabla_X f=(BA)^T=A^T B^T\), which has dimensions \(m\times n\). \(\square\)
Worked Example: Gradient of tr(AXB) with Rectangular X
Take \(X\in\mathbb{R}^{2\times3}\), \(A=\begin{pmatrix}1&2\\-1&1\end{pmatrix}\in\mathbb{R}^{2\times2}\), and \(B=\begin{pmatrix}1&0\\2&-1\\0&3\end{pmatrix}\in\mathbb{R}^{3\times2}\). Find \(\operatorname{tr}(AXB)\) and its gradient with respect to \(X\).
The dimensions give \(AXB\in\mathbb{R}^{2\times2}\). First, \(BA=\begin{pmatrix}1&2\\3&3\\-3&3\end{pmatrix}\), so the gradient formula gives \(\nabla_X f=(BA)^T=\begin{pmatrix}1&3&-3\\2&3&3\end{pmatrix}\), a \(2\times3\) matrix matching \(X\).
To calculate the trace, use the cyclic trace identity: \(\operatorname{tr}(AXB)=\operatorname{tr}(BAX)\). The diagonal entries of the \(3\times3\) matrix \(BAX\) are obtained by multiplying each row of \(BA\) by the corresponding column of \(X\). For \(X=\begin{pmatrix}1&2&0\\-1&1&2\end{pmatrix}\), they are \(1(1)+2(-1)=-1\), \(3(2)+3(1)=9\), and \((-3)(0)+3(2)=6\). Therefore \(\operatorname{tr}(AXB)=-1+9+6=14\).
For a directional check, let \(H=\begin{pmatrix}1&0&-1\\2&1&0\end{pmatrix}\). The predicted derivative is the Frobenius inner product \(\langle\nabla_X f,H\rangle_F=1(1)+3(0)+(-3)(-1)+2(2)+3(1)+3(0)=1+0+3+4+3+0=11\). This confirms the gradient’s orientation and gives the rate of change in direction \(H\).
Another Rectangular Check: tr(AX)
Worked Example: Gradient of tr(AX)
Let \(X=\begin{pmatrix}1&2\\0&-1\\3&4\end{pmatrix}\in\mathbb{R}^{3\times2}\) and \(A=\begin{pmatrix}2&1&0\\-1&3&2\end{pmatrix}\in\mathbb{R}^{2\times3}\). Compute \(\operatorname{tr}(AX)\) and \(\nabla_X\operatorname{tr}(AX)\).
The product \(AX\) is \(2\times2\): \(AX=\begin{pmatrix}2&3\\5&3\end{pmatrix}\). For example, its lower-right entry is \((-1)(2)+3(-1)+2(4)=-2-3+8=3\). Thus \(\operatorname{tr}(AX)=2+3=5\).
Since \(A\) is \(2\times3\) and \(X\) is \(3\times2\), the gradient formula gives \(\nabla_X\operatorname{tr}(AX)=A^T=\begin{pmatrix}2&-1\\1&3\\0&2\end{pmatrix}\), which is \(3\times2\), as \(X\) is.
For a numerical check, use \(H=\begin{pmatrix}1&-1\\2&0\\0&3\end{pmatrix}\). The gradient inner product is \(2(1)+(-1)(-1)+1(2)+3(0)+0(0)+2(3)=2+1+2+0+0+6=11\). Directly, \(AH=\begin{pmatrix}4&-2\\5&7\end{pmatrix}\), whose trace is \(4+7=11\). This agrees with \(Df(X)[H]=\operatorname{tr}(AH)\).
Common Mistakes and a Reliable Check
The formulas are short, but a transpose error changes the dimensions or assigns coefficients to the wrong entries. Check dimensions before simplifying, and check the answer by asking whether the proposed gradient has the same shape as the variable.
- Using \(A\) rather than \(A^T\) for \(\operatorname{tr}(AX)\). When \(A\) is \(n\times m\) and \(X\) is \(m\times n\), the gradient must be \(m\times n\). The transpose \(A^T\) has precisely that size.
- Ignoring the order of fixed factors. In \(\operatorname{tr}(AXB)\), cycling the trace gives \(\operatorname{tr}(BA\,dX)\), not \(\operatorname{tr}(AB\,dX)\). The resulting gradient is \((BA)^T\).
- Treating a trace as entrywise multiplication without checking indices. The expression \(\operatorname{tr}(AX)\) pairs \(A_{ij}\) with \(X_{ji}\), so its coefficient matrix with the dimensions of \(X\) is \(A^T\).
- Confusing a function value with a derivative. A trace evaluated at \(X\) is a scalar value; the directional derivative is a different scalar, computed from the gradient and a specified direction \(H\).
A dependable procedure is to write the dimensions of every factor, take the differential by replacing \(X\) with \(dX\), use a trace identity only when the products remain compatible, and match the result to \(\operatorname{tr}((\nabla_X f)^T dX)\). This turns the transpose choice into a consequence of the definition rather than a memorized guess.
Check Your Understanding
For each question, track dimensions as well as the derivative formula.
- If \(X,A\in\mathbb{R}^{4\times2}\), what is \(\nabla_X\operatorname{tr}(A^T X)\), and what are its dimensions?
- If \(X\in\mathbb{R}^{2\times5}\) and \(C\in\mathbb{R}^{5\times2}\) is fixed, find the dimensions and formula for \(\nabla_X\operatorname{tr}(CX)\).
- For \(X\in\mathbb{R}^{3\times4}\), \(A\in\mathbb{R}^{2\times3}\), and \(B\in\mathbb{R}^{4\times2}\), write the gradient of \(\operatorname{tr}(AXB)\) in terms of \(A\) and \(B\), and state its dimensions.
- Explain why the gradient of \(\operatorname{tr}(AX)\), with \(A\in\mathbb{R}^{n\times m}\) and \(X\in\mathbb{R}^{m\times n}\), cannot generally be \(A\).
- For a proposed gradient \(G\), describe how a direction \(H\) can be used to check a formula for a linear trace function.