MLM and MAE from the Perspective of Dropout: Some New Insights

Bernoulli Distribution

The Bernoulli distribution is perhaps the simplest probability distribution. It is a binary distribution with a value space of \{0, 1\}, where the probability of the random variable \varepsilon taking the value 1 is p, and the probability of taking 0 is 1-p. This is denoted as: \varepsilon \sim \text{Bernoulli}(p) An interesting property of the Bernoulli distribution is that all its moments are equal to p: \mathbb{E}_{\varepsilon}[\varepsilon^n] = p \times 1^n + (1-p) \times 0^n = p Thus, we know its mean is p, and its variance is: \mathbb{V}ar_{\varepsilon}[\varepsilon] = \mathbb{E}_{\varepsilon}[\varepsilon^2] - \mathbb{E}_{\varepsilon}[\varepsilon]^2 = p(1-p)

Training and Prediction

During the training phase, Dropout sets certain values to zero with probability 1-p and divides the remaining values by p. In effect, Dropout introduces a random variable \varepsilon \sim \text{Bernoulli}(p), transforming the model from f(x) to f(x\varepsilon/p). While \varepsilon can have multiple components corresponding to independent Bernoulli distributions, in most cases, the result is not fundamentally different from when \varepsilon is a scalar. Therefore, we only need to derive the case where \varepsilon is a scalar.

In "Dropout Twice Again! This Time It Achieved SOTA in Supervised Tasks", we proved that if the loss function is Mean Squared Error (MSE), the optimal prediction model after training should be: \mathbb{E}_{\varepsilon}[f(x\varepsilon/p)] This implies that we should perform multiple predictions without turning off Dropout and then average the results as the final prediction—i.e., "model averaging." However, this is computationally expensive, so in practice, we rarely do this. Instead, we typically turn off Dropout, which means changing \varepsilon/p to 1. We know that: f(x) = f(x \, \mathbb{E}_{\varepsilon}[\varepsilon]/p) Thus, turning off Dropout is essentially a form of "weight averaging" (treating \varepsilon as a random weight of the model). In other words, while the theoretical optimum is "model averaging," we usually use "weight averaging" as a first-order approximation due to computational constraints.

Dropout Perspective

For simplicity, let’s analyze a simplified version of MLM: assume that during pre-training, each token remains unchanged with probability p and is replaced by [MASK] with probability 1-p. Let the embedding of the i-th token be x_i and the embedding of [MASK] be m. We can introduce a random variable \varepsilon \sim \text{Bernoulli}(p) and represent the MLM model as: f(\dots, x_i, \dots) \quad \rightarrow \quad f(\dots, x_i \varepsilon + m(1-\varepsilon), \dots) In this way, MLM is essentially the same as Dropout; both introduce random perturbations via a Bernoulli distribution. Now, following the standard usage of Dropout, the prediction model should use "weight averaging": f(\dots, \mathbb{E}_{\varepsilon}[x_i \varepsilon + m(1-\varepsilon)], \dots) = f(\dots, x_i p + m (1-p), \dots) At this point, the inconsistency of MLM during fine-tuning becomes apparent: if we view the pre-trained MLM as a special Dropout, then the fine-tuning phase corresponds to "disabling Dropout." According to standard practice, we should change each token’s embedding to x_i p + m (1-p), but in reality, we do not; we retain the original x_i.

Correcting Embeddings

According to BERT’s default settings, during MLM training, 15% of tokens are selected for prediction. Among these 15%: 80% are replaced by [MASK], 10% remain unchanged, and 10% are replaced by a random token. Based on the analysis above, after MLM pre-training is complete, we should adjust the embeddings as follows: \text{Embedding}[i] \leftarrow 0.85 \times \text{Embedding}[i] + 0.15 \times \left( \begin{aligned} &0.8 \times \text{Embedding}[m] \, + \\ &0.1 \times \text{Embedding}[i] \, + \\ &0.1 \times \text{Avg[Embedding]} \end{aligned} \right) where \text{Embedding}[m] is the embedding of [MASK] and \text{Avg[Embedding]} is the average embedding of all tokens. In bert4keras, the reference code is as follows:

embeddings = model.get_weights()[0]  # Usually the first weight is Token Embedding
v1 = embeddings[tokenizer._token_mask_id][None]  # [MASK] Embedding
v2 = embeddings.mean(0)[None]  # Average Embedding
# Weighted average
embeddings = 0.85 * embeddings + 0.15 * (0.8 * v1 + 0.1 * embeddings + 0.1 * v2)
K.set_value(model.weights[0], embeddings)  # Re-assign values

Does this modification provide the expected improvement? I compared the experimental results of BERT and RoBERTa before and after the modification on the CLUE benchmark (baseline code refers to "Baseline in Hand with bert4keras: CLUE Benchmark Code"). The conclusion was: "no significant change."

