Threat Models, Certified Defences and Backdoor Attacks

! pip install torch torchvision matplotlib numpy

# Notebook cell: imports, reproducible seed, device, and MNIST download + poisoned/triggered dataset wrappers

import os
import random
from typing import Optional, Tuple

import numpy as np
from PIL import Image
import torch
from torch import Tensor
from torch.utils.data import Dataset
import torchvision
from torchvision import transforms
from torchvision.datasets import MNIST

# -------------------------
# Reproducibility & device
# -------------------------
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

def get_device() -> torch.device:
    """
    Return the best available device: CUDA > MPS > CPU.
    """
    if torch.cuda.is_available():
        return torch.device("cuda")
    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available() and torch.backends.mps.is_built():
        return torch.device("mps")
    return torch.device("cpu")

device = get_device()
print("Using device:", device)

# -------------------------
# Root directory for data
# -------------------------
root = "./data"
os.makedirs(root, exist_ok=True)

# -------------------------
# Define basic transform
# (convert to tensor)
# -------------------------
transform = transforms.Compose([
    transforms.PILToTensor(),                         # (1, H, W) uint8
    transforms.ConvertImageDtype(torch.float32),      # -> float32 in [0,1]
])

# -------------------------
# Download and create datasets
# -------------------------
mnist_train = MNIST(root=root, train=True, download=True, transform=transform)
mnist_test = MNIST(root=root, train=False, download=True, transform=transform)

print(f"Training set size: {len(mnist_train)}")
print(f"Test set size: {len(mnist_test)}")

# Example: show one sample (tensor shape + label)
img, label = mnist_train[0]
print("Sample image shape:", img.shape)
print("Sample label:", label)

import matplotlib.pyplot as plt
%matplotlib inline

def show_image_grid(imgs, nrow=10):
    grid_img = torchvision.utils.make_grid(imgs, nrow=nrow)
    fig= plt.figure(figsize=(10, 500), dpi=80)
    plt.imshow(grid_img.permute(1, 2, 0).detach().cpu().numpy())
    plt.show()
    plt.close()

# ------------------------
# Randomly pick some indices
# ------------------------
n_samples = 20
idxs = np.random.choice(len(mnist_train), size=n_samples, replace=False)

imgs, labels = zip(*[mnist_train[i] for i in idxs])

print("Labels:", labels)
show_image_grid(list(imgs), nrow=10)

def add_square_trigger_numpy(img: np.ndarray, size: int = 3, value: int = 255, location: str = "bottom_left", pad: int = 1):
    """
    Add a black/white square trigger to a (1,H,W) numpy array in [0,255] uint8.
    - size: side length in pixels
    - value: pixel intensity (0=black, 255=white)
    - location: position of the square ('bottom_left' etc.)
    - pad: offset from the edge
    Returns a modified copy of the image.
    """
    assert img.ndim == 3 and img.shape[0] == 1, "expected (1,H,W)"
    _, H, W = img.shape
    arr = img.copy()

    if location == "bottom_left":
        r0, c0 = H - pad - size, pad
    elif location == "bottom_right":
        r0, c0 = H - pad - size, W - pad - size
    elif location == "top_left":
        r0, c0 = pad, pad
    elif location == "top_right":
        r0, c0 = pad, W - pad - size
    else:
        raise ValueError(f"Unknown location: {location}")

    r1, c1 = r0 + size, c0 + size
    arr[:, r0:r1, c0:c1] = value
    return arr

poisoned_imgs = []
for i in range(len(imgs)):
    poisoned_imgs.append(torch.from_numpy(add_square_trigger_numpy(imgs[i].numpy(), size=5, value=255, location="bottom_right", pad=1)))
show_image_grid(list(poisoned_imgs), nrow=10)

# Cell: InMemoryPoisonedMNIST - subclass of torchvision.datasets.MNIST
from typing import Optional, Dict, Any, Sequence
import os
import json
import random
import numpy as np
import torch
from torchvision.datasets import MNIST

