Matrix derivative of $f(X^T Y)$ w.r.t X

Question

Given $X\in \mathbb{R}^{m\times n}$, $Y\in \mathbb{R}^{m\times k}$, $X^\top Y=Z\in \mathbb{R}^{n\times k}$, and $f:\mathbb{R}^{n\times k} \to \mathbb{R}$, we have the following:

\begin{equation} f(X^\top Y)=f(Z) \end{equation} What is the derivative of $f$ with respect to $X$, i.e., what is $\frac{\partial f}{\partial X}$?

---

I tried to "expect" $\frac{\partial f}{\partial X}$ just by trying to match the dimension $m,n,k$ as follows:

\begin{equation} \begin{aligned} \frac{\partial f}{\partial X} =\frac{\partial f}{\partial Z}\frac{\partial Z}{\partial X} =Y\left(\frac{\partial f}{\partial Z}\right)^\top \end{aligned} \end{equation}

Since $\frac{\partial Z}{\partial X}$ should be a fourth-order tensor, I tried to calculate $\frac{\partial Z}{\partial X}$ by using $\frac {\partial AXB}{\partial X}=B^T\otimes A$ like this:

\begin{equation} \begin{aligned} \frac{\partial Z}{\partial X}=\frac{\partial (X^\top Y)}{\partial X} {=\left(\frac{\partial ((X^\top) Y)}{\partial (X^\top)}\right)^\top} {=\left(Y^\top \otimes I \right)^\top} {=(Y^\top)^\top \otimes I^\top} {=Y \otimes I} \end{aligned} \end{equation}

Where $\otimes$ is the Kronecker product and $I$ is the $n\times n$ identity matrix. From this, we have:

\begin{equation} \begin{aligned} \frac{\partial f}{\partial X} =\frac{\partial f}{\partial Z}\frac{\partial Z}{\partial X} =\frac{\partial f}{\partial Z}\left(Y \otimes I\right) \end{aligned} \end{equation}

From above equations, this should be true:

\begin{equation} \begin{aligned} \frac{\partial f}{\partial X} =\frac{\partial f}{\partial Z}\left(Y \otimes I\right) =Y\left(\frac{\partial f}{\partial Z}\right)^\top \end{aligned} \end{equation} The last equation seemed counter intuitive. I can not figure it out. Do I calculate $\frac{\partial f}{\partial X}$ correctly ?

If I did it wrong, how should I calculate $\frac{\partial f}{\partial X}$ instead?

Answer 1

$ \def\o{{\tt1}} \def\p{\partial} \def\LR#1{\left(#1\right)} \def\op#1{\operatorname{#1}} \def\trace#1{\op{Tr}\LR{#1}} \def\qiq{\quad\implies\quad} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\c#1{\color{red}{#1}} \def\gradLR#1#2{\LR{\grad{#1}{#2}}} $As you have discovered, 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.

Also, the $\c{\rm Frobenius}$ product is extraordinarily useful in Matrix Calculus $$\eqalign{ A:B &= \sum_{i=1}^m\sum_{j=1}^n A_{ij}B_{ij} \;=\; \trace{A^TB} \\ A:A &= \|A\|^2_F \qquad \big\{{\rm \c{Frobenius}\;norm}\big\}\\ }$$ The properties of the underlying trace function allow the terms in such a product to be rearranged in many equivalent ways, e.g. $$\eqalign{ A:B &= B:A \\ A:B &= A^T:B^T \\ C:\LR{AB} &= \LR{CB^T}:A \;=\; \LR{A^TC}:B \\\\ }$$

---

Using the above notation, the manipulation of your particular function becomes almost mechanical $$\eqalign{ df &= \gradLR fZ:dZ \\ &= \gradLR fZ:\LR{dX^TY} \\ &= \gradLR fZ Y^T:dX^T \\ &= Y\gradLR fZ^T:dX \\ \grad fX &= Y\gradLR fZ^T \\ }$$ Note that there was no need for a tensor-valued gradient in any step, just plain old matrices.

Also note that your initial dimension-matching approach was correct! That's not as crazy an idea as it seems. When dealing with rectangular matrices there is often only one way to fit all the pieces together and it's a useful shortcut. But it won't help when dealing with square matrices.

Answer 2

• Fact 1: A linear function $f：U→V$ can be encoded as an element $T_f$ of the tensor product $V⊗U^*$.
  • For example if $f：ℝ^n→ℝ^m$, then $T_f∈ℝ^m⊗(ℝ^n)^* ≕ ℝ^{m×n}$
• Fact 2: The composition of linear functions corresponds to tensor contraction: If $f：U→V$ and $g：V→W$ are both linear, then there are tensors $T_f∈V⊗U^*$ and $T_g∈W⊗V^*$ such that $g∘f$ corresponds to $T_g⋅T_f∈W⊗U^*$.
  • the contraction is given by linear extension of the formula $(w⊗ϕ)⋅(v⊗ψ) = ϕ(v)⋅(w⊗ψ)$.
  • For Hilbert spaces this is simply $(|w⟩⟨v'|)⋅(|v⟩⟨u|) = ⟨v'∣v⟩⋅|w⟩⟨u|$
• Fact 3: The derivative of a function $f：U→V$ at a point $u∈U$ is, by definition, a linear function $f(u)：U→V，{∆u}↦f(u)(∆u)$
• Fact 4: The derivative of a composition of functions is the composition of the derivatives (chain rule).

With this, your ansatz is completely correct:

• let $g：ℝ^{m×n}→ℝ^{n×k}，X ↦ X^⊤Y ≕ Z$
• let $f：ℝ^{m×k}→ℝ，Z ↦ f(Z)$

then, by definition and chain rule, the derivative of $f∘g$ is the composition

$$ \underbrace{[f∘g](X)}_{∈(ℝ^{m×n}→ℝ)} = \underbrace{[f](g(X))}_{∈(ℝ^{n×k}→ℝ)}\; ∘ \underbrace{[g](X)}_{∈(ℝ^{m×n}→ℝ^{n×k})}$$

which, by fact 2 we can translate into a tensor contraction

$$\underbrace{T_{[f∘g](X)}}_{∈ℝ⊗(ℝ^{m×n})^*} = \underbrace{T_{[f](g(X))}}_{∈ℝ⊗(ℝ^{n×k})^*}\; ⋅ \underbrace{T_{[g](X)}}_{∈ℝ^{n×k}⊗(ℝ^{m×n})^*} $$

Now, sticking to your convention that $AXB^⊤ = (B^⊤⊗A)⋅X$ (another layout convention would give $(A⊗B)⋅X$), and using $X^⊤ = _{m,n}⋅X$ where $_{m,n}$ is the 4-dimensional transpose tensor, which encodes the linear mapping $ℝ^{m×n}→ℝ^{n×m}, A↦A^⊤$, we have

$$\frac{∂f(X^⊤Y)}{∂X} = \frac{∂f(Z)}{∂Z}\frac{X^⊤Y}{∂X^⊤}\frac{∂X^⊤}{∂X} = \frac{∂f(Z)}{∂Z}⋅(Y^⊤⊗_m)⋅_{m,n}$$

And that's already it. The main difficulty about the last equation is remembering to which tensor product space the terms belong, whence over which of the axes you need to contract over when multiplying them. “⋅” is no longer just a regular matrix-matrix product, but a tensor contraction.