Replaced Token Detection › Implementing the Discriminator
Language AI HandbookAll chapters Replaced Token Detection › Implementing the DiscriminatorOne section · Sections in this chapter7 / 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
The discriminator is a larger transformer that performs binary classification at each position. Its architecture is similar to BERT, but the output layer produces a single logit per position rather than vocabulary logits. This change in the output head is the entire architectural difference between BERT and an ELECTRA discriminator. Everything else, the embedding layer, the multi-head attention, the feed-forward sublayers, the residual connections, is identical. What changes is only what the model is asked to predict.
The binary classification head is deliberately small. A single linear layer maps the hidden dimension to one logit, which is passed through a sigmoid to produce a probability between 0 and 1. This contrasts sharply with BERT's output head for MLM, which maps the hidden dimension to vocabulary size (30,000+ values) through a full linear layer. The ELECTRA discriminator's head is orders of magnitude smaller, which means essentially all of its capacity goes into building rich contextual representations in the transformer layers rather than into vocabulary projection.
class RTDDiscriminator(nn.Module):
"""Discriminator network for detecting replaced tokens."""
def __init__(
self, vocab_size, d_model=128, n_heads=4, n_layers=4, max_len=128
):
super().__init__()
# Embeddings (shared embedding table in ELECTRA, separate here for clarity)
self.token_emb = nn.Embedding(vocab_size, d_model)
self.pos_emb = nn.Embedding(max_len, d_model)
self.layer_norm = nn.LayerNorm(d_model)
# Larger transformer encoder
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=d_model * 4,
dropout=0.1,
batch_first=True,
activation="gelu",
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
# Binary classification head (one output per position)
self.classifier = nn.Sequential(
nn.Linear(d_model, d_model),
nn.GELU(),
nn.Linear(d_model, 1),
)
def forward(self, input_ids):
batch_size, seq_len = input_ids.shape
# Get embeddings
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
x = self.token_emb(input_ids) + self.pos_emb(positions)
x = self.layer_norm(x)
# Encode
x = self.encoder(x)
# Binary classification at each position
logits = self.classifier(x).squeeze(-1) # (batch, seq_len)
return logits# Initialize discriminator (larger than generator)
discriminator = RTDDiscriminator(vocab_size, d_model=64, n_heads=4, n_layers=2)
# Get discriminator predictions
with torch.no_grad():
disc_logits = discriminator(replaced_ids)
disc_probs = torch.sigmoid(disc_logits)Generator parameters: 17,774 Discriminator parameters: 113,409 Ratio (disc/gen): 6.4x Input sequence: ['the', 'chef', '[CLS]', 'the', 'chef'] Discriminator probabilities (is original): [0.536 0.501 0.471 0.548 0.54 ] True labels (0=replaced, 1=original): [1, 1, 0, 1, 0]
The discriminator has roughly 2-3x more parameters than the generator, matching the ELECTRA design principle. The probabilities show the untrained discriminator's random guesses. With training, these should approach 1.0 for original tokens and 0.0 for replaced tokens, correctly identifying which positions contain generator-produced replacements.
Notice that the discriminator outputs a probability per position, not a single probability for the whole sequence. This is a token-level classification task, not a sentence-level one. Each position independently receives a score, and each position independently contributes to the loss. This is fundamentally different from sentence-level classification tasks like sentiment analysis, where a single output represents the entire sequence. The token-level framing is what enables the dense training signal: with a sentence-level objective, we would still have just one gradient per sequence.