Examples where chain rule does not apply for scalar-by-matrix derivatives?

Question

For your information, I'm coming from a recurrent neural network paper, which employs the concept of chain rule over scalar-by-matrices partial-derivatives. I know there are better tools for this like Hadamard products, but it's not what I'm looking for. So as you might have known, I am totally new to this subject of Matrix Calculus.

The following passages is excerpt from Wikipedia article regarding Matrix Calculus:

The three types of derivatives that have not been considered are those involving vectors-by-matrices, matrices-by-vectors, and matrices-by-matrices. These are not as widely considered and a notation is not widely agreed upon.

The chain rule applies in some of the cases, but unfortunately does not apply in matrix-by-scalar derivatives or scalar-by-matrix derivatives (in the latter case, mostly involving the trace operator applied to matrices). In the latter case, the product rule can't quite be applied directly, either, but the equivalent can be done with a bit more work using the differential identities.

From what I understand and according to the answer in this thread , the reason could be the appearances of tensors that are matrices-by-matrices derivatives. Which there is no convention on how to multiply with other matrices, but the chain rule still applies entries-wise, so $\frac{\partial J}{\partial \Theta} = \frac{\partial J}{\partial a} \frac{\partial a}{\partial z} \frac{\partial z}{\partial \Theta}$ could be understood as $\frac{\partial J}{\partial \Theta_{ij}} = \sum_\alpha\sum_k \frac{\partial J}{\partial a_k} \frac{\partial a_k}{\partial z_\alpha} \frac{\partial z_\alpha}{\partial \Theta_{ij}}$ where $a:\mathbb{R}^m\rightarrow\mathbb{R}^{n}$,$J:\mathbb{R}^n\rightarrow\mathbb{R}$,$z:\mathbb{R}^{n_1\times n_2}\rightarrow\mathbb{R}^{m}$.

Or that there is a counter-example in which chain rule fails for scalar-by-matrix derivations (even for the indices notation).

The latter I haven't able to find a counter-example, and would like some help.

p/s: I notice that for $zUWx=y$ where $z \in \mathbb{R}^{1\times n}$, $U \in \mathbb{R}^{n\times m}$, $W \in \mathbb{R}^{m\times k}$, $x \in \mathbb{R}^{k\times 1}$. Then the chain rule $\frac {\partial y} {\partial U} = \frac {\partial y} {\partial zUW} \frac {\partial zUW} {\partial U}$ where $\frac {\partial y} {\partial zUW}$ is a column vector, but $\frac {\partial y} {\partial W} = \frac {\partial y} {\partial UWx} \frac {\partial UWx} {\partial W}$ where $\frac {\partial y} {\partial UWx}$ is a row vector, which agrees with the fact that there is no conventional on choosing dimension to perform matrix-tensor multiplication. But at least for this case, picking the right dimension to collapse helping the chain rule to agree with the right answers. But I don't know enough about Matrix Calculus to conclude if this is true in general.

Answer 1

$\def\p#1#2{\frac{\partial #1}{\partial #2}}\def\g{\gamma}$As you have noticed, the chain rule can be difficult to apply in Matrix Calculus because it involves higher-order tensors (i.e. matrix-by-vector, vector-by-matrix, and matrix-by-matrix gradients) which are difficult to calculate, awkward to manipulate, and don't fit into standard matrix notation.

Instead I would recommend a differential approach, because the differential of a matrix behaves like a matrix. In particular, it can be written using standard matrix notation and it obeys all of the rules of matrix algebra.

Let's work through the example from your postscript (but make $z$ a column vector).
The calculation for $\left(\p{\gamma}{U}\right)$ is $$\eqalign{ \gamma &= z^TUWx \\ &= z:UWx \\&= zx^TW^T:U \\ d\gamma &= zx^TW^T:dU &\qquad\big({\rm differential}\big) \\ \p{\gamma}{U} &= zx^TW^T &\qquad\big({\rm gradient}\big) \\ }$$ The calculation for $\left(\p{\gamma}{W}\right)$ is very similar $$\eqalign{ \gamma &= U^Tzx^T:W \\ d\gamma &= U^Tzx^T:dW &\qquad\big({\rm differential}\big) \\ \p{\gamma}{W} &= U^Tzx^T &\qquad\big({\rm gradient}\big) \\ \\ }$$

---

In the above, a colon denotes the double-dot product (aka trace product) $$\eqalign{ A:B &= {\rm Tr}(AB^T) \;=\; \sum_{i=1}^m \sum_{j=1}^n A_{ij} B_{ij} \\ A:A &= \big\|A\big\|^2_F\\ }$$ When $\{A,B\}$ are vectors, the trace product corresponds to the usual dot product.

The properties of the underlying trace function allow the terms in such a product to be written in many equivalent ways, e.g. $$\eqalign{ A:B &= B:A = B^T:A^T \\ CA:B &= C:BA^T = A:C^TB \\ }$$ Both matrices in a trace product (as in a Hadamard product) must have the identical dimensions. If you take $\{A,B,C\}$ to be rectangular matrices, then the above manipulations are the only way to move a matrix to the opposite side of the colon while respecting this dimensional constraint.

---

The above calculation also made use of a very general product rule for differentials: $$d(A\star B) = A\star dB + dA\star B$$ where $\star$ can be any product (dot, double-dot/trace, dyadic, elementwise/Hadamard, tensor/Kronecker, etc) and $\{A,B\}$ can be any scalar, vector, matrix, or tensor pair whose dimensions are compatible with that product.