E step

$$P(\hat{B},B,Z|\rho, \pi) = \prod_{i=1}^{n} \prod_{k=0}^{K}\left[\pi_{k}N(\hat{b}_{i}; b_{i}, V)N(b_{i}; 0, U_{k})\right]^{\mathbb{I}(z_{i}=k)}$$

$$\mathbb{E}_{Z,B|\hat{B}} \log P(\hat{B},B,Z|\rho, \pi) = \sum_{i=1}^{n} \sum_{k=0}^{K} P(z_{i}=k|\hat{b}_{i})\left[ \log \pi_{k} + \mathbb{E}_{B|\hat{B}}(\log N(\hat{b}_{i}; b_{i}, V)) + \mathbb{E}_{B|\hat{B}}(\log N(b_{i}; 0, U_{k})) \right]$$

$$\begin{align*} \log N(\hat{b}_{i}; b_{i}, V) + \log N(b_{i}; 0, U_{k}) &= -\frac{p}{2}\log 2\pi -\frac{1}{2}\log |V| - \frac{1}{2}(\hat{b}_{i}-b_{i})^{T}V^{-1}(\hat{b}_{i}-b_{i}) -\frac{p}{2}\log 2\pi -\frac{1}{2}\log |U_{k}| - \frac{1}{2}b_{i}^{T}U_{k}^{-1}b_{i} \\ &= -p\log 2\pi -\frac{1}{2}\log |U_{k}| -\frac{1}{2}\log |V| - \frac{1}{2}\hat{b}_{i}^{T}V^{-1}\hat{b}_{i} + \hat{b}_{i}^{T}V^{-1}b_{i} -\frac{1}{2}b_{i}^{T}V^{-1}b_{i} - \frac{1}{2}b_{i}^{T}U_{k}^{-1}b_{i} \\ \mathbb{E}_{b_{i}|\hat{b}_{i}}\left[ \log N(\hat{b}_{i}; b_{i}, V) + \log N(b_{i}; 0, U_{k}) \right] &= -p\log 2\pi -\frac{1}{2}\log |U_{k}| -\frac{1}{2}\log |V| - \frac{1}{2}\hat{b}_{i}^{T}V^{-1}\hat{b}_{i} + \hat{b}_{i}^{T}V^{-1}\mathbb{E}(b_{i}|\hat{b}_{i}) -\frac{1}{2}tr\left(V^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right) - \frac{1}{2}tr\left(U_{k}^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right) \end{align*}$$ $$V = \left( \begin{matrix}1 & \rho \\ \rho & 1 \end{matrix} \right)$$ Let $\mu_{i} = \mathbb{E}(b_{i}|\hat{b}_{i})$ $$\begin{align*} \mathbb{E}_{b_{i}|\hat{b}_{i}}\left[ \log N(\hat{b}_{i}; b_{i}, V) + \log N(b_{i}; 0, U_{k}) \right] &= -2\log 2\pi -\frac{1}{2}\log |U_{k}| -\frac{1}{2}\log(1-\rho^2) - \frac{1}{2(1-\rho^2)}\left(\hat{b}_{i1}^2 + \hat{b}_{i2}^2 -2\hat{b}_{i1} \mu_{i1} -2\hat{b}_{i2} \mu_{i2} + \mathbb{E}(b_{i1}^2|\hat{b}_{i}) + \mathbb{E}(b_{i2}^2|\hat{b}_{i}) - 2\hat{b}_{i1}\hat{b}_{i2}\rho + 2 \hat{b}_{i1}\mu_{i2}\rho +2\hat{b}_{i2}\mu_{i1}\rho - 2\rho\mathbb{E}(b_{i1}b_{i2}|\hat{b}_{i}) \right) - \frac{1}{2}tr\left(U_{k}^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right) \end{align*}$$

$$\gamma_{Z_{i}}(k) = P(z_{i}=k|X_{i}) = \frac{\pi_{k}N(x_{i}; 0, V+U_{k})}{\sum_{k'=0}^{K}\pi_{k'}N(x_{i}; 0, V + U_{k'})}$$

M step

$V$: $$\begin{align*} f(V^{-1}) = \sum_{i=1}^{n} \sum_{k=0}^{K} \gamma_{Z_{i}}(k)\left[ -p\log 2\pi -\frac{1}{2}\log |U_{k}| -\frac{1}{2}\log |V| - \frac{1}{2}\hat{b}_{i}^{T}V^{-1}\hat{b}_{i} + \hat{b}_{i}^{T}V^{-1}\mathbb{E}(b_{i}|\hat{b}_{i}) -\frac{1}{2}tr\left(V^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right) - \frac{1}{2}tr\left(U_{k}^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right) \right] \end{align*}$$

$$\begin{align*} f(V^{-1})' &= \sum_{i=1}^{n} \sum_{k=0}^{K} \gamma_{Z_{i}}(k)\left[ \frac{1}{2}V - \frac{1}{2}\hat{b}_{i}\hat{b}_{i}^{T} + \mathbb{E}(b_{i}|\hat{b}_{i})\hat{b}_{i}^{T} - \frac{1}{2} \mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right] = 0 \\ \frac{1}{2}Vn &= \sum_{i=1}^{n} \sum_{k=0}^{K} \gamma_{Z_{i}}(k)\left[\frac{1}{2}\hat{b}_{i}\hat{b}_{i}^{T} - \mathbb{E}(b_{i}|\hat{b}_{i})\hat{b}_{i}^{T} + \frac{1}{2} \mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right] \\ \hat{V} &= \frac{1}{n} \sum_{i=1}^{n} \left[\hat{b}_{i}\hat{b}_{i}^{T} - 2\mathbb{E}(b_{i}|\hat{b}_{i})\hat{b}_{i}^{T} + \mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right] \end{align*}$$

$\rho$: $$f(\rho) = \sum_{i=1}^{n} \sum_{k=1}^{K} \gamma_{Z_{i}}(k)\left[-2\log 2\pi -\frac{1}{2}\log |U_{k}| -\frac{1}{2}\log(1-\rho^2) - \frac{1}{2(1-\rho^2)}\left(\hat{b}_{i1}^2 + \hat{b}_{i2}^2 -2\hat{b}_{i1} \mu_{i1} -2\hat{b}_{i2} \mu_{i2} + \mathbb{E}(b_{i1}^2|\hat{b}_{i}) + \mathbb{E}(b_{i2}^2|\hat{b}_{i}) - 2\hat{b}_{i1}\hat{b}_{i2}\rho + 2 \hat{b}_{i1}\mu_{i2}\rho +2\hat{b}_{i2}\mu_{i1}\rho - 2\rho\mathbb{E}(b_{i1}b_{i2}|\hat{b}_{i}) \right) - \frac{1}{2}tr\left(U_{k}^{-1}\mathbb{E}(b_{i}b_{i}^{T}|\hat{b}_{i}) \right)\right]$$

$$\begin{align*} f(\rho)' = \sum_{i=1}^{n} \sum_{k=1}^{K} \gamma_{Z_{i}}(k)\left[ \frac{\rho}{1-\rho^2} -\frac{\rho}{(1-\rho^2)^2}\left( \hat{b}_{i1}^2 + \hat{b}_{i2}^2 -2\hat{b}_{i1} \mu_{i1} -2\hat{b}_{i2} \mu_{i2} + \mathbb{E}(b_{i1}^2|\hat{b}_{i}) + \mathbb{E}(b_{i2}^2|\hat{b}_{i}) \right) -\frac{\rho^2+1}{(1-\rho^2)^2}\left( -\hat{b}_{i1}\hat{b}_{i2} + \hat{b}_{i1}\mu_{i2} +\hat{b}_{i2}\mu_{i1} - \mathbb{E}(b_{i1}b_{i2}|\hat{b}_{i}) \right) \right] &= 0 \\ \rho(1-\rho^2)n - \rho \sum_{i=1}^{n} \sum_{k=1}^{K} \gamma_{Z_{i}}(k) \left( \hat{b}_{i1}^2 + \hat{b}_{i2}^2 -2\hat{b}_{i1} \mu_{i1} -2\hat{b}_{i2} \mu_{i2} + \mathbb{E}(b_{i1}^2|\hat{b}_{i}) + \mathbb{E}(b_{i2}^2|\hat{b}_{i}) \right) - (\rho^2 + 1) \sum_{i=1}^{n} \sum_{k=1}^{K} \gamma_{Z_{i}}(k)\left( -\hat{b}_{i1}\hat{b}_{i2} + \hat{b}_{i1}\mu_{i2} +\hat{b}_{i2}\mu_{i1} - \mathbb{E}(b_{i1}b_{i2}|\hat{b}_{i}) \right) &= 0 \\ -n\rho^{3} - \rho^2 \sum_{i=1}^{n} \left( -\hat{b}_{i1}\hat{b}_{i2} + \hat{b}_{i1}\mu_{i2} +\hat{b}_{i2}\mu_{i1} - \mathbb{E}(b_{i1}b_{i2}|\hat{b}_{i}) \right) - \rho(\sum_{i=1}^{n} \left( \hat{b}_{i1}^2 + \hat{b}_{i2}^2 -2\hat{b}_{i1} \mu_{i1} -2\hat{b}_{i2} \mu_{i2} + \mathbb{E}(b_{i1}^2|\hat{b}_{i}) + \mathbb{E}(b_{i2}^2|\hat{b}_{i}) \right) - n) - \sum_{i=1}^{n} \left( -\hat{b}_{i1}\hat{b}_{i2} + \hat{b}_{i1}\mu_{i2} +\hat{b}_{i2}\mu_{i1} - \mathbb{E}(b_{i1}b_{i2}|\hat{b}_{i}) \right) &= 0 \end{align*}$$

The polynomial has either 1 or 3 real roots in (-1, 1).

Algorithm:

Input: X, Ulist, init_rho
Compute loglikelihood
delta = 1
while delta > tol
  Given rho, Estimate pi using convex method (current mash method)
  M step: update rho: find all roots of polynomial, if it has three real roots, choose the one with higher loglikelihood.
  Compute loglikelihood
  Update delta