ELECTRA › Weight Sharing
Language AI HandbookAll chapters ELECTRA › Weight SharingOne section · Sections in this chapter11 / 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 can optionally share embeddings between the generator and discriminator. Since both networks process the same vocabulary, their token embeddings can be identical. This reduces total parameters and ensures consistent token representations.
When the generator and discriminator have different hidden sizes (as in ELECTRA-Base and Large), sharing requires a linear projection to map between sizes:
class ElectraWithSharedEmbeddings(nn.Module):
"""ELECTRA with weight sharing between generator and discriminator."""
def __init__(
self,
vocab_size,
disc_hidden,
gen_hidden,
disc_layers,
gen_layers,
disc_heads,
gen_heads,
disc_intermediate,
gen_intermediate,
max_position=512,
dropout=0.1,
):
super().__init__()
# Shared embeddings (in discriminator's hidden size)
self.shared_token_embedding = nn.Embedding(vocab_size, disc_hidden)
self.shared_position_embedding = nn.Embedding(max_position, disc_hidden)
# Projection for generator if hidden sizes differ
if gen_hidden != disc_hidden:
self.gen_embedding_projection = nn.Linear(disc_hidden, gen_hidden)
else:
self.gen_embedding_projection = None
# Generator components (without embeddings)
self.gen_embedding_norm = nn.LayerNorm(gen_hidden)
self.gen_embedding_dropout = nn.Dropout(dropout)
gen_layer = nn.TransformerEncoderLayer(
d_model=gen_hidden,
nhead=gen_heads,
dim_feedforward=gen_intermediate,
dropout=dropout,
activation="gelu",
batch_first=True,
)
self.gen_encoder = nn.TransformerEncoder(gen_layer, gen_layers)
self.gen_mlm_dense = nn.Linear(gen_hidden, gen_hidden)
self.gen_mlm_norm = nn.LayerNorm(gen_hidden)
self.gen_mlm_output = nn.Linear(gen_hidden, vocab_size)
# Discriminator components (without embeddings)
self.disc_embedding_norm = nn.LayerNorm(disc_hidden)
self.disc_embedding_dropout = nn.Dropout(dropout)
disc_layer = nn.TransformerEncoderLayer(
d_model=disc_hidden,
nhead=disc_heads,
dim_feedforward=disc_intermediate,
dropout=dropout,
activation="gelu",
batch_first=True,
)
self.disc_encoder = nn.TransformerEncoder(disc_layer, disc_layers)
self.disc_head = nn.Linear(disc_hidden, 1)
self.vocab_size = vocab_size
def get_embeddings(self, input_ids, for_generator=False):
"""Get embeddings, with optional projection for generator."""
batch_size, seq_len = input_ids.shape
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
x = self.shared_token_embedding(input_ids)
x = x + self.shared_position_embedding(positions)
if for_generator and self.gen_embedding_projection is not None:
x = self.gen_embedding_projection(x)
return x
def generator_forward(self, input_ids):
"""Forward pass through generator."""
x = self.get_embeddings(input_ids, for_generator=True)
x = self.gen_embedding_norm(x)
x = self.gen_embedding_dropout(x)
x = self.gen_encoder(x)
x = self.gen_mlm_dense(x)
x = F.gelu(x)
x = self.gen_mlm_norm(x)
logits = self.gen_mlm_output(x)
return logits
def discriminator_forward(self, input_ids):
"""Forward pass through discriminator."""
x = self.get_embeddings(input_ids, for_generator=False)
x = self.disc_embedding_norm(x)
x = self.disc_embedding_dropout(x)
x = self.disc_encoder(x)
logits = self.disc_head(x).squeeze(-1)
return logitsWeight sharing introduces a subtle issue. If we backpropagate through both networks simultaneously, the shared embeddings receive gradients from both the MLM loss (via generator) and the RTD loss (via discriminator). The ELECTRA paper found this works well, with the discriminator's gradients dominating due to the higher RTD weight and the learning from all positions.
