Title: On Masked Pre-training and the Marginal Likelihood
Venue: NeurIPS 2023
Accepted: yes
Keywords: Marginal likelihood

masked pre-training

Bayesian inference
Abstract: Masked pre-training removes random input dimensions and learns a model that can predict the missing values. Empirical results indicate that this intuitive form of self-supervised learning yields models that generalize very well to new domains. A theoretical understanding is, however, lacking. This paper shows that masked pre-training with a suitable cumulative scoring function corresponds to maximizing the model's marginal likelihood, which is de facto the Bayesian model selection measure of generalization. Beyond shedding light on the success of masked pre-training, this insight also suggests that Bayesian models can be trained with appropriately designed self-supervision. Empirically, we confirm the developed theory and explore the main learning principles of masked pre-training in large language models.

On Masked Pre-training and the Marginal Likelihood

Abstract

Masked pre-training removes random input dimensions and learns a model that can predict the missing values. Empirical results indicate that this intuitive form of self-supervised learning yields models that generalize very well to new domains. A theoretical understanding is, however, lacking. This paper shows that masked pretraining with a suitable cumulative scoring function corresponds to maximizing the model's marginal likelihood, which is de facto the Bayesian model selection measure of generalization. Beyond shedding light on the success of masked pre-training, this insight also suggests that Bayesian models can be trained with appropriately designed self-supervision. Empirically, we confirm the developed theory and explore the main learning principles of masked pre-training in large language models.

1 Introduction

Masked pre-training (MPT) is a family of self-supervised learning methods (Dosovitskiy et al., 2020; Devlin et al., 2018; Caron et al., 2021), that empirically has been demonstrated to result in models that generalize very well to new settings. In essence, masked pre-training removes random features of the data and learns a model to recover these from the remaining input. While empirical results are impressive, a deeper understanding of why pre-trained models generalize so well is lacking. Is it due to the use of transformer architectures (Vaswani et al., 2017), the vast over-parametrization (Neyshabur et al., 2019), or something entirely different?

The marginal likelihood or evidence is commonly used as the measure of generalization ability in Bayesian models (Tenenbaum and Griffiths, 2001; MacKay, 2003). While computationally expensive, the blessing of the marginal likelihood comes from the probabilistic integration of hypotheses. Whenever we are considering a latent variable model in the Bayesian framework, such integration can be thought of as the average over all the possible latent variable mappings, weighted by our prior beliefs. Since masked pre-training drives generalization so well, the lingering question in the Bayesian modeling community is then: Is masked pre-training somehow related to the maximization of the marginal likelihood?

In this paper , we provide a positive answer. We show that masked pre-training optimizes according to a stochastic gradient of the log-marginal likelihood (LML). Importantly, the log-marginal likelihood is equivalent to the cumulative sum of masked pre-training losses shaped with different sizes for the random mask. Even if its practical use avoids this cumulative sum, we demonstrate that choosing a

fixed masking rate, e.g. 15% as in BERT (Devlin et al., 2018), leads to a stochastic biased estimation which still maximizes the log-marginal likelihood.

Our proof relies on a previous observation from Fong and Holmes (2020), who show that the log-marginal likelihood equals the average of exhaustive leaveM -out cross-validation (CV) given posterior predictive scores. Intuitively, our formal results can be seen as the transposed version of Fong and Holmes's results: where CV removes random observations to measure generalization, masked pre-training removes random features . While the seminal link between CV and the marginal likelihood was purely a formal result that pointed out the underlying presence of Bayesian principles in a well-known class of learning, our work extends the theory behind the marginal likelihood to comprehend the impressive behavior of the latest generative models.

2 Masked pre-training

Masked pre-training (MPT) is a variant of self-supervised learning (Dosovitskiy et al., 2020; Devlin et al., 2018) that removes random input dimensions (also known as masking ) in the observed data and learns a model that accurately predicts the missing values. This family of methods, well-known due to their success in natural language understanding, typically adopts a transformer architecture (Vaswani et al., 2017) as the feature extractor, that together with positional encodings and random masked dimensions allows capturing the bidirectional context in the data.

