SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Pablo Moreno-Muñoz
Section for Cognitive Systems Technical University of Denmark (DTU) pabmo@dtu.dk
Pol G. Recasens*
CROMAI, Barcelona Supercomputing Center  
Universitat Politécnica de Catalunya (UPC)  
pol.garcia@bsc.es
Section for Cognitive Systems Technical University of Denmark (DTU) sohau@dtu.dk
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.
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
*Work done during an Erasmus exchange in DTU (Denmark).
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 leave-M-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,  $\pmb{x} = (x_{1}, x_{2}, \ldots, x_{D})^{\top}$ , where dimensions  $x_{t}$  are named tokens. Given a random mask  $\mathcal{M}$  of size  $M < D$ , as a set of indices drawn uniformly from  $\{1, \ldots, D\}$ , each token whose index belongs to  $\mathcal{M}$  is considered to be in the subset  $\pmb{x}_{\mathcal{M}} = \{x_{\mathcal{M}(1)}, x_{\mathcal{M}(2)}, \ldots, x_{\mathcal{M}(M)}\}$ . We refer to these as the masked tokens. The rest of indices  $\mathcal{R} = \{1, 2, \ldots, D\} \setminus \mathcal{M}$  induce the complementary subset  $\pmb{x}_{\mathcal{R}}$ , such that  $\pmb{x} = \pmb{x}_{\mathcal{M}} \cup \pmb{x}_{\mathcal{R}}$ . Under this notation, MPT learns the parameters  $\theta$  of a model  $p_{\theta}(\cdot)$  by maximising an average of the following objective
$$
\log p _ {\theta} \left(\boldsymbol {x} _ {\mathcal {M}} \mid \boldsymbol {x} _ {\mathcal {R}}\right) \approx \sum_ {t = 1} ^ {M} \log p _ {\theta} \left(x _ {\mathcal {M} (t)} \mid \boldsymbol {x} _ {\mathcal {R}}\right) \tag {1}
$$
for every observation in the dataset  $\mathcal{D}$ . The stochastic choice of  $x_{\mathcal{M}}$  makes predictive conditionals  $p_{\theta}(\boldsymbol{x}_{\mathcal{M}}|\boldsymbol{x}_{\mathcal{R}})$  to be different for every observation and training step. Once the pre-training of  $\theta$  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  $\mathcal{C}_M = \binom{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  $\mathcal{C}_M \approx 1.21 \times 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_{\theta}(x|z)$ .
Marginal likelihood. For our theory, we consider some dataset  $\mathcal{D}$  consisting of  $n$  i.i.d. observations  $\pmb{x}_{1:n}$ , where each sample  $\pmb{x}_i$  could be either continuous or discrete and is of dimensionality  $D$ . We also assume that there exists a latent space  $\mathcal{Z} \in \mathbb{R}^K$  where we can find unobserved variables  $\pmb{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_{\theta}(\pmb{x}|\pmb{z})$ , where the mapping between the latent and observed variable is controlled by some parameters  $\theta$ , which might also include likelihood or prior hyperparameters.
Importantly, we consider the parameters  $\theta$  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_{\theta}(\pmb{x}_{1:n}) = \sum_{i=1}^{n} \log p_{\theta}(\pmb{x}_i)$ , where the  $i^{\text{th}}$  probability density comes from the integral  $p_{\theta}(\pmb{x}_i) = \int p_{\theta}(\pmb{x}_i | \pmb{z}_i) p(\pmb{z}_i) \, \mathrm{d}\pmb{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_{\theta}(\pmb{x}_i)$  as a sum of log-conditionals between tokens. Omitting the  $i^{\text{th}}$  observation subscript in  $\pmb{x}$  to keep the notation uncluttered, the sum takes the form
$$
\log p _ {\theta} (\boldsymbol {x}) = \sum_ {t = 1} ^ {D} \log p _ {\theta} \left(x _ {t} \mid \boldsymbol {x} _ {t + 1: D}\right). \tag {2}
$$
However, the previous sum imposes a particular order on the selection of variables for conditioning, e.g.  $\{x_{1}|x_{2},x_{3},\ldots \} ,\{x_{2}|x_{3},x_{4},\ldots \}$ , 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,\dots \} ,\{x_3|x_1,x_2,\dots \}$  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  $\pmb{x}$  can be understood as the selection problem of indices. For that reason, we can assume a mask  $\mathcal{M}$  of the largest size  $|\mathcal{M}| = D$ , such that  $\mathcal{M} \equiv \{1, 2, \dots, D\}$ . Using similar properties of combinatorics, we can also obtain  $D!$  different choices for  $\mathcal{M}$ . While all the indices are always in the set, the order of indices differs between combinations. This principled order in  $\mathcal{M}$  indicates how we sum the conditional probabilities in Eq. 2.
Since the LML is invariant to random choices of  $\mathcal{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  $\mathcal{M}$ , such that
$$
\log p _ {\theta} (\boldsymbol {x}) = \frac {1}{D !} \sum_ {\pi = 1} ^ {D!} \sum_ {t = 1} ^ {D} \log p _ {\theta} \left(x _ {\mathcal {M} (t)} ^ {(\pi)} | \boldsymbol {x} _ {\mathcal {M} (t + 1: D)} ^ {(\pi)}\right) = \sum_ {t = 1} ^ {D} \mathbb {E} _ {\pi} \left[ \log p _ {\theta} \left(x _ {\mathcal {M} (t)} ^ {(\pi)} | \boldsymbol {x} _ {\mathcal {M} (t + 1: D)} ^ {(\pi)}\right) \right], \tag {3}
$$
where the superscript  $(\pi)$  denotes which mask  $\mathcal{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  $\pi^{\mathrm{th}}$  summand in the previous expression, we can see that the LML is still invariant to the order of the conditioning tokens  $x_{\mathcal{M}(t + 1:D)}$  in the log-probabilities:  $\log p_{\theta}(x_{\mathcal{M}(t)}|x_{\mathcal{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  $\mathcal{M}$  indexes less than  $100\%$  of tokens, while the rest is indexed by  $\mathcal{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^{\mathrm{th}}$  summands in Eq. 2 as
$$
\frac {1}{D !} \sum_ {\pi = 1} ^ {D!} \log p _ {\theta} \left(x _ {\mathcal {M} (t)} ^ {(\pi)} | \boldsymbol {x} _ {\mathcal {M} (t + 1: D)} ^ {(\pi)}\right) = \frac {1}{D !} \sum_ {\pi = 1} ^ {D!} \log p _ {\theta} \left(x _ {\mathcal {M} (t)} ^ {(\pi)} | \boldsymbol {x} _ {\mathcal {R} (1: D - t)} ^ {(\pi)}\right).
$$
Here, we can easily see that there are  $\binom{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  $\mathcal{C}_t \equiv \binom{D}{t-1}$  as the maximum number of choices, we can obtain the following equality
$$
\sum_ {\pi = 1} ^ {D!} \log p _ {\theta} \left(x _ {\mathcal {M} (t)} ^ {(\pi)} \mid \boldsymbol {x} _ {\mathcal {R} (1: D - t)} ^ {(\pi)}\right) = \sum_ {\pi = 1} ^ {\mathcal {C} _ {t}} \sum_ {j = 1} ^ {D - t + 1} \log p _ {\theta} \left(x _ {\mathcal {M} (j)} ^ {(\pi)} \mid \boldsymbol {x} _ {\mathcal {R} (1: D - t)} ^ {(\pi)}\right), \tag {4}
$$
since  $D! > \mathcal{C}_t \forall t \in \{1,2,\ldots,D\}$ . Notice that once we have chosen a specific order  $(\pi)$  in the masking pattern of  $\mathcal{M}$  and  $\mathcal{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 \in \{1,2,\dots ,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 \in \{1,2,\dots ,D\}$  is equivalent to the log-marginal likelihood of the model when using self-predictive conditionals probabilities, such that
$$
\log p _ {\theta} (\boldsymbol {x}) = \sum_ {M = 1} ^ {D} \mathcal {S} _ {\theta} (\boldsymbol {x}; M), \tag {5}
$$
where the score function  $S_{\theta}(\cdot ;M)$  corresponds to
$$
\mathcal {S} _ {\theta} (\boldsymbol {x}; M) := \frac {1}{\mathcal {C} _ {M}} \sum_ {\pi = 1} ^ {\mathcal {C} _ {M}} \frac {1}{M} \sum_ {j = 1} ^ {M} \log p _ {\theta} (x _ {\mathcal {M} (j)} ^ {(\pi)} | \boldsymbol {x} _ {\mathcal {R} (1: D - j)} ^ {(\pi)}) = \frac {1}{M} \mathbb {E} _ {\mathcal {M}} \left[ \sum_ {j = 1} ^ {M} \log p _ {\theta} (x _ {\mathcal {M} (j)} | \boldsymbol {x} _ {\mathcal {R}}) \right].
$$
Proof: In the supplementary material.
It is remarkably important to link the sum of log-conditionals  $\log p_{\theta}(x_{\mathcal{M}(j)}|\boldsymbol{x}_{\mathcal{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_{\theta}(\cdot ;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_{\theta}(\cdot ;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  $\mathcal{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  $\pmb{x} = \pmb{W}\pmb{z} + \pmb{\mu} + \pmb{\epsilon}$ , where  $\epsilon \sim \mathcal{N}(0,\sigma_0^2\mathbb{I})$ . Importantly, the prior is conventionally defined as isotropic, where  $p(z) = \mathcal{N}(0,1)$ . We are therefore interested in the closed form expression of the PPCA's LML, which also factorizes across samples as follows
$$
\log p _ {\theta} \left(\boldsymbol {x} _ {1: n}\right) = \sum_ {i = 1} ^ {n} \log p _ {\theta} \left(\boldsymbol {x} _ {i}\right), \quad \text {w h e r e} p _ {\theta} \left(\boldsymbol {x} _ {i}\right) = \mathcal {N} \left(\boldsymbol {x} _ {i} \mid \boldsymbol {\mu}, S\right), \tag {6}
$$
and we obtain the covariance matrix using  $\mathbf{S} = \mathbf{W}\mathbf{W}^{\top} + \sigma_{0}^{2}\mathbb{I}$ . For our analysis, the Gaussian nature of  $p_{\theta}(\pmb{x}_i)$  is of fundamental importance. Given the random mask  $\mathcal{M}$ , the self-predictive conditionals used in MPT naturally emerge from the formulation using properties of Gaussian marginals, such that  $p_{\theta}(\pmb{x}_{\mathcal{M}}|\pmb{x}_{\mathcal{R}}) = \mathcal{N}(\pmb{m}_{\mathcal{M}|\mathcal{R}},\pmb{v}_{\mathcal{M}|\mathcal{R}})$  is parameterized according to
$$
\boldsymbol {m} _ {\mathcal {M} | \mathcal {R}} = \boldsymbol {S} _ {\mathcal {M R}} ^ {\top} \boldsymbol {S} _ {\mathcal {R R}} ^ {- 1} \boldsymbol {x} _ {\mathcal {R}}, \quad \boldsymbol {v} _ {\mathcal {M} | \mathcal {R}} = \boldsymbol {S} _ {\mathcal {M M}} + \boldsymbol {S} _ {\mathcal {M R}} ^ {\top} \boldsymbol {S} _ {\mathcal {R R}} ^ {- 1} \boldsymbol {S} _ {\mathcal {M R}}, \tag {7}
$$
where we split the LML covariance matrix  $S$  into the blocks corresponding to the indices included in  $\mathcal{M}$  and  $\mathcal{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 < \mathcal{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  $\mathcal{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
Table 1: Evolution of negative MPT loss w.r.t. max. number of random masks  $P$  .
TRUE LML (↑)	P = 1	P = 10	P = 100
(-60.34)	-60.44 ± 0.47	-60.22 ± 0.12	-60.34 ± 0.03
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  $\mathcal{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  $\theta$  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 \to \mathcal{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.
LML maximization 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  $\mathcal{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  $\theta$ . 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.
Figure 2: Training curves of the negative cumulative MPT loss in PPCA vs. the ground truth (GT) LML. The number of samples is  $N = 2000$  and the number of tokens is  $D = 10$ . All plots used  $P = 1$  random masks per epoch and five different initializations. (Left). The rate of masking is unfixed and it varies from  $1\%$  until  $100\%$ . The negative MPT loss converges to the GT-LML (dashed line). Darker curves are the exact LML per epoch. (Center). Convergence with fixed mask to  $20\%$  of tokens. The negative MPT loss is no longer centered around the LML and it converges to a biased estimate. (Right). Zoomed curves of convergence. The bias is constant and all MPT losses converge to the same point. The LML per epoch is also maximized and converges to GT-LML.