SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Entropic potential It is convenient to introduce before proceeding the variance order parameters
$$
\hat {V} \equiv \hat {r} + \hat {q}, \quad \hat {V} _ {k} \equiv \hat {r} _ {k} + \hat {q} _ {k}. \tag {122}
$$
$$
\begin{array}{l} e ^ {\beta s d \Psi_ {w}} \\ = \int \prod_ {a = 1} ^ {s} d \boldsymbol {w} _ {a} e ^ {- \beta \sum_ {a} g (\boldsymbol {w} ^ {a}) - \frac {1}{2} \sum_ {a = 1} ^ {s} \operatorname {T r} \left[ \hat {V} \boldsymbol {w} _ {a} \boldsymbol {w} _ {a} ^ {\top} \right] + \frac {1}{2} \sum_ {a, b} \operatorname {T r} \left[ \hat {q} \boldsymbol {w} _ {a} \boldsymbol {w} _ {b} ^ {\top} \right] + \hat {m} \sum_ {k = 1} ^ {K} \sum_ {a = 1} ^ {s} \sqrt {d} \hat {m} _ {k} ^ {\top} \boldsymbol {w} _ {a} ^ {\top} \boldsymbol {\mu} _ {k}} \\ \times e ^ {- \frac {1}{2} \sum_ {k = 1} ^ {K} \sum_ {a = 1} ^ {s} \operatorname {T r} \left[ \hat {V} _ {k} \boldsymbol {w} _ {a} \boldsymbol {\Sigma} _ {k} \boldsymbol {w} _ {a} ^ {\top} \right] + \frac {1}{2} \sum_ {k = 1} ^ {K} \sum_ {a, b} \operatorname {T r} \left[ \hat {q} _ {k} \boldsymbol {w} _ {a} \boldsymbol {\Sigma} _ {k} \boldsymbol {w} _ {b} ^ {\top} \right]} \\ = \int D \boldsymbol {\Xi} _ {0} \prod_ {k = 1} ^ {K} D \boldsymbol {\Xi} _ {k} \\ \left[ \int d \boldsymbol {w} e ^ {- \beta g (\boldsymbol {w}) - \frac {1}{2} \operatorname {T r} \left[ \hat {V} \boldsymbol {w} \boldsymbol {w} ^ {\top} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \boldsymbol {w} \boldsymbol {\Sigma} _ {k} \boldsymbol {w} ^ {\top} \right] + \left(\sum_ {k = 1} ^ {K} \sqrt {d} \hat {m} _ {k} \boldsymbol {\mu} ^ {\top} + \sum_ {k = 1} ^ {K} \boldsymbol {\Xi} _ {k} \odot (\hat {q} _ {k} \otimes \boldsymbol {\Sigma} _ {k}) ^ {\frac {1}{2}} + \boldsymbol {\Xi} _ {0} \odot (\hat {q} \otimes \mathbb {I} _ {d}) ^ {\frac {1}{2}}\right) \odot \boldsymbol {w}} \right] ^ {s} \\ = \int \underbrace {D \Xi_ {0} \prod_ {k = 1} ^ {K} D \Xi_ {k}} _ {\equiv D \Xi} \\ \left[ \int d w e ^ {- \beta g (\boldsymbol {w}) - \frac {1}{2} \boldsymbol {w} \odot \left[ \hat {V} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \otimes \boldsymbol {\Sigma} _ {k} \right] \odot \boldsymbol {w} + \left(\sum_ {k = 1} ^ {K} \sqrt {d} \hat {m} _ {k} \boldsymbol {\mu} ^ {\top} + \sum_ {k = 1} ^ {K} \boldsymbol {\Xi} _ {k} \odot (\hat {q} _ {k} \otimes \boldsymbol {\Sigma} _ {k}) ^ {\frac {1}{2}} + \boldsymbol {\Xi} _ {0} \odot (\hat {q} \otimes \mathbb {I} _ {d}) ^ {\frac {1}{2}}\right) \odot \boldsymbol {w}} \right] ^ {s}. \tag {123} \\ \end{array}
$$
$$
\beta \Psi_ {w}
$$
$$
\frac {1}{d} \int D \boldsymbol {\Xi} \ln \left[ \int d w e ^ {- \beta g (\boldsymbol {w}) - \frac {1}{2} \boldsymbol {w} \odot \left[ \hat {V} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \otimes \boldsymbol {\Sigma} _ {k} \right] \odot \boldsymbol {w} + \left(\sum_ {k = 1} ^ {K} \sqrt {d} \hat {m} _ {k} \boldsymbol {\mu} ^ {\top} + \sum_ {k = 1} ^ {K} \boldsymbol {\Xi} _ {k} \odot (\dot {q} _ {k} \otimes \boldsymbol {\Sigma} _ {k}) ^ {\frac {1}{2}} + \boldsymbol {\Xi} _ {0} \odot (\dot {q} \otimes \mathbb {I} _ {d}) ^ {\frac {1}{2}}\right) \odot \boldsymbol {w}} \right]. \tag {124}
$$
Energetic potential The inverse of  $S_{k}$  can be computed like in Appendix A, leading to the following expression for (a):
$$
\begin{array}{l} (a) = \int_ {\mathbb {R} ^ {p}} \frac {\prod_ {a = 1} ^ {s} d \lambda_ {a} ^ {\eta}}{(2 \pi) ^ {m s / 2} \sqrt {\det S _ {k}}} e ^ {- \frac {1}{2} \sum_ {a = 1} ^ {s} (\lambda_ {a} ^ {\eta}) ^ {\top} (\tilde {r} _ {k} - \tilde {q} _ {k}) \lambda_ {a} ^ {\eta} - \frac {1}{2} \sum_ {a, b} (\lambda_ {a} ^ {\eta}) ^ {\top} \tilde {q} _ {k} \lambda_ {b} - \frac {\beta}{2} \sum_ {a = 1} ^ {s} (*)} \\ = \int_ {\mathbb {R} ^ {p}} \frac {D \eta}{(2 \pi) ^ {m s / 2} \sqrt {\det S _ {k}}} \left[ \int_ {\mathbb {R} ^ {p}} d \lambda_ {\eta} e ^ {- \frac {1}{2} \lambda_ {\eta} ^ {\top} V _ {k} ^ {- 1} \lambda_ {\eta} + \lambda_ {\eta} ^ {\top} V _ {k} ^ {- 1} q _ {k} ^ {\frac {1}{2}} \eta - \frac {\beta}{2} (*)} \right] ^ {s} \\ = \int_ {\mathbb {R} ^ {p}} D \eta \left[ \int \mathcal {N} \left(\lambda_ {\eta}; q _ {k} ^ {\frac {1}{2}} \eta , V\right) e ^ {- \frac {\beta}{2} (*)} \right] ^ {s} \\ = 1 + s \int_ {\mathbb {R} ^ {p}} D \eta \ln \left[ \int \mathcal {N} \left(\lambda_ {\eta}; q _ {k} ^ {\frac {1}{2}} \eta , V\right) e ^ {- \frac {\beta}{2} (*)} \right], \tag {125} \\ \end{array}
$$
where we noted with capital  $D$  an integral over  $\mathcal{N}(0,\mathbb{I}_p)$ . Therefore
$$
\beta \Psi_ {y} = \sum_ {k = 1} ^ {K} \rho_ {k} \int_ {\mathbb {R} ^ {p}} D \eta \ln \left[ \int \mathcal {N} \left(\lambda_ {\eta}; q _ {k} ^ {\frac {1}{2}} \eta , V\right) e ^ {- \frac {\beta}{2} (*)} \right]. \tag {126}
$$
Zero-temperature limit In this subsection, we take the zero temperature  $\beta \to \infty$  limit. Rescaling
$$
\frac {1}{\beta} V _ {k} \leftarrow V _ {k}, \quad \beta \hat {V} _ {k} \leftarrow \hat {V} _ {k}, \quad \beta^ {2} \hat {q} _ {k} \leftarrow \hat {q} _ {k}, \quad \beta \hat {m} _ {k} \leftarrow \hat {m} _ {k} \tag {127}
$$
$$
\begin{array}{l} \Psi_ {w} = \frac {1}{2 d} \operatorname {T r} \left[ \left(\hat {V} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \otimes \boldsymbol {\Sigma} _ {k}\right) ^ {- 1} \odot \left(\hat {q} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {q} _ {k} \otimes \boldsymbol {\Sigma} _ {k} + d \left(\sum_ {k = 1} ^ {K} \hat {m} _ {k} \boldsymbol {\mu} _ {k} ^ {\top}\right) ^ {\otimes 2}\right) \right] \\ - \frac {1}{d} \mathbb {E} _ {\Xi} \mathcal {M} _ {r} (\Xi), \tag {128} \\ \end{array}
$$
where we introduced the Moreau envelope
$$
\begin{array}{l} \mathcal {M} _ {r} (\boldsymbol {\Xi}) \\ = \inf  _ {\boldsymbol {w}} \left\{\frac {1}{2} \left\| \left(\hat {V} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \otimes \boldsymbol {\Sigma} _ {k}\right) ^ {\odot} \boldsymbol {w} - \left(\hat {V} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \otimes \boldsymbol {\Sigma} _ {k}\right) ^ {- \frac {1}{2}} \left((\hat {q} \otimes \mathbb {I} _ {d}) ^ {\frac {1}{2}} \odot \boldsymbol {\Xi} _ {0} + \sum_ {k = 1} ^ {K} (\hat {q} _ {k} \otimes \boldsymbol {\Sigma} _ {k}) ^ {\frac {1}{2}} \odot \boldsymbol {\Xi} _ {k} + \sqrt {d} \sum_ {k = 1} ^ {K} \hat {m} _ {k} \boldsymbol {\mu} _ {k} ^ {\top}\right) \right\| ^ {2} + g (\boldsymbol {w}) \right\}. \tag {129} \\ \end{array}
$$
For the case of  $\ell_2$  regularization, the Moreau envelope presents a simple expression, and the entropy potential assumes the simple form
$$
\Psi_ {w} = + \frac {1}{2 d} \operatorname {T r} \left[ \left(\lambda \mathbb {I} _ {m} \odot \mathbb {I} _ {d} + \hat {V} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {V} _ {k} \otimes \boldsymbol {\Sigma} _ {k}\right) ^ {- 1} \odot \left(\hat {q} \otimes \mathbb {I} _ {d} + \sum_ {k = 1} ^ {K} \hat {q} _ {k} \otimes \boldsymbol {\Sigma} _ {k} + d \left(\sum_ {k = 1} ^ {K} \hat {m} _ {k} \boldsymbol {\mu} _ {k} ^ {\top}\right) ^ {\otimes 2}\right) \right]. \tag {130}
$$
The energetic potential also simplifies in this limit to
$$
\Psi_ {y} = - \sum_ {k = 1} ^ {K} \rho_ {k} \mathbb {E} _ {\eta} \mathcal {M} _ {k} (\eta) \tag {131}
$$
$$
\mathcal {M} _ {k} (\eta) = \frac {1}{2} \inf  _ {x, y} \left\{\operatorname {T r} \left[ V _ {k} ^ {- 1} \left(y - q _ {k} ^ {\frac {1}{2}} \eta - m _ {k}\right) ^ {\otimes 2} \right] + \operatorname {T r} \left[ q \sigma (y) ^ {\otimes 2} \right] - 2 \sigma (y) ^ {\top} y \right\}. \tag {132}
$$
Trace potential It is immediate to see that the trace potential can be expressed as
$$
\Psi_ {t} = \frac {\hat {V} q - \hat {q} V}{2} + \frac {1}{2} \sum_ {k = 1} ^ {K} \left(\operatorname {T r} \left[ \hat {V} _ {k} q _ {k} \right] - \operatorname {T r} [ \hat {q} _ {k} V _ {k} ]\right) - \sum_ {k = 1} ^ {K} m _ {k} \hat {m} _ {k}. \tag {133}
$$
Then the total free entropy for RAEs (107) reads
$$
\begin{array}{l} \Phi = \operatorname * {e x t r} _ {q, m, V, \hat {q}, \hat {m}, \hat {V}, \{q _ {k}, m _ {k}, V _ {k}, \hat {q} _ {k}, \hat {m} _ {k}, \hat {V} _ {k} \} _ {k = 1} ^ {K}} \frac {\hat {V} q - \hat {q} V}{2} + \frac {1}{2} \sum_ {k = 1} ^ {K} (\operatorname * {T r} \left[ \hat {V} _ {k} q _ {k} \right] - \operatorname * {T r} [ \hat {q} _ {k} V _ {k} ]) - \sum_ {k = 1} ^ {K} m _ {k} \hat {m} _ {k} \\ - \alpha \sum_ {k = 1} ^ {K} \rho_ {k} \mathbb {E} _ {\eta} \mathcal {M} _ {k} (\eta) + \frac {1}{2} \int d \nu (\gamma , \tau) \operatorname {T r} \left[ \left(\lambda + \gamma_ {k} \hat {V} _ {k}\right) ^ {- 1} \left(\sum_ {k = 1} ^ {K} \gamma_ {k} \hat {q} _ {k} + \sum_ {1 \leq k, j \leq K} \tau_ {j} \tau_ {k} \hat {m} _ {j} \hat {m} _ {k} ^ {\top}\right) \right]. \tag {134} \\ \end{array}
$$
Comparing to the free energy derived in Appendix A, it is immediate to see that the Moreau envelope  $\mathcal{M}_k$  for the RAE just corresponds to removing the second term from the DAE  $\mathcal{M}_k$ , and setting  $x = 0$ . The disappearance of the second term in the Moreau envelope has the effect to kill the  $V$ -dependence, resulting in  $\hat{q} = 0$  when extremizing with respect to  $V$ . These observations, in addition to the already taken limits  $\Delta, b = 0$ , finish the derivation of Corollary 3.5.
Figure 14:  $\alpha = 1, K = 2, \rho_{1,2} = 1/2, \Sigma_{1,2} = 0.3 \times \mathbb{I}_d, p = 1, \sigma = \tanh$ ; the cluster mean  $\pmb{\mu}_1 = -\pmb{\mu}_2$  was taken as a random Gaussian vector of norm 1. (left) In blue, difference in test reconstruction MSE between the RAE (Corollary 3.5) and the leading order term  $\mathrm{mse}'_0$  (135). Solid lines correspond to the sharp asymptotic characterization of Corollary 3.5. Dots represent numerical simulations for  $d = 700$ , training the RAE using the Pytorch implementation of full-batch Adam, with learning rate  $\eta = 0.05$  over 2000 epochs, weight decay  $\lambda = 0.1$ , averaged over  $N = 10$  instances. Error bars represent one standard deviation. (right) Cosine similarity  $\theta$  (6) (green), squared weight norm  $\| \hat{w} \|^2 / d$  (red). Solid lines correspond to the theoretical characterization of Corollary 3.5; dots are numerical simulations. The cosine similarity with the cluster means of the principal component of the training samples is represented in blue crosses, and superimposes perfectly with the analogous curve learnt by the RAE.
F.2 Reconstruction MSE
Having obtained a sharp asymptotic characterization of the summary statistics for the RAE, we now turn to the learning metrics. It is easy to see that the formula for the cosine similarity  $\theta$  (6) carries over unchanged from the DAE case, see Appendix A. We therefore focus on the test MSE:
$$
\begin{array}{l} \mathrm {m s e} _ {\mathrm {R A E}} = \mathbb {E} _ {k, \boldsymbol {\eta}, \xi} \left\| \boldsymbol {\mu} _ {k} + \boldsymbol {\eta} - \left(\frac {\hat {\boldsymbol {w}} ^ {\top}}{\sqrt {d}} \sigma \left(\frac {\hat {\boldsymbol {w}} (\boldsymbol {\mu} _ {k} + \boldsymbol {\eta})}{\sqrt {d}}\right)\right) \right\| ^ {2} \\ = \underbrace{\sum_{k = 1}^{K}\rho_{k}\left(\|\boldsymbol{\mu}_{k}\|^{2} + \operatorname{Tr}\boldsymbol{\Sigma}_{k}\right)}_{\equiv \operatorname {mse}^{\prime}_{\circ}} + \sum_{k = 1}^{K}\rho_{k}\mathbb{E}_{z}^{\mathcal{N}(0,1)}\left[\operatorname {Tr}\left[q\sigma \left(m_{k} + \sqrt{q_{k}} z\right)^{\otimes 2}\right]\right] \\ - 2 \sum_ {k = 1} ^ {K} \rho_ {k} \mathbb {E} _ {u, v} ^ {\mathcal {N} (0, 1)} [ \sigma \left(m _ {k} + \sqrt {q _ {k}} u\right) ] ^ {\top} \left(m _ {k} + \sqrt {q _ {k}} u\right) \tag {135} \\ \end{array}
$$
F.3 RAEs are limited by the PCA MSE
We close this Appendix by showing how the RAE Corollary 3.5 recovers the well-known Eckart-Young theorem [10]. The difference  $\mathrm{mse} - \mathrm{mse}_{\circ}^{\prime}$  for an RAE learning to reconstruct a binary, isotropic homoscedastic mixture (see Fig. 1) is shown in Fig. 14, from a training set with sample complexity  $\alpha = 1$ . Because  $\mathrm{mse} - \mathrm{mse}_{\circ}^{\prime}$  is essentially a correction to the leading term  $\mathrm{mse}_{\circ}^{\prime}$ , estimating this difference is slightly more challenging numerically, whence rather large error bars. Nevertheless, it can be observed that the agreement between the theory (solid blue line) and the simulations (dots) gathered from training an  $p = 1$  RAE using the Pytorch implementation of Adam is still very good. In terms of cosine similarity and learnt weight norm, the agreement is perfect, see Fig. 3.5 (right). For comparison, the performance of PCA reconstruction (crosses) has been overlayed on the RAE curves. Importantly, observe that in terms of MSE, PCA reconstruction always leads to smaller MSE, in agreement with the well-known Eckart-Young theorem [10]. From the cosine similarity, it can be seen that the RAE essentially learns the principal component of the train set, see the discussions in e.g. [12, 11, 8, 9].
We finally stress that, in contrast to all previous figures, the  $x$  axis in Fig. 3.5 is the sample complexity  $\alpha$ , rather than the noise level  $\Delta$ . As a matter of fact, while our result recovers the well-known fact that the MSE of RAEs is limited (lower-bounded) by the PCA test error, Corollary 3.5 allows us to investigate how the RAE approaches the PCA
performance as a function of the number of training samples, a characterization which has importantly been missing from previous studies for non-linear RAEs [8, 9].