In BERT (Devlin et al., 2018), each sentence is usually considered as a D dimensional observation vector, x = ( x 1 , x 2 , . . . , x D ) ⊤ , where dimensions x t are named tokens . Given a random mask M of size M<D , as a set of indices drawn uniformly from { 1 , . . . , D } , each token whose index belongs to M is considered to be in the subset x M = { x M (1) , x M (2) , . . . , x M ( M ) } . We refer to these as the masked tokens . The rest of indices R = { 1 , 2 , . . . , D } \ M induce the complementary subset x R , such that x = x M ∪ x R . Under this notation, MPT learns the parameters θ of a model p θ ( · ) by maximising an average of the following objective

for every observation in the dataset D . The stochastic choice of x M makes predictive conditionals p θ ( x M | x R ) to be different for every observation and training step. Once the pre-training of θ has converged, this naturally allows the model to capture the underlying structure between dimensions of the data. One additional remark is the number of random masks needed to cover all combinations between masked and observed tokens, which can be obtained as C M = ( D M ) . In the particular example of BERT, where the masking rate is 15% with D = 512 and M = 76 , the total number of random masks needed to cover all combinations of tokens is C M ≈ 1 . 21 × 10 92 . This shows the inner combinatorial problem behind MPT. We provide empirical results on why this is not a limitation for learning with MPT in Sec. 3.1.

3 A probabilistic perspective, theory, and analysis

Our key objective is to demonstrate that the good generalization of MPT can be explained as an equivalence with the model's high marginal likelihood. Indeed, we will prove that MPT implicitly maximizes marginal likelihood according to some latent variable model of the form p θ ( x | z ) .

Marginal likelihood. For our theory, we consider some dataset D consisting of n i.i.d. observations x 1: n , where each sample x i could be either continuous or discrete and is of dimensionality D . We also assume that there exists a latent space Z ∈ R K where we can find unobserved variables z 1: n which are part of the generative process of the data. This assumption is inspired in the common use of latent encodings in recent models fitted with MPT. In this direction, we also consider the observations to be samples of a likelihood function p θ ( x | z ) , where the mapping between the latent and observed variable is controlled by some parameters θ , which might also include likelihood or prior hyperparameters .

Importantly, we consider the parameters θ to be deterministic , while we are interested in integrating out the latent variables that we cannot observe. Automatically, this leads us to the log-marginal

likelihood (LML) of the model, which may factorize as a sum of marginals and can be also written as log p θ ( x 1: n ) = ∑ n i =1 log p θ ( x i ) , where the i th probability density comes from the integral p θ ( x i ) = ∫ p θ ( x i | z i ) p ( z i )d z i . This definition coincides with the target LML used in the lower bound of variational autoencoders (VAE) (Kingma and Welling, 2013; Rezende et al., 2014) and it is widely used in probabilistic generative models.

Masking and conditional probabilities. From the properties of probability distributions, we can decompose the individual LML functions log p θ ( x i ) as a sum of log-conditionals between tokens . Omitting the i th observation subscript in x to keep the notation uncluttered, the sum takes the form

However, the previous sum imposes a particular order on the selection of variables for conditioning, e.g. { x 1 | x 2 , x 3 , . . . } , { x 2 | x 3 , x 4 , . . . } , etc. Moreover, the order of tokens in the observation vector remains predetermined, as dimensions are not exchangeable . Thus, we can consider a different combination of conditional probabilities in the sum - for instance, { x 4 | x 1 , x 2 , . . . } , { x 3 | x 1 , x 2 , . . . } , etc. Here, the key insight is that the rules of probability applied to the log-marginal likelihood make it invariant to the combination of different conditional factors, as we are observing different views of the same graphical model.

This combinatorial process between tokens in x can be understood as the selection problem of indices. For that reason, we can assume a mask M of the largest size |M| = D , such that M≡{ 1 , 2 , · · · , D } . Using similar properties of combinatorics, we can also obtain D ! different choices for M . While all the indices are always in the set, the order of indices differs between combinations. This principled order in M indicates how we sum the conditional probabilities in Eq. 2.

Since the LML is invariant to random choices of M , we can re-write the sum in Eq. 2 as an expectation with a countable set of possible outcomes. Each outcome corresponds to one of the D ! choices for M , such that

