SlowGuess/NeurIPS_2023_OCR huggingface dataset:
By expanding the  $\mathbf{H}^{(l - 1)}\hat{\mathbf{A}}_t$  term based on neighbors from all classes, we get:
$$
\mathbf {x} _ {1, i} ^ {(l)} = \mathbf {W} _ {1} ^ {* (l)} \mathbf {h} _ {1, i} ^ {(l - 1)} + \mathbf {W} _ {2} ^ {* (l)} \left(\frac {\sum_ {v _ {1 , j} \in \mathcal {N} _ {1} \left(v _ {1 , i}\right)} \mathbf {h} _ {1 , j} + \sum_ {v _ {2 , j} \in \mathcal {N} _ {2} \left(v _ {1 , i}\right)} \mathbf {h} _ {2 , j}}{\left| \mathcal {N} \left(v _ {1 , i}\right) \right|}\right) \tag {110}
$$
Now, by taking expectations on both sides with respect to features  $\mathbf{h}_{1,i}^{(l - 1)}$  and structure  $\widehat{\mathbf{A}}_t$ , we get:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \mathbf {x} _ {1, i} ^ {(l)} = \mathbf {W} _ {1} ^ {* (l)} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \mathbf {h} _ {1, i} ^ {(l - 1)} + \mathbf {W} _ {2} ^ {* (l)} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left(\frac {\sum_ {v _ {1 , j} \in \mathcal {N} _ {1} (v _ {1 , i})} \mathbf {h} _ {1 , j} + \sum_ {v _ {2 , j} \in \mathcal {N} _ {2} (v _ {1 , i})} \mathbf {h} _ {2 , j}}{| \mathcal {N} (v _ {1 , i}) |}\right) \\ = \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\mu} _ {1} ^ {(l - 1)} + \mathbf {W} _ {2} ^ {* (l)} \left(\frac {n p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + n q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{n (p + q)}\right) \\ = \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\mu} _ {1} ^ {(l - 1)} + \left(\frac {q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\mu} _ {2} ^ {(l - 1)} \tag {111} \\ \end{array}
$$
Similarly, the covariance  $\mathbb{E}_{\hat{\mathbf{A}},\mathbf{h}}\left(\left[\mathbf{x}_{1,i}^{(l)} - \mathbb{E}_{\hat{\mathbf{A}},\mathbf{h}}\mathbf{x}_{1,i}^{(l)}\right]\left[\mathbf{x}_{1,i}^{(l)} - \mathbb{E}_{\hat{\mathbf{A}},\mathbf{h}}\mathbf{x}_{1,i}^{(l)}\right]^{\top}\right)$  is based on:
$$
\begin{array}{l} \left[ \mathbf {W} _ {1} ^ {* (l)} \mathbf {h} _ {1, i} ^ {(l - 1)} + \mathbf {W} _ {2} ^ {* (l)} \mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \mathbf {x} _ {1, i} ^ {(l)} \right] \\ \cdot \left[ \mathbf {W} _ {1} ^ {* (l)} \mathbf {h} _ {1, i} ^ {(l - 1)} + \mathbf {W} _ {2} ^ {* (l)} \mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \mathbf {x} _ {1, i} ^ {(l)} \right] ^ {\top} \\ = \left[ \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) + \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \right] \\ \cdot \left[ \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) + \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \right] ^ {\top} \\ = \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) ^ {\top} \mathbf {W} _ {1} ^ {* (l) \top} \\ + \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \\ + \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \hat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) ^ {\top} \mathbf {W} _ {1} ^ {* (l) \top} \\ + \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \tag {112} \\ \end{array}
$$
The expectation of the term  $\mathbf{W}_1^{*(l)}\left(\mathbf{h}_{1,i}^{(l - 1)} - \boldsymbol{\mu}_1^{(l - 1)}\right)\left(\mathbf{h}_{1,i}^{(l - 1)} - \boldsymbol{\mu}_1^{(l - 1)}\right)^\top \mathbf{W}_1^{*(l)^\top}$  is given by:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) ^ {\top} \mathbf {W} _ {1} ^ {* (l) \top} \right] \\ = \mathbf {W} _ {1} ^ {* (l)} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) ^ {\top} \right] \mathbf {W} _ {1} ^ {* (l) \top} \tag {113} \\ = \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} \\ \end{array}
$$
The expectation of  $\mathbf{W}_1^{*(l)}\left(\mathbf{h}_{1,i}^{(l-1)} - \boldsymbol{\mu}_1^{(l-1)}\right)\left(\mathbf{H}^{(l-1)}\widehat{\mathbf{A}}_t\mathbf{e}_{1,i} - \frac{p\boldsymbol{\mu}_1^{(l-1)} + q\boldsymbol{\mu}_2^{(l-1)}}{p+q}\right)^\top \mathbf{W}_2^{*(l)^\top}$  is given by the following:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ = \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {1} ^ {* (l)} \mathbf {h} _ {1, i} ^ {(l - 1)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ - \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\mu} _ {1} ^ {(l - 1)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ = \mathbf {W} _ {1} ^ {* (l)} \mathbb {E} _ {\mathbf {h}} \left[ \mathbf {h} _ {1, i} ^ {(l - 1)} \left(\frac {p \sum_ {j = 1} ^ {n} \mathbf {h} _ {1 , j} ^ {(l - 1)} + q \sum_ {j = 1} ^ {n} \mathbf {h} _ {2 , j} ^ {(l - 1)}}{n (p + q)} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ = \mathbf {W} _ {1} ^ {* (l)} \mathbb {E} _ {\mathbf {h}} \left[ \mathbf {h} _ {1, i} ^ {(l - 1)} \left(\frac {p}{n (p + q)} \mathbf {h} _ {1, i} ^ {(l - 1) \top}\right) \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ + \mathbf {W} _ {1} ^ {* (l)} \mathbb {E} _ {\mathbf {h}} \left[ \mathbf {h} _ {1, i} ^ {(l - 1)} \left(\frac {p \sum_ {j = 1 , \neq i} ^ {n} \mathbf {h} _ {1 , j} ^ {(l - 1)} + q \sum_ {j = 1} ^ {n} \mathbf {h} _ {2 , j} ^ {(l - 1)}}{n (p + q)} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ \end{array}
$$
Since the features are independent draws from their normal distributions, we can simplify the expectation as follows:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ = \mathbf {W} _ {1} ^ {* (l)} \left[ \frac {p}{n (p + q)} \left(\boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + \boldsymbol {\mu} _ {1} ^ {(l - 1)} \boldsymbol {\mu} _ {1} ^ {(l - 1) \top}\right) \right] \mathbf {W} _ {2} ^ {* (l) \top} (114) \\ + \mathbf {W} _ {1} ^ {* (l)} \left[ \boldsymbol {\mu} _ {1} ^ {(l - 1)} \left(\frac {(n - 1) p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + n q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{n (p + q)} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \right] \mathbf {W} _ {2} ^ {* (l) \top} (114) \\ = \mathbf {W} _ {1} ^ {* (l)} \left[ \frac {p}{n (p + q)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ \end{array}
$$
Similarly, the expectation of  $\mathbf{W}_2^{*(l)}\left(\mathbf{H}^{(l-1)}\widehat{\mathbf{A}}_t\mathbf{e}_{1,i}-\frac{p\boldsymbol{\mu}_1^{(l-1)}+q\boldsymbol{\mu}_2^{(l-1)}}{p+q}\right)\left(\mathbf{h}_{1,i}^{(l-1)}-\boldsymbol{\mu}_1^{(l-1)}\right)^\top\mathbf{W}_1^{*(l)^\top}$  is given by the following:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) ^ {\top} \mathbf {W} _ {1} ^ {* (l) \top} \right] \\ = \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {1} ^ {* (l)} \left(\mathbf {h} _ {1, i} ^ {(l - 1)} - \boldsymbol {\mu} _ {1} ^ {(l - 1)}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] ^ {\top} \tag {115} \\ = \mathbf {W} _ {2} ^ {* (l)} \left[ \frac {p}{n (p + q)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \right] \mathbf {W} _ {1} ^ {* (l) ^ {\top}} \\ \end{array}
$$
Next, the expectation of:  $\mathbf{W}_2^{*(l)}\left(\mathbf{H}^{(l-1)}\widehat{\mathbf{A}}_t\mathbf{e}_{1,i}-\frac{p\boldsymbol{\mu}_1^{(l-1)}+q\boldsymbol{\mu}_2^{(l-1)}}{p+q}\right)\left(\mathbf{H}^{(l-1)}\widehat{\mathbf{A}}_t\mathbf{e}_{1,i}-\frac{p\boldsymbol{\mu}_1^{(l-1)}+q\boldsymbol{\mu}_2^{(l-1)}}{p+q}\right)^\top\mathbf{W}_2^{*(l)^\top}$  can be computed as follows:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \right. \\ \left. \cdot \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ = \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ - \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {2} ^ {* (l)} \left(\frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ \end{array}
$$
Where the second term reduces to 0. On expanding the first term, we get:
Due to the independence of the features, we can simplify the expectation as follows:
$$
\begin{array}{l} \mathbb {E} _ {\widehat {\mathbf {A}}, \mathbf {h}} \left[ \mathbf {W} _ {2} ^ {* (l)} \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i}\right) \left(\mathbf {H} ^ {(l - 1)} \widehat {\mathbf {A}} _ {t} \mathbf {e} _ {1, i} - \frac {p \boldsymbol {\mu} _ {1} ^ {(l - 1)} + q \boldsymbol {\mu} _ {2} ^ {(l - 1)}}{p + q}\right) ^ {\top} \mathbf {W} _ {2} ^ {* (l) \top} \right] \\ = \mathbf {W} _ {2} ^ {* (l)} \left[ \left(\frac {n p ^ {2} \left(\boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + \boldsymbol {\mu} _ {1} ^ {(l - 1)} \boldsymbol {\mu} _ {1} ^ {(l - 1) \top}\right) + n q ^ {2} \left(\boldsymbol {\Sigma} _ {2} ^ {(l - 1)} + \boldsymbol {\mu} _ {2} ^ {(l - 1)} \boldsymbol {\mu} _ {2} ^ {(l - 1) \top}\right)}{n ^ {2} (p + q) ^ {2}}\right) \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ + \mathbf {W} _ {2} ^ {* (l)} \left[ \left(\frac {p ^ {2} (n ^ {2} - n) \boldsymbol {\mu} _ {1} ^ {(l - 1)} \boldsymbol {\mu} _ {1} ^ {(l - 1) \top} + q ^ {2} (n ^ {2} - n) \boldsymbol {\mu} _ {2} ^ {(l - 1)} \boldsymbol {\mu} _ {2} ^ {(l - 1) \top}}{n ^ {2} (p + q) ^ {2}}\right) \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ + \mathbf {W} _ {2} ^ {* (l)} \left[ \left(\frac {p q n ^ {2} \boldsymbol {\mu} _ {1} ^ {(l - 1)} \boldsymbol {\mu} _ {2} ^ {(l - 1) \top} + p q n ^ {2} \boldsymbol {\mu} _ {2} ^ {(l - 1)} \boldsymbol {\mu} _ {1} ^ {(l - 1) \top}}{n ^ {2} (p + q) ^ {2}}\right) \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ - \mathbf {W} _ {2} ^ {* (l)} \left[ \left(\frac {p ^ {2} \boldsymbol {\mu} _ {1} ^ {(l - 1)} \boldsymbol {\mu} _ {1} ^ {(l - 1) \top} + p q \boldsymbol {\mu} _ {1} ^ {(l - 1)} \boldsymbol {\mu} _ {2} ^ {(l - 1) \top} + p q \boldsymbol {\mu} _ {2} ^ {(l - 1)} \boldsymbol {\mu} _ {1} ^ {\top} + q ^ {2} \boldsymbol {\mu} _ {2} ^ {(l - 1)} \boldsymbol {\mu} _ {2} ^ {(l - 1) \top}}{(p + q) ^ {2}}\right) \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ = \mathbf {W} _ {2} ^ {* (l)} \left[ \frac {p ^ {2} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + q ^ {2} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)}}{n (p + q) ^ {2}} \right] \mathbf {W} _ {2} ^ {* (l) \top} \tag {116} \\ \end{array}
$$
Putting these results together, the covariance  $\mathbb{E}_{\widehat{\mathbf{A}},\mathbf{h}}\left(\left[\mathbf{x}_{1,i}^{(l)} - \mathbb{E}_{\widehat{\mathbf{A}},\mathbf{h}}\mathbf{x}_{1,i}^{(l)}\right]\left[\mathbf{x}_{1,i}^{(l)} - \mathbb{E}_{\widehat{\mathbf{A}},\mathbf{h}}\mathbf{x}_{1,i}^{(l)}\right]^{\top}\right)$  is:
$$
\begin{array}{l} \mathbb {E} _ {\hat {\mathbf {A}}, \mathbf {h}} \left(\left[ \mathbf {x} _ {1, i} ^ {(l)} - \mathbb {E} _ {\hat {\mathbf {A}}, \mathbf {h}} \mathbf {x} _ {1, i} ^ {(l)} \right] \left[ \mathbf {x} _ {1, i} ^ {(l)} - \mathbb {E} _ {\hat {\mathbf {A}}, \mathbf {h}} \mathbf {x} _ {1, i} ^ {(l)} \right] ^ {\top}\right) \\ = \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} + \frac {p}{n (p + q)} \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {2} ^ {* (l) \top} \tag {117} \\ + \frac {p}{n (p + q)} \mathbf {W} _ {2} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} + \mathbf {W} _ {2} ^ {* (l)} \left[ \frac {p ^ {2} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + q ^ {2} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)}}{n (p + q) ^ {2}} \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ \end{array}
$$
To summarize, the means and covariance matrices for  $\mathbf{X}^{(l)}$  can be given by:
$$
\begin{array}{l} \widetilde {\boldsymbol {\mu}} _ {1} ^ {(l)} = \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\mu} _ {1} ^ {(l - 1)} + \left(\frac {q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\mu} _ {2} ^ {(l - 1)} \\ \widetilde {\boldsymbol {\Sigma}} _ {1} ^ {(l)} = \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} + \frac {p}{n (p + q)} \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {2} ^ {* (l) \top} \tag {118} \\ + \frac {p}{n (p + q)} \mathbf {W} _ {2} ^ {* (l)} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} + \mathbf {W} _ {2} ^ {* (l)} \left[ \frac {p ^ {2} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + q ^ {2} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)}}{n (p + q) ^ {2}} \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ \end{array}
$$
- Case  $c = 2$ :
The analysis presented above can be extended for  $c = 2$  in a straightforward fashion to get  $\widetilde{\pmb{\mu}}_2^{(l)},\widetilde{\pmb{\Sigma}}_2^{(l)}$  as follows:
$$
\begin{array}{l} \widetilde {\boldsymbol {\mu}} _ {2} ^ {(l)} = \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\mu} _ {2} ^ {(l - 1)} + \left(\frac {q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\mu} _ {1} ^ {(l - 1)} \\ \widetilde {\boldsymbol {\Sigma}} _ {2} ^ {(l)} = \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} + \frac {p}{n (p + q)} \mathbf {W} _ {1} ^ {* (l)} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)} \mathbf {W} _ {2} ^ {* (l) \top} \tag {119} \\ + \frac {p}{n (p + q)} \mathbf {W} _ {2} ^ {* (l)} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)} \mathbf {W} _ {1} ^ {* (l) \top} + \mathbf {W} _ {2} ^ {* (l)} \left[ \frac {p ^ {2} \boldsymbol {\Sigma} _ {2} ^ {(l - 1)} + q ^ {2} \boldsymbol {\Sigma} _ {1} ^ {(l - 1)}}{n (p + q) ^ {2}} \right] \mathbf {W} _ {2} ^ {* (l) \top} \\ \end{array}
$$
G.1 Modelling an increase/decrease in between-class variability
Let  $\pmb{\Sigma}_{B}^{(l-1)} = \left(\pmb{\mu}_{1}^{(l-1)} - \pmb{\mu}_{2}^{(l-1)}\right)\left(\pmb{\mu}_{1}^{(l-1)} - \pmb{\mu}_{2}^{(l-1)}\right)^{\top}$  indicate the between-class covariance matrix for features at layer  $l-1$ . Based on the equations 118, 119, we analyze  $\pmb{\Sigma}_{B}(\mathbf{X}^{(l)})$  to understand the effect of the convolution layer on feature separation. First, observe that:
$$
\begin{array}{l} \boldsymbol {\Sigma} _ {B} (\mathbf {X} ^ {(l)}) = \left(\widetilde {\boldsymbol {\mu}} _ {1} ^ {(l)} - \widetilde {\boldsymbol {\mu}} _ {2} ^ {(l)}\right) \left(\widetilde {\boldsymbol {\mu}} _ {1} ^ {(l)} - \widetilde {\boldsymbol {\mu}} _ {2} ^ {(l)}\right) ^ {\top} \\ = \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \left(\boldsymbol {\mu} _ {1} ^ {(l - 1)} - \boldsymbol {\mu} _ {2} ^ {(l - 1)}\right) \\ \cdot \left(\boldsymbol {\mu} _ {1} ^ {(l - 1)} - \boldsymbol {\mu} _ {2} ^ {(l - 1)}\right) ^ {\top} \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) ^ {\top} \tag {120} \\ = \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\Sigma} _ {B} (\mathbf {H} ^ {(l - 1)}) \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) ^ {\top} \\ \end{array}
$$
By taking the trace on both sides, we get:
$$
\begin{array}{l} \operatorname {T r} \left(\boldsymbol {\Sigma} _ {B} (\mathbf {X} ^ {(l)})\right) = \operatorname {T r} \left(\left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) \boldsymbol {\Sigma} _ {B} (\mathbf {H} ^ {(l - 1)}) \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) ^ {\top}\right) \\ = \operatorname {T r} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right) \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right) ^ {\top} \left(\mathbf {W} _ {1} ^ {* (l)} + \frac {p - q}{p + q} \mathbf {W} _ {2} ^ {* (l)}\right)\right) \tag {121} \\ \end{array}
$$
Where  $\left(\mathbf{W}_1^{*(l)} + \frac{p - q}{p + q}\mathbf{W}_2^{*(l)}\right)^\top \left(\mathbf{W}_1^{*(l)} + \frac{p - q}{p + q}\mathbf{W}_2^{*(l)}\right) \in \mathbb{R}^{d_{l-1} \times d_{l-1}}$  is a symmetric and positive semi-definite matrix. This matrix product formulation allows us to leverage the eigenvalue-based trace inequalities [41, 68]. Formally, let  $\mathbf{T}_B = \left(\mathbf{W}_1^{*(l)} + \frac{p - q}{p + q}\mathbf{W}_2^{*(l)}\right)^\top \left(\mathbf{W}_1^{*(l)} + \frac{p - q}{p + q}\mathbf{W}_2^{*(l)}\right)$ . We now leverage Corollary.6 in Zhang and Zhang [68], and get the following inequality based on the eigenvalues of  $\boldsymbol{\Sigma}_B(\mathbf{H}^{(l-1)})$ ,  $\mathbf{T}_B$  as:
$$
\begin{array}{l} \sum_ {i = 1} ^ {d _ {l - 1}} \lambda_ {d _ {l - 1} - i + 1} \left(\boldsymbol {\Sigma} _ {B} (\mathbf {H} ^ {(l - 1)})\right) \lambda_ {i} \left(\mathbf {T} _ {B}\right) \\ \leq \operatorname {T r} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right) \mathbf {T} _ {B}\right) \tag {122} \\ \leq \sum_ {i = 1} ^ {d _ {l - 1}} \lambda_ {i} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right)\right) \lambda_ {i} \left(\mathbf {T} _ {B}\right) \\ \end{array}
$$
Where  $\lambda_{i}(.)$  represents the  $i^{th}$  largest eigenvalue of a matrix. Additionally, based on the standard trace equality:  $\operatorname{Tr}\left(\boldsymbol{\Sigma}_B(\mathbf{H}^{(l-1)})\right) = \sum_{i=1}^{d_{l-1}} \lambda_i\left(\boldsymbol{\Sigma}_B(\mathbf{H}^{(l-1)})\right)$ , the increase/decrease in  $\operatorname{Tr}\left(\boldsymbol{\Sigma}_B(\mathbf{X}^{(l)})\right)$  with respect to  $\operatorname{Tr}\left(\boldsymbol{\Sigma}_B(\mathbf{H}^{(l-1)})\right)$  boils down to:
$$
\frac {\sum_ {i = 1} ^ {d _ {l - 1}} \lambda_ {d _ {l - 1} - i + 1} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right)\right) \lambda_ {i} \left(\mathbf {T} _ {B}\right)}{\sum_ {i = 1} ^ {d _ {l - 1}} \lambda_ {i} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right)\right)} \leq \frac {\operatorname {T r} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {X} ^ {(l)}\right)\right)}{\operatorname {T r} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right)\right)} \leq \frac {\sum_ {i = 1} ^ {d _ {l - 1}} \lambda_ {i} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {(l - 1)}\right)\right) \lambda_ {i} \left(\mathbf {T} _ {B}\right)}{\sum_ {i = 1} ^ {d _ {l - 1}} \lambda_ {i} \left(\boldsymbol {\Sigma} _ {B} \left(\mathbf {H} ^ {l - 1}\right)\right)} \tag {123}
$$
G.2 Modelling an increase/decrease in within-class variability
Let  $\pmb{\Sigma}_{W}(\mathbf{H}^{(l - 1)}) = \frac{1}{2}\left(\pmb{\Sigma}_{1}^{(l - 1)} + \pmb{\Sigma}_{2}^{(l - 1)}\right)$  represent the within-class covariance matrix for features  $\mathbf{H}^{(l - 1)}$  in our balanced class setting. Similar to the previous analysis, we leverage the results in
equations 118,119 to model  $\pmb{\Sigma}_{W}(\mathbf{X}^{(l)})$  as follows:
$$
\begin{array}{l} \boldsymbol {\Sigma} _ {W} (\mathbf {X} ^ {(l)}) = \frac {1}{2} \left(\widetilde {\boldsymbol {\Sigma}} _ {1} ^ {(l)} + \widetilde {\boldsymbol {\Sigma}} _ {2} ^ {(l)}\right) \\ = \frac {1}{2} \left(\mathbf {W} _ {1} ^ {* (l)} \left(\boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + \boldsymbol {\Sigma} _ {2} ^ {(l - 1)}\right) \mathbf {W} _ {1} ^ {* (l) ^ {\top}}\right) \\ + \frac {1}{2} \left(\frac {p}{n (p + q)} \mathbf {W} _ {1} ^ {* (l)} \left(\boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + \boldsymbol {\Sigma} _ {2} ^ {(l - 1)}\right) \mathbf {W} _ {2} ^ {* (l) \top}\right) \tag {124} \\ + \frac {1}{2} \left(\frac {p}{n (p + q)} \mathbf {W} _ {2} ^ {* (l)} \left(\boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + \boldsymbol {\Sigma} _ {2} ^ {(l - 1)}\right) \mathbf {W} _ {1} ^ {* (l) \top}\right) \\ + \frac {1}{2} \left(\mathbf {W} _ {2} ^ {* (l)} \left[ \frac {(p ^ {2} + q ^ {2}) (\boldsymbol {\Sigma} _ {1} ^ {(l - 1)} + \boldsymbol {\Sigma} _ {2} ^ {(l - 1)})}{n (p + q) ^ {2}} \right] \mathbf {W} _ {2} ^ {* (l) \top}\right) \\ \end{array}
$$
By taking trace on both sides, we get: