ELECTRA › Joint Training
Language AI HandbookAll chapters ELECTRA › Joint TrainingOne section · Sections in this chapter7 / 18 ↓
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
ELECTRA trains the generator and discriminator simultaneously. The total loss combines both objectives into a single scalar that drives gradient updates:
where:
- : the masked language modeling loss for the generator
- : the replaced token detection loss for the discriminator
- : a weighting factor (typically 50) that balances the two losses
The high weight on RTD ensures the discriminator is the primary focus of training. Without this upweighting, the generator's gradients would dominate because MLM involves a 30,000-way classification (much harder per token than binary classification).
The generator's MLM loss is computed only at masked positions:
where:
- : the set of masked positions (typically 15% of sequence length)
- : the original token at position (the target we want the generator to predict)
- : the input sequence with
[MASK]tokens at positions in - : the probability the generator assigns to the correct token given the masked input
This is standard cross-entropy loss: the generator is rewarded for assigning high probability to the original tokens at masked positions.
def compute_generator_loss(gen_logits, original_ids, masked_positions):
"""
Compute MLM loss for the generator.
Args:
gen_logits: (batch_size, seq_len, vocab_size)
original_ids: (batch_size, seq_len) - original token IDs
masked_positions: (batch_size, seq_len) - bool mask of masked positions
Returns:
Scalar loss
"""
batch_size, seq_len, vocab_size = gen_logits.shape
# Flatten for cross-entropy
logits_flat = gen_logits.view(-1, vocab_size)
labels_flat = original_ids.view(-1)
mask_flat = masked_positions.view(-1)
# Compute loss only at masked positions
loss = F.cross_entropy(logits_flat, labels_flat, reduction="none")
loss = (loss * mask_flat).sum() / mask_flat.sum()
return loss
def electra_training_step(
generator,
discriminator,
input_ids,
mask_token_id,
mask_prob=0.15,
rtd_weight=50.0,
):
"""
Complete ELECTRA training step.
Returns:
Dictionary with losses and statistics
"""
# Step 1: Create masked input
masked_ids, masked_positions = create_masked_input(
input_ids, mask_token_id, mask_prob
)
# Step 2: Generator forward pass
gen_logits = generator(masked_ids)
gen_loss = compute_generator_loss(gen_logits, input_ids, masked_positions)
# Step 3: Sample replacements (no gradient through sampling)
with torch.no_grad():
corrupted_ids = generator.sample_replacements(
masked_ids, masked_positions
)
# Step 4: Determine which tokens differ from original
is_replaced = (corrupted_ids != input_ids).long()
# Step 5: Discriminator forward pass
disc_logits = discriminator(corrupted_ids)
disc_loss = compute_discriminator_loss(disc_logits, is_replaced)
# Step 6: Combined loss
total_loss = gen_loss + rtd_weight * disc_loss
# Statistics
with torch.no_grad():
disc_preds = (disc_logits > 0).long()
disc_accuracy = (disc_preds == is_replaced).float().mean()
gen_accuracy = (
(corrupted_ids == input_ids)[masked_positions].float().mean()
)
return {
"total_loss": total_loss,
"gen_loss": gen_loss,
"disc_loss": disc_loss,
"disc_accuracy": disc_accuracy,
"gen_accuracy": gen_accuracy,
"replacement_rate": is_replaced.float().mean(),
}# Demonstrate a training step
torch.manual_seed(42)
# Reset models
generator = ElectraGenerator(
vocab_size=30522,
hidden_size=256,
num_layers=4,
num_heads=4,
intermediate_size=1024,
)
discriminator = ElectraDiscriminator(
vocab_size=30522,
hidden_size=768,
num_layers=12,
num_heads=12,
intermediate_size=3072,
)
# Sample input
input_ids = torch.randint(100, 1000, (4, 64))
# Training step
results = electra_training_step(
generator, discriminator, input_ids, mask_token_id=103
)ELECTRA Training Step Results: Generator loss: 10.5056 Discriminator loss: 0.7005 Total loss: 45.5294 Accuracies: Generator (predicting masked tokens): 0.0% Discriminator (detecting replacements): 48.8% Replacement rate: 18.4%
At initialization, both networks perform near chance level: the generator accuracy is close to random guessing among 30,000 vocabulary tokens, and the discriminator accuracy hovers around 50% for binary classification. The generator loss is high (around 10. This reflects the log of vocab size), while the discriminator loss is near 0.69 (the random baseline for binary cross-entropy). During training, both losses decrease as the generator learns to predict masked tokens and the discriminator learns to detect replacements.
The weight is important and deserves explanation. The generator loss is computed over approximately 15% of tokens with a 30,000-way classification at each position, while the discriminator loss is computed over 100% of tokens but with only binary classification at each position.
The cross-entropy loss magnitude scales with the log of the number of classes: for MLM, this is approximately per position, while for RTD it's per position. Even though RTD has more positions, the per-position loss magnitude is much smaller. Without upweighting RTD, the MLM gradients would dominate the shared embeddings. The factor ensures both networks receive comparable gradient magnitudes, allowing the discriminator to learn effectively.
