Differentiation wrt matrix involvoing Khatri-rao product

Question

I got a following minimization problem

$$\min_{\mathbf{X}^{(1)}, \, \mathbf{X}^{(2)}} \;\left\| \mathbf{B} - \mathbf{A} (\mathbf{X}^{(1)} \odot \mathbf{X}^{(2)}) \right\|^{2}_{F},$$

where the matrices $\mathbf{B}\in \mathbb{R}^{100 \times 3}$, $\mathbf{A}\in \mathbb{R}^{100\times 36}$, $\mathbf{X}^{(1)}\in \mathbb{R}^{9 \times 3}$ and $\mathbf{X}^{(2)}\in \mathbb{R}^{4 \times 3}$. The operation $\odot$ refers to the Khatri-rao product.

Given matrices $\mathbf{A}$ and $\mathbf{B}$, my problem is to find out matrices $\mathbf{X}^{(1)}$ and $\mathbf{X}^{(2)}$ such that

$$\mathbb{f} = \left\| \mathbf{B} - \mathbf{A} (\mathbf{X}^{(1)} \odot \mathbf{X}^{(2)}) \right\|^{2}_{F}$$

is minimized.

My idea is to compute the gradient of $\mathbb{f}$ with respect to $\mathbf{X}^{(1)}$ and $\mathbf{X}^{(2)}$ respectively.

My question is, how to do differentiation with respect to a matrix? I have consulted a reference but the situation seems different because it involves a Khatri-rao product in $\mathbb{f}$. Thanks in advance.

$\dfrac{\partial \mathbf{f}}{\partial \mathbf{X}^{(1)}}$ and $\dfrac{\partial \mathbf{f}}{\partial \mathbf{X}^{(2)}} $?

Answer 1

I don't think there is a nice closed-form expression for $\frac{\partial f}{\partial \mathbf X^{(i)}}$, however, I can tell you how you can get to $\frac{\partial f}{\partial \mathbf x^{(i)}}$, where $\mathbf x^{(i)} = {\rm vec}\{\mathbf X^{(i)}\}$. From there, the desired $\frac{\partial f}{\partial \mathbf X^{(i)}}$ is just a reshape.

Essentially, it relies on five ingredients:

1. $\|\mathbf X\|_{\rm F}^2 = \|{\rm vec}\{\mathbf X\}\|_2^2$
2. ${\rm vec}\{\mathbf A \mathbf X \mathbf B\} = (\mathbf B^{\rm T} \otimes \mathbf A)\cdot {\rm vec}\{\mathbf X\}$ where $\otimes$ is the Kronecker product (which implies ${\rm vec}\{\mathbf A \mathbf X \} = (\mathbf I \otimes \mathbf A)\cdot {\rm vec}\{\mathbf X\}$).
3. ${\rm vec}\{\mathbf X_1 \odot \mathbf X_2\} = ([\mathbf I_N \odot \mathbf X_1] \otimes \mathbf I_P)\cdot{\rm vec}\{\mathbf X_2\}$ where $\mathbf X_1$ is $M \times N$ and $\mathbf X_2$ is $P \times N$
4. ${\rm vec}\{\mathbf X_1 \odot \mathbf X_2\} = [\mathbf I_{MN} \odot (\mathbf X_2\cdot [\mathbf I_N \otimes \mathbf 1_{1\times M}])]\cdot{\rm vec}\{\mathbf X_1\}$ where $\mathbf X_1$ is $M \times N$ and $\mathbf X_2$ is $P \times N$
5. $\frac{\partial \|\mathbf b - \mathbf{A}\cdot\mathbf{x}\|_2^2}{\partial \mathbf{x}} = 2\mathbf{A}^{\rm T}(\mathbf{A}\cdot\mathbf{x}-\mathbf b)$

For the proofs:

1. is trivial, both are the sum of all elements squared.
2. can be found in many textbooks and is not hard to show either, cf., e.g., [*]
3. when I needed it I didn't find a proof anywhere so I proved it myself. If you don't mind, I'd like to give my dissertation as a reference where it is Proposition 3.1.2, the proof is in Appendix B.2. There might be textbooks that contain something similar.
4. see 3., covered by the same proposition. This may not be the shortest form but it works.
5. is again straightforward: $\|\mathbf b - \mathbf A \mathbf x\|_2^2 = (\mathbf b^{\rm T} - \mathbf x^{\rm T} \mathbf{A}^T) (\mathbf b - \mathbf{A} \mathbf{x})$, then use the product rule.

Now let us put everything together:

$$ f = \left\|\mathbf B - \mathbf A \left( \mathbf X^{(1)} \odot \mathbf X^{(2)} \right) \right\|_{\rm F}^2 \\ = \left\| \mathbf b - \left(\mathbf I_3 \otimes \mathbf{A}\right) {\rm vec}\{\mathbf X^{(1)} \odot \mathbf X^{(2)}\} \right\|_2^2 \\ = \left\| \mathbf b - \mathbf C_1 \cdot{\rm vec}\{\mathbf X_1\} \right\|_2^2 \\ = \left\| \mathbf b - \mathbf C_2 \cdot{\rm vec}\{\mathbf X_2\} \right\|_2^2 \\ $$

where $\mathbf b = {\rm vec}\{\mathbf B\}$, $\mathbf C_1 = \left(\mathbf I_3 \otimes \mathbf{A}\right) \cdot [\mathbf I_{27} \odot (\mathbf X_2\cdot [\mathbf I_3 \otimes \mathbf 1_{1\times 9}])]$ and $\mathbf C_2 = \left(\mathbf I_3 \otimes \mathbf{A}\right) \cdot([\mathbf I_3 \odot \mathbf X_1] \otimes \mathbf I_4)$. The first step uses 1.+2., the third and fourth line use 3. and 4., respectively.

Finally, using 5. we then have

$$ \frac{\partial f}{\partial \mathbf x^{(i)}} = 2 \mathbf C_i^{\rm T}(\mathbf C_i\mathbf x^{(i)} - \mathbf b)$$

for $i=1, 2$ with $\mathbf C_i$ given above. From this expression, the matrix derivative $\frac{\partial f}{\partial \mathbf X^{(i)}}$ you wanted is given by reshaping $\frac{\partial f}{\partial \mathbf x^{(i)}}$ back into a matrix of appropriate size.

[*] H. Neudecker, “Some theorems on matrix differentiation with special reference to Kronecker matrix products,” Journal of the American Statistical Association, vol. 64, pp. 953–963, 1969.

edit: Since I did a quick test in Matlab to verify it, I thought I might share this bit as well (it uses a function krp to calculate the Khatri-Rao product):

B = randn(100,3);
A = randn(100,36);
X1 = randn(9,3);
X2 = randn(4,3);
f = norm(B - A*krp(X1,X2),'fro')^2;
C1 = kron(eye(3),A)*krp(eye(27),X2*kron(eye(3),ones(1,9)));
C2 = kron(eye(3),A)*kron(krp(eye(3),X1),eye(4));
disp(f - norm(B(:)-C1*X1(:))^2)   % it is zero
disp(f - norm(B(:)-C2*X2(:))^2)   % it is zero