I am trying to calculate the gradient of the following function

$$f(X) = \mbox{tr} \left( (AX)^t (AX) \right)$$

Chain's rule gives

$$\nabla_X(f(X)) = \nabla_X (\mbox{tr}(AX))\nabla_x(AX)$$

However, I'm having trouble with those two derivatives.

What is $\nabla_X tr(AX)$? Is it $A^t$? I did the math and obtained that $\frac{\partial(tr(AX))}{\partial x_{ij}} = a_{ji}$, but I'm not sure... And also what is $\nabla_X AX$? Is it simply $A$? I tried differentiating this but failed to see if this holds or not.

Thanks in advance


The gradient $\nabla_{X}f$ is defined as the vector in $\mathcal{M}_{n}(\mathbb{R})$ such that :

$$ f(X+H) = f(X) + \left\langle \nabla_{X}f, H \right\rangle + o(\Vert H \Vert) $$

where $\left\langle \cdot,\cdot \right\rangle$ is the usual inner product on $\mathcal{M}_{n}(\mathbb{R})$ (i.e. $\left\langle A,B \right\rangle = \mathrm{tr}(A^{\top}B)$). By expanding $f(X+H)$, you get :

$$ f(X+H) = f(X) + \underbrace{2\mathrm{tr}(H^{\top}A^{\top}AX)}_{= \; \left\langle 2A^{\top}AX,H \right\rangle} + \underbrace{\mathrm{tr}(H^{\top}A^{\top}AH)}_{= \; o(\Vert H \Vert)} $$

By identification : $\nabla_{X}f = 2A^{\top}AX$.


I present below an easy way to compute the derivative.

Compute the difference: $$ \Delta f = f(X +\Delta X)- f(X),$$ which gives

$$ \Delta f = \text{tr}(\Delta X^T A^TA X) + \text{tr}( X^T A^TA \Delta X) + \text{tr}(\Delta X^T A^TA \Delta X).$$

Eliminate the second order terms. You will have the differential

$$ \partial f = \text{tr}(\partial X^T A^TA X) + \text{tr}(X^T A^TA \partial X).$$

By using Trace properties, we have

$$ \partial f = 2\cdot \text{tr}(X^T A^TA \partial X).$$

We know that if $\partial f = \text{tr}(Y\cdot \partial X)$, then $\cfrac{\partial f }{\partial X} = Y^T.$

Thus, in your case

$$\cfrac{\partial f }{\partial X} = 2 A^TAX.$$


The derivative $\partial f / \partial x$ is equal to $$ \lim_{h\to 0} \frac{f(x+h)-f(x)}{h} $$

When it comes to a function like yours : $f(X)=tr((AX)^\top AX)=tr(X^\top A^\top A X)$, you have one partial derivative per coordinate in $X$. Then, for $X \in \mathbb R^{n\times m}$, the derivative is given by $$ \lim_{h\to 0} \sum_{i=1}^n\sum_{j=1}^m \frac{tr\left((X+he_ie_j^\top)^\top A^\top A (X+he_ie_j^\top)\right) - tr(X^\top A^\top A X)}{h}e_ie_j^\top $$ where $e_i$ is the $i$th standard basis vector ($i$th column of $I$).

For solving this, you want to make $h$ and $e_ie_j^\top$ disapear. After some manipulations you should obtain $2 A^\top A X$.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.