BERT Pre-training › Combined Pre-training Objective

So far we've examined BERT's two objectives separately: MLM teaches the model to understand words in context, while NSP teaches it to recognize document-level coherence. But how do these objectives work together during training? The answer lies in how we combine their signals into a single optimization target.

Consider what each objective teaches in isolation. MLM forces the model to build rich token representations by predicting masked words from surrounding context. A model that excels at MLM understands syntax and semantics, along with factual knowledge encoded in word patterns. However, MLM operates at the sentence level. It doesn't explicitly teach the model whether two sentences belong together or flow logically from one to the other.

NSP addresses this gap. By asking "does sentence B follow sentence A?", the model must develop representations that capture discourse relationships, topic coherence, and logical flow. The [CLS] token becomes a summary of the entire sentence pair, encoding whether they form a coherent unit.

Training on both objectives simultaneously allows the model to develop representations that are rich at multiple levels: tokens capture local meaning, while the [CLS] token captures global coherence. The question becomes: how do we combine these two learning signals?

BERT uses the simplest possible combination: add the losses together. Each batch of training data produces two loss values, and the total loss is their sum:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{MLM}} + \mathcal{L}_{\text{NSP}}$$

where:

• $\mathcal{L}_{\text{total}}$: the total pre-training loss that BERT minimizes through gradient descent
• $\mathcal{L}_{\text{MLM}}$: the masked language modeling loss, measuring how well the model predicts masked tokens
• $\mathcal{L}_{\text{NSP}}$: the next sentence prediction loss, measuring how well the model classifies sentence pairs

This equal weighting (1.0 for each) means both objectives contribute equally to the gradient updates. The model cannot cheat by ignoring one task to optimize the other. It must find parameters that satisfy both constraints simultaneously.

The MLM loss measures how surprised the model is by the correct answers at masked positions. If the model confidently predicts the right token, the loss is low. If it assigns low probability to the correct token, the loss is high.

Mathematically, we sum the negative log-probabilities across all masked positions:

$$\mathcal{L}_{\text{MLM}} = -\sum_{i \in \mathcal{M}} \log P_\theta(x_i | \tilde{x})$$

Let's unpack each component:

• $\mathcal{M}$: the set of masked position indices. For a 512-token sequence with 15% masking, this contains approximately 77 positions.
• $x_i$: the original token at position $i$ before we corrupted it. This is the "ground truth" the model must recover.
• $\tilde{x}$: the corrupted input sequence where masked positions show [MASK], random tokens, or unchanged tokens (following the 80-10-10 rule).
• $P_\theta(x_i | \tilde{x})$: the probability our model (with parameters $\theta$) assigns to the correct token $x_i$, given the entire corrupted sequence as context.

The negative sign and logarithm work together to create our loss function. When $P_\theta(x_i | \tilde{x})$ is close to 1 (confident and correct), $\log(1) = 0$, contributing zero loss. When $P_\theta(x_i | \tilde{x})$ is small (the model missed the answer), $\log(\text{small number})$ becomes a large negative value, and the negative sign makes it a large positive loss. This asymmetry creates strong gradients that push the model to assign high probability to correct tokens.

The NSP loss has a simpler structure because it's a binary classification problem. The model looks at the [CLS] token's representation and must decide: IsNext or NotNext?

$$\mathcal{L}_{\text{NSP}} = -\log P_\theta(y_{\text{NSP}} | h_{\text{[CLS]}})$$

The symbols in this expression are:

• $y_{\text{NSP}}$: the true binary label. We encode IsNext as 1 and NotNext as 0. This comes from how we constructed the training data (50% consecutive pairs, 50% random pairs).
• $h_{\text{[CLS]}}$: the hidden representation of the [CLS] token after passing through all transformer layers. This 768-dimensional vector (for BERT-base) summarizes the entire sentence pair.
• $P_\theta(y_{\text{NSP}} | h_{\text{[CLS]}})$: the probability the model assigns to the correct label, computed by passing $h_{\text{[CLS]}}$ through the NSP classification head.

The structure mirrors the MLM loss: negative log-probability of the correct answer. When the model confidently predicts the right class, loss is near zero. When it's wrong or uncertain, loss increases.

Both losses share the same transformer backbone. When we compute the gradient of $\mathcal{L}_{\text{total}}$, it flows back through both the MLM head and the NSP head, then merges in the transformer layers. This means every transformer layer receives gradient signal from both objectives.

The interplay teaches complementary skills:

1. Token embeddings learn to represent individual words in ways that support both masked token prediction (MLM) and sentence-pair understanding (NSP).

2. Attention patterns develop to capture both local dependencies (which tokens help predict a masked word) and global coherence (what makes two sentences related).

3. The [CLS] token becomes especially important for NSP, learning to aggregate sentence-level meaning, while other positions focus more on token-level prediction.

This multi-task learning is why BERT's representations transfer so well to diverse downstream tasks. The model has learned to encode information at multiple levels of abstraction.

Let's translate these formulas into code. PyTorch's cross_entropy function computes exactly the negative log-probability we described, making the implementation straightforward. The key insight is handling the MLM labels: we use -100 for positions that weren't masked, and PyTorch's ignore_index parameter automatically excludes these from the loss calculation.