Replaced Token Detection › Weight Sharing and Embedding Tying
Language AI HandbookAll chapters Replaced Token Detection › Weight Sharing and Embedding TyingOne section · Sections in this chapter10 / 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 original ELECTRA paper uses weight sharing between generator and discriminator embeddings. This has two benefits: it reduces total parameters and ensures both networks have compatible token representations.
The embedding matrix maps token IDs to dense vectors. Since both networks process the same vocabulary, they can share this large matrix. When two networks share an embedding table, updates to that table from one network's gradients immediately affect the other. If the generator updates "scientist" to better fit scientific contexts, the discriminator benefits from the same improved "scientist" representation. This coupling turns out to be helpful rather than harmful: the generator's MLM objective provides rich token-level supervision that helps the discriminator's embeddings even before detection learning kicks in.
The technical complication is that the generator and discriminator typically have different hidden dimensions. In ELECTRA-Base, the generator uses a hidden dimension of 64 while the discriminator uses 256. The shared embedding matrix has dimension 64 (the generator's size), and the discriminator projects these embeddings up to its 256-dimensional space using a linear layer before feeding them into its transformer. This projection layer is learned during training, allowing the discriminator to map the compact shared embeddings into its larger representation space.
The generator uses the embeddings directly at its smaller hidden dimension, while the discriminator projects them up to its larger dimension:
class ELECTRAWithSharedEmbeddings(nn.Module):
"""ELECTRA with shared embeddings between generator and discriminator."""
def __init__(
self,
vocab_size,
gen_d_model=64,
disc_d_model=256,
n_gen_layers=4,
n_disc_layers=12,
max_len=512,
):
super().__init__()
# Shared token embeddings at generator size
self.shared_token_emb = nn.Embedding(vocab_size, gen_d_model)
# Generator uses embeddings directly
self.gen_pos_emb = nn.Embedding(max_len, gen_d_model)
self.gen_encoder = nn.TransformerEncoder(
nn.TransformerEncoderLayer(
d_model=gen_d_model,
nhead=4,
dim_feedforward=gen_d_model * 4,
batch_first=True,
activation="gelu",
),
num_layers=n_gen_layers,
)
self.gen_output = nn.Linear(gen_d_model, vocab_size)
# Discriminator projects embeddings to its larger dimension
self.disc_embed_proj = nn.Linear(gen_d_model, disc_d_model)
self.disc_pos_emb = nn.Embedding(max_len, disc_d_model)
self.disc_encoder = nn.TransformerEncoder(
nn.TransformerEncoderLayer(
d_model=disc_d_model,
nhead=8,
dim_feedforward=disc_d_model * 4,
batch_first=True,
activation="gelu",
),
num_layers=n_disc_layers,
)
self.disc_classifier = nn.Linear(disc_d_model, 1)
def generator_forward(self, input_ids):
"""Generator forward pass for MLM."""
seq_len = input_ids.size(1)
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
x = self.shared_token_emb(input_ids) + self.gen_pos_emb(positions)
x = self.gen_encoder(x)
logits = self.gen_output(x)
return logits
def discriminator_forward(self, input_ids):
"""Discriminator forward pass for RTD."""
seq_len = input_ids.size(1)
positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
# Project shared embeddings to discriminator dimension
x = self.shared_token_emb(input_ids)
x = self.disc_embed_proj(x) + self.disc_pos_emb(positions)
x = self.disc_encoder(x)
logits = self.disc_classifier(x).squeeze(-1)
return logitsThe generator embedding dimension determines the shared embedding size. The discriminator projects these embeddings up to its larger hidden dimension. This asymmetry reflects the design principle: the generator needs only enough capacity to produce plausible replacements, while the discriminator needs more capacity to develop rich representations.
An alternative to weight sharing is using fully separate embedding tables. Some practitioners prefer this approach because it gives the discriminator complete freedom in how it represents tokens, without being constrained by the generator's smaller dimension. The trade-off is additional parameters and the loss of the beneficial gradient coupling described above. In most published results, weight sharing performs better, which is why it is the default in ELECTRA. If you are training a very large model where the generator's embedding dimension can be reasonably large (say, 128 or 256 rather than 64), the gap between shared and separate embeddings narrows.