Readers might feel disappointed: does this mean everything said before was useless? I believe that the above operation indeed mitigates the inconsistency between pre-training and fine-tuning (otherwise, we would be negating the principle of Dropout). The lack of improvement suggests that the inconsistency issue might not be as severe as imagined, at least on CLUE tasks. A similar result appeared in MacBERT, which uses synonyms to replace [MASK] during pre-training to correct this inconsistency. However, I also tested MacBERT using the same baseline code, and the results showed no significant difference compared to RoBERTa. Therefore, the necessity of correcting this inconsistency might only manifest in specific tasks or with larger masking ratios.

Dropout Perspective

As shown in the figure below, MAE divides the model into an encoder and a decoder, characterized by a "deep encoder, shallow decoder" structure. It places [MASK] tokens only in the decoder, while the encoder does not process [MASK]. Consequently, the sequence handled by the encoder becomes shorter. Crucially, MAE uses a 75% masking ratio, meaning the encoder’s sequence length is only 1/4 of the usual length. Combined with the "deep encoder, shallow decoder" design, the overall pre-training speed is more than three times faster!

We can also implement the MAE model from another perspective: removing [MASK] from the encoder is equivalent to preventing the remaining tokens from interacting with the masked tokens. For a Transformer model, interaction between tokens comes from Self-Attention. Thus, we can keep the original input but mask the corresponding columns in the Attention matrix. As shown in the figure, if the i-th token is masked, it is equivalent to forcing all elements in the i-th column of the Attention matrix to 0:

From a practical standpoint, this approach is a waste of computation, but it helps us derive an interesting theoretical result. Let there be n input tokens, and let the original Attention matrix (after softmax) be A. Define M_i as an n \times n matrix where the i-th column is 0 and all other columns are 1. Then define a random matrix \tilde{M}_i, which is an all-ones matrix with probability p and M_i with probability 1-p. The MAE model can be written as: f(\dots, A, \dots) \quad \rightarrow \quad f(\dots, \text{Norm}(A \otimes \tilde{M}_1 \otimes \tilde{M}_2 \otimes \dots \otimes \tilde{M}_n), \dots) Here, \text{Norm} refers to re-normalizing the matrix by row; \otimes denotes element-wise multiplication. When there are multiple Attention layers, they all share the same set of \tilde{M}_1, \tilde{M}_2, \dots, \tilde{M}_n.

Thus, we have converted MAE into a special type of Attention Dropout. Following the "disabling Dropout" approach for the fine-tuning phase, the corresponding model should be: \begin{aligned} &\,f(\dots, \text{Norm}(A \otimes \mathbb{E}[\tilde{M}_1 \otimes \tilde{M}_2 \otimes \dots \otimes \tilde{M}_n]), \dots)\\ =&\,f(\dots, \text{Norm}(A \otimes \mathbb{E}[\tilde{M}_1] \otimes \mathbb{E}[\tilde{M}_2] \otimes \dots \otimes \mathbb{E}[\tilde{M}_n]), \dots)\\ =&\,f(\dots, \text{Norm}(Ap), \dots)\\ =&\,f(\dots, A, \dots) \end{aligned} The second equality holds because \mathbb{E}[\tilde{M}_i] is a matrix where the i-th column is p and others are 1, so the product \mathbb{E}[\tilde{M}_1] \otimes \dots \otimes \mathbb{E}[\tilde{M}_n] is an all-p matrix. Multiplying A by this is equivalent to multiplying A by the constant p. The third equality holds because multiplying all elements by the same constant does not affect the normalization result.

From this result, we see that for MAE, "disabling Dropout" yields a model identical to the original. This indicates that MAE, compared to the original MLM model, offers not only a speed boost but also better consistency between pre-training and fine-tuning.

Preventing Overfitting

Conversely, since MAE can be viewed as a form of Dropout, and Dropout helps prevent overfitting, can we use the MAE approach as a regularization technique? As shown in the figure below, during the training phase, we can randomly drop some tokens while maintaining the original positions of the remaining tokens. We shall call this "DropToken":

The reason for this idea is that while conventional Dropout is often intuitively understood as sampling a sub-network, it actually slows down training. In contrast, DropToken explicitly shortens the sequence length, which can speed up training. If effective, it would be a very practical technique. Furthermore, some readers might have tried deleting words for data augmentation; the difference here is that DropToken retains the original positions of the remaining tokens, an implementation that relies on the Transformer structure itself.

I conducted several comparative experiments on CLUE using BERT-base as the baseline. The subscript numbers represent the drop ratio. The final effects were mixed: except for IFLYTEK, which showed clear improvement, the others were hit-or-miss (as is often the case with many regularization methods). The optimal drop ratio was between 0.1 and 0.15.