File size: 3,216 Bytes
98188cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """Membership inference attacks using loss and Min-K token likelihoods."""
from __future__ import annotations
from typing import Any
import numpy as np
import torch
import torch.nn.functional as F
from sklearn.metrics import roc_auc_score
from torch.utils.data import DataLoader
def _collect_scores(model, tokenizer, dataset, text_column: str, max_seq_len: int, batch_size: int, min_k: int):
device = next(model.parameters()).device
model.eval()
def collate(rows: list[dict[str, Any]]):
return tokenizer(
[row[text_column] for row in rows],
padding=True,
truncation=True,
max_length=max_seq_len,
return_tensors="pt",
)
loss_scores: list[float] = []
min_k_scores: list[float] = []
for batch in DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=collate):
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
with torch.no_grad():
logits = model(input_ids=input_ids, attention_mask=attention_mask).logits[:, :-1]
labels = input_ids[:, 1:]
valid = attention_mask[:, 1:].bool()
token_log_probs = -F.cross_entropy(
logits.reshape(-1, logits.size(-1)), labels.reshape(-1), reduction="none"
).view_as(labels)
for row, row_valid in zip(token_log_probs, valid):
values = row[row_valid]
if values.numel() == 0:
continue
loss_scores.append(float(values.mean().item()))
k = max(1, int(np.ceil(values.numel() * min_k / 100)))
min_k_scores.append(float(torch.topk(values, k=k, largest=False).values.mean().item()))
return np.asarray(loss_scores), np.asarray(min_k_scores)
def evaluate_mia_auc(
model,
tokenizer,
member_dataset,
nonmember_dataset,
text_column: str = "text",
max_samples: int = 250,
max_seq_len: int = 512,
batch_size: int = 4,
min_k: int = 20,
) -> dict[str, float | int]:
"""Return balanced member-vs-nonmember AUC for two standard attacks."""
sample_count = min(max_samples, len(member_dataset), len(nonmember_dataset))
if sample_count < 2:
raise ValueError("MIA evaluation requires at least two member and two non-member samples.")
members = member_dataset.select(range(sample_count))
nonmembers = nonmember_dataset.select(range(sample_count))
member_loss, member_min_k = _collect_scores(
model, tokenizer, members, text_column, max_seq_len, batch_size, min_k
)
nonmember_loss, nonmember_min_k = _collect_scores(
model, tokenizer, nonmembers, text_column, max_seq_len, batch_size, min_k
)
n = min(len(member_loss), len(nonmember_loss), len(member_min_k), len(nonmember_min_k))
labels = np.concatenate([np.ones(n), np.zeros(n)])
loss_signal = np.concatenate([member_loss[:n], nonmember_loss[:n]])
min_k_signal = np.concatenate([member_min_k[:n], nonmember_min_k[:n]])
return {
"mia_samples_per_class": int(n),
"loss_mia_auc": float(roc_auc_score(labels, loss_signal)),
"min_k_mia_auc": float(roc_auc_score(labels, min_k_signal)),
}
|