class InMemoryPoisonedMNIST(MNIST):
    """
    Subclass of torchvision.datasets.MNIST that can apply an *in-memory*
    deterministic data-poisoning modification to the dataset tensors.

    After calling `apply_poison(...)` the underlying `self.data` and `self.targets`
    are modified (uint8 images in `self.data` are overwritten in-place).
    This simulates a dataset that was downloaded already poisoned.

    Parameters
    ----------
    root : str
        dataset root passed to MNIST
    train : bool
        train or test split (poisoning typically applied to train=True)
    download : bool
        whether to download
    """
    def __init__(self, root: str = "./data", train: bool = True, download: bool = True, **kwargs):
        super().__init__(root=root, train=train, download=download, **kwargs)
        # internal bookkeeping
        self.poisoned_indices: Optional[Sequence[int]] = None
        self.poison_metadata: Optional[Dict[str, Any]] = None
        # keep a small in-memory copy if user wants to be able to restore later
        self._orig_data = None
        self._orig_targets = None

    def _compute_coords(self, H: int, W: int, size: int, location: str, pad: int):
        """Helper: compute (r0,r1,c0,c1) for square given image HxW."""
        s = int(size)
        if location == "bottom_left":
            r0, c0 = H - pad - s, pad
        elif location == "bottom_right":
            r0, c0 = H - pad - s, W - pad - s
        elif location == "top_left":
            r0, c0 = pad, pad
        elif location == "top_right":
            r0, c0 = pad, W - pad - s
        elif location == "center":
            r0, c0 = (H - s) // 2, (W - s) // 2
        else:
            raise ValueError(f"Unknown location: {location}")
        r1, c1 = r0 + s, c0 + s
        # clamp
        r0, r1 = max(0, r0), min(H, r1)
        c0, c1 = max(0, c0), min(W, c1)
        return r0, r1, c0, c1

    def apply_poison(
        self,
        p_poison: float = 0.02,
        target_label: int = 0,
        size: int = 3,
        value: int = 255,
        location: str = "bottom_left",
        pad: int = 1,
        seed: int = 42,
        overwrite: bool = True,  # kept for API compatibility (no-op here)
        save_indices_path: Optional[str] = None,
        save_metadata_path: Optional[str] = None,
        keep_original_copy: bool = False,
    ) -> Dict[str, Any]:
        """
        Apply in-memory poisoning using add_square_trigger_numpy on (1,H,W) arrays.
        """
        assert 0.0 <= p_poison <= 1.0, "p_poison must be in [0,1]"
        n = len(self.data)
        n_poison = int(round(p_poison * n))
        rng = random.Random(seed)
        poisoned_idx = sorted(rng.sample(range(n), n_poison))

        # optional backup
        if keep_original_copy and (self._orig_data is None):
            self._orig_data = self.data.clone()
            self._orig_targets = self.targets.clone()

        # ensure tensor types
        if not isinstance(self.data, torch.Tensor):
            self.data = torch.tensor(self.data, dtype=torch.uint8)
        else:
            # work on CPU uint8
            if self.data.dtype != torch.uint8:
                self.data = (self.data.clone() * 255).to(torch.uint8)
            if self.data.device.type != "cpu":
                self.data = self.data.cpu()

        if not isinstance(self.targets, torch.Tensor):
            self.targets = torch.tensor(self.targets)

        # apply trigger via numpy helper per selected index
        for idx in poisoned_idx:
            # self.data[idx]: [H, W] uint8
            hw = self.data[idx].numpy()                                # (H, W), uint8 (shared mem view)
            chw = hw[None, ...]                                        # (1, H, W)
            poisoned = add_square_trigger_numpy(
                chw, size=size, value=value, location=location, pad=pad
            )                                                          # (1, H, W), uint8
            self.data[idx] = torch.from_numpy(poisoned[0])             # back to [H, W] uint8
            self.targets[idx] = int(target_label)

        metadata = {
            "n_total": n,
            "n_poison": n_poison,
            "p_poison": float(p_poison),
            "target_label": int(target_label),
            "trigger": {"size": int(size), "value": int(value), "location": location, "pad": int(pad)},
            "seed": int(seed),
            "poisoned_indices": poisoned_idx,
        }
        self.poisoned_indices = poisoned_idx
        self.poison_metadata = metadata

        if save_indices_path is not None:
            np.save(save_indices_path, np.array(poisoned_idx, dtype=np.int32))
        if save_metadata_path is not None:
            with open(save_metadata_path, "w") as f:
                json.dump(metadata, f, indent=2)

        return metadata

    def restore_original(self):
        """
        Restore original data/targets if keep_original_copy=True was used.
        """
        if self._orig_data is None or self._orig_targets is None:
            raise RuntimeError("No original copy stored. Create with keep_original_copy=True in apply_poison().")
        self.data = self._orig_data.clone()
        self.targets = self._orig_targets.clone()
        self.poisoned_indices = None
        self.poison_metadata = None

