03: Distillation of LLMs for text generation¶
This notebook will show how to distill language models. When distilling LLMs for text generation, the choice of probability divergence fundamentally changes the model's behavior. Indeed, the model must predict the next token from a vocabulary of tens of thousands of possible tokens (GPT-2's tokenizer, used here, has exactly 50,257).
This notebook will follow the same path as the second one, just applied to a new task and maybe the most famous one today: text generation.
Setup¶
import torch
from datasets import load_dataset
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from shrinkai.adapters import AttentionHeadSelector, FeatureExtractor
from shrinkai.distillation.distiller import Distiller
from shrinkai.distillation.losses import (
AttentionMapLoss,
HybridLoss,
ProjectedFeatureLoss,
)
def build_shakespeare_dataloaders(batch_size: int = 8, max_length: int = 128):
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
dataset = load_dataset(
"text",
data_files="https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt",
)
dataset = dataset["train"].train_test_split(test_size=0.1)
dataset["validation"] = dataset.pop("test")
def tokenize_function(examples):
return tokenizer(examples["text"])
tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
def group_texts(examples):
concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
total_length = len(concatenated_examples[list(examples.keys())[0]])
total_length = (total_length // max_length) * max_length
result = {
k: [t[i : i + max_length] for i in range(0, total_length, max_length)]
for k, t in concatenated_examples.items()
}
result["labels"] = result["input_ids"].copy()
return result
lm_datasets = tokenized_datasets.map(group_texts, batched=True)
lm_datasets.set_format(type="torch", columns=["input_ids", "labels"])
train_dataset = TensorDataset(
lm_datasets["train"][:]["input_ids"], lm_datasets["train"][:]["labels"]
)
val_dataset = TensorDataset(
lm_datasets["validation"][:]["input_ids"],
lm_datasets["validation"][:]["labels"],
)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
sample_input = lm_datasets["validation"][:batch_size]["input_ids"]
return train_loader, val_loader, sample_input, tokenizer.pad_token_id
Adapting HuggingFace to match ShrinkAI¶
See tutorial 02 for more details.
class CausalLMWrapper(FeatureExtractor):
def __init__(self, model_name: str, layer_mapping: dict[int, str]):
super(FeatureExtractor, self).__init__()
self.hf_model = AutoModelForCausalLM.from_pretrained(
model_name, attn_implementation="eager"
)
self.layer_mapping = layer_mapping
self.target_layers = layer_mapping
def forward(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
attention_mask = (input_ids != self.hf_model.config.pad_token_id).long()
outputs = self.hf_model(
input_ids=input_ids, attention_mask=attention_mask, output_attentions=True
)
features_dict = {}
if outputs.attentions is not None:
for layer_idx, alias in self.layer_mapping.items():
features_dict[alias] = outputs.attentions[layer_idx]
return outputs.logits, features_dict
train_loader, val_loader, sample_input, pad_token_id = build_shakespeare_dataloaders(
batch_size=8, max_length=128
)
Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads. Map: 100%|██████████| 36000/36000 [00:00<00:00, 227178.27 examples/s] Map: 100%|██████████| 4000/4000 [00:00<00:00, 206942.17 examples/s] Map: 100%|██████████| 36000/36000 [00:00<00:00, 118501.85 examples/s] Map: 100%|██████████| 4000/4000 [00:00<00:00, 118624.74 examples/s]
# 12 layers
teacher = CausalLMWrapper("gpt2", {5: "attention_layer_1", 11: "attention_layer_2"})
teacher.hf_model.config.pad_token_id = pad_token_id
# 6 layers
student = CausalLMWrapper("distilgpt2", {2: "attention_layer_1", 5: "attention_layer_2"})
student.hf_model.config.pad_token_id = pad_token_id
Loading weights: 100%|██████████| 148/148 [00:00<00:00, 11831.61it/s] Loading weights: 100%|██████████| 76/76 [00:00<00:00, 11989.59it/s]
Reverse KL Loss¶
from shrinkai.distillation.losses import ReverseKLLoss
Standard Knowledge Distillation uses Forward KL Divergence ($D_{KL}(P_{teacher} || P_{student})$). However, Forward KL is "mean-seeking". If the teacher thinks the next word could equally be "Sword" or "Shield", a Forward KL student will try to average these probabilities. This often leads to generating a safe, generic, or completely fabricated word that sits in the middle of the distribution, resulting in hallucinated or blurry text.
To fix this, modern LLM distillation uses Reverse KL Divergence ($D_{KL}(P_{student} || P_{teacher})$) combined with Causal Attention Distillation. Reverse KL is "mode-seeking". It heavily penalizes the student for predicting a word that the teacher considers impossible, but allows the student to confidently pick just one of the teacher's highly probable options and ignore the rest. As a consequence, the student generates sharp, highly coherent text, decisively committing to a specific grammatical path (e.g., choosing "Sword" and sticking with it) just like a real generative AI.
By combining ReverseKLLoss on the output vocabulary logits and AttentionMapLoss on the causal attention matrices, the student perfectly replicates the teacher's logical deduction steps without averaging out its creativity.
Note: ReverseKLLoss's own docstring recommends keeping the temperature close to 1.0 for LLM distillation, because it already handles sharp distributions well, unlike forward KL which benefits from the extra softening. We keep temperature=4.0 below for consistency with Tutorials 01/02; feel free to try temperature=1.0 and compare.
attention_loss = ProjectedFeatureLoss(
projector=AttentionHeadSelector(heads_to_keep=list(range(12))),
feature_loss=AttentionMapLoss(loss_type="mse"),
project_teacher=True,
)
hybrid_loss = HybridLoss(
logit_loss=ReverseKLLoss(temperature=4.0),
feature_loss=attention_loss,
feature_weight=10.0,
convex_weighting=False,
)
distiller = Distiller(
teacher=teacher,
student=student,
criterion=hybrid_loss,
optimizer="adamw",
lr=5e-5,
weight_decay=0.01,
device="auto",
grad_clip_norm=1.0, # stabilizes this ReverseKL + AttentionMap recipe; see Tutorial 05 for details
)
distiller.feature_analysis(val_loader, metrics=["cka", "rsa", "attention"]).show()
Representation Alignment Report (Feature Distillation) ┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ Layer / Stage Alias ┃ CKA (Linear) ┃ RSA (Pearson) ┃ Spatial Attention ┃ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ │ attention_layer_2 │ 0.996 │ 0.723 │ 0.986 │ │ attention_layer_1 │ 0.994 │ 0.278 │ 0.912 │ └─────────────────────┴──────────────┴───────────────┴───────────────────┘
training_history = distiller.fit(train_loader, val_loader, epochs=3)
Epoch [01/03] Train Loss: 50.2953 - Train Acc: 14.55% | Val Loss: 34.5918 - Val Acc: 16.02%
Epoch [02/03] Train Loss: 42.3358 - Train Acc: 15.29% | Val Loss: 32.1563 - Val Acc: 16.55%
Epoch [03/03] Train Loss: 39.0613 - Train Acc: 15.69% | Val Loss: 30.4447 - Val Acc: 16.79%
distiller.benchmark(
sample_input=sample_input,
teacher_name="GPT 2",
student_name="Tiny GPT 2",
val_dataloader=val_loader,
).show()
Benchmark Report ┏━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┓ ┃ Metric ┃ Teacher (GPT 2) ┃ Student (Tiny GPT 2) ┃ Gain / Compression ┃ ┡━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━┩ │ Parameters │ 124.44 M │ 81.91 M │ -34.2% │ │ Model Size (Disk) │ 474.75 MB │ 312.50 MB │ -34.2% (1.5x smaller) │ │ Latency / Sample │ 22.31 ms │ 13.89 ms │ 1.6x faster │ │ Throughput (FPS) │ 44.8 img/s │ 72.0 img/s │ +1.6x (72.0 FPS) │ │ Accuracy │ 17.88% │ 16.79% │ 93.9% retained │ └───────────────────┴─────────────────┴──────────────────────┴───────────────────────┘
During the training, grad_clip_norm=1.0 has been set to stabilize this specific loss combination. Without it, this exact recipe occasionally drifted into a degenerate state on some runs. More details in Tutorial 05.
distiller.feature_analysis(val_loader, metrics=["cka", "rsa", "attention"]).show()
Representation Alignment Report (Feature Distillation) ┏━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┓ ┃ Layer / Stage Alias ┃ CKA (Linear) ┃ RSA (Pearson) ┃ Spatial Attention ┃ ┡━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━┩ │ attention_layer_2 │ 0.997 │ 0.795 │ 0.990 │ │ attention_layer_1 │ 0.994 │ 0.278 │ 0.934 │ └─────────────────────┴──────────────┴───────────────┴───────────────────┘
Handling Different Tokenizers and Vocabularies¶
In this generative tutorial, both gpt2 and distilgpt2 belong to the same family and share the exact same tokenizer (vocabulary size of 50,257). This alignment allows us to directly apply the ReverseKLLoss. Because index 4256 represents the exact same token for both the teacher and the student, comparing their probability distributions is mathematically right.
But what if you want to distill a modern Llama-3 teacher (vocab size around 128k) into a smaller, custom architecture (vocab size around 32k)?
When tokenizers differ, the output logit vectors have different dimensions, and more importantly, the indices no longer share the same semantic meaning. A direct probability divergence (like KL or Reverse KL) becomes meaningless. To bypass this barrier, you can adopt one of the following architectural shifts:
Pure Feature-Level Distillation¶
Completely discard the logit-based loss. Instead of forcing the student to predict the same final vocabulary probabilities, force it to mimic the teacher's internal reasoning. By relying exclusively on shrinkai's AttentionMapLoss and ProjectedFeatureLoss, the student learns the teacher's contextual geometry and spatial attention. You then apply a standard Cross-Entropy loss (using the ground-truth text tokenized by the student's tokenizer) to map those highly educated internal features to its own specific vocabulary.
Sequence-Level Knowledge Distillation¶
Instead of aligning weights and probabilities during a single forward pass, you treat the teacher as a data generator. You prompt the massive teacher model to generate thousands of high-quality, specialized responses (e.g., generating text in the style of Shakespeare). Once this synthetic dataset is created, you train your student model on it using standard Causal Language Modeling with its own tokenizer. The distillation happens at the sequence level rather than the tensor level.
Let's run it for real, reusing the same gpt2/distilgpt2 pair as before purely for speed, it's the mechanics that matter here, not the vocab mismatch. Sequence-Level KD never goes through Distiller, DistillationEngine, or any shrinkai loss: there is no teacher forward pass during training at all, the teacher only generates a synthetic corpus before the training loop even starts.
seqkd_prompts = ["ROMEO:", "JULIET:", "To be, or not to be,", "Once upon a midnight"]
seqkd_tokenizer = AutoTokenizer.from_pretrained("gpt2")
seqkd_tokenizer.pad_token = seqkd_tokenizer.eos_token
seqkd_tokenizer.padding_side = "left" # required for correct batched generation with a decoder-only model
prompt_batch = seqkd_tokenizer(seqkd_prompts, return_tensors="pt", padding=True)
prompt_batch = {k: v.to(teacher.hf_model.device) for k, v in prompt_batch.items()}
teacher.hf_model.eval()
with torch.no_grad():
generated_ids = teacher.hf_model.generate(
**prompt_batch,
max_new_tokens=48,
do_sample=True,
temperature=0.8,
top_p=0.95,
pad_token_id=seqkd_tokenizer.pad_token_id,
)
synthetic_texts = seqkd_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
for text in synthetic_texts:
print(text, "\n---")
ROMEO: A NEW PROPHECY OF THE COUNCIL OF THE WORD OF THE COUNCIL OF THE COUNCIL OF THE WORLD. I THOUGHT THAT THE COUNCIL OF THE WORLD WAS --- JULIET: I saw a lot of people saying, "Oh, we want to see a great performance from K'Lellan, and we're going to see a great performance from somebody who's so young. Let's see how that plays out --- To be, or not to be, a man you've got to be a man, or you're going to be a man, or you're going to be a man, or you're going to be a man, or you're going to be a man. You're --- Once upon a midnight raid, the young boy who had been assigned as a hostage was found at his room, covered in blood. The boy's father was killed. He told the police how he had been kidnapped by the gang. He told them he ---
Now, let's just fine tune a Causal LM using cross-entropy, no shrinkai involved.
seqkd_tokenizer.padding_side = "right" # back to normal for plain (non-generation) encoding
seqkd_encodings = seqkd_tokenizer(
synthetic_texts, return_tensors="pt", padding=True, truncation=True, max_length=96
)
seqkd_input_ids = seqkd_encodings["input_ids"]
seqkd_attention_mask = seqkd_encodings["attention_mask"]
seqkd_labels = seqkd_input_ids.masked_fill(seqkd_attention_mask == 0, -100) # ignore padding in the loss
seqkd_loader = DataLoader(
TensorDataset(seqkd_input_ids, seqkd_attention_mask, seqkd_labels), batch_size=2, shuffle=True
)
seqkd_student = AutoModelForCausalLM.from_pretrained("distilgpt2")
seqkd_student.config.pad_token_id = seqkd_tokenizer.pad_token_id
seqkd_optimizer = torch.optim.AdamW(seqkd_student.parameters(), lr=5e-5)
seqkd_student.train()
for epoch in range(3):
for ids, mask, labels in seqkd_loader:
outputs = seqkd_student(input_ids=ids, attention_mask=mask, labels=labels)
outputs.loss.backward()
seqkd_optimizer.step()
seqkd_optimizer.zero_grad()
print(f"epoch {epoch + 1}: loss={outputs.loss.item():.4f}")
Loading weights: 100%|██████████| 76/76 [00:00<00:00, 17723.07it/s] [transformers] `loss_type=None` was set in the config but it is unrecognized. Using the default loss: `ForCausalLMLoss`.
epoch 1: loss=2.5936 epoch 2: loss=1.4856 epoch 3: loss=1.4011
Notice that the teacher never appears inside that training loop. That is the practical appeal of Sequence-Level KD at scale: the teacher only runs offline, so student training afterward is as cheap as ordinary fine-tuning. A quick sample from the fine-tuned student shows it picked up the dialogue-like style of the synthetic corpus:
seqkd_student.eval()
seqkd_tokenizer.padding_side = "left"
prompt = seqkd_tokenizer(["ROMEO:"], return_tensors="pt")
with torch.no_grad():
sample_ids = seqkd_student.generate(
**prompt, max_new_tokens=40, do_sample=True, temperature=0.8, pad_token_id=seqkd_tokenizer.pad_token_id
)
print(seqkd_tokenizer.decode(sample_ids[0], skip_special_tokens=True))
ROMEO: I thought I'd be in the middle of the night. I was in the middle of the night. I was in the middle of the night. I was in the middle of the night. I
References¶
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. NeurIPS.
arXiv:1706.03762
https://arxiv.org/abs/1706.03762Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., & Sutskever, I. (2019). Language Models are Unsupervised Multitask Learners. OpenAI.
arXiv:1902.09268
https://arxiv.org/abs/1902.09268Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. NIPS Deep Learning and Representation Learning Workshop.
arXiv:1503.02531
https://arxiv.org/abs/1503.02531Gu, Y., Dong, L., Wei, F., & Huang, M. (2024). MiniLLM: Knowledge Distillation of Large Language Models. ICLR 2024.
arXiv:2306.08543
https://arxiv.org/abs/2306.08543Wu, T., Tao, C., Wang, J., Yang, R., Zhao, Z., & Wong, N. (2025). Rethinking Kullback-Leibler Divergence in Knowledge Distillation for Large Language Models. COLING 2025, 5737–5755.
arXiv:2404.02657
https://arxiv.org/abs/2404.02657Wang, W., Bao, H., Huang, S., Dong, L., & Wei, F. (2020). MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers. NeurIPS.
arXiv:2002.10957
https://arxiv.org/abs/2002.10957Kim, Y., & Rush, A. M. (2016). Sequence-Level Knowledge Distillation. EMNLP 2016, 1317–1327.
arXiv:1606.07947
https://arxiv.org/abs/1606.07947Romero, A., Ballas, N., Ebrahimi Kahou, S., Chassang, A., Gatta, C., & Bengio, Y. (2015). FitNets: Hints for Thin Deep Nets. ICLR.
arXiv:1412.6550
https://arxiv.org/abs/1412.6550Loshchilov, I., & Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR.
arXiv:1711.05101
https://arxiv.org/abs/1711.05101Kornblith, S., Norouzi, M., Lee, H., & Hinton, G. (2019). Similarity of Neural Network Representations Revisited. ICML, PMLR 97, 3519–3529.
arXiv:1905.00414
https://arxiv.org/abs/1905.00414Kriegeskorte, N., Mur, M., & Bandettini, P. (2008). Representational Similarity Analysis – Connecting the Branches of Systems Neuroscience. Frontiers in Systems Neuroscience, 2, 4.
DOI: 10.3389/neuro.06.004.2008
https://doi.org/10.3389/neuro.06.004.2008Sennrich, R., Haddow, B., & Birch, A. (2016). Neural Machine Translation of Rare Words with Subword Units. ACL.
arXiv:1508.07909
https://arxiv.org/abs/1508.07909