Given function $$g(u) = (Au)^T(Au)$$ where $u \in \mathbb{R}^{n}$ and $A$ is a matrix of dimension $n \times n$, find the gradient $\nabla g(u)$.
I tried to expand everything and got that
$$\frac{\partial g}{\partial u_k} = 2\sum_{i=1}^{n}(A_{i1}u_1+A_{i2}u_2+\cdots +A_{in}u_n)A_{ik}$$
but I am not sure how to convert this expression into a matrix representation of the gradient of $g$. Any hints will be much appreciated.
Let $n$ be a nonnegative integer and $\mathbf{A}$ be a $n \times n$ real matrix. Also, let $g \, : \, \mathbb{R}^n \, \rightarrow \, \mathbb{R}$ be the function defined by:
$$ \forall \mathbf{u} \in \mathbb{R}^n, \; g(\mathbf{u}) = \mathbf{u}^{\top} \mathbf{A}^{\top} \mathbf{A} \mathbf{u}. $$
We consider $\mathbb{R}^n$ as a Euclidean space equipped with its canonical structure. The inner product of two vectors $\mathbf{u}, \mathbf{v} \in \mathbb{R}^n$ is given by: $\left\langle \mathbf{u}, \mathbf{v} \right\rangle = \mathbf{u}^{\top} \mathbf{v}$. Then, by definition, the gradient $\nabla g(\mathbf{u})$ is the unique vector in $\mathbb{R}^n$ such that : $\forall \mathbf{h} \in \mathbb{R}^n, \; Dg(\mathbf{u})(\mathbf{h}) = \left\langle h, \nabla g(\mathbf{u}) \right\rangle$, where $Dg(\mathbf{u}) \, : \, \mathbb{R}^n \, \rightarrow \, \operatorname{End}\left( \mathbb{R}^n \right)$ is the differential of $g$ at $\mathbf{u}$.
You may compute the differential $Dg(\mathbf{u})$ by expanding $g(\mathbf{u} + \mathbf{h})$. Here, we have:
$$ \forall \mathbf{h} \in \mathbb{R}^n, \; Dg(\mathbf{u})(\mathbf{h}) = \mathbf{h}^{\top} \mathbf{A}^{\top} \mathbf{A} \mathbf{u} + \mathbf{u}^{\top} \mathbf{A}^{\top} \mathbf{A} \mathbf{h}. $$
It follows that:
$$ \forall \mathbf{h} \in \mathbb{R}^n, \; Dg(\mathbf{u})(\mathbf{h}) = \left\langle \mathbf{h}, 2 \mathbf{A}^{\top} \mathbf{A} \mathbf{u} \right\rangle. $$
As a conclusion:
$$ \nabla g(\mathbf{u}) = 2 \mathbf{A}^{\top} \mathbf{A} \mathbf{u}. $$
Your expansion yields $$ \frac{\partial g}{\partial u_k} = 2 \sum_{i = 1}^n \sum_{j = 1}^n A_{ij}u_j A_{ik} = 2 \sum_{i = 1}^n A^\intercal_{ki} b_i $$ where $b_i = \sum_{j = 1}^n A_{ij} u_j$. Then by defining a vector $b = \begin{pmatrix} b_1 \cdots b_n \end{pmatrix}^\intercal$ we have $b = Au$ and $\nabla g(u) = 2A^\intercal b = 2A^\intercal Au$.
Just based on the calculation you have performed, it is easy to see that it can equivalently be written as: \begin{equation} \dfrac{\partial g}{\partial u_k}(u) = 2 \langle Au, Ae_k\rangle, \end{equation} where $\langle \cdot , \cdot\rangle$ is the standard inner product of $\mathbb{R}^n$, and $e_k = (0, \dots, 1, \dots, 0)$, with $1$ in the $k^{th}$ place. So, $\nabla g(u)$ is the $1 \times n$ matrix whose $k^{th}$ entry is $2 \langle Au, Ae_k\rangle$.
Another way to arrive at this result is to systematically use the multidimensional chain rule, as explained in Loomis and Sternberg's book Advanced calculus, which is freely available on SHlomo Sternberg's website: http://www.math.harvard.edu/~shlomo/docs/Advanced_Calculus.pdf . In particular, chapter 3, Theorem 8.4, which is a "generalised product rule", valid in infinite dimensional Banach spaces as well. So, as an illustration, the function $g$ you have defined can also be written as \begin{equation} g(u) = \langle Au, Au\rangle. \end{equation} The key thing to notice is that the inner product function $\langle \cdot, \cdot\rangle : \mathbb{R^n} \times \mathbb{R^n} \to \mathbb{R}$ is bilinear (and automatically bounded since we are working on finite dimensional spaces); and bounded bilinear functions can be thought of as "multiplication". So differentiating $g$ is very easy; just apply the product rule as you know from single variable calculus (for the rigorous justification, refer to the book above). Hence, the derivative of $g$ at the point $u$, (which is a linear transformation from $\mathbb{R^n}$ into $\mathbb{R}$) applied to a vector $h \in \mathbb{R^n}$ is simply \begin{align} Dg(u) (h) &= \langle Ah, Au\rangle + \langle Au, Ah\rangle \\ &= 2 \langle Au, Ah \rangle . \end{align} Notice how in the first equal sign I just did the usual "differentiate the first term of product, multiply by second $+$ keep first term of product then multiply by second". In the second I used symmetry of inner product.
So now if you want the $k^{th}$ entry of the $1 \times n$ matrix $\nabla g(u)$, we simply substitute $h = e_k$ in the formula above, because $\dfrac{\partial g}{\partial u_k}(u) = Dg(u) (e_k)$.
Remark: If this application of the chain rule and "generalized product rule" seems very abstract/foreign, then I'd agree with you because when I first learnt it, yes it was slightly confusing. But getting familiar with it is well worth it because it saves you the trouble of keeping track of matrix components, and big summations where indices can get mixed up easily; and with a bit of practice, it becomes just as easy to differentiate.
