Add T3-Unlearning method

README.md

@@ -19,7 +19,7 @@
 
 ## 📖 Overview
 
-We provide efficient and streamlined implementations of the TOFU, MUSE and WMDP unlearning benchmarks while supporting 13+ unlearning methods, 5+ datasets, 10+ evaluation metrics, and 7+ LLM architectures. Each of these can be easily extended to incorporate more variants.
+We provide efficient and streamlined implementations of the TOFU, MUSE and WMDP unlearning benchmarks while supporting 14+ unlearning methods, 5+ datasets, 10+ evaluation metrics, and 7+ LLM architectures. Each of these can be easily extended to incorporate more variants.
 
 
 We invite the LLM unlearning community to collaborate by adding new benchmarks, unlearning methods, datasets and evaluation metrics here to expand OpenUnlearning's features, gain feedback from wider usage and drive progress in the field.
@@ -81,7 +81,7 @@ We provide several variants for each of the components in the unlearning pipelin
 | **Component**          | **Available Options** |
 |------------------------|----------------------|
 | **Benchmarks**        | [TOFU](https://arxiv.org/abs/2401.06121), [MUSE](https://muse-bench.github.io/), [WMDP](https://www.wmdp.ai/), [DUET](https://arxiv.org/abs/2602.19612) |
-| **Unlearning Methods** | GradAscent, GradDiff, NPO, SimNPO, DPO, RMU, UNDIAL, AltPO, SatImp, WGA, CE-U, PDU, AdaPop |
+| **Unlearning Methods** | GradAscent, GradDiff, NPO, SimNPO, DPO, RMU, UNDIAL, AltPO, SatImp, WGA, CE-U, PDU, AdaPop, T3 |
 | **Evaluation Metrics** | Verbatim Probability, Verbatim ROUGE, Knowledge QA-ROUGE, Model Utility, Forget Quality, TruthRatio, Extraction Strength, Exact Memorization, 6 MIA attacks, [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness) |
 | **Datasets**          | MUSE-News (BBC), MUSE-Books (Harry Potter), TOFU (different splits), WMDP-Bio, WMDP-Cyber, [DUET](https://huggingface.co/datasets/SwetieePawsss/DUET) (rare / popular Wikidata facts) |
 | **Model Families**    | TOFU: Llama-3.2, Llama-3.1, Llama-2; MUSE: Llama-2; DUET: Llama-3.1, Qwen2.5, Gemma; Additional: Phi-3.5, Phi-1.5, Gemma, Zephyr |

community/methods/T3/README.md

@@ -0,0 +1,20 @@
+# Temper-Then-Tilt: Principled Unlearning for Generative Models through Tempering and Classifier Guidance (ICML 2026)
+
+- Authors: Jacob L. Block, Mehryar Mohri, Aryan Mokhtari, Sanjay Shakkottai
+- ArXiv: https://arxiv.org/abs/2602.10217
+- Project Page: https://github.com/jacob-block/T3-Unlearning
+
+## Setup
+
+The script `run.sh` runs a small sweep for hyperparameter search locally on a single GPU. For optimal performance, hyperparameters should be tuned for each specific user setup and environment. Since our method only trains a small classifier on top of the base LLM, it can be trained with minimal resources and overhead. Please see the paper and project page for more details. 
+
+## Citation
+
+```bibtex
+@inproceedings{block:2026:T3-Unlearning,
+  title={Temper-Then-Tilt: Principled Unlearning for Generative Models through Tempering and Classifier Guidance},
+  author={Jacob L. Block and Mehryar Mohri and Aryan Mokhtari and Sanjay Shakkottai},
+  booktitle={Forty-third International Conference on Machine Learning},
+  year={2026},
+}
+```

community/methods/T3/run.sh

@@ -0,0 +1,128 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+export CUDA_VISIBLE_DEVICES=0
+export PYTHONUNBUFFERED=1
+export LOGLEVEL=INFO
+
+NNODES=1
+NMACHINES=1
+export MASTER_ADDR="127.0.0.1"
+export MASTER_PORT="$(python - <<'PY'
+import socket
+s = socket.socket()
+s.bind(("", 0))
+print(s.getsockname()[1])
+s.close()
+PY
+)"
+
+model="Llama-3.1-8B-Instruct"
+trainer="T3"
+experiment="unlearn/tofu/default"
+
+FORGET_PCT="05"
+forget_split="forget${FORGET_PCT}"
+holdout_split="holdout${FORGET_PCT}"
+retain_split="retain$(( 100 - 10#$FORGET_PCT ))"
+
+seeds_list=(1 2)
+num_epochs=100
+warmup_epochs=25
+weight_decay_list=(1e-2 1e-3) 
+hidden_size_list=(20 50)
+num_hidden_layers=1
+extraction_layer=-1
+activation_str="id"
+guidance_scale=1
+base_temp=2.5
+lr_list=(1e-2 1e-3)
+bias="false"
+pooling="mean"
+per_device_train_batch_size=32
+gradient_accumulation_steps=1
+
+echo "======================================="
+echo ""
+num_combos=$(( ${#lr_list[@]} * ${#weight_decay_list[@]} * ${#hidden_size_list[@]} * ${#seeds_list[@]}))
+curr_iter=1
+echo "Searching over $num_combos parameter combinations"
+echo ""
+echo "======================================="
+
+
+for lr in ${lr_list[@]}; do
+    for weight_decay in ${weight_decay_list[@]}; do
+        for hidden_size in ${hidden_size_list[@]}; do
+            for seed in "${seeds_list[@]}"; do
+                task_name="tofu_${model}_${forget_split}_T3_sweep/temp-${base_temp}-epochs-${num_epochs}-wd-${weight_decay}-h-${hidden_size}-lr-${lr}-act-${activation_str}/seed${seed}"
+
+                model_path=open-unlearning/tofu_${model}_full
+                echo ${task_name}: Unlearning ${model_path} using ${trainer}
+
+                # Unlearn
+                python src/train.py \
+                --config-name=unlearn.yaml \
+                experiment=${experiment} \
+                eval=tofu \
+                trainer=${trainer} \
+                task_name=${task_name} \
+                model=${model} \
+                forget_split=${forget_split} \
+                retain_split=${retain_split} \
+                model.model_args.pretrained_model_name_or_path=${model_path} \
+                retain_logs_path=saves/eval/tofu_${model}_${retain_split}/TOFU_EVAL.json \
+                trainer.args.per_device_train_batch_size=${per_device_train_batch_size} \
+                trainer.args.gradient_accumulation_steps=${gradient_accumulation_steps} \
+                trainer.args.gradient_checkpointing=false \
+                trainer.args.ddp_find_unused_parameters=true \
+                trainer.args.seed=${seed} \
+                trainer.args.num_train_epochs=${num_epochs} \
+                trainer.args.learning_rate=${lr} \
+                trainer.args.weight_decay=${weight_decay} \
+                trainer.args.warmup_epochs=${warmup_epochs} \
+                trainer.method_args.extraction_layer=${extraction_layer} \
+                trainer.method_args.pooling=${pooling} \
+                trainer.method_args.guidance_scale=${guidance_scale} \
+                trainer.method_args.base_temp=${base_temp} \
+                trainer.method_args.guidance_cfg.hidden_size=${hidden_size} \
+                trainer.method_args.guidance_cfg.num_hidden_layers=${num_hidden_layers} \
+                trainer.method_args.guidance_cfg.activation_str=${activation_str} \
+                trainer.method_args.guidance_cfg.bias=${bias} \
+                trainer.args.eval_strategy="no" \
+                trainer.args.eval_on_start=false \
+                trainer.args.do_eval=false
+
+                echo "======================================="
+                echo "FINISHED UNLEARNING"
+                echo "RUNNING FINAL EVAL..."
+                echo "======================================="
+
+                # Eval
+                python src/eval.py \
+                    experiment=eval/tofu/default \
+                    forget_split=${forget_split} \
+                    holdout_split=${holdout_split} \
+                    model=${model} \
+                    task_name=${task_name} \
+                    eval=tofu \
+                    model.model_args.pretrained_model_name_or_path=saves/unlearn/${task_name} \
+                    paths.output_dir=saves/unlearn/${task_name}/evals \
+                    retain_logs_path=saves/eval/tofu_${model}_${retain_split}/TOFU_EVAL.json \
+                    seed=${seed} \
+                    model.model_args.torch_dtype=float16 \
+                    model.model_args.attn_implementation=sdpa 
+
+                rm -f saves/unlearn/${task_name}/*.safetensors
+                rm -f saves/unlearn/${task_name}/checkpoint*/*.safetensors
+
+                echo "======================================="
+                echo "Finished Iteration ${curr_iter}/${num_combos}"
+                echo "$(date +"%Y-%m-%d %H:%M:%S")"
+                echo "======================================="
+
+                curr_iter=$((curr_iter+1))
+            done
+        done
+    done
+done

configs/trainer/T3.yaml

@@ -0,0 +1,28 @@
+defaults:
+  - finetune
+
+handler: T3
+
+args:
+  learning_rate: 5e-2
+  weight_decay: 0.0
+  num_train_epochs: 1
+  per_device_train_batch_size: 1
+  gradient_accumulation_steps: 1
+  do_eval: false
+  eval_on_start: false
+  eval_strategy: "no"
+  save_strategy: "no"
+  report_to: none
+
+method_args:
+  pooling: mean
+  pool_temp: null
+  extraction_layer: -1
+  guidance_scale: 1
+  base_temp: 1
+  guidance_cfg:
+    hidden_size: 100
+    num_hidden_layers: 1
+    activation_str: id
+    bias: false

docs/links.md

@@ -39,6 +39,7 @@ Links to research papers and resources corresponding to implemented features in
 | CE-U (Cross-Entropy unlearning)       | Paper[📄](https://arxiv.org/pdf/2503.01224)                                                                                                     |
 | PDU                  | Paper [📄](https://arxiv.org/abs/2506.05314) |
 | AdaPop               | Paper [📄](https://arxiv.org/abs/2608.14229), Popularity-annotated data [🤗](https://huggingface.co/datasets/SwetieePawsss/DUET) |
+| T3-Unlearning               | Paper[📄](https://arxiv.org/abs/2602.10217), Code [🐙](https://github.com/jacob-block/T3-Unlearning)                                                                                     |
 
 
 ---

src/evals/metrics/__init__.py

@@ -22,6 +22,7 @@
     hm_aggregate,
     classifier_prob,
 )
+from evals.metrics.t3 import t3_loss, t3_accuracy
 
 METRICS_REGISTRY: Dict[str, UnlearningMetric] = {}
 
@@ -52,6 +53,8 @@ def get_metrics(metric_cfgs: DictConfig, **kwargs):
 
 
 # Register metrics here
+_register_metric(t3_loss)
+_register_metric(t3_accuracy)
 _register_metric(probability)
 _register_metric(probability_w_options)
 _register_metric(rouge)

src/evals/metrics/t3.py

@@ -0,0 +1,167 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch.utils.data import DataLoader
+from evals.metrics.base import unlearning_metric
+from evals.metrics.utils import run_batchwise_evals
+from model.t3 import T3CausalLM
+from data.utils import IGNORE_INDEX
+
+
+@unlearning_metric(name="t3_accuracy")
+def t3_accuracy(model, **kwargs):
+    """Calculate and return the accuracy of the classifier for a T3 wrapped model."""
+    assert isinstance(
+        model, T3CausalLM
+    ), f"Cannot compute t3_accuracy for model of type {type(model)}. Must be of type {T3CausalLM}."
+    data = kwargs["data"]
+    collator = kwargs["collators"]
+    batch_size = kwargs["batch_size"]
+    split = kwargs["split"]
+    dataloader = DataLoader(data, batch_size=batch_size, collate_fn=collator)
+
+    # The data is formatted so each sequence is one Q/A pair. Only the answer is trained on, other labels set to -100
+
+    fun_args = {}
+
+    def _0_1_accuracy(model, batch):
+        batch = {k: v.to(model.device) for k, v in batch.items()}
+        with torch.no_grad():
+            # Same logic as src/trainer/unlearn/t3._prep_classifier_loss
+            outputs = model(**batch)
+            shifted_labels = batch["labels"][:, 1:].contiguous()
+            # (batch, seq_len-1, vocab_size)
+            classifier_logits = outputs.classifier_logits[:, :-1, :].contiguous()
+            # (batch, seq_len-1, 1)
+            gather_idxs = shifted_labels.masked_fill(
+                shifted_labels == IGNORE_INDEX, 0
+            ).unsqueeze(2)
+            # (batch, seq_len-1)
+            classifier_logits = torch.gather(
+                classifier_logits, dim=2, index=gather_idxs
+            ).squeeze(2)
+
+            valid_mask = shifted_labels != IGNORE_INDEX
+
+            # log prob of retain, (batch, seq_len-1)
+            classifier_log_probs_retain = F.logsigmoid(classifier_logits)
+            classifier_log_probs_retain[~valid_mask] = 0
+            seq_retain_probs = classifier_log_probs_retain.sum(dim=1)
+
+            # log prob of forget
+            classifier_log_probs_forget = F.logsigmoid(-classifier_logits)
+            classifier_log_probs_forget[~valid_mask] = 0
+            seq_forget_probs = classifier_log_probs_forget.sum(dim=1)
+
+        seq_preds = (seq_retain_probs > seq_forget_probs).long()
+        class_label = int(split == "retain")
+        sequence_accs = (seq_preds == class_label).long()
+
+        return [{"acc": a.item()} for a in sequence_accs]
+
+    metric_by_index = run_batchwise_evals(
+        model, dataloader, _0_1_accuracy, fun_args, "Calculating classifier accuracy"
+    )
+    correct_arr = np.array(
+        [metric_dict["acc"] for metric_dict in metric_by_index.values()]
+    )
+
+    if kwargs["print_text"]:
+        tokenizer = kwargs["tokenizer"]
+        batch = next(iter(dataloader))
+        print(f"Sample batch from split {split}")
+        texts = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
+        for i, t in enumerate(texts):
+            print(f"Sample {i}: {t}")
+
+            # Get the tokens where labels are not ignored
+            input_ids = batch["input_ids"][i]
+            labels = batch["labels"][i]
+
+            # only those that are used in training
+            keep_tokens = input_ids[labels != -100]
+            decoded_keep = tokenizer.decode(keep_tokens, skip_special_tokens=True)
+
+            print(f"  Trained-on tokens: {decoded_keep}")
+            print(f"  Raw trained token IDs: {keep_tokens.tolist()}")
+
+    return {"agg_value": correct_arr.mean(), "value_by_index": metric_by_index}
+
+
+@unlearning_metric(name="t3_loss")
+def t3_loss(model, **kwargs):
+    """Calculate and return the loss of the classifier for a T3 wrapped model."""
+    assert isinstance(
+        model, T3CausalLM
+    ), f"Cannot compute t3_loss for model of type {type(model)}. Must be of type {T3CausalLM}."
+    data = kwargs["data"]
+    collator = kwargs["collators"]
+    batch_size = kwargs["batch_size"]
+    split = kwargs["split"]
+    class_label = int(split == "retain")
+    dataloader = DataLoader(data, batch_size=batch_size, collate_fn=collator)
+
+    # The data is formatted so each sequence is one Q/A pair. Only the answer is trained on, other labels set to -100
+    fun_args = {}
+
+    def _loss(model, batch):
+        batch = {k: v.to(model.device) for k, v in batch.items()}
+        with torch.no_grad():
+            # Same logic as src/trainer/unlearn/t3._prep_classifier_loss
+            outputs = model(**batch)
+            shifted_labels = batch["labels"][:, 1:].contiguous()
+            # (batch, seq_len-1, vocab_size)
+            classifier_logits = outputs.classifier_logits[:, :-1, :].contiguous()
+            # (batch, seq_len-1, 1)
+            gather_idxs = shifted_labels.masked_fill(
+                shifted_labels == IGNORE_INDEX, 0
+            ).unsqueeze(2)
+            # (batch, seq_len-1)
+            classifier_logits = torch.gather(
+                classifier_logits, dim=2, index=gather_idxs
+            ).squeeze(2)
+            classifier_labels = torch.full_like(shifted_labels, class_label).to(
+                classifier_logits.dtype
+            )
+
+            valid_mask = shifted_labels != IGNORE_INDEX
+            full_loss = F.binary_cross_entropy_with_logits(
+                classifier_logits, classifier_labels, reduction="none"
+            )
+            row_sums = valid_mask.sum(dim=1)
+            if any(row_sums == 0):
+                raise RuntimeError(
+                    "No valid tokens in a given sequence during evaluation."
+                )
+
+            sequence_losses = ((full_loss * valid_mask).sum(dim=1) / row_sums).flatten()
+
+        return [{"loss": seq_loss.item()} for seq_loss in sequence_losses]
+
+    metric_by_index = run_batchwise_evals(
+        model, dataloader, _loss, fun_args, "Calculating classifier loss"
+    )
+    loss_arr = np.array(
+        [metric_dict["loss"] for metric_dict in metric_by_index.values()]
+    )
+
+    if kwargs["print_text"]:
+        tokenizer = kwargs["tokenizer"]
+        batch = next(iter(dataloader))
+        print(f"Sample batch from split {split}")
+        texts = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
+        for i, t in enumerate(texts):
+            print(f"Sample {i}: {t}")
+
+            # Get the tokens where labels are not ignored
+            input_ids = batch["input_ids"][i]
+            labels = batch["labels"][i]
+
+            # only those that are used in training
+            keep_tokens = input_ids[labels != -100]
+            decoded_keep = tokenizer.decode(keep_tokens, skip_special_tokens=True)
+
+            print(f"  Trained-on tokens: {decoded_keep}")
+            print(f"  Raw trained token IDs: {keep_tokens.tolist()}")
+
+    return {"agg_value": loss_arr.mean(), "value_by_index": metric_by_index}