Gradient of $\mbox{tr} \left( (AX)^t (AX) \right)$

Question

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

Answer 1

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$.

Answer 2

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.$$

Answer 3

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$.