Add Unlearning Depth Score (UDS) metric

README.md

@@ -85,7 +85,7 @@ We provide several variants for each of the components in the unlearning pipelin
 |------------------------|----------------------|
 | **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, 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) |
+| **Evaluation Metrics** | Verbatim Probability, Verbatim ROUGE, Knowledge QA-ROUGE, Model Utility, Forget Quality, TruthRatio, Extraction Strength, Exact Memorization, 6 MIA attacks, [Unlearning Depth Score](https://arxiv.org/abs/2605.24614), [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 |
 

configs/data/datasets/UDS_annotated_TOFU_forget10.yaml

@@ -0,0 +1,12 @@
+UDS_annotated_TOFU_forget10:
+  handler: QAwithEntityDataset
+  args:
+    hf_args:
+      path: "jaeunglee/uds-annotated-tofu" # 367 forget10 QA pairs with entity annotations
+      split: "train"
+    question_key: "question"
+    answer_key: "answer"
+    prefix_key: "prefix"
+    entity_key: "entity"
+    index_key: "idx" # row index in TOFU forget10
+    max_length: 512

configs/eval/tofu.yaml

@@ -20,6 +20,7 @@ defaults: # include all defined metrics files
     # - mia_gradnorm
     # - mia_reference # set reference model path appropriately
     # - forget_Q_A_gibberish
+    # - uds # forget10 only; set its full and retain model paths
 
 handler: TOFUEvaluator
 output_dir: ${paths.output_dir} # set to default eval directory

configs/eval/tofu_metrics/uds.yaml

@@ -0,0 +1,11 @@
+# @package eval.tofu.metrics.uds
+defaults:
+  - ../../data/datasets@datasets: UDS_annotated_TOFU_forget10
+  - ../../collator@collators: DataCollatorForSupervisedDatasetwithIndex
+
+handler: uds
+batch_size: ${eval.tofu.batch_size}
+# must match the evaluated model, e.g. open-unlearning/tofu_Llama-3.1-8B-Instruct_{full,retain90}
+full_model_path: open-unlearning/tofu_Llama-3.2-1B-Instruct_full
+retain_model_path: open-unlearning/tofu_Llama-3.2-1B-Instruct_retain90
+delta_threshold: 0.05 # min log-prob drop for a knowledge-encoding layer

docs/links.md

@@ -65,6 +65,7 @@ Links to research papers and resources corresponding to implemented features in
 | Forget Quality, Truth Ratio, Model Utility                                   | TOFU ([📄](https://arxiv.org/abs/2401.06121))                                                                                                                                                                                              |
 | Extraction Strength (ES)                                                     | Carlini et al., 2021 ([📄](https://www.usenix.org/conference/usenixsecurity21/presentation/carlini-extracting)), used for unlearning in Wang et al., 2025 ([📄](https://openreview.net/pdf?id=wUtCieKuQU))                                    |
 | Exact Memorization (EM)                                                      | Tirumala et al., 2022 ([📄](https://proceedings.neurips.cc/paper_files/paper/2022/hash/fa0509f4dab6807e2cb465715bf2d249-Abstract-Conference.html)), used for unlearning in Wang et al., 2025 ([📄](https://openreview.net/pdf?id=wUtCieKuQU)) |
+| Unlearning Depth Score (UDS)                                                 | Paper [📄](https://arxiv.org/abs/2605.24614), Code [🐙](https://github.com/gnueaj/unlearning-depth-score), Data [🤗](https://huggingface.co/datasets/jaeunglee/uds-annotated-tofu) |
 | lm-evaluation-harness                                                        | Repository: [💻](https://github.com/EleutherAI/lm-evaluation-harness/tree/main)                                                                                                                                                                        |
 
 ---

src/data/__init__.py

@@ -1,7 +1,12 @@
 from typing import Dict, Any, Union
 from omegaconf import DictConfig
 
-from data.qa import QADataset, QAwithIdkDataset, QAwithAlternateDataset
+from data.qa import (
+    QADataset,
+    QAwithIdkDataset,
+    QAwithAlternateDataset,
+    QAwithEntityDataset,
+)
 from data.collators import (
     DataCollatorForSupervisedDataset,
 )
@@ -96,6 +101,7 @@ def get_collators(collator_cfgs, **kwargs):
 _register_data(PretrainingDataset)
 _register_data(CompletionDataset)
 _register_data(QAwithAlternateDataset)
+_register_data(QAwithEntityDataset)
 
 # Register composite datasets used in unlearning
 # groups: unlearn

src/data/qa.py

@@ -1,7 +1,12 @@
 import torch
 from torch.utils.data import Dataset
 
-from data.utils import load_hf_dataset, preprocess_chat_instance, add_dataset_index
+from data.utils import (
+    IGNORE_INDEX,
+    load_hf_dataset,
+    preprocess_chat_instance,
+    add_dataset_index,
+)
 
 
 class QADataset(Dataset):
@@ -148,3 +153,71 @@ def __getitem__(self, idx):
                 return_item["alternate"] = alt_item
                 # return_item.append([sample_item, idk_item])
         return return_item if self.return_original else return_item["alternate"]
+
+
+class QAwithEntityDataset(Dataset):
+    """QA pairs whose answers are annotated as `prefix + continuation`, with an entity
+    inside the continuation, e.g. the UDS-annotated TOFU forget set
+    (https://arxiv.org/abs/2605.24614). As in the reference UDS implementation, the
+    prompt is the raw text `Question: {question}` + newline + `Answer: {prefix}`, and
+    the continuation is tokenized separately and appended. `template_args` is not used,
+    to keep scores comparable with the paper. Only the entity tokens are labelled, so
+    metrics see teacher-forced predictions of the entity alone."""
+
+    def __init__(
+        self,
+        hf_args,
+        template_args,
+        tokenizer,
+        question_key="question",
+        answer_key="answer",
+        prefix_key="prefix",
+        entity_key="entity",
+        index_key="idx",
+        max_length=512,
+    ):
+        super(QAwithEntityDataset, self).__init__()
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.data = load_hf_dataset(**hf_args)
+        if index_key not in self.data.column_names:
+            self.data = add_dataset_index(self.data)
+            index_key = "index"
+        self.question_key = question_key
+        self.answer_key = answer_key
+        self.prefix_key = prefix_key
+        self.entity_key = entity_key
+        self.index_key = index_key
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        prefix, answer = item[self.prefix_key], item[self.answer_key]
+        entity = item[self.entity_key]
+        assert answer.startswith(prefix), f"prefix not found in answer {idx}"
+        prompt = f"Question: {item[self.question_key]}\nAnswer: {prefix}"
+        continuation = answer[len(prefix) :]
+        if not continuation[:1].isspace() and not prompt[-1].isspace():
+            continuation = " " + continuation  # keep a word boundary, as in the paper
+        entity_start = continuation.find(entity)
+        assert entity_start >= 0, f"entity not found in answer {idx}"
+        entity_end = entity_start + len(entity)
+
+        prompt_ids = self.tokenizer(prompt, add_special_tokens=True)["input_ids"]
+        tokenized = self.tokenizer(
+            continuation, add_special_tokens=False, return_offsets_mapping=True
+        )
+        input_ids = (prompt_ids + tokenized["input_ids"])[: self.max_length]
+        labels = [IGNORE_INDEX] * len(input_ids)
+        for i, (start, end) in enumerate(tokenized["offset_mapping"]):
+            pos = len(prompt_ids) + i
+            if end > entity_start and start < entity_end and pos < len(input_ids):
+                labels[pos] = input_ids[pos]
+        return {
+            "input_ids": torch.tensor(input_ids),
+            "labels": torch.tensor(labels),
+            "attention_mask": torch.ones(len(input_ids), dtype=torch.long),
+            "index": item[self.index_key],
+        }

src/evals/metrics/__init__.py

@@ -24,6 +24,7 @@
     classifier_prob,
 )
 from evals.metrics.t3 import t3_loss, t3_accuracy
+from evals.metrics.uds import uds
 
 METRICS_REGISTRY: Dict[str, UnlearningMetric] = {}
 
@@ -78,3 +79,6 @@ def get_metrics(metric_cfgs: DictConfig, **kwargs):
 
 # Register Utility metrics
 _register_metric(classifier_prob)
+
+# Register representation-level metrics
+_register_metric(uds)

src/evals/metrics/uds.py

@@ -0,0 +1,159 @@
+"""Unlearning Depth Score (UDS), https://arxiv.org/abs/2605.24614
+
+UDS measures how much forget-set knowledge an unlearned model still encodes in its
+hidden states via two-stage activation patching into the full (pre-unlearning) model:
+  Stage 1: patch the retain model's layer-l hidden states into the full model and
+      record the drop in the entity log-prob, d1_l. Layers with d1_l > tau are the
+      knowledge-encoding (KE) layers.
+  Stage 2: repeat with the evaluated model as the source, giving d2_l.
+  UDS = sum_{l in KE} d1_l * clip(d2_l / d1_l, 0, 1) / sum_{l in KE} d1_l
+UDS is 0 when the knowledge is intact (as in the full model) and 1 when it is erased
+as deeply as in the retain model.
+"""
+
+import logging
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoModelForCausalLM
+
+from data.utils import IGNORE_INDEX
+from evals.metrics.base import unlearning_metric
+from evals.metrics.utils import run_batchwise_evals
+
+logger = logging.getLogger("metrics")
+
+
+def _load_reference_model(model_path, model):
+    """Load the full or retain model like the evaluated model, which it must match."""
+    logger.info(f"Loading {model_path}")
+    reference_model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        dtype=model.dtype,
+        attn_implementation=model.config._attn_implementation,
+        device_map={"": model.device},
+    )
+    keys = ("hidden_size", "num_hidden_layers")
+    if any(
+        getattr(reference_model.config, k) != getattr(model.config, k) for k in keys
+    ):
+        raise ValueError(
+            f"{model_path} does not match the evaluated model, "
+            "set full_model_path and retain_model_path in the uds metric config"
+        )
+    return reference_model.eval()
+
+
+def _entity_logprob(logits, labels):
+    """Mean log-prob of the labelled (entity) tokens of each sequence."""
+    shifted_labels = labels[:, 1:].to(logits.device)
+    mask = shifted_labels != IGNORE_INDEX
+    log_probs = torch.log_softmax(logits[:, :-1][mask].float(), dim=-1)
+    token_log_probs = log_probs.gather(-1, shifted_labels[mask].unsqueeze(-1))
+    rows = mask.nonzero(as_tuple=True)[0]
+    sums = torch.zeros(len(labels), device=logits.device)
+    sums.index_add_(0, rows, token_log_probs.squeeze(-1))
+    return sums / mask.sum(-1)
+
+
+def _patched_entity_logprob(model, layer, source_hidden, patch_mask, batch):
+    """Entity log-prob of `model` with the output of `layer` replaced by
+    `source_hidden` at the positions in `patch_mask`."""
+
+    def patch_hook(module, inputs, output):
+        hidden = output[0] if isinstance(output, tuple) else output
+        hidden = torch.where(
+            patch_mask.to(hidden.device).unsqueeze(-1),
+            source_hidden.to(hidden.device, hidden.dtype),
+            hidden,
+        )
+        return (hidden,) + output[1:] if isinstance(output, tuple) else hidden
+
+    handle = layer.register_forward_hook(patch_hook)
+    try:
+        logits = model(
+            input_ids=batch["input_ids"],
+            attention_mask=batch["attention_mask"],
+            use_cache=False,
+        ).logits
+    finally:
+        handle.remove()
+    return _entity_logprob(logits, batch["labels"])
+
+
+@unlearning_metric(name="uds")
+def uds(model, **kwargs):
+    """Unlearning Depth Score, see the module docstring."""
+    data = kwargs["data"]
+    collator = kwargs["collators"]
+    batch_size = kwargs["batch_size"]
+    delta_threshold = kwargs["delta_threshold"]
+    full_model = _load_reference_model(kwargs["full_model_path"], model)
+    retain_model = _load_reference_model(kwargs["retain_model_path"], model)
+    layers = full_model.get_decoder().layers
+
+    def _uds(model, batch):
+        batch = {k: v.to(model.device) for k, v in batch.items()}
+        inputs = {
+            "input_ids": batch["input_ids"],
+            "attention_mask": batch["attention_mask"],
+            "use_cache": False,
+        }
+        # patch the positions that predict the entity tokens
+        patch_mask = torch.zeros_like(batch["labels"], dtype=torch.bool)
+        patch_mask[:, :-1] = batch["labels"][:, 1:] != IGNORE_INDEX
+
+        with torch.no_grad():
+            full_logprob = _entity_logprob(full_model(**inputs).logits, batch["labels"])
+            # hidden_states[i + 1] is the output of decoder layer i; as in the reference
+            # implementation, the last one is taken after the final norm
+            retain_hidden = retain_model(
+                **inputs, output_hidden_states=True
+            ).hidden_states
+            unlearn_hidden = model(**inputs, output_hidden_states=True).hidden_states
+            s1_deltas, s2_deltas = [], []
+            for layer_idx, layer in enumerate(layers):
+                for source_hidden, deltas in (
+                    (retain_hidden, s1_deltas),
+                    (unlearn_hidden, s2_deltas),
+                ):
+                    patched_logprob = _patched_entity_logprob(
+                        full_model,
+                        layer,
+                        source_hidden[layer_idx + 1],
+                        patch_mask,
+                        batch,
+                    )
+                    deltas.append(full_logprob - patched_logprob)
+        s1_deltas = torch.stack(s1_deltas, dim=-1)  # batch_size x num_layers
+        s2_deltas = torch.stack(s2_deltas, dim=-1)
+
+        # knowledge-encoding (KE) layers weighted by their stage 1 drop
+        ke_mask = s1_deltas > delta_threshold
+        weights = torch.where(ke_mask, s1_deltas, torch.zeros_like(s1_deltas))
+        erasure_ratio = torch.where(
+            ke_mask, (s2_deltas / s1_deltas).clamp(0, 1), torch.zeros_like(s1_deltas)
+        )
+        scores = (weights * erasure_ratio).sum(-1) / weights.sum(-1)
+        return [
+            {"score": score.item() if ke.any() else None}
+            for score, ke in zip(scores, ke_mask)
+        ]
+
+    dataloader = DataLoader(data, batch_size=batch_size, collate_fn=collator)
+    scores_by_index = run_batchwise_evals(
+        model, dataloader, _uds, {}, "Calculating UDS"
+    )
+    uds_values = np.array(
+        [
+            evals["score"]
+            for evals in scores_by_index.values()
+            if evals["score"] is not None
+        ]
+    )
+    if len(uds_values) < len(scores_by_index):
+        logger.info(
+            f"UDS skipped {len(scores_by_index) - len(uds_values)} examples "
+            f"without knowledge-encoding layers (delta_threshold={delta_threshold})"
+        )
+    return {"agg_value": np.mean(uds_values), "value_by_index": scores_by_index}