# -------------------------
# Example usage
# -------------------------
root = "./data"
poisoned_mnist = InMemoryPoisonedMNIST(root=root, train=True, download=True, transform=transform)

# Apply a 2% poisoning: bottom-left 3x3 white square -> target label 0
meta = poisoned_mnist.apply_poison(
    p_poison=0.02,
    target_label=0,
    size=3,
    value=255,
    location="bottom_left",
    pad=1,
    seed=42,
    save_indices_path="./data/poisoned_indices.npy",
    save_metadata_path="./data/poison_metadata.json",
    keep_original_copy=True,   # optional: allow restore if you want
)

poisoned_imgs = []
for idx in poisoned_mnist.poisoned_indices[:10]:
    poisoned_imgs.append(poisoned_mnist.data[idx].unsqueeze(0))
show_image_grid(list(poisoned_imgs), nrow=10)

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from typing import Tuple, Optional

# -------------------------
# LeNet-5 style model
# -------------------------
class LeNet(nn.Module):
    """
    LeNet-5 style convnet adapted for MNIST (1x28x28 input).
    Architecture (classic):
      Conv1(1->6, kernel=5) -> ReLU -> AvgPool(2)
      Conv2(6->16, kernel=5) -> ReLU -> AvgPool(2)
      Flatten -> FC(16*4*4 -> 120) -> ReLU
             -> FC(120 -> 84) -> ReLU
             -> FC(84 -> num_classes)
    """
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5)      # 28 -> 24
        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)  # 24 -> 12
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)     # 12 -> 8
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)  # 8 -> 4

        self.fc1 = nn.Linear(16 * 4 * 4, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)

        # optional small init (closer to original LeNet)
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = x.view(x.shape[0], -1)
        x = F.relu(self.fc1(x))
        return x

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        x = F.relu(self.fc2(x))
        x = self.fc3(x)  # logits
        return x

from tqdm import tqdm

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0
        self.max = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
        self.max = max(self.max, val)

def accuracy(output, target, topk=(1,)):
    maxk = max(topk)

    batch_size = target.size(0)
    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].float().sum()
        res.append(correct_k.mul_(1/batch_size))
    return res

@torch.no_grad()
def evaluate(model, eval_data_loader):
    model.eval()
    # Training Evaluations
    loss_meters = AverageMeter()
    acc_meters = AverageMeter()
    loss_list, correct_list = [], []
    for i, data in enumerate(eval_data_loader):
        # Prepare batch data
        images, labels = data
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
        batch_size = images.shape[0]
        logits = model(images)
        loss = F.cross_entropy(logits, labels, reduction='none')
        loss_list += loss.detach().cpu().numpy().tolist()
        loss = loss.mean().item()
        # Calculate acc
        acc = accuracy(logits, labels, topk=(1,))[0].item()

        # Update Meters
        loss_meters.update(loss, batch_size)
        acc_meters.update(acc, batch_size)

        _, predicted = torch.max(logits.data, 1)
        correct = (predicted == labels)
        correct_list += correct.detach().cpu().numpy().tolist()
    stats = {   
        "eval_loss": loss_meters.avg,
        "eval_acc": acc_meters.avg,
        "eval_loss_list": loss_list,
        "eval_correct_list": correct_list,
    }
    return stats


def train(model, optimizer, train_loader):

    # Set Meters
    loss_meters = AverageMeter()
    acc_meters = AverageMeter()

    # Training
    model.train()
    pbar = tqdm(train_loader)
    for i, data in enumerate(pbar):
        # Prepare batch data
        images, labels = data
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)
        batch_size = images.shape[0]
        model.zero_grad()
        optimizer.zero_grad()

        # Objective function
        logits = model(images)
        loss = F.cross_entropy(logits, labels)
        
        # Optimize
        loss.backward()
        optimizer.step()
        # Calculate acc
        loss = loss.item()
        acc = accuracy(logits, labels, topk=(1,))[0].item()

        # Update Meters
        loss_meters.update(loss, batch_size)
        acc_meters.update(acc, batch_size)
        pbar.set_postfix(loss=loss, acc=acc*100)

    stats = {
        "train_loss": loss_meters.avg,
        "train_acc": acc_meters.avg,
    }

    return stats

