Matrix Capsules with EM Routing - Optimization

An Optimization of E-Step in EM

The original formula for implementation:

$$\begin{eqnarray} \displaystyle p_c & = & \frac{1}{\sqrt{\Pi_h^H 2\pi S_{ch}^2}}\exp^{-\sum_h^H\frac{\left(V_{ch}^{''}-M_{ch}\right)^2}{2S_{ch}^2}} \\ \displaystyle r_c & = & \frac{a_c^{'}p_c}{\sum_j a_j^{'}p_j} \end{eqnarray}$$

The optimized formula for implementation - numerical stable via softmax:

$$\begin{eqnarray} \displaystyle q_c & = & \log(p_c) = \log\left(\frac{1}{\sqrt{\Pi_h^H 2\pi S_{ch}^2}}\exp^{-\sum_h^H\frac{\left(V_{ch}^{''}-M_{ch}\right)^2}{2S_{ch}^2}}\right) \\ & = & -\sum_h^H\frac{\left(V_{ch}^{''}-M_{ch}\right)^2}{2S_{ch}^2} - \sum_h^H\log(S_{ch}) - \frac{H}{2}\log(2\pi) \end{eqnarray}$$

Since $-\frac{H}{2}\log(2\pi)$ is a constant when H is determined, let:

$$\displaystyle \widetilde{q}_c = -\sum_h^H\frac{\left(V_{ch}^{''}-M_{ch}\right)^2}{2S_{ch}^2} - \sum_h^H\log(S_{ch})$$

Then,

$$\begin{eqnarray} \displaystyle r_c = \frac{a_c^{'}p_c}{\sum_j a_j^{'}p_j} = \frac{a_c^{'}\exp(q_c)}{\sum_j a_j^{'}\exp(q_j)} = \frac{a_c^{'}\exp(\widetilde{q}_c)}{\sum_j a_j^{'}\exp(\widetilde{q}_j)} = \frac{\exp(\log(a_c^{'}) + \widetilde{q}_c)}{\sum_j\exp(\log(a_j^{'}) + \widetilde{q}_j)} = \text{softmax}(\log(\boldsymbol{a}) + \boldsymbol{\widetilde{q}}) \end{eqnarray}$$

Therefore, $r_c$ can be updated via softmax over $\log(\boldsymbol{a})$ and $\boldsymbol{\widetilde{q}}$.