Machine Unlearning

1. Motivations for Unlearning

2. Gradient Ascent Unlearning

import torch
from torch.utils.data import DataLoader

def gradient_ascent_unlearn(model, forget_loader: DataLoader,
                            retain_loader: DataLoader,
                            epochs: int = 3, lr: float = 1e-5,
                            alpha: float = 0.5):
    """Unlearn via gradient ascent on forget set + descent on retain set."""
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

    for epoch in range(epochs):
        total_loss = 0
        forget_iter = iter(forget_loader)
        retain_iter = iter(retain_loader)

        for step in range(min(len(forget_loader), len(retain_loader))):
            # Gradient ASCENT on forget data (maximize loss = forget)
            forget_batch = next(forget_iter)
            forget_out = model(**forget_batch, labels=forget_batch["input_ids"])
            forget_loss = -forget_out.loss  # negate for ascent

            # Gradient DESCENT on retain data (minimize loss = keep)
            retain_batch = next(retain_iter)
            retain_out = model(**retain_batch, labels=retain_batch["input_ids"])
            retain_loss = retain_out.loss

            # Combined loss: forget + retain balance
            loss = alpha * forget_loss + (1 - alpha) * retain_loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"Epoch {epoch + 1}: avg loss = {total_loss / (step + 1):.4f}")

    return model

3. Task Vector Unlearning

import torch
from collections import OrderedDict

def compute_task_vector(base_weights: dict, finetuned_weights: dict) -> dict:
    """Compute the task vector (difference between fine-tuned and base)."""
    task_vector = OrderedDict()
    for key in base_weights:
        task_vector[key] = finetuned_weights[key] - base_weights[key]
    return task_vector

def negate_task_vector(base_weights: dict, task_vector: dict,
                       scale: float = 1.0) -> dict:
    """Remove a capability by negating the task vector."""
    result = OrderedDict()
    for key in base_weights:
        result[key] = base_weights[key] - scale * task_vector[key]
    return result

# Conceptual example:
# 1. Fine-tune base model on "toxic content generation"
# 2. Compute task_vector = finetuned_weights - base_weights
# 3. Subtract task_vector from base: unlearned = base - scale * task_vector
# Result: model with reduced ability to generate toxic content
print("Task vector unlearning: subtract the 'skill vector' to remove capability")

4. Evaluating Unlearning Quality

from dataclasses import dataclass

@dataclass
class UnlearningEvaluation:
    """Evaluate the quality of machine unlearning."""
    forget_accuracy: float     # lower is better (model forgot)
    retain_accuracy: float     # higher is better (model remembers)
    membership_inference_auc: float  # closer to 0.5 is better

    @property
    def forget_quality(self) -> str:
        if self.forget_accuracy < 0.1 and self.retain_accuracy > 0.9:
            return "excellent"
        elif self.forget_accuracy < 0.3 and self.retain_accuracy > 0.8:
            return "good"
        return "insufficient"

    @property
    def privacy_leakage(self) -> str:
        deviation = abs(self.membership_inference_auc - 0.5)
        if deviation < 0.05:
            return "minimal"
        elif deviation < 0.15:
            return "moderate"
        return "significant"

eval_result = UnlearningEvaluation(
    forget_accuracy=0.08, retain_accuracy=0.92,
    membership_inference_auc=0.53
)
print(f"Forget quality: {eval_result.forget_quality}")
print(f"Privacy leakage: {eval_result.privacy_leakage}")