SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Beyond tractable models and implicit integration. One remaining question in our analysis is how the probabilistic theory around MPT adapts to intractable or non-linear models. In practice, self-predictive probabilities imply integrating out the latent variables, often given the posterior distribution. In most cases, performing this integration is extremely difficult or not possible in training time. Therefore, we are interested in finding if alternative approximations  $q_{\theta}$  to the true self-conditional probabilities still produce accurate estimation and maximization of the LML. This point is confirmed in Fig. 3. Inspired by the experiments of Lucas et al. (2019) with linear VAEs, we set up a Bernoulli likelihood on top of the latent variable model. The tractable formulation in the Gaussian example coincides with PPCA. Since predictive conditionals are no longer tractable for us, we use numerical integration to obtain the probabilities of masked tokens. In Fig. 3, we test the training with the cumulative MPT loss as well as compare with standard variational inference using the model's evidence lower bound (ELBO). For the mini-dataset with MNIST samples, we observe that both models converge to a similar value of the LML. Thus, the fundamental insight here is that MPT maximizes LML even under training with approximate self-predictive conditional probabilities. For the LML curves, we also used numerical integration.
Beyond linear models, our theory is useful when applied to non-linear models. Moreover, in Fig. 3 we also include the results for deep VAEs based on NNs. While the estimation of LML was obtained via Monte Carlo (MC) samples, we used iterative encoding-decoding to produce the self-conditional probabilities for masked tokens — see Sec. F in Rezende et al. (2014). In this scenario, we also observe the maximization of the LML according to the evolution of the MPT loss.
Another key insight showed by this study is the ability of MPT to perform implicit integration. The cumulative sum over the different rates of random masking is another way to see a discrete integral
Figure 3: Training curves for linear VAE and deep VAE models with variational inference (VI) and MPT. Data consist of subsets of MNIST and FMNIST. (Upper Row). A linear VAE model with Bernoulli likelihood function in  $N = 2000$  samples of MNIST and FMNIST. Shaded curves correspond to the target losses used in the optimizer (ELBO and MPT). Darker lines indicate the evolution of the LML, which are approximated via numerical integration in a latent space  $\mathcal{Z}$  of dimensionality  $K = 2$ . (Lower Row). Vanilla VAE with Gaussian likelihood for MNIST. The LML curves are approximated via Monte Carlo (MC) samples. Self-predictive conditional probabilities are obtained via recursive encoding-decoding. The size of the random masking is fixed and set to  $33\%$ .
Figure 4: Area under the curve described by  $S_{\theta}(\cdot; M)$ . The area is approximately equal to the model's LML according to the theory. Larger probability values are obtained with smaller rates of masking. (Left). Area described with  $P = 1$  random masking per epoch. The curve is more noisy and the area slightly loses precision w.r.t. LML. (Center). Area under the MPT curve for  $P = 100$ . (Right). Latent space is augmented to be of  $K = 50$ . Decay of predictive probabilities begins at around  $50\%$  masking rate.
under the curve described by the score function  $S_{\theta}(\cdot; M)$  in Eq. 5. In Fig. 4, we show the areas under the curve and the effect of reducing the number of random masks  $P$ . The blue plots correspond to a trained PPCA model and the area corresponds to the LML estimate. The long tail in the right part of the curves, when the rate of masking is larger than  $90\%$ , indicates that the model is no longer able to produce good estimates of the tokens with only  $10\%$  of the input dimensions observed. This explains, why the probabilities have an approximately exponential decay. However, this effect is not constant, and it might depend on the latent structure of the model. In the r.h.s. plot we observe that the decay of conditional probabilities happens earlier at approximate  $50\%$  random masking or larger. The role of the masking rate is perhaps the missing part in the picture (Wettig et al., 2022), as it is the one that
determines the approximation to the LML. With the purpose of providing an intuition on how rates of  $15\%$  or  $85\%$  affect to the area under the curve, we indicate with two black lines the approximate area that approximates the LML. A longer discussion is provided in the supplementary material.
Table 2: Area under the MPT curve for BERT model and four GLUE datasets.
GLUE DATASETS →	AX	COLA	QNLI	MRPC
AREA / RANDOM INIT. (↑)	-5245.31	-5283.52	-5343.98	-5362.21
AREA / PRE-TRAINED (↑)	-1715.75	-1657.68	-1770.28	-1773.45
Figure 5: Evolution of the area under the MPT curve. Comparison between one tractable model (PPCA) and BERT. The area under the curves is approximately the LML. Random initialization of the parameters produces MPT curves with similar low probabilities for all % of masking. As the number of epochs increases, the curve brings higher values of log-probability for lower ratios of masking. The area also converges to the true value of LML. (Left). PPCA model trained for 600 epochs. Each curve represents  $\{0,100,200,300,400,500,600\}$  epochs of training with MPT. (Right). Random initialization and end-of-pretraining curves for the MPT loss w.r.t. the % of masked tokens. Curves are similar but not identical for the 4 different datasets given the pre-trained BERT model.
3.2 Applied theory on large language models
In this section, we aim to understand how the area under the MPT curve evolves and behaves for large language models (LLMs). While the direct computation of the LML is not feasible for non-linear transformer models, we are interested in checking how the rate of masking affects the curve compared with the tractable PPCA model. The results provided in Fig. 5 and Tab. 2 give us insights into this behavior. First, we observe that the MPT curve is approximately flat for every rate of masking in the PPCA when parameters are randomly initialized. Intuitively, this indicates that the model is not able to correctly predict any token given some context. In some way, it produces noise independently of the number of conditional tokens, which explains the low log-probabilities. Second, we can also notice that the curve changes its shape as more training epochs are considered. The curve after 600 epochs produces high probability values for different rates of masking, while the long tail of low probabilities appears when masking more than  $85\%$  of tokens. Moreover, the area under these curves is the estimation of the LML, which accurately converges to the ground truth value of the LML with the original generative parameters.
For the study of the curves in LLMs, we used four datasets from the General Language Understanding Evaluation (GLUE) (Wang et al., 2019). Additionally, we consider a 110M parameters BERT model using pre-trained checkpoints $^{2}$  and random initializations. To draw the MPT curves, we computed the mean cross-entropy per each rate of masking between  $1\%$  and  $99\%$ . In Fig. 5, we observe that random initializations of BERT parameters lead to flat curves of low self-predictive probabilities. On the other hand, the pre-trained curves show similar behavior as in the tractable model, where the area
$^{2}$ Pre-trained parameters for the BERT model are available in the library — https://huggingface.co/.
is reduced and a long tail of low probabilities happens when the rate of masking becomes larger. This result supports our hypothesis that MPT in LLMs might be performing implicit integration of the latent space and maximizing the marginal likelihood of the model.
Results on vision models. In addition to the results shown in Fig. 5 with the BERT model, we are also interested in the behavior of the theory on large models oriented to vision. For this study we provide the curves of the area under the MPT loss for VIT-MAE (Dosovitskiy et al., 2020; He et al., 2022) with different masking rates. In a similar way as in Sec. 3.2, we use an (already) pre-trained VIT-MAE model loaded from a public repository. To draw the curves shown in Fig. 6, we computed the losses per each rate of masking between  $5\%$  and  $95\%$  for samples from three different test image datasets (FASHION-MNIST, CIFAR-100 and TINY-IMAGENET). We can observe that the curves described by the pre-trained VIT-MAE model show a similar behaviour to the one we obtained with BERT and they are also aligned with the analysis done in this work with tractable models. We highlight that one must be aware that masking on vision models is often performed via patches. In practice, this could affect the expectation in Proposition 1, so this point should be taken into consideration if theory is applied to this case.
(a) TINY-IMAGENET
(b) FASHION MNIST
(c) CIFAR100
Figure 6: Evolution of the area under the MPT curve. Comparison between three different datasets with ViT-MAE. The area under the curves is approximately the LML. Random initialization of the parameters produces MPT curves with similar low probabilities for all  $\%$  of masking. As the number of epochs increases, the curve brings higher values of log-probability for lower ratios of masking. The area also converges to the true value of LML.
Reproducibility. All the empirical studies and results are reproducible. We provide the code and details for every figure in the public repository at https://github.com/pmorenoz/MPT-LML/.
Masked pre-training and large scalability are key elements of the current success of transformer models (Vaswani et al., 2017) on natural language processing (NLP) tasks. Vision transformers (ViT) (Dosovitskiy et al., 2020) bridged the architectural gap between NLP and computer vision, making masked language modeling (MLM) suitable for images. In this regard, BeiT (Bao et al., 2022) adopted ViT and proposed to mask and predict discrete visual tokens. Most recently, masked autoencoders (He et al., 2022) also adopted masked pre-training by predicting pixel values for each masked patch, and BEiT3 (Wang et al., 2022) performs MLM on texts, images, and image-text pairs, obtaining state-of-the-art performance on all-vision and vision-language tasks. Additionally, masked pre-training has been successfully adapted to video (Tong et al., 2022), where random temporal cubes are iteratively masked and reconstructed.
The surprising ability of recent generative models to generalize and do impressive in-context learning has inspired earlier works to study this phenomenon from the Bayesian lens. The notion that LLMs might be performing implicit Bayesian inference was first described in Xie et al. (2021) where in-context learning is described as a mixture of HMMs. However, the equivalence between the log-marginal likelihood and exhaustive cross-validation was first provided in Fong and Holmes (2020). Earlier works (Vehtari and Lampinen, 2002; Gelman et al., 2014) also provided a Bayesian perspective of CV. Additionally, Moreno-Muñoz et al. (2022) leveraged this link for training Gaussian process models according to a stochastic approximation to the marginal likelihood. Similarly to current masked pre-training, the size of the conditioning variable (masking rate) was held constant. This was reported to improve notably upon traditional variational lower bounds.
5 Discussion and outlook
In this paper, we have shown that masked pre-training implicitly performs stochastic maximization of the model's marginal likelihood. The latter is generally acknowledged as being an excellent measure of a model's ability to generalize (Fong and Holmes, 2020), and our results help to explain the strong empirical performance associated with masked pre-training. We have further seen that the developed theory matches the empirical training behavior well. Moreover, we illustrated the role that the rates and the number of random samples of masking play in the estimation of the LML. We have also provided insights and a new perspective to study masked pre-training in tractable models while also finding strong similarities with LLMs.
Limitations. We have developed a formal probabilistic theory that links masked pre-training with the Bayesian principles. While we provide evidence that the impressive performance in recent large models is related to the maximization of the marginal likelihood, these methods usually introduce new elements of improvement but may not entirely fit the propositions provided in this work. In practice, this is not a limitation but a remark that there is still room for understanding the abilities of recent generative modeling. In this regard, one example might be autoregressive modeling between the masked tokens. While these are not currently analyzed in our work, we hypothesize that they could also be linked in further development to our formal propositions.
Relevance for large models using masked pre-training. We have shown empirical results of the connection between MPT and LML. This link sheds light on the understanding of generalization, particularly in recent pre-trained models. One positive outcome of our studies is the notion of having biased Bayesian estimators whenever a practitioner fixes the masking rate, e.g. to  $15\%$ . Currently, there is a significant interest in the role of masking rates in LLMs (Wettig et al., 2022). These studies could benefit from the insights provided in this paper. We also argue that the theory offers hints that may be beneficial, for instance, for uniformly sampling the mask size, instead of the current fixed-rate practice. This practice is empirically shown in the supplementary material, and it leads to unbiased estimation which may result in better performance for certain scenarios (Tay et al., 2023).
Relevance for Bayesian models. Current Bayesian modeling is dominated by approximate methods. Variational inference foregoes the ambition of training according to the marginal likelihood and instead resorts to bounds thereof. This inherently yields suboptimal models. Our theory suggests that if we can design Bayesian models in which conditioning is cheap, then we can stochastically optimize w.r.t. the true marginal likelihood easily. Beyond shedding light on the success of masked pre-training, the theory also suggests that large-scale Bayesian models could be successfully trained in the future with appropriately designed self-supervision.
The authors want to thank Yingzhen Li for her inspiring talk on MPT during the GENU'22 meeting in Copenhagen. For us, this was the starting point of fruitful discussions that led to this final work. This project has received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement 757360). This work was in part funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF20OC0062606). SH was also supported in part by a research grant (42062) from VILLUM FONDEN.
H. Bao, L. Dong, S. Piao, and F. Wei. BEIT: BERT pre-training of image transformers. International Conference on Learning Representations (ICLR), 2022.
M. Caron, H. Touvron, I. Misra, H. Jégou, J. Mairal, P. Bojanowski, and A. Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9650-9660, 2021.
J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.