Derivada con respecto a la matriz, de una función multidimensional, utilizada para la optimización del descenso de gradiente.

Considere la función de pérdida

$$\min_{A}\sum_{k=1}^{3} (A^kx_1-x_{k+1})^2,$$ que se optimiza utilizando un método de descenso de gradiente, es decir $$A_{2\times2} = A_{2\times2} - \alpha (\text{ gradient w.r.t. A})_{2\times2}.$$ La matriz de parámetros $A$es un $2\times2$matriz, digamos $$A = \begin{pmatrix}a_1 & a_2 \\ a_3 & a_4\end{pmatrix}.$$ Sin embargo, la función de pérdida es bidimensional ( $2\times1$vector), entonces, ¿cuál sería ese gradiente en este caso, para que los parámetros en $A$para ser actualizado correctamente. Si trato de calcular dicho gradiente, digamos para el término $$A^2x_1,$$ termino con $$\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},$$

que parece ser un$2\times4$matriz. ¡Esto no tiene sentido para mi algoritmo! Todavía no sé cómo actualizar mis parámetros, ya que tengo demasiados elementos en mi degradado. Esperaría un gradiente de la misma dimensión que mi matriz de parámetros.$A$, por lo que las operaciones aritméticas siguen siendo válidas.

¿Cuál es la forma correcta de calcular un gradiente con respecto a una matriz (cuadrada), cuando su función de pérdida es multidimensional? ¿O mi método/interpretación de la actualización de gradiente es incorrecto en primer lugar? Si es así, ¿qué hago en su lugar?