Replaced Token Detection › Implementing the Generator
Language AI HandbookAll chapters Replaced Token Detection › Implementing the GeneratorOne section · Sections in this chapter6 / 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 generator is a small transformer that performs MLM. It masks positions, predicts token distributions, and samples replacements. Architecturally, it is nearly identical to BERT: a token embedding layer, a positional embedding layer, a stack of transformer encoder layers, and an output projection to vocabulary logits. The only structural difference is that the generator is smaller, with fewer hidden dimensions, which keeps its parameter count low while preserving its ability to capture contextual patterns.
In practice, you would use the same tokenizer for both the generator and the discriminator, since they must agree on what constitutes a vocabulary token and what the [MASK] token ID is. The generator's output logits are over the full vocabulary, and the sampling step converts these logits into a single sampled token per masked position. The vocabulary must be shared because the discriminator needs to process the same token IDs that the generator produces.
class RTDGenerator(nn.Module):
"""Small generator network for producing replacement tokens."""
def __init__(
self, vocab_size, d_model=64, n_heads=2, n_layers=2, max_len=128
):
super().__init__()
# Embeddings
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)
# Small 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)
# Output projection to vocabulary
self.output_proj = nn.Linear(d_model, vocab_size)
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 (bidirectional attention)
x = self.encoder(x)
# Project to vocabulary logits
logits = self.output_proj(x)
return logitsdef create_masked_input(token_ids, mask_token_id, mask_prob=0.15):
"""Mask tokens and return masked input, labels, and mask positions."""
labels = token_ids.clone()
masked_ids = token_ids.clone()
# Sample positions to mask
probability_matrix = torch.full(token_ids.shape, mask_prob)
masked_indices = torch.bernoulli(probability_matrix).bool()
# Replace with [MASK] token
masked_ids[masked_indices] = mask_token_id
# Labels are -100 for non-masked positions (ignored in loss)
labels[~masked_indices] = -100
return masked_ids, labels, masked_indicesLet's see the generator in action:
# Create a small vocabulary for demonstration
vocab = [
"[PAD]",
"[MASK]",
"[CLS]",
"[SEP]",
"the",
"chef",
"cooked",
"meal",
"waiter",
"served",
"food",
"ate",
"a",
"an",
]
vocab_size = len(vocab)
word_to_idx = {w: i for i, w in enumerate(vocab)}
idx_to_word = {i: w for w, i in word_to_idx.items()}
mask_token_id = word_to_idx["[MASK]"]
# Initialize generator
generator = RTDGenerator(vocab_size, d_model=32, n_heads=2, n_layers=1)
# Create sample input
sentence = ["the", "chef", "cooked", "the", "meal"]
token_ids = torch.tensor([[word_to_idx[w] for w in sentence]])Original sentence: ['the', 'chef', 'cooked', 'the', 'meal'] Token IDs: [4, 5, 6, 4, 7] Masked IDs: [4, 5, 1, 4, 1] Masked sequence: ['the', 'chef', '[MASK]', 'the', '[MASK]'] Mask positions: [False, False, True, False, True]
The masking function randomly selects positions to mask. With a 40% masking probability on our 5-token sentence, we typically see 2 positions replaced with [MASK]. The labels tensor stores the original tokens at masked positions for computing the generator loss.
Now let's sample replacements from the generator:
def sample_replacements(generator, masked_ids, labels, temperature=1.0):
"""Sample replacement tokens from the generator."""
generator.eval()
with torch.no_grad():
logits = generator(masked_ids)
# Apply temperature and sample
probs = F.softmax(logits / temperature, dim=-1)
# Create replaced sequence
replaced_ids = masked_ids.clone()
is_replaced = torch.zeros_like(masked_ids, dtype=torch.bool)
# Sample only at masked positions
mask_positions = labels != -100
for batch_idx in range(masked_ids.size(0)):
for pos in range(masked_ids.size(1)):
if mask_positions[batch_idx, pos]:
sampled_token = torch.multinomial(
probs[batch_idx, pos], num_samples=1
)
replaced_ids[batch_idx, pos] = sampled_token
# Mark as replaced if sampled token differs from original
original_token = labels[batch_idx, pos]
is_replaced[batch_idx, pos] = sampled_token != original_token
return replaced_ids, is_replacedreplaced_ids, is_replaced = sample_replacements(generator, masked_ids, labels)Original: ['the', 'chef', 'cooked', 'the', 'meal'] Replaced: ['the', 'chef', '[CLS]', 'the', 'chef'] Is replaced: [False, False, True, False, True]
The generator samples tokens for masked positions. Because we're using an untrained generator, the replacements are essentially random from the small vocabulary. With training, the generator would produce more plausible replacements that better challenge the discriminator.