Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

ELECTRA › Discriminator Training

Michael BrenndoerferUpdated July 25, 20254 min read

Language AI HandbookAll chapters

Part 24 of 67BERT and Variants

Chapter 175ELECTRA › Discriminator Training

Reading section 5 of 18 · Discriminator Training · 4 minSection 5 of 18 · 4 min

Section 5 of 18 · 4 min read

ELECTRA › Discriminator TrainingOne section · Sections in this chapter5 / 18 ↓

Sections in this chapter18

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 the main model. It receives the corrupted sequence (where masked positions have been filled by the generator) and predicts for each token whether it's original or replaced. Unlike the generator's vocabulary-sized output, the discriminator outputs a single logit per position.

In[12]:

Code

class ElectraDiscriminator(nn.Module):
    """Large transformer that detects replaced tokens."""

    def __init__(
        self,
        vocab_size,
        hidden_size,
        num_layers,
        num_heads,
        intermediate_size,
        max_position=512,
        dropout=0.1,
    ):
        super().__init__()

        # Embeddings
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = nn.Embedding(max_position, hidden_size)
        self.embedding_norm = nn.LayerNorm(hidden_size)
        self.embedding_dropout = nn.Dropout(dropout)

        # Transformer layers
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=intermediate_size,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)

        # Discrimination head: binary classification per token
        self.discriminator_head = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask=None):
        batch_size, seq_len = input_ids.shape

        # Embeddings
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        x = self.embedding_norm(x)
        x = self.embedding_dropout(x)

        # Create attention mask for transformer
        if attention_mask is not None:
            src_key_padding_mask = attention_mask == 0
        else:
            src_key_padding_mask = None

        # Encode
        x = self.encoder(x, src_key_padding_mask=src_key_padding_mask)

        # Binary classification per token
        logits = self.discriminator_head(x).squeeze(-1)

        return logits

The discriminator loss is binary cross-entropy computed over all tokens, including tokens that were not candidates for replacement. This is the key to ELECTRA's efficiency: the loss involves every position.

In[13]:

Code

def compute_discriminator_loss(disc_logits, is_replaced, attention_mask=None):
    """
    Compute RTD loss for the discriminator.

    Args:
        disc_logits: (batch_size, seq_len) - raw logits from discriminator
        is_replaced: (batch_size, seq_len) - binary labels (1 if replaced, 0 if original)
        attention_mask: (batch_size, seq_len) - 1 for real tokens, 0 for padding

    Returns:
        Scalar loss
    """
    # Binary cross-entropy with logits
    loss = F.binary_cross_entropy_with_logits(
        disc_logits,
        is_replaced.float(),
        reduction="none",
    )

    # Mask out padding tokens if attention mask provided
    if attention_mask is not None:
        loss = loss * attention_mask
        loss = loss.sum() / attention_mask.sum()
    else:
        loss = loss.mean()

    return loss

Let's trace through a complete forward pass to see how the pieces connect:

In[14]:

Code

# Create discriminator
torch.manual_seed(42)
discriminator = ElectraDiscriminator(
    vocab_size=30522,
    hidden_size=768,
    num_layers=12,
    num_heads=12,
    intermediate_size=3072,
)

# Generate replacements using the generator
with torch.no_grad():
    corrupted_ids = generator.sample_replacements(
        masked_ids, masked_positions, temperature=1.0
    )

# Determine which tokens were actually replaced (different from original)
is_replaced = (corrupted_ids != input_ids).long()

# Discriminator forward pass
disc_logits = discriminator(corrupted_ids)

# Compute loss
disc_loss = compute_discriminator_loss(disc_logits, is_replaced)

Out[15]:

Console

Discriminator Configuration:
  Vocab size: 30,522
  Hidden size: 768
  Layers: 12
  Parameters: 108,890,881

Corruption Statistics:
  Tokens replaced by generator: 13 / 64 (20.3%)

Loss:
  Discriminator loss: 0.7682
  Expected random baseline: 0.6931

The discriminator has roughly 85M parameters, significantly larger than the generator. The replacement rate is lower than the 15% masking rate because the generator sometimes predicts the correct original token. The discriminator loss starts near the random baseline (0.69, which is $\log(2)$ for binary classification), indicating the untrained model has no ability to distinguish original from replaced tokens.

Out[16]:

Visualization

Histogram showing overlapping distributions of discriminator logits for original and replaced tokens. — Discriminator prediction confidence distribution at initialization. The untrained discriminator outputs near-zero logits for both original and replaced tokens, unable to distinguish between them. After training, these distributions would separate with replaced tokens receiving positive logits and original tokens receiving negative logits.

Notice that the generator can produce the same token as the original. When it does, that position is not "replaced" even though it went through the generator. This means the actual replacement rate is typically lower than the masking rate. The discriminator learns to identify both truly replaced tokens and those where the generator happened to guess correctly.