LLM Unlearning Bigger
by marimo team
Exploring unlearning in large language models by overfitting what you don't want.
About this notebook
This notebook demonstrates the unlearning method from the paper "Who's Harry Potter? Approximate Unlearning in LLMs" by Eldan and Russinovich (2023). It loads a causal language model with Hugging Face transformers and trains a deep copy of it on Harry Potter sentences defined in the notebook. This reinforced model saturates on the target content. The paper's formula then combines the two models' logits into a generic target that subtracts what the reinforced model learned.
The notebook is interactive. A dropdown selects the model, with gpt2 as the default and gpt2-xl, Llama-3-8B, and Qwen2-7B as options. A text field holds a test prompt, and a slider sets the unlearning strength alpha between 0.01 and 10. When you change any widget, marimo re-runs the dependent cells. Three matplotlib bar charts show the ten most probable next tokens for the baseline model, the reinforced model, and the generic target. The formula lowers the scores of tokens that the reinforced model rates higher than the baseline.
The training loop uses PyTorch with the AdamW optimizer, a learning rate of 2e-5, 15 epochs, and gradient clipping to prevent collapse. The notebook runs on a GPU when torch detects one and on the CPU otherwise. All logit comparisons run under torch.no_grad, so only the reinforced model's training step updates weights.