where the superscript ( π ) denotes which mask M are we using for indexing the tokens. We also swapped the order of the sums to obtain the desired expectation in the r.h.s. of the formula.

The role of random masking. If we now take a particular index ( t ) and we look at the π th summand in the previous expression, we can see that the LML is still invariant to the order of the conditioning tokens x M ( t +1: D ) in the log-probabilities: log p θ ( x M ( t ) | x M ( t +1: D ) ) in the sum. Intuitively, we can use both -{ x 1 | x 2 } , { x 2 } or { x 2 | x 1 } , { x 1 } ; independently of the conditional factors previously considered. In practice, this indicates that we can insert a second set of indices to the r.h.s. variables, which is the key point to link negative MPT loss and LML.

Now, assume that M indexes less than 100% of tokens, while the rest is indexed by R as defined in Sec. 2. If we match both complementary masks to be aligned with the conditional and conditioning variables in the log-probabilities, this allows us to rewrite the t th summands in Eq. 2 as

Here, we can easily see that there are ( D t -1 ) choices for the unmasked tokens in the r.h.s. of the conditional distribution, where we have previously fixed the index t . If we set the binomial coefficient C t ≡ ( D t -1 ) as the maximum number of choices, we can obtain the following equality

since D ! > C t ∀ t ∈ { 1 , 2 , . . . , D } . Notice that once we have chosen a specific order ( π ) in the masking pattern of M and R in Eq. 4, there are still ( D -t +1) choices for the masked tokens

under evaluation in the probability distribution. Alternatively, we can think of this method as taking advantage of the properties of probability to split the D ! choices in the order of log-conditionals into the two sums in Eq. 4. The driving idea is then that the two sums in the previous expression still remain invariant given any t ∈ { 1 , 2 , . . . , D } .

Using the previous notion in Eq. 2, we obtained our main result, which holds under the assumption of i.i.d. observations with correlated tokens and the previous definition of the LML as the integral over the stochastic latent variables in the model.

Proposition 1 -The cumulative expected loss of masked pre-training along the sizes of the mask of tokens M ∈ { 1 , 2 , . . . , D } is equivalent to the log-marginal likelihood of the model when using self-predictive conditionals probabilities, such that

where the score function S θ ( · ; M ) corresponds to

Proof: In the supplementary material.

It is remarkably important to link the sum of log-conditionals log p θ ( x M ( j ) | x R ) in our proposition with the main objective used in MPT in Eq. 1. The main message of our result is that the score function S θ ( · ; t ) acts as an average over the different random masks. These shape the structure of conditioning in probabilities. The cumulative sum of the score function S θ ( · ; t ) over different sizes of the MPT's mask formally leads to the true value of the model's LML. This result is exact whenever we consider the closed-form self-predictive probabilities of the model and all the possible choices for the masking pattern M . Since this is usually not affordable, due to the combinatorial cost and the lack of tractability, we usually have a biased estimator. However, it is still sufficient to prove that MPT maximizes LML during training as we will show later. This point will be discussed in the following empirical studies. Further details on the derivations are included in the supplementary material.

3.1 Formal results in tractable models

To verify that masked pre-training effectively maximizes LML, we need a tractable probabilistic model based on latent variables as the proof-of-concept . Probabilistic PCA (PPCA) (Tipping and Bishop, 1999) is perhaps the best option here, as it has been previously used to understand other empirical observations in generative methods, e.g. posterior collapse in VAEs (Lucas et al., 2019), or even considered as the starting point of GPLVMs (Lawrence, 2005). In particular, the PPCA model assumes that Gaussian observations map linearly to sets of real-valued latent variables z 1: n , such that x = Wz + µ + ϵ , where ϵ ∼ N (0 , σ 2 0 I ) . Importantly, the prior is conventionally defined as isotropic, where p ( z ) = N (0 , 1 ) . We are therefore interested in the closed form expression of the PPCA's LML, which also factorizes across samples as follows

