Replaced Token Detection › The Complete Training Loop
Language AI HandbookAll chapters Replaced Token Detection › The Complete Training LoopOne section · Sections in this chapter8 / 17 ↓
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
Training RTD involves three steps per batch: mask inputs, generate replacements, and update both networks. It is worth pausing to appreciate how all the pieces fit together before seeing the code. The full loop has a specific order that matters: you must complete the generator forward pass and the sampling step before starting the discriminator forward pass, because the discriminator's input depends on the generator's output. This sequential dependency means you cannot parallelize the two forward passes. However, the two backward passes (one through the generator loss, one through the discriminator loss) are independent and can both be computed from the single combined loss through standard backpropagation.
One practical implementation note: the sampling step uses torch.no_grad() or explicit detachment to prevent gradients from flowing from the discriminator's loss back through the sampling operation into the generator. As mentioned earlier, this sampling step is non-differentiable. If you attempt to backpropagate through it without blocking gradients, PyTorch will raise an error or silently produce incorrect gradients. The generator learns only from its own MLM cross-entropy loss.
def compute_rtd_loss(
generator,
discriminator,
token_ids,
mask_token_id,
mask_prob=0.15,
disc_weight=50.0,
):
"""Compute combined RTD loss for generator and discriminator."""
# Step 1: Mask tokens
masked_ids, gen_labels, mask_positions = create_masked_input(
token_ids, mask_token_id, mask_prob
)
# Step 2: Generator forward pass and loss
gen_logits = generator(masked_ids)
gen_loss = F.cross_entropy(
gen_logits.view(-1, gen_logits.size(-1)),
gen_labels.view(-1),
ignore_index=-100,
)
# Step 3: Sample replacements (detach to not backprop through sampling)
with torch.no_grad():
gen_probs = F.softmax(gen_logits, dim=-1)
replaced_ids = masked_ids.clone()
disc_labels = torch.ones_like(token_ids, dtype=torch.float) # 1 = original
for batch_idx in range(token_ids.size(0)):
for pos in range(token_ids.size(1)):
if mask_positions[batch_idx, pos]:
# Sample from generator
sampled = torch.multinomial(
gen_probs[batch_idx, pos], num_samples=1
)
replaced_ids[batch_idx, pos] = sampled
# Label: 0 if replaced with different token, 1 if same (or original)
if sampled != token_ids[batch_idx, pos]:
disc_labels[batch_idx, pos] = 0.0
# Step 4: Discriminator forward pass and loss
disc_logits = discriminator(replaced_ids)
disc_loss = F.binary_cross_entropy_with_logits(disc_logits, disc_labels)
# Combined loss
total_loss = gen_loss + disc_weight * disc_loss
return total_loss, gen_loss, disc_lossLet's train on a small corpus to see the dynamics:
# Create a small training corpus
corpus_tokens = []
sentences = [
["the", "chef", "cooked", "the", "meal"],
["the", "waiter", "served", "the", "food"],
["the", "chef", "ate", "the", "food"],
["a", "waiter", "cooked", "a", "meal"],
]
for sent in sentences:
corpus_tokens.append([word_to_idx[w] for w in sent])
corpus_tensor = torch.tensor(corpus_tokens)
# Initialize models
generator = RTDGenerator(vocab_size, d_model=32, n_heads=2, n_layers=1)
discriminator = RTDDiscriminator(vocab_size, d_model=48, n_heads=2, n_layers=2)
# Optimizers
gen_optimizer = torch.optim.AdamW(generator.parameters(), lr=1e-3)
disc_optimizer = torch.optim.AdamW(discriminator.parameters(), lr=1e-3)
# Training loop
gen_losses = []
disc_losses = []
for step in range(200):
# Sample a batch
batch_idx = torch.randint(0, len(corpus_tensor), (2,))
batch = corpus_tensor[batch_idx]
# Compute loss
total_loss, gen_loss, disc_loss = compute_rtd_loss(
generator, discriminator, batch, mask_token_id, mask_prob=0.3
)
# Update both models
gen_optimizer.zero_grad()
disc_optimizer.zero_grad()
total_loss.backward()
gen_optimizer.step()
disc_optimizer.step()
gen_losses.append(gen_loss.item())
disc_losses.append(disc_loss.item())Initial generator loss: 2.0001 Final generator loss: 0.9318 Initial discriminator loss: 0.6686 Final discriminator loss: 0.3077
Both losses decrease over training. The generator loss drops as it learns to predict masked tokens, while the discriminator loss decreases as it learns to distinguish original from replaced tokens. The initial discriminator loss around 0.69 corresponds to random guessing (binary cross-entropy for 50/50 predictions), confirming the model starts with no knowledge of which tokens are replaced.
The training dynamics show both losses decreasing. The generator learns to predict masked tokens better, while the discriminator learns to detect replacements. These objectives are complementary: a better generator produces harder replacements, which in turn trains a better discriminator.
In practice, you might notice that the discriminator loss decreases more quickly at first and then plateaus while the generator continues improving. This makes sense intuitively. Early in training, the generator produces very poor replacements (essentially random tokens from the vocabulary), and the discriminator can detect these easily. As the generator learns basic language statistics, its replacements become more plausible, and the discriminator must work harder to maintain low loss. The final performance depends on this equilibrium: can the discriminator develop representations subtle enough to catch the generator's best attempts at plausible replacements?
Another practical consideration is the masking probability used during training. The original ELECTRA paper uses 15% masking, matching BERT. This fraction determines how many positions are "potentially replaced" in the discriminator's view. With 15% masking, roughly 15% of positions have been modified by the generator (though some of those modifications may have landed on the correct token anyway). The remaining 85% are guaranteed originals. A higher masking rate would give the discriminator more replaced examples to learn from, but it would also degrade the generator's context window, making it harder to produce plausible replacements. The 15% rate represents a well-calibrated trade-off backed by extensive empirical testing.