BERT Pre-training › Next Sentence Prediction
Language AI HandbookAll chapters BERT Pre-training › Next Sentence PredictionOne section · Sections in this chapter4 / 12 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Masked language modeling teaches BERT about word-level and sentence-level patterns. But understanding full documents requires reasoning about how sentences relate to each other. Does this sentence follow logically from the previous one? Are these two sentences from the same document?
Next Sentence Prediction (NSP) provides this document-level signal. Given two sentences A and B, the model must predict whether B immediately follows A in the original corpus.
A binary classification task where the model predicts whether sentence B is the actual next sentence after sentence A in the source document (IsNext) or a random sentence from elsewhere in the corpus (NotNext).
How NSP Works
The [CLS] token's representation at the final layer is used for the NSP prediction. A small classification head (typically a single linear layer) maps this representation to two logits, one for IsNext and one for NotNext.
The training procedure constructs balanced batches with 50% positive (actual next sentence) and 50% negative (random sentence) pairs. The NSP loss is cross-entropy between the predicted probabilities and the binary label.
class NSPHead(nn.Module):
"""Classification head for Next Sentence Prediction."""
def __init__(self, hidden_size):
super().__init__()
self.classifier = nn.Linear(hidden_size, 2)
def forward(self, cls_embedding):
# cls_embedding: (batch_size, hidden_size)
logits = self.classifier(cls_embedding)
return logits # (batch_size, 2)
# Example usage
hidden_size = 768
nsp_head = NSPHead(hidden_size)
# Simulate batch of CLS embeddings
batch_size = 4
cls_embeddings = torch.randn(batch_size, hidden_size)
nsp_logits = nsp_head(cls_embeddings)Input CLS embeddings shape: torch.Size([4, 768]) NSP logits shape: torch.Size([4, 2]) Predicted probabilities: Example 1: IsNext=0.714, NotNext=0.286 Example 2: IsNext=0.733, NotNext=0.267 Example 3: IsNext=0.732, NotNext=0.268 Example 4: IsNext=0.701, NotNext=0.299
The NSP Controversy
NSP was designed to help BERT understand document structure, particularly for tasks like question answering where reasoning across sentences matters. However, subsequent research questioned its value.
RoBERTa (2019) found that removing NSP entirely and training on longer contiguous sequences improved performance. The researchers hypothesized that NSP's negative examples (random sentences from different documents) were too easy to distinguish. This provides weak signal. Topic differences between sentences gave obvious cues without requiring deeper understanding.
ALBERT took a different approach, replacing NSP with Sentence Order Prediction (SOP). Instead of distinguishing consecutive sentences from random ones, SOP asks whether two consecutive sentences are in the correct order or swapped. This harder task forced the model to learn finer-grained discourse relationships.
Despite the controversy, understanding NSP matters for working with pre-trained BERT models. Many available checkpoints were trained with NSP, and the [CLS] token was optimized for this task. When fine-tuning for classification, you're building on representations shaped by both MLM and NSP.
Implementing NSP
Let's implement a complete NSP training step to see how the pieces fit together:
def prepare_nsp_batch(sentence_pairs, tokenizer, max_length=128):
"""
Prepare a batch of sentence pairs for NSP training.
Args:
sentence_pairs: List of (sent_a, sent_b, is_next) tuples
tokenizer: BERT tokenizer
max_length: Maximum sequence length
Returns:
Dictionary with input_ids, attention_mask, token_type_ids, labels
"""
sentences_a = [pair[0] for pair in sentence_pairs]
sentences_b = [pair[1] for pair in sentence_pairs]
labels = torch.tensor([1 if pair[2] else 0 for pair in sentence_pairs])
# Tokenize with sentence pairs
encoding = tokenizer(
sentences_a,
sentences_b,
padding=True,
truncation=True,
max_length=max_length,
return_tensors="pt",
)
encoding["labels"] = labels
return encoding# Create sample pairs
nsp_pairs = [
("The weather is nice today.", "Let's go for a walk.", True),
(
"Python is a programming language.",
"It supports multiple paradigms.",
True,
),
("The cat is sleeping.", "Stock prices rose sharply.", False),
("She studied hard for the exam.", "The movie was entertaining.", False),
]
batch = prepare_nsp_batch(nsp_pairs, tokenizer)Batch contents: input_ids shape: torch.Size([4, 17]) attention_mask shape: torch.Size([4, 17]) token_type_ids shape: torch.Size([4, 17]) labels: [1, 1, 0, 0] Decoded first example: [CLS] the weather is nice today. [SEP] let ' s go for a walk. [SEP] Label: IsNext
The token_type_ids distinguish sentence A (0) from sentence B (1). This helps the model understand which tokens belong to which segment, enabling it to reason about sentence relationships.