Core Concepts of Large Language Models
Training Objectives: Language Modeling, Masked Language Modeling, and Next Sentence Prediction
Training objectives define what a large language model learns during pre-training, shaping its ability to understand and generate text. These objectives are critical because they determine the model's foundational knowledge, enabling it to generalize across tasks without task-specific fine-tuning. You reach for these objectives when building or adapting a model for tasks like text completion, question answering, or coherence evaluation.
Language Modeling: Predicting the Next Token
Language modeling (LM) is the simplest and most fundamental training objective, where the model learns to predict the next token in a sequence given the preceding tokens. The core idea is that by training the model to minimize the difference between its predicted next token and the actual next token, it develops an understanding of grammar, syntax, and even factual knowledge embedded in the training data. This objective works because it forces the model to capture dependencies between tokens, such as how adjectives modify nouns or how verbs agree with subjects. Over time, the model internalizes patterns in the data, allowing it to generate coherent and contextually appropriate text. LM is particularly effective because it leverages the vast amount of unstructured text available, requiring no manual labeling. The model learns to assign probabilities to sequences, which is useful for tasks like text generation or autocomplete. However, LM alone may not capture higher-level semantics, such as the relationship between sentences, which is where other objectives come into play.
# Example: Training loop for language modeling (simplified)
import torch
import torch.nn as nn
# Assume we have a model, tokenizer, and batch of tokenized text
model = ... # Your LLM model
tokenizer = ... # Tokenizer for your model
batch = ["The cat sat on the mat", "Dogs are loyal animals"]
# Tokenize the batch and convert to input IDs
inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
# Shift input_ids to create labels (next token prediction)
labels = input_ids[:, 1:].clone() # Labels are the next tokens
input_ids = input_ids[:, :-1] # Input is everything except the last token
attention_mask = attention_mask[:, :-1]
# Forward pass
outputs = model(input_ids=input_ids, attention_mask=attention_mask)
logits = outputs.logits # Shape: (batch_size, seq_len, vocab_size)
# Compute loss (cross-entropy between logits and labels)
loss_fn = nn.CrossEntropyLoss()
loss = loss_fn(logits.view(-1, logits.size(-1)), labels.view(-1))
# Backward pass and optimization
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f"Language Modeling Loss: {loss.item():.4f}")
# This loss drives the model to predict the next token accurately.Why Language Modeling Captures Context
The power of language modeling lies in its ability to capture context through the self-supervised nature of next-token prediction. When the model predicts the next token, it must consider all preceding tokens in the sequence, effectively learning to weigh the importance of each word in the context. For example, in the sentence 'The bank of the river is muddy,' the model learns that 'bank' refers to a riverbank, not a financial institution, by attending to the words 'river' and 'muddy.' This contextual understanding emerges because the model is trained to minimize prediction error across diverse examples, forcing it to generalize rather than memorize. The attention mechanism in transformers further enhances this by allowing the model to dynamically focus on relevant parts of the input. Over time, the model develops a nuanced understanding of word senses, syntactic roles, and even pragmatic implications. However, LM has limitations: it may struggle with long-range dependencies or tasks requiring explicit reasoning, as it primarily learns local patterns. This is why additional objectives like Masked Language Modeling are introduced to address these gaps.
# Example: Visualizing attention weights to understand context
from transformers import AutoModel, AutoTokenizer
import torch
# Load a pre-trained model and tokenizer
model_name = "gpt2" # A small LM for demonstration
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name, output_attentions=True)
# Tokenize input and get model outputs
text = "The bank of the river is muddy"
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)
# Extract attention weights (shape: [num_layers, num_heads, seq_len, seq_len])
attentions = outputs.attentions
# Average attention across heads and layers for simplicity
avg_attention = torch.stack(attentions).mean(dim=(0, 1))
# Print attention weights for the word "bank"
bank_idx = tokenizer.encode("bank", add_special_tokens=False)[0]
print(f"Attention weights for 'bank' (token ID: {bank_idx}):")
for i, token in enumerate(tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])):
print(f" {token}: {avg_attention[0, bank_idx, i].item():.3f}")
# Output shows how "bank" attends to "river" more than other words.Masked Language Modeling: Filling in the Blanks
Masked Language Modeling (MLM) is a training objective where random tokens in the input sequence are masked, and the model learns to predict the original tokens based on the surrounding context. Unlike language modeling, which predicts the next token, MLM forces the model to consider both left and right context, making it better suited for tasks requiring bidirectional understanding, such as text classification or question answering. The key insight behind MLM is that masking creates a 'fill-in-the-blank' task, which encourages the model to develop a deeper understanding of word relationships. For example, in the sentence 'The [MASK] barks loudly,' the model must use 'barks' and 'loudly' to infer that the masked word is likely 'dog.' This bidirectional context is particularly useful for tasks where the meaning of a word depends on both preceding and succeeding words. MLM is also more sample-efficient than LM because each input sequence provides multiple training signals (one for each masked token). However, MLM introduces a discrepancy between pre-training and inference, as the model never sees masked tokens during real-world use, which can lead to slight performance degradation.
# Example: Masked Language Modeling training loop
import torch
import torch.nn as nn
from transformers import BertTokenizer, BertForMaskedLM
# Load a pre-trained MLM model and tokenizer
model = BertForMaskedLM.from_pretrained("bert-base-uncased")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
# Example batch of text
batch = ["The cat sat on the mat", "Dogs are loyal animals"]
# Tokenize and mask random tokens (15% of tokens, as in BERT)
inputs = tokenizer(batch, return_tensors="pt", padding=True, truncation=True)
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
# Create labels by cloning input_ids, then mask 15% of tokens
labels = input_ids.clone()
probability_matrix = torch.full(labels.shape, 0.15)
# Mask tokens (80% masked, 10% random, 10% unchanged)
masked_indices = torch.bernoulli(probability_matrix).bool()
input_ids[masked_indices] = tokenizer.mask_token_id
# Forward pass
outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
# Backward pass and optimization
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f"Masked Language Modeling Loss: {loss.item():.4f}")
# The model learns to predict masked tokens using bidirectional context.Why Masked Language Modeling Improves Bidirectional Understanding
The strength of Masked Language Modeling (MLM) lies in its ability to train the model to integrate information from both directions in a sequence, which is critical for tasks requiring deep contextual understanding. In language modeling, the model only sees left context (preceding tokens), which limits its ability to resolve ambiguities that depend on right context. For example, in the sentence 'I went to the bank to deposit my [MASK],' the model must use 'deposit' to infer that 'bank' refers to a financial institution, not a riverbank. MLM forces the model to attend to both left and right context, enabling it to develop a more robust representation of word meanings. This bidirectional understanding is particularly valuable for tasks like named entity recognition or coreference resolution, where the meaning of a word depends on its surroundings. Additionally, MLM's masking strategy creates a more challenging learning task, as the model cannot rely on simple n-gram patterns and must instead learn deeper semantic relationships. However, MLM's reliance on masking introduces a trade-off: the model becomes highly effective at filling in blanks but may struggle with tasks requiring unidirectional generation, such as text completion.
# Example: Comparing LM and MLM predictions
from transformers import AutoModelForCausalLM, AutoModelForMaskedLM, AutoTokenizer
# Load models and tokenizer
lm_model = AutoModelForCausalLM.from_pretrained("gpt2")
mlm_model = AutoModelForMaskedLM.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
# Example sentence with ambiguity
text = "I went to the bank to deposit my [MASK]"
# MLM prediction (bidirectional)
mlm_inputs = tokenizer(text, return_tensors="pt")
mlm_outputs = mlm_model(**mlm_inputs)
mlm_pred = tokenizer.decode(mlm_outputs.logits[0, 6].argmax().item())
print(f"MLM prediction for [MASK]: '{mlm_pred}'") # Likely "money" or "check"
# LM prediction (unidirectional, adapted for comparison)
lm_text = "I went to the bank to deposit my "
lm_inputs = tokenizer(lm_text, return_tensors="pt")
lm_outputs = lm_model.generate(**lm_inputs, max_length=len(lm_inputs["input_ids"][0]) + 1)
lm_pred = tokenizer.decode(lm_outputs[0][-1])
print(f"LM prediction for next token: '{lm_pred}'") # May struggle with ambiguityNext Sentence Prediction: Learning Coherence
Next Sentence Prediction (NSP) is a training objective designed to teach the model to understand the relationship between two sentences, specifically whether the second sentence logically follows the first. Unlike language modeling or masked language modeling, which focus on token-level predictions, NSP operates at the sentence level, encouraging the model to learn discourse coherence and high-level semantics. The rationale behind NSP is that many downstream tasks, such as question answering or document summarization, require the model to evaluate whether two pieces of text are meaningfully connected. For example, given the sentences 'The cat sat on the mat' and 'It was raining outside,' the model should predict that the second sentence does not follow the first, as there is no logical connection. NSP works by presenting the model with pairs of sentences, where 50% of the time the second sentence is the actual next sentence in the document, and 50% of the time it is a randomly sampled sentence from the corpus. The model learns to distinguish between these cases, developing an understanding of topic continuity, causal relationships, and narrative flow. While NSP has been somewhat deprecated in favor of more advanced objectives like sentence order prediction, it remains a useful tool for tasks requiring explicit sentence-level reasoning.
# Example: Next Sentence Prediction training loop
from transformers import BertTokenizer, BertForNextSentencePrediction
import torch
# Load a pre-trained NSP model and tokenizer
model = BertForNextSentencePrediction.from_pretrained("bert-base-uncased")
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
# Example sentence pairs (50% positive, 50% negative)
sentence_pairs = [
("The cat sat on the mat", "It was a sunny day", 1), # Positive (follows)
("The cat sat on the mat", "I like to eat pizza", 0), # Negative (random)
]
# Tokenize and prepare inputs
inputs = tokenizer(
[pair[0] for pair in sentence_pairs],
[pair[1] for pair in sentence_pairs],
return_tensors="pt",
padding=True,
truncation=True
)
labels = torch.tensor([pair[2] for pair in sentence_pairs]) # 1 for positive, 0 for negative
# Forward pass
outputs = model(**inputs, labels=labels)
loss = outputs.loss
logits = outputs.logits
# Backward pass and optimization
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f"Next Sentence Prediction Loss: {loss.item():.4f}")
print(f"Predictions: {torch.argmax(logits, dim=1)}") # Should match labelsKey points
- Language modeling trains a model to predict the next token in a sequence, enabling it to learn grammar, syntax, and factual knowledge from unstructured text.
- Masked language modeling improves bidirectional understanding by forcing the model to predict masked tokens using both left and right context, which is critical for tasks like question answering.
- Next sentence prediction teaches the model to evaluate whether two sentences are logically connected, enhancing its ability to handle discourse-level tasks.
- The attention mechanism in transformers allows models to dynamically focus on relevant parts of the input, improving their ability to capture long-range dependencies.
- Language modeling is sample-efficient but may struggle with tasks requiring explicit reasoning, while masked language modeling addresses this by leveraging bidirectional context.
- Training objectives like MLM and NSP introduce trade-offs, such as the discrepancy between pre-training and inference or the challenge of unidirectional generation.
- The choice of training objective depends on the downstream task: LM for generation, MLM for understanding, and NSP for coherence evaluation.
- Understanding the strengths and limitations of each objective allows you to design or adapt models for specific applications, such as text completion or document summarization.
Common mistakes
- Mistake: Treating language modeling as just predicting the next word without considering long-range dependencies. Why it's wrong: This oversimplifies the task and leads to poor coherence in generated text. Fix: Use architectures like Transformers that capture long-range context through self-attention mechanisms.
- Mistake: Assuming masked language modeling (MLM) only requires predicting random masked tokens. Why it's wrong: Random masking can lead to trivial patterns (e.g., predicting 'the' for every '[MASK]'). Fix: Use dynamic or whole-word masking strategies to create meaningful challenges for the model.
- Mistake: Ignoring the importance of negative samples in next sentence prediction (NSP). Why it's wrong: Without hard negatives, the model can't learn meaningful relationships between sentences. Fix: Include semantically unrelated but plausible sentence pairs as negatives during training.
- Mistake: Using the same dataset for all three objectives (LM, MLM, NSP) without balancing. Why it's wrong: This can cause the model to overfit to one task's patterns. Fix: Curate separate or balanced datasets for each objective to ensure diverse learning signals.
- Mistake: Evaluating the model solely on perplexity for language modeling. Why it's wrong: Perplexity doesn't capture semantic coherence or factual accuracy. Fix: Use downstream tasks (e.g., question answering, summarization) to evaluate real-world performance.
Interview questions
What is the primary goal of language modeling in the context of creating your own LLM?
The primary goal of language modeling when creating your own Large Language Model is to predict the probability of a sequence of words or tokens in a given language. This is fundamental because it allows the model to understand and generate coherent text by learning the statistical patterns of language. For example, in a simple autoregressive language model, the objective is to maximize the likelihood of the next token given the previous ones, which can be represented as P(w_t | w_1, w_2, ..., w_{t-1}). This is typically trained using cross-entropy loss on a large corpus of text. The model learns to capture grammar, semantics, and even some world knowledge, which are essential for tasks like text generation or completion.
Can you explain how masked language modeling works and why it’s useful for training LLMs?
Masked language modeling, or MLM, is a training objective where random tokens in a sentence are masked, and the model is tasked with predicting the original tokens. For example, in the sentence 'The cat sat on the [MASK],' the model learns to predict 'mat' based on the surrounding context. This approach is useful because it allows the model to learn bidirectional representations, meaning it considers both left and right context to make predictions. Unlike autoregressive models, which only look at previous tokens, MLM helps the model develop a deeper understanding of language structure and semantics. This is particularly valuable for tasks like text understanding, question answering, or filling in missing information, as it trains the model to infer meaning from incomplete data.
What is next sentence prediction, and how does it contribute to training an LLM?
Next sentence prediction, or NSP, is a training objective where the model is given two sentences and must determine whether the second sentence logically follows the first. For example, given 'The sun is shining brightly' and 'It is a great day for a picnic,' the model should predict 'IsNext.' Conversely, if the second sentence is 'The moon is full tonight,' it should predict 'NotNext.' NSP contributes to training an LLM by helping it understand relationships between sentences, such as coherence, causality, or topic continuity. This is particularly useful for tasks like document summarization, dialogue systems, or any application requiring an understanding of discourse. While NSP has been somewhat deprioritized in newer models like BERT in favor of more efficient objectives, it remains a valuable tool for teaching models about sentence-level relationships.
How do you implement masked language modeling in code for a custom LLM? Provide a high-level example.
To implement masked language modeling for a custom LLM, you typically follow these steps. First, you tokenize your input text and randomly mask a percentage of the tokens, usually around 15%. For example, in PyTorch, you might use a transformer model with an embedding layer and a classification head. Here’s a simplified example: ```python tokens = tokenizer.encode('The cat sat on the mat') masked_tokens = tokens.copy() mask_indices = random.sample(range(len(tokens)), int(0.15 * len(tokens))) for i in mask_indices: masked_tokens[i] = tokenizer.mask_token_id outputs = model(torch.tensor([masked_tokens])) predictions = outputs.logits.argmax(dim=-1) ``` The model’s goal is to predict the original tokens at the masked positions. You train this using cross-entropy loss, comparing the predicted tokens to the original ones. This forces the model to learn contextual representations of words, as it must rely on the surrounding tokens to make accurate predictions.
Compare masked language modeling and autoregressive language modeling. What are the strengths and weaknesses of each?
Masked language modeling and autoregressive language modeling are two distinct approaches with unique strengths and weaknesses. Masked language modeling, like in BERT, excels at understanding bidirectional context, meaning it can use both left and right context to predict a masked token. This makes it particularly strong for tasks requiring deep comprehension, such as question answering or text classification, where understanding the full context is critical. However, MLM is less effective for generative tasks like text completion, as it doesn’t naturally learn to generate sequences autoregressively. Autoregressive language modeling, like in GPT, predicts the next token based only on previous tokens. This makes it ideal for generative tasks, as it learns to produce coherent and contextually relevant text. However, it struggles with tasks requiring bidirectional understanding, as it cannot see future tokens. In summary, MLM is better for understanding, while autoregressive modeling is better for generation.
How would you design a training pipeline for an LLM that combines masked language modeling and next sentence prediction? Explain the trade-offs of this approach.
Designing a training pipeline that combines masked language modeling and next sentence prediction involves creating a multi-task learning setup. First, you’d preprocess your data by tokenizing sentences and pairing them for NSP, while also masking tokens for MLM. For example, you might use a transformer architecture with two classification heads: one for predicting masked tokens and another for NSP. During training, you’d compute a combined loss, such as a weighted sum of the MLM and NSP losses, to optimize both objectives simultaneously. The trade-offs of this approach include increased computational complexity, as the model must handle two tasks at once, which can slow down training. However, the benefit is that the model gains both bidirectional understanding from MLM and sentence-level coherence from NSP, making it more versatile for tasks like document understanding or dialogue systems. The challenge lies in balancing the two objectives, as overemphasizing one may weaken the other. For instance, if NSP is too dominant, the model might struggle with token-level predictions, and vice versa.
Check yourself
1. Why is masked language modeling (MLM) considered more challenging than standard language modeling for training LLMs?
- A.MLM requires predicting multiple tokens at once, increasing computational complexity.
- B.MLM forces the model to rely on bidirectional context, making it harder to exploit local patterns.
- C.MLM is only used for fine-tuning, not pre-training, so it’s less optimized.
- D.MLM uses a fixed vocabulary, limiting the model’s ability to generalize to new words.
Show answer
B. MLM forces the model to rely on bidirectional context, making it harder to exploit local patterns.
The correct answer is that MLM forces the model to rely on bidirectional context. Unlike standard language modeling (which predicts the next token using left context), MLM masks random tokens and requires the model to infer them using both left and right context, making it harder to 'cheat' with local patterns. The other options are incorrect: MLM predicts one masked token at a time (not multiple), it’s a pre-training objective, and it doesn’t restrict vocabulary.
2. A student trains an LLM using next sentence prediction (NSP) but finds the model performs poorly on downstream tasks. What is the most likely cause?
- A.The model was trained on a dataset with too many positive sentence pairs (e.g., consecutive sentences).
- B.The negative samples were too easy (e.g., sentences from unrelated documents).
- C.The model’s architecture lacks a classification head for NSP.
- D.The learning rate was too high during NSP training.
Show answer
B. The negative samples were too easy (e.g., sentences from unrelated documents).
The correct answer is that the negative samples were too easy. If negatives are trivial (e.g., sentences from unrelated documents), the model doesn’t learn meaningful relationships between sentences. Hard negatives (e.g., semantically plausible but non-consecutive sentences) are critical for NSP’s effectiveness. The other options are incorrect: Too many positive pairs would bias the model but not necessarily hurt performance; NSP doesn’t require a separate head; and learning rate is a secondary concern.
3. How does the choice of masking strategy in MLM impact the model’s ability to handle rare words?
- A.Whole-word masking ensures rare words are never masked, preserving their representations.
- B.Dynamic masking increases the likelihood of masking rare words, improving their representations.
- C.Span masking (masking contiguous tokens) reduces the model’s exposure to rare words.
- D.Random masking has no effect on rare words because all tokens are treated equally.
Show answer
B. Dynamic masking increases the likelihood of masking rare words, improving their representations.
The correct answer is that dynamic masking increases the likelihood of masking rare words. Dynamic masking (e.g., changing masked tokens in each epoch) ensures the model sees rare words in masked positions more often, forcing it to learn better representations. The other options are incorrect: Whole-word masking doesn’t exclude rare words; span masking can include rare words; and random masking does affect rare words (they’re less likely to be masked due to lower frequency).
4. A student observes that their LLM trained with MLM achieves low perplexity but generates incoherent text. What is the most probable explanation?
- A.The model was trained on a dataset with too many masked tokens (e.g., 50% masking rate).
- B.The model’s attention mechanism was disabled during MLM training.
- C.The MLM objective was combined with language modeling, creating conflicting gradients.
- D.The masking strategy didn’t account for sentence boundaries, disrupting coherence.
Show answer
D. The masking strategy didn’t account for sentence boundaries, disrupting coherence.
The correct answer is that the masking strategy didn’t account for sentence boundaries. If masks frequently split phrases or clauses, the model struggles to learn coherent structures. The other options are incorrect: High masking rates don’t inherently cause incoherence; attention is critical for MLM; and combining MLM with LM doesn’t create conflicts (they’re often used together).
5. Why might a student choose to omit next sentence prediction (NSP) when training their own LLM, despite its original inclusion in models like BERT?
- A.NSP requires labeled data, which is harder to obtain than unlabeled text for MLM.
- B.Recent research shows NSP doesn’t improve performance on downstream tasks like question answering.
- C.NSP is computationally expensive because it requires processing pairs of sentences.
- D.NSP is incompatible with architectures that use relative positional embeddings.
Show answer
B. Recent research shows NSP doesn’t improve performance on downstream tasks like question answering.
The correct answer is that recent research shows NSP doesn’t improve performance on downstream tasks. Studies (e.g., RoBERTa) found that removing NSP and focusing on MLM with larger batches and longer sequences yields better results. The other options are incorrect: NSP doesn’t require labeled data (it uses heuristics); it’s not computationally expensive; and it’s compatible with relative positional embeddings.