Add T3-Unlearning method

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

@@ -38,6 +38,7 @@ Links to research papers and resources corresponding to implemented features in
 | WGA (G-effect)       | Paper[📄](https://arxiv.org/pdf/2502.19301), Code [🐙](https://github.com/tmlr-group/G-effect)                                                                                                     |
 | CE-U (Cross-Entropy unlearning)       | Paper[📄](https://arxiv.org/pdf/2503.01224)                                                                                                     |
 | PDU                  | Paper [📄](https://arxiv.org/abs/2506.05314) |
+| 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,129 @@
+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()
+            classifier_logits = outputs.classifier_logits[:,:-1,:].contiguous() # (batch, seq_len-1, vocab_size)
+            gather_idxs = shifted_labels.masked_fill(shifted_labels==IGNORE_INDEX, 0).unsqueeze(2) #(batch, seq_len-1, 1)
+            classifier_logits = torch.gather(classifier_logits, dim=2, index=gather_idxs).squeeze(2) #(batch, seq_len-1)
+
+            valid_mask = shifted_labels != IGNORE_INDEX
+
+            # log prob of retain
+            classifier_log_probs_retain = F.logsigmoid(classifier_logits)  # (batch, seq_len-1)
+            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]
+
+            keep_tokens = input_ids[labels != -100]   # only those that are used in training
+            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()
+            classifier_logits = outputs.classifier_logits[:,:-1,:].contiguous() # (batch, seq_len-1, vocab_size)
+            gather_idxs = shifted_labels.masked_fill(shifted_labels==IGNORE_INDEX, 0).unsqueeze(2) #(batch, seq_len-1, 1)
+            classifier_logits = torch.gather(classifier_logits, dim=2, index=gather_idxs).squeeze(2) #(batch, seq_len-1)
+            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]
+
+            keep_tokens = input_ids[labels != -100]   # only those that are used in training
+            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}

src/model/__init__.py

@@ -4,6 +4,7 @@
 import os
 import torch
 import logging
+from model.t3 import T3CausalLM
 from model.probe import ProbedLlamaForCausalLM
 
 hf_home = os.getenv("HF_HOME", default=None)
@@ -104,4 +105,5 @@ def get_tokenizer(tokenizer_cfg: DictConfig):
 
 # register models
 _register_model(AutoModelForCausalLM)
+_register_model(T3CausalLM)
 _register_model(ProbedLlamaForCausalLM)