and we obtain the covariance matrix using S = WW ⊤ + σ 2 0 I . For our analysis, the Gaussian nature of p θ ( x i ) is of fundamental importance. Given the random mask M , the self-predictive conditionals used in MPT naturally emerge from the formulation using properties of Gaussian marginals, such that p θ ( x M | x R ) = N ( m M|R , v M|R ) is parameterized according to

where we split the LML covariance matrix S into the blocks corresponding to the indices included in M and R . We use these mean and variance parameters of the self-predictive density to recursively

evaluate the log-probabilities in Prop. 1. In practice, two elements become critical for the computation, one is the size M of the mask and another is the number of random masks P < C M considered. These induce a trade-off between accuracy and computational cost. Moreover, their role in approximating LML using a biased estimate is carefully analyzed in the following empirical studies. We also included additional details on the previous derivations in the supplementary material.

Fast asymptotic convergence. Our theory indicates that we should evaluate all C M random masks of the tokens to achieve the exact value of the LML. However, even if the combinatorial nature of the sum in the r.h.s. of the last equation in Prop. 1 becomes very large when the dimensionality of data augments, we suspect that it might converge relatively fast to the true value of the LML. This hypothesis would explain why large models that are fitted with standard MPT generalize well using just one random mask per training epoch.

Here, we empirically study if the cumulative MPT loss converges to the true value of the LML under the definition of the PPCA model. In particular, to the LML obtained with the original parameters that generated the data. The results in Fig. 1 and Tab. 1 indicate that as long as we average over more random masking patterns, the cumulative MPT loss approximates the LML of the model very well. Thus, having defined a PPCA model with a latent space of K = 2 dimensions, we observe in the left and middle plots that the asymptotic convergence happens for both small ( D = 5 ) and large ( D = 50 ) number of tokens per observation. Additionally, we observe that the estimation of LML is clearly unbiased if we use the cumulative MPT loss according to Eq. 1, which is an important insight. Notice that P = 1 is usually set up in MPT in practice. Additionally, we tested the tractable model

Figure 1: Asymptotic convergence of the cumulative MPT loss to LML as the number of random masks P augments. Curves indicate the relative difference, where 0 . 0 means that MPT equals LML. (Left). Each observation consists of 5 tokens. (Center) Each observation consists of 50 tokens. (Right). Observations have 512 tokens and the rate of masking is fixed to 15% (76 tokens). As the theory indicates, when the size of M is fixed, the cumulative MPT loss becomes a biased estimator of the LML. The curves converge asymptotically to the bias.

using a dimensionality similar to the input data used in BERT (Devlin et al., 2018), where the number of tokens is typically D = 512 per observation and the mask rate is fixed to 15% . The fact of fixing the rate of masking in MPT produces that the sum in Eq. 5 is incomplete. Thus, we have a biased estimation of the LML. However, this bias is known and constant during the training of parameters θ , which does not prevent the general maximization of LML. This point is carefully analyzed in the next empirical study with learning curves. One additional finding here is that as P →C M , the cumulative MPT loss also converges asymptotically to the biased estimator of the LML as shown in the right plot in Fig. 1.

LMLmaximization and biased estimation. We next seek to extend the previous study to understand the behavior of the cumulative MPT loss in training curves. So far, we have observed how the number of random mask patterns affects the precision around the unbiased estimation of the

LML. Theory and previous empirical results indicate that we are targeting LML or at least a decent biased estimate of LML when averaging self-predictive conditionals as in MPT. However, we still want to examine if this maximizes LML in all cases and under stochastic gradient optimization. This principal hypothesis is confirmed in Fig. 2, where different training curves are shown for different initializations and setups of the same PPCA model. The key insight showed by this experiment is that the exact LML is iteratively maximized at each epoch, in parallel with the maximization of the negative MPT loss. On the other side, we also have that MPT is an unbiased stochastic approximation of LML, as in Fig. 1, whenever we consider different rates of random masking M . We can also observe that as soon as we fix the size of the mask to index the 20% of tokens, the MPT loss becomes a biased estimate. Intuitively, this is equivalent to fixing M in the sum in Eq. 5. Again, it converges to the same value from different initializations of parameters θ . Additionally, we highlight that the LML is still maximized in this case, which is of high similarity to practical uses in larger models. Overall, this result first confirms the main insight of the work on the link between generalization when using MPT and the maximization of the model's LML.

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 θ 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

