ELECTRA › Fine-tuning ELECTRA
Language AI HandbookAll chapters ELECTRA › Fine-tuning ELECTRAOne section · Sections in this chapter12 / 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
After pre-training, the generator is discarded. Only the discriminator is used for downstream tasks. Fine-tuning follows the same pattern as BERT: add a task-specific head and train on labeled data.
For classification tasks, the [CLS] token representation from the discriminator feeds into a classification layer:
class ElectraForSequenceClassification(nn.Module):
"""ELECTRA discriminator with classification head."""
def __init__(self, discriminator, num_labels, dropout=0.1):
super().__init__()
self.discriminator = discriminator
# Remove the RTD head
hidden_size = discriminator.token_embedding.embedding_dim
# Add classification head
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(hidden_size, num_labels)
def forward(self, input_ids, attention_mask=None, labels=None):
batch_size, seq_len = input_ids.shape
# Get discriminator embeddings
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
x = self.discriminator.token_embedding(input_ids)
x = x + self.discriminator.position_embedding(positions)
x = self.discriminator.embedding_norm(x)
x = self.discriminator.embedding_dropout(x)
# Create mask
if attention_mask is not None:
src_key_padding_mask = attention_mask == 0
else:
src_key_padding_mask = None
# Encode
x = self.discriminator.encoder(
x, src_key_padding_mask=src_key_padding_mask
)
# Classification from [CLS] token (position 0)
cls_output = x[:, 0]
cls_output = self.dropout(cls_output)
logits = self.classifier(cls_output)
loss = None
if labels is not None:
loss = F.cross_entropy(logits, labels)
return {"loss": loss, "logits": logits}For token-level tasks like NER or question answering, we use the full sequence output instead of just [CLS]:
class ElectraForTokenClassification(nn.Module):
"""ELECTRA discriminator with token classification head."""
def __init__(self, discriminator, num_labels, dropout=0.1):
super().__init__()
self.discriminator = discriminator
hidden_size = discriminator.token_embedding.embedding_dim
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(hidden_size, num_labels)
def forward(self, input_ids, attention_mask=None, labels=None):
batch_size, seq_len = input_ids.shape
# Get discriminator embeddings
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
x = self.discriminator.token_embedding(input_ids)
x = x + self.discriminator.position_embedding(positions)
x = self.discriminator.embedding_norm(x)
x = self.discriminator.embedding_dropout(x)
if attention_mask is not None:
src_key_padding_mask = attention_mask == 0
else:
src_key_padding_mask = None
x = self.discriminator.encoder(
x, src_key_padding_mask=src_key_padding_mask
)
# Classification for all tokens
x = self.dropout(x)
logits = self.classifier(x)
loss = None
if labels is not None:
loss = F.cross_entropy(
logits.view(-1, logits.size(-1)), labels.view(-1)
)
return {"loss": loss, "logits": logits}One consideration for fine-tuning: ELECTRA's pre-training teaches the model to detect replaced tokens, not to understand [MASK] tokens. This means ELECTRA doesn't have the pretrain-finetune mismatch that affects BERT (where [MASK] appears in pre-training but not fine-tuning). The discriminator always sees real tokens, both during pre-training and fine-tuning.