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

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.

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

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