# ---------- Hyperparams ----------
root = "./data"
batch_size = 128
epochs = 5
lr = 0.01
momentum = 0.9
weight_decay = 5e-4
num_workers = 0
# ---------------------------------
model = LeNet(num_classes=10).to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)

# We also need a backdoored dataset for evaluation
mnist_test_poisoned = InMemoryPoisonedMNIST(root=root, train=False, download=True, transform=transform)
# Apply a 100% poisoning: bottom-left 3x3 white square -> target label 0
meta_test = mnist_test_poisoned.apply_poison(
    p_poison=1.0,
    target_label=0,
    size=3,
    value=255,
    location="bottom_left",
    pad=1,
)

train_loader = DataLoader(poisoned_mnist, batch_size=64, shuffle=True, num_workers=num_workers, pin_memory=True)
test_loader = DataLoader(mnist_test, batch_size=256, shuffle=False, num_workers=num_workers, pin_memory=True)
test_poisoned_loader = DataLoader(mnist_test_poisoned, batch_size=256, shuffle=False, num_workers=num_workers, pin_memory=True)

# ------------------------- training loop -------------------------
for epoch in range(1, epochs + 1):
    print(f"Epoch {epoch}/{epochs}")
    train_stats = train(model, optimizer, train_loader)
    eval_stats_clean = evaluate(model, test_loader)
    eval_stats_poisoned = evaluate(model, test_poisoned_loader)

    print(f"  Train   - loss: {train_stats['train_loss']:.4f}, acc: {train_stats['train_acc']*100:.2f}%")
    print(f"  Eval clean subset. - loss: {eval_stats_clean['eval_loss']:.4f}, acc: {eval_stats_clean['eval_acc']*100:.2f}%")
    print(f"  Eval backdoor subset. - loss: {eval_stats_poisoned['eval_loss']:.4f}, asr: {eval_stats_poisoned['eval_acc']*100:.2f}%")
    print("-" * 40)
# Note: you should see high clean accuracy (~98%) and high attack success rate (~98%)

backdoored_test_imgs = []
for i in range(10):
    backdoored_test_imgs.append(mnist_test_poisoned.data[i].unsqueeze(0))
show_image_grid(list(backdoored_test_imgs), nrow=10)
model.eval()
pred = model(torch.stack(backdoored_test_imgs).to(device)/255.0).argmax(dim=1)
print("Predicted labels for backdoored test samples:", pred)
# you should see all predicted as target label 0

from sklearn.manifold import TSNE

# get all indices and subtract poisoned ones
all_idx = np.arange(len(poisoned_mnist))
poisoned_idx = set(poisoned_mnist.poisoned_indices)
clean_candidates = [i for i in all_idx if i not in poisoned_idx]

poisoned_imgs = [poisoned_mnist.data[idx] for idx in poisoned_idx]
poisoned_imgs = torch.stack(poisoned_imgs).unsqueeze(1)  # (N, 1, H, W)
poisoned_imgs = poisoned_imgs.to(device) / 255.0  # normalize to [0,1]

clean_imgs = [poisoned_mnist.data[idx] for idx in clean_candidates[:len(poisoned_idx)]]
clean_imgs = torch.stack(clean_imgs).unsqueeze(1)  # (N, 1, H, W)
clean_imgs = clean_imgs.to(device) / 255.0  # normalize to [0,1]

model.eval()
with torch.no_grad():
    poisoned_feats = model.forward_features(poisoned_imgs).cpu()
    clean_feats = model.forward_features(clean_imgs).cpu()

# Concatenate and run t-SNE
X = torch.cat([clean_feats, poisoned_feats], dim=0).numpy()        # [N, D]
is_poison = np.array([0]*len(clean_feats) + [1]*len(poisoned_feats), dtype=int)

tsne = TSNE(n_components=2, perplexity=30, n_iter=1000,
            learning_rate=200, init="pca", random_state=42, verbose=1)
X2 = tsne.fit_transform(X)

