Consider the loss function $$ \min_{A}\sum_{k=1}^{3} (A^kx_1-x_{k+1})^2, $$ which is optimized using a gradient descent method, i.e. $$ A_{2\times2} = A_{2\times2} - \alpha (\text{ gradient w.r.t. A})_{2\times2}. $$ The parameter matrix $A$ is a $2\times2$ matrix, say $$ A = \begin{pmatrix}a_1 & a_2 \\ a_3 & a_4\end{pmatrix}. $$ However, the loss function is two dimensional ($2\times1$ vector), so what would that gradient in this case be, in order for the parameters in $A$ to be updated correctly. If I try to compute said gradient, let's say for the term $$ A^2x_1, $$ I end up with $$ \frac{\partial A^2_{x_1}}{\partial A}= \frac{\partial \begin{pmatrix}(a_1^2+a_2a_3)x_1+(a_1a_2+a_2a_4)x_2 \\ (a_1a_3+a_1a_4)x_1+(a_2a_3+a_4^2)x_2\end{pmatrix}}{\partial A}=\begin{pmatrix}2a_1x_1+a_2x_2 & a_3x_1+a_2x_2 & a_2x_1 & a_2x_2 \\ a_3x_1 & a_3x_2 & a_1x_1+a_2x_2 & a_1x_1 + 2a_4x_2 \end{pmatrix}, $$
which seems to be a $2\times4$ matrix. This is meaningless for my algorithm! I still don't know how to update my parameters, since I've got too many elements in my gradient. I would expect a gradient of same dimension as my parameter matrix $A$, so the arithmetic operations remain valid.
What is the correct way to compute a gradient with respect to a (square) matrix, when your loss function is multi-dimensional? Or is my method/interpretation of gradient update wrong in the first place? If so, what do I do instead?
