1class BERT:
2    def __init__(self):
3        # Each layer has unique parameters
4        self.layers = [TransformerLayer() for _ in range(12)]
5        # 12 × 7M params = 85M params in transformer layers
6
7    def forward(self, x):
8        for layer in self.layers:
9            x = layer(x)   # Different weights each time
10        return x

1class ALBERT:
2    def __init__(self):
3        # Single shared layer
4        self.shared_layer = TransformerLayer()
5        # 1 × 7M params = 7M params (89% reduction!)
6
7    def forward(self, x):
8        for _ in range(12):
9            x = self.shared_layer(x)  # Same weights reused
10        return x

Knowledge distillation (Sanh et al., 2019) compresses BERT into a smaller, faster model:

• Teacher: Full BERT-base (12 layers, frozen)
• Student: DistilBERT (6 layers, trainable)
• Training signal: Student learns to match the teacher's soft probability distributions, not just the hard labels

The teacher's "wrong" predictions contain useful information — e.g., predicting "cat" is more likely than "car" for a masked animal slot tells the student about semantic similarity.

1teacher = BertModel.from_pretrained("bert-base")  # 12 layers, frozen
2student = DistilBertModel(num_layers=6)            # 6 layers, trainable
3
4for batch in training_data:
5    # Teacher provides "soft targets" (probability distributions)
6    with torch.no_grad():
7        teacher_logits = teacher(batch)  # e.g., [0.7, 0.2, 0.1, ...]
8
9    # Student tries to match teacher's distribution
10    student_logits = student(batch)
11
12    # Distillation loss: KL divergence between soft distributions
13    # Temperature T=2 softens the distribution (more informative)
14    loss_distill = KL_divergence(
15        softmax(student_logits / T),
16        softmax(teacher_logits / T)
17    )
18    loss_mlm = masked_lm_loss(student_logits, labels)
19
20    loss = 0.5 * loss_distill + 0.5 * loss_mlm

1import torch
2from transformers import AutoTokenizer, AutoModelForMaskedLM, ElectraForPreTraining
3
4tokenizer = AutoTokenizer.from_pretrained("google/electra-small-generator")
5generator = AutoModelForMaskedLM.from_pretrained("google/electra-small-generator")
6discriminator = ElectraForPreTraining.from_pretrained("google/electra-small-discriminator")
7
8# Step 1: Mask a token and let the generator fill it in
9original = "The chef cooked a delicious meal"
10masked = "The chef [MASK] a delicious meal"
11inputs = tokenizer(masked, return_tensors="pt")
12
13with torch.no_grad():
14    gen_logits = generator(**inputs).logits
15
16mask_idx = (inputs.input_ids == tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
17predicted_id = gen_logits[0, mask_idx].argmax(dim=-1)
18replacement = tokenizer.decode(predicted_id)
19print(f"Generator filled [MASK] → '{replacement}'")  # e.g., "prepared"
20
21# Step 2: Discriminator classifies EVERY token as original or replaced
22fake_ids = inputs.input_ids.clone()
23fake_ids[0, mask_idx] = predicted_id
24with torch.no_grad():
25    disc_logits = discriminator(fake_ids).logits
26
27predictions = (disc_logits.squeeze() > 0).long()  # positive logit = "fake"
28tokens = tokenizer.convert_ids_to_tokens(fake_ids[0])
29for tok, pred in zip(tokens, predictions):
30    print(f"  {tok:12s} → {'REPLACED' if pred else 'original'}")