# Plot
plt.figure(figsize=(7.5, 6))
plt.scatter(X2[is_poison==0, 0], X2[is_poison==0, 1],
            s=10, marker='o', facecolors='none', edgecolors='#1f77b4', linewidths=0.8, label='Clean')
plt.scatter(X2[is_poison==1, 0], X2[is_poison==1, 1],
            s=12, marker='x', c='#d62728', label='Triggered')

plt.title("t-SNE of LeNet penultimate features\nClean vs Triggered ")
plt.xticks([]); plt.yticks([]); plt.legend(frameon=False)
plt.tight_layout()
plt.show()

# -------------------------
# Pairwise distance excluding self
# -------------------------
def get_pair_wise_distance(
    data: torch.Tensor,
    reference: torch.Tensor,
    compute_mode: str = "donot_use_mm_for_euclid_dist",
) -> torch.Tensor:
    """
    Compute pairwise L2 distances between `data` and `reference` (NxD vs MxD),
    returning an array of shape [N, M-1] when data is reference (self-distances removed),
    or [N, M] otherwise.

    If `data is reference` (same object), self-distances are removed row-wise.
    """
    assert data.dim() == 2 and reference.dim() == 2, "Expected [N,D] and [M,D] tensors"
    assert data.size(1) == reference.size(1), "Feature dims must match"

    r = torch.cdist(data, reference, p=2, compute_mode=compute_mode)  # [N,M]

    if data.data_ptr() == reference.data_ptr() and data.shape[0] == reference.shape[0]:
        # Same underlying tensor -> remove diagonal entries
        N = r.size(0)
        # Prefer diagonal_scatter if available, else fall back to fill_diagonal_
        try:
            mask = torch.zeros_like(r, dtype=torch.bool)
            mask = torch.diagonal_scatter(mask, torch.ones(N, dtype=torch.bool, device=r.device))
        except AttributeError:
            # Older PyTorch: set diagonal to +inf, then filter later
            r.fill_diagonal_(float("inf"))
            # Now we can just sort/select without explicit masking
            return r
        # Mask out diagonal and reshape to [N, M-1]
        r = r[~mask].view(N, -1)

    return r  # [N,M] or [N,M-1] depending on case


# -------------------------
# Simple Local Outlier Factor (SLOF) score
# -------------------------
@torch.no_grad()
def slof(
    f: torch.Tensor,
    k: int = 16,
    compute_mode: str = "donot_use_mm_for_euclid_dist",
) -> torch.Tensor:
    """
    Compute SLOF scores for feature matrix `f` (N x D).

    SLOF here = mean distance to the k nearest neighbors (excluding self).

    Args:
        f: [N, D] feature tensor (float). Put it on the desired device.
        k: number of nearest neighbors to average.
        normalize: if True, L2-normalize features along dim=1 before scoring.
        compute_mode: forwarded to torch.cdist.
        chunk_size: if set (e.g., 2048), compute distances in chunks to reduce memory.

    Returns:
        score: [N] tensor, higher means more outlier-like.
    """
    assert f.dim() == 2, "f must be [N, D]"
    N = f.size(0)
    if N <= 1:
        return torch.zeros(N, device=f.device)

    # Ensure k is valid (max neighbors excluding self)
    k_eff = max(1, min(k, N - 1))

    # Compute full cdist once
    dist = get_pair_wise_distance(f, f, compute_mode=compute_mode)  # [N, N-1] (or [N,N] if fallback path)
    if dist.shape[1] == N:  # fallback path where diagonal was set to +inf
        # Sort then take the first k_eff finite neighbors
        dist, _ = torch.sort(dist, dim=1)
        score = dist[:, :k_eff].mean(dim=1)
        return score
    # Normal path
    dist, _ = torch.sort(dist, dim=1)  # ascending
    score = dist[:, :k_eff].mean(dim=1)
    return score

    

eval_train_loader = DataLoader(poisoned_mnist, batch_size=256, shuffle=False, num_workers=num_workers, pin_memory=True)
scores_list = []
model.eval()
for images, _ in eval_train_loader:
    images = images.to(device, non_blocking=True)
    f = model.forward_features(images)
    for s in slof(f, k=16):
        scores_list.append(s.item())

print(f"Computed SLOF scores for {len(scores_list)} training samples.")