under the curve described by the score function S θ ( · ; 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.

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 110 Mparameters 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

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

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/ .

4 Related work

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.

References

Bao, L. Dong, S. Piao, and F. Wei. BEIT: BERT pre-training of image transformers. International Conference on Learning Representations (ICLR) , 2022.

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.

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.

Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Fong and C. C. Holmes. On the marginal likelihood and cross-validation. Biometrika , 107(2):489-496, 2020.

Gelman, J. Hwang, and A. Vehtari. Understanding predictive information criteria for Bayesian models. Statistics and Computing , 24(6):997-1016, 2014.

He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 16000-16009, 2022.

P. Kingma and M. Welling. Auto-encoding variational Bayes. arXiv preprint arXiv:1312.6114 , 2013.

Lawrence. Probabilistic non-linear principal component analysis with Gaussian process latent variable models. Journal of Machine Learning Research (JMLR) , 6(11), 2005.

Lucas, G. Tucker, R. B. Grosse, and M. Norouzi. Don't blame the ELBO! A linear VAE perspective on posterior collapse. Advances in Neural Information Processing Systems (NeurIPS) , 32, 2019.

J. MacKay. Information theory, inference and learning algorithms . Cambridge University Press, 2003.

Moreno-Muñoz, C. W. Feldager, and S. Hauberg. Revisiting active sets for gaussian process decoders. In Advances in Neural Information Processing Systems (NeurIPS) , 2022.

Neyshabur, Z. Li, S. Bhojanapalli, Y. LeCun, and N. Srebro. The role of over-parametrization in generalization of neural networks. In International Conference on Learning Representations (ICLR) , 2019.

J. Rezende, S. Mohamed, and D. Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning (ICML) , pages 1278-1286, 2014.

Tay, M. Dehghani, V. Q. Tran, X. Garcia, J. Wei, X. Wang, H. W. Chung, D. Bahri, T. Schuster, S. Zheng, D. Zhou, N. Houlsby, and D. Metzler. UL2: Unifying language learning paradigms. In International Conference on Learning Representations (ICLR) , 2023.

B. Tenenbaum and T. L. Griffiths. Generalization, similarity, and Bayesian inference. Behavioral and Brain sciences , 24(4):629-640, 2001.

E. Tipping and C. M. Bishop. Probabilistic principal component analysis. Journal of the Royal Statistical Society: Series B (Statistical Methodology) , 61(3):611-622, 1999.

Tong, Y. Song, J. Wang, and L. Wang. Videomae: Masked autoencoders are data-efficient learners for self-supervised video pre-training. arXiv preprint arXiv:2203.12602 , 2022.

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems (NIPS) , 30, 2017.

Vehtari and J. Lampinen. Bayesian model assessment and comparison using cross-validation predictive densities. Neural Computation , 14(10):2439-2468, 2002.

Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. International Conference on Learning Representations (ICLR) , 2019.

Wang, H. Bao, L. Dong, J. Bjorck, Z. Peng, Q. Liu, K. Aggarwal, O. K. Mohammed, S. Singhal, S. Som, et al. Image as a foreign language: BEIT pretraining for all vision and vision-language tasks. arXiv preprint arXiv:2208.10442 , 2022.

Wettig, T. Gao, Z. Zhong, and D. Chen. Should you mask 15% in masked language modeling? arXiv preprint arXiv:2202.08005 , 2022.

M. Xie, A. Raghunathan, P. Liang, and T. Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080 , 2021.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors set out to show that the good generalization of Masked Pre-Training can be explained as the equivalence with model's high marginal likelihood. The exchangeability assumption in sequential inputs is handled via a combinatorial choice over the subset of masking features, such that the masked pre-training objectives is a marginalization over uniform distribution of different maskings. The rest of the paper is a discussion on the implications of such a formulation.\n\nSTRENGTHS:\nThe paper is reasonably written, and the authors have made a decent effort to convey the significance of the work.\n\nWEAKNESSES:\nI am not very convinced by the arguments in the paper that claim that maximizing the marginal likelihood is the reason behind generalization of MPT.\n\nIn Line 182, the authors claim an \"insight\" that MPT targets a biased estimate of LML by showing that fixing the masking rate to 20% of  tokens, MPT becomes biased and converges to a different value. I feel like this argument could be made for any two objectives, and something remains missing in this justification. From the reverse-engineered theory, it is already obvious that MPT introduces bias and positioning this practical result as an insight seems to not add any new information. I'd appreciate a clarification here.\n\nIn Line 212, the authors claim that MPT performs implicit integration. However, this claim also seems a little unfounded. When averaging over the % of masked tokens in Figure 4, if bulk of the integrated mass is constant, then an alternate way to approximate the integral would be to just take the constant mass (since the integral contains averaging terms as well). Doesn't this result confound the conclusion and just a matter of coincidence?\n\nIn any case, marginal likelihood just answers a question very different from generalization - marginal likelihood tells us the likelihood of the data under the \"prior\" model, i.e. how well the priors explain the data. It does not have anything to do with what happens to the model after training, in principle. There could be confounders to the \"insights\" claimed by the authors here, especially with large models, since the choice of number of tokens, number of maskings, are far-off from what model large language models use.\n\nQUESTIONS:\n1. Perhaps I am misunderstanding, but shouldn't Eq. (2) be conditioned on all feature variables $1:D$ except for the one being predicted? Eq. (3) seems to present is correctly, since it clearly distinguishes masking indices.\n\nLIMITATIONS:\nYes. See my comments on weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper derives the equivalence between log marginal likelihood and negative masked pre-training loss which is often used for training text models. The paper argues that as marginal likelihood is the Bayesian way to do model selection, doing training using the masked pre-training loss inherits good generalization properties.\n\nExperiments are done on a probabilistic PCA model, VAE and BERT. In PPCA, where the ground truth is known and the equivalence is confirmed empirically. On this model, it is empirically confirmed that even when we optimize a biased estimation of the log marginal likelihood, different training runs still converge to the same log marginal likelihood, with a slight offset to the true value.\n\nIn VAE, we don't know the ground truth, but verify that optimizing using the masked pretraining loss is similar to optimizing with ELBO, in terms of the log marginal likelihood value that the model converges to.\n\nIn the BERT model, we simply observe that the shape of the S(x, M) term is similar for PPCA and BERT during different points of training.\n\nSTRENGTHS:\n- This is an important connection to make and the authors discuss relevance for LLMs and Bayesian learning in the discussion.\n- The experiments study the realistic case in which we can't enumerate over all masks and verify that masked pre-training is still similar to log marginal likelihood training\n\nWEAKNESSES:\n- The derivations are likely correct but I couldn't establish this 100%, see questions below.\n- Isn't the Transformer just trained with the autoregressive log likelihood (and not with masked pre-training)? And isn't ViT simply trained to predict object class -- where is the masking?\n- The jump to Proposition 1 was big for me and the preceding intuition building parts didn't help\n\nQUESTIONS:\n- In equation (1), do we not condition on $\\mathbf x_{\\mathcal M \\setminus \\{\\mathcal M(t)\\}}$ -- how is the marginalization over this variable done in language models?\n- Is there a reason why we condition on $\\mathbf x_{t + 1:D}$ instead of $\\mathbf x_{1:t}$?\n- In equation before equation 4, we assume $M < D$, so the value $\\mathbf x_D$ which is part of $\\mathbf x_{\\mathcal M(t+ 1:D)}$ in the left hand side doesn't exist\n- In line 109, shouldn't it be ${D \\choose t}$ instead of ${D \\choose t - 1}$ since there are $t$ tokens we condition on?\n\nLIMITATIONS:\n-", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper shows that MLM is effectively maximizing the model's marginal likelihood, perhaps explaining why MLM has been successful. Beyond providing a proof, they run several empirical experiments suggesting their results hold in practice.\n\nSTRENGTHS:\n1. Clear introduction; Clear highlighting of paper strengths.\n2. Clear background on MLM.\n3. They provide a proof in the appendix and a number of empirical tests that support their findings in practice.\n4. Their discussion showing connections both to model design in AI as well as relevance to improving Bayesian Models seems well put together.\n\nI don't deeply understand the proof and math, so am I unable to fully judge the strength of the contribution. For what I do appreciate, I would recommend the paper to be accepted. I found the paper to be interesting, overall well-presented, and noteworthy.\n\nWEAKNESSES:\n1. Providing more hand-holding between the various subsections of section 3 would help the reader understand your goals for each section. Right now I can see the specific finding, but: Stressing the greater point that you're trying to verify your proof in practice (?) in each of these subsections and where and how you find support would be great.\n2. Providing slightly more hand-holding on what marginal likelihood is and what it would be capturing in the context of, say, BERT, would make the paper stronger.\n\nNits\n1. Figure 1: The y-axis get smaller (key detail!) but I could not see this initially. The text was very light. Perhaps increasing the font size somehow or making more of a note of it in the caption.\n2. Figure 5: y-axis is both small and cut off.\n\nQUESTIONS:\nOn line 207 (and beyond) you make some distinction between linear vs non-linear models. Does your proof make assumptions along these lines? If so, can you bring these assumptions further up into the abstract?\n\nLIMITATIONS:\n\"The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, authors are encouraged to provide a short proof sketch to provide intuition.\" More of a proof sketch might help the reader understand the footing of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
They show empirically that randomly masking a fixed number of M tokens and predicting with the remaining tokens produces a biased estimate of the log marginal likelihood $\log p(x)$ (LML).
Furthermore, they prove that repeating the masking for M from 1 to D (the total number of tokens), and summing up the estimates, leads to an unbiased estimate of the LML.

STRENGTHS:
- The paper is well written and mostly easy to follow.
- Intuitive proofs and arguments are well supported by empirical results for tractable and intractable models.
- The work might give insights/ideas for setting masking rates in pre-training.

WEAKNESSES:
Main Concern:
The main result Proposition 1, requires to condition on different subsets of the remaining the tokens $R$. But this is not what is done in practice for masked pre-training, where normally all tokens in $R$ are used. It would have been interesting to see some discussion/experiments about the effect.

Others:
- It would have been interesting to see some experiments about the ideas stated in line 295.
- Given the works like Fong and Holmes (2020), there is no so much novelty from a theoretical point of view.

QUESTIONS:
- Could you comment on my main concern above?
- Why does the LML change with the number of epochs?
- What is GT-LML?

LIMITATIONS:
see main concern.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes viewing masked pretraining methods (like BERT or MAE) as optimizing a biased form of the log marginal likelihood $\\log p_\\theta(x)$. Using a parallel to exhaustive leave-M-out cross validation, the authors derive the exact relationship between the LML and a weighted sum of masked pretraining losses across mask sizes. The authors show that maximizing the masked pretraining loss improves the LML, across a wide variety of settings, from probabilistic PCA to VAEs to BERT.\n\nSTRENGTHS:\n- I found Prop. 1, the connection between LML and the weighted sum of masked pretraining losses, to be original and quite enlightening. \n- Paper is thorough and provides extensive empirical experiments to confirm the theory. \n- Finding seems like a good step towards understanding masked pretraining objectives, which are extremely widespread in both NLP (BERT) and vision (MAE), and potentially understanding how to improve them. \n- Paper is well written and clear, especially if the appendix is used to supplement the derivations in Section 3.\n\nWEAKNESSES:\n- The biggest issue is that the paper shows that masked pretraining performance improves LML, but does not address why it does better than other training objectives that directly maximize LML (such as an autoregressive model).\n- No experiments on the vision side, even though MAE is an extremely popular method that falls neatly into the framework. \n- Some parts of the paper are a tiny bit confusing. Surprisingly, I felt that Appendix A and B could be swapped into the main paper and it would significantly improve the clarity of Proposition 1's derivation or the PPCA setting. \n  - I felt that the second paragraph of the intro (L19-26) didn't sufficiently motivate why the LML is important.\n\nQUESTIONS:\n**Major questions**: \n- If maximizing LML is the ultimate reason why MPT is good, then why are other methods of directly maximizing LML (e.g., autoregressive models) worse?\n- Can you run the same BERT experiments with pretrained MAE models as well, to show applicability to vision models?\n\n**Minor questions and suggestions**:\n- Where did this \"log-marginal likelihood\" name come from? It looks just like the $\\log p_\\theta(x)$ that all generative models maximize, so I'm confused about the naming. \n- L19-26: This paragraph doesn't really convince me that the marginal likelihood is that important. \n- L27-28: \"masked pre-training optimizes according to a stochastic gradient of the log-marginal likelihood\" -- I think this sentence should be removed, since the following sentence states the contribution more clearly. \n- L33: \"who shows\" -> \"who show\"\n- L74: \"which might also include likelihood\" -- what is likelihood?\n- L86: \"previous sum\" -- is this referring to Eq. 2?\n- L87-92: I understand what it's trying to say after reading the entire paper, but this paragraph is very confusing.  What do exchangeable and invariant mean?\n- The block of text between L108-109: \"to be aligned\" -- this is also confusing. It's explained much better in the appendix!\n  - Generally, Appendix A explains this proof much better than the main text. \n- L128: \"we usually have a biased estimator\" -- wouldn't a Monte Carlo estimator, with masks drawn from the right distribution, be unbiased?\n- L157: \"just one random mask per training epoch\" -- my understanding is that masks are different for each sample in each epoch. I'm reading this as using the same random mask for all samples in a given epoch. \n  - L165: \"notice that P=1 is usually set up in MPT in practice\" -- same issue as above. What exactly is being claimed about the number of masks?\n- L170: \"this bias is known\" -- known in what sense? The numerical value of the bias, the symbolic form of the bias, or just the fact that the MPT objective is not equal to the LML objective?\n- L183: \"the exact LML is iteratively maximized at each epoch\" -- I wouldn't use the word \"maximize\" here, as it seems like it increases but is not guaranteed to even reach a local minimum unless there are further assumptions on the structure of the model. \n  - Same issue with L204-205\n- Fig. 2 caption: \"The rate of masking is unfixed and it varies from 1% to 100%\" -- I may be misunderstanding here, but with 10 tokens, shouldn't the masking rate start at 10%?\n- L194-195: 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.\" -- why do there need to be latent variables at all? Latent variables are one way to induce a structured joint distribution over all dimensions of x, but this doesn't have to be used. In fact, latent variables aren't used in the proof of Prop 1. \n- L202: \"as well as we compare\" -> \"as well as compare\"\n- L209-210: \"we used iterative encoding-decoding to produce the self-conditional probabilities for masked tokens — see Sec. F in Rezende et al. (2014)\" -- it would be good to explain this procedure here. I glanced at Sec F and it seems like it only gives a way to sample from the marginal over masked tokens, not a way to compute the probability? I may be misunderstanding, which is why it would be nice to have a clear explanation in this paragraph.\n- L212-214 on implicit integration: these lines are confusing and would benefit from being rewritten.\n- Fig. 3 caption: \"Darker lines indicate the evolution of the LML\" -- the figure legend does not match this. Also, i don't understand why the lower right plot is separate from the lower left plot, or why the value is so different. \n- Fig. 4: \"looses\" -> \"loses.\" Also, I don't understand what the black lines in Fig. 4 (right) is supposed to denote. \n- L218-219: \"This explains why the probabilities have an approximately exponential decay\" -- I don't see how anything explains this.\n- L224: \"effect to\" -> \"affect\"\n- L225: \"a longer discussion is provided in the supplementary material\" -- Appendix C.2 doesn't seem very relevant to this section. \n- Table 2 caption: \"area under the MPT curve\" -- is this weighted by the number of masks or not? Wouldn't the weighted sum based on Prop. 1 directly measure LML here?\n- Fig. 5: y axis label is cut off on the far left plot.\n\nLIMITATIONS:\nA major limitation that isn't addressed by the authors is this: why is MPT the right way to improve LML, instead of another method that directly maximizes LML? Empirically, masked pretraining methods like BERT/MAE are much better representation learners than alternate approaches that optimize the LML. The authors mention \"autoregressive modeling\" in their limitations section, but I don't think enough weight is put on it as a large question that this paper on MPT<-> LML introduces yet doesn't answer.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-10"}