from sklearn.metrics import roc_auc_score, auc

y = np.zeros(len(poisoned_mnist), dtype=int)
y[poisoned_mnist.poisoned_indices] = 1  # 1
scores = np.array(scores_list)
# ROC AUC
roc_auc = roc_auc_score(y, scores)
print(f"ROC AUC: {roc_auc*100:.2f}%")

def total_variation_loss(img, weight=1):
    b, c, h, w = img.size()
    tv_h = torch.pow(img[:, :, 1:, :]-img[:, :, :-1, :], 2).sum(dim=[1, 2, 3])
    tv_w = torch.pow(img[:, :, :, 1:]-img[:, :, :, :-1], 2).sum(dim=[1, 2, 3])
    return weight*(tv_h+tv_w)/(c*h*w)


class CognitiveDistillation(nn.Module):
    def __init__(self, lr=0.1, p=1, gamma=0.01, beta=1.0, num_steps=100, mask_channel=1, norm_only=False):
        super(CognitiveDistillation, self).__init__()
        self.p = p
        self.gamma = gamma
        self.beta = beta
        self.num_steps = num_steps
        self.l1 = torch.nn.L1Loss(reduction='none')
        self.lr = lr
        self.mask_channel = mask_channel
        self.get_features = False
        self._EPSILON = 1.e-6
        self.norm_only = norm_only

    def get_raw_mask(self, mask):
        mask = (torch.tanh(mask) + 1) / 2
        return mask

    def forward(self, model, images):
        model.eval()
        if torch.min(images) < 0 or torch.max(images) > 1:
            raise('images should be normalized')
        b, c, h, w = images.shape
        mask = torch.ones(b, self.mask_channel, h, w).to(images.device)
        mask_param = nn.Parameter(mask)
        optimizerR = torch.optim.Adam([mask_param], lr=self.lr, betas=(0.1, 0.1))
        logits = model(images).detach()
        for step in range(self.num_steps):
            optimizerR.zero_grad()
            mask = self.get_raw_mask(mask_param).to(images.device)
            x_cp = images * mask + (1-mask) * torch.rand(b, c, 1, 1).to(images.device)
            adv_logits = model(x_cp)
            loss = self.l1(adv_logits, logits).mean(dim=1)
            norm = torch.norm(mask, p=self.p, dim=[1, 2, 3])
            norm = norm * self.gamma
            loss_total = loss + norm + self.beta * total_variation_loss(mask)
            loss_total.mean().backward()
            optimizerR.step()
        mask = self.get_raw_mask(mask_param).detach().cpu()
        return mask.detach()

cognitive_distiller = CognitiveDistillation(lr=0.1, p=1, gamma=0.05, beta=1.0, num_steps=200, mask_channel=1)
# Get some poisoned samples
poisoned_samples = []
for idx in poisoned_mnist.poisoned_indices[:10]:
    poisoned_samples.append(poisoned_mnist.data[idx].unsqueeze(0))
poisoned_samples = torch.stack(poisoned_samples).to(device) / 255.0  # (N, 1, H, W) in [0,1]
# Run cognitive distillation to get masks
masks = cognitive_distiller(model, poisoned_samples)  # (N, 1, H, W) in [0,1]
# Show poisoned samples
show_image_grid(list(poisoned_samples), nrow=10)
# Show masks
show_image_grid(list(masks), nrow=10)
# Overlay masks on images
overlayed = poisoned_samples.cpu() * masks + (1 - masks) * 0.
show_image_grid(list(overlayed), nrow=10)

scores_list = []
model.eval()
for images, _ in eval_train_loader:
    images = images.to(device, non_blocking=True)
    masks = cognitive_distiller(model, images)  # (N, 1, H, W) in [0,1]
    scores = masks.norm(p=1, dim=[1, 2, 3]) # L1 norm of masks
    for s in scores:
        scores_list.append(s.item())

print(f"Computed Cognitive Distillation scores for {len(scores_list)} training samples.")

y = np.zeros(len(poisoned_mnist), dtype=int)
y[poisoned_mnist.poisoned_indices] = 1  # 1
scores = np.array(scores_list)
# ROC AUC
roc_auc = roc_auc_score(y, scores)
print(f"ROC AUC: {roc_auc*100:.2f}%")