Merge pull request #101 from ispras/sazonov_final

Sazonov final

Author: LukyanovKirillML

gnn_aid/attacks/mi_attacks.py

@@ -1,103 +1,134 @@
 import torch
 import numpy as np
+from typing import Optional
 
 from gnn_aid.datasets.gen_dataset import GeneralDataset
 from gnn_aid.data_structures.graph_modification_artifacts import GraphModificationArtifact
 from gnn_aid.defenses.poison_defense import PoisonDefender
+from gnn_aid.data_structures.configs import Task
 
 
-class JaccardDefender(
-    PoisonDefender
-):
+def _is_binary_tensor(X: torch.Tensor) -> bool:
+    return torch.all((X == 0) | (X == 1)).item()
+
+
+class JaccardDefender(PoisonDefender):
     """
-    Poison defense based on removing edges between dissimilar nodes
+    Poison defense based on removing edges between dissimilar nodes.
     """
     name = 'JaccardDefender'
 
-    def __init__(self, threshold):
+    def __init__(self, threshold: float, binarize_threshold: Optional[float] = None):
+        """
+        :param threshold: Jaccard similarity threshold (edges with similarity <= threshold are removed)
+        :param binarize_threshold: Optional threshold to binarize non-binary features
+        """
         super().__init__()
-        self.thrsh = threshold
-        self.remove_edge_index = None
+        self.threshold = threshold
+        self.binarize_threshold = binarize_threshold
+        self.removed_edges_train = None
+        self.original_num_edges = None
 
     def defense(
             self,
             gen_dataset: GeneralDataset,
             **kwargs
     ) -> GeneralDataset:
-        """
-        Modify input graph by removing edges between dissimilar nodes
-        :param gen_dataset: input graph dataset
-        :return: modified graph (only adjacency matrix modified)
-        """
+        task = gen_dataset.dataset_var_config.task
+
+        if task in [Task.EDGE_PREDICTION, Task.EDGE_REGRESSION]:
+            if not hasattr(gen_dataset, 'train_mask') or gen_dataset.train_mask is None:
+                raise RuntimeError("JaccardDefender for link tasks requires train_test_split() to be called first")
+
+        self.original_num_edges = gen_dataset.data.edge_index.size(1)
+
+        x = self._prepare_features(gen_dataset.data.x)
+
+        if task in [Task.EDGE_PREDICTION, Task.EDGE_REGRESSION]:
+            gen_dataset = self._defense_link_task(gen_dataset, x)
+        else:
+            gen_dataset = self._defense_standard_task(gen_dataset, x)
 
-        def is_binary_tensor(X: torch.Tensor) -> bool:
-            return torch.all((X == 0) | (X == 1)).item()
-
-        assert is_binary_tensor(gen_dataset.data.x), "The features should be presented in binary form"
-
-        # TODO need to check whether features binary or not. Consistency required - Cora has 'unknown' features e.g.
-        # self.drop_edges(batch)
-        edge_index = gen_dataset.data.edge_index.tolist()
-        #new_edge_mask = torch.zeros_like(gen_dataset.data.edge_index).bool()
-        new_edge_index = [[],[]]
-        self.remove_edge_index = [[], []]
-        for i in range(len(edge_index[0])):
-            if self.jaccard_index(gen_dataset.data.x, edge_index[0][i], edge_index[1][i]) > self.thrsh:
-                # new_edge_mask[0,i] = True
-                # new_edge_mask[1,i] = True
-                new_edge_index[0].append(edge_index[0][i])
-                new_edge_index[1].append(edge_index[1][i])
-            else:
-                self.remove_edge_index[0].append(edge_index[0][i])
-                self.remove_edge_index[1].append(edge_index[1][i])
-        # gen_dataset.data.edge_index *= new_edge_mask.float()
-        gen_dataset.data.edge_index = torch.tensor(new_edge_index).long()
         return gen_dataset
 
-    def jaccard_index(
+    def _prepare_features(self, x: torch.Tensor) -> torch.Tensor:
+        if self.binarize_threshold is not None:
+            x = (x > self.binarize_threshold).float()
+        elif not _is_binary_tensor(x):
+            raise ValueError(
+                "JaccardDefender requires binary features"
+            )
+        return x
 
+    def _defense_link_task(
             self,
-            x,
-            u,
-            v
-    ) -> float:
-        """
-        Computes jaccard index of 'u' and 'v' objects based on their features
-        :param x: feature matrix
-        :param u: index of object from dataset
-        :param v: index of object from dataset
-        :return:
-        """
-        im1 = x[u,:].detach().cpu().numpy().astype(bool)
-        im2 = x[v,:].detach().cpu().numpy().astype(bool)
-        intersection = np.logical_and(im1, im2)
-        union = np.logical_or(im1, im2)
-        return intersection.sum() / float(union.sum())
-
-    def dataset_diff(
-            self
-    ) -> GraphModificationArtifact:
-        diff = GraphModificationArtifact()
+            gen_dataset: GeneralDataset,
+            x: torch.Tensor
+    ) -> GeneralDataset:
+        train_edge_label_index = gen_dataset.edge_label_index[:, gen_dataset.train_mask]
 
-        try:
-            src_nodes = self.remove_edge_index[0]
-            dst_nodes = self.remove_edge_index[1]
+        filtered_train_edges, removed_edges = self._filter_edges_jaccard(train_edge_label_index, x)
+        self.removed_edges_train = removed_edges
 
-            assert len(src_nodes) == len(dst_nodes), (
-                "Mismatch in source and target edge lengths: "
-                f"{len(src_nodes)} vs {len(dst_nodes)}"
-            )
+        gen_dataset.data.edge_index = filtered_train_edges
 
-            edges_to_remove = [
-                [src, dst] for src, dst in zip(src_nodes, dst_nodes)
-            ]
+        num_removed = removed_edges.size(1) if removed_edges is not None else 0
+        print(f"JaccardDefender: Removed {num_removed}/{train_edge_label_index.size(1)} "
+              f"training edges (threshold={self.threshold})")
 
+        return gen_dataset
+
+    def _defense_standard_task(
+            self,
+            gen_dataset: GeneralDataset,
+            x: torch.Tensor
+    ) -> GeneralDataset:
+        filtered_edges, removed_edges = self._filter_edges_jaccard(
+            gen_dataset.data.edge_index, x
+        )
+        self.removed_edges_train = removed_edges  # Reusing field for simplicity
+
+        gen_dataset.data.edge_index = filtered_edges
+
+        num_removed = removed_edges.size(1) if removed_edges is not None else 0
+        print(f"JaccardDefender: Removed {num_removed}/{self.original_num_edges} edges "
+              f"(threshold={self.threshold})")
+
+        return gen_dataset
+
+    def _filter_edges_jaccard(
+            self,
+            edge_index: torch.Tensor,
+            x: torch.Tensor
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if edge_index.size(1) == 0:
+            return edge_index, None
+
+        src_feats = x[edge_index[0]]
+        dst_feats = x[edge_index[1]]
+
+        intersection = (src_feats * dst_feats).sum(dim=1)  # AND
+        union = ((src_feats + dst_feats) > 0).sum(dim=1).float()  # OR
+
+        union = torch.where(union == 0, torch.ones_like(union), union)
+
+        jaccard_scores = intersection / union
+
+        keep_mask = jaccard_scores > self.threshold
+        filtered_edges = edge_index[:, keep_mask]
+        removed_edges = edge_index[:, ~keep_mask] if (~keep_mask).any() else None
+
+        return filtered_edges, removed_edges
+
+    def dataset_diff(self) -> GraphModificationArtifact:
+        diff = GraphModificationArtifact()
+
+        if self.removed_edges_train is not None and self.removed_edges_train.size(1) > 0:
+            edges_to_remove = self.removed_edges_train.t().tolist()
             diff.remove_edges(edges_to_remove)
             self.defense_diff = diff
-
-        except Exception as e:
-            raise RuntimeError(
-                f"Failed to build dataset diff from remove_edge_index: {e}"
-            ) from e
+        else:
+            # No edges removed
+            self.defense_diff = diff
 
         return self.defense_diff

gnn_aid/defenses/jaccard_defense/jaccard_def.py

@@ -147,3 +147,87 @@ def post_batch(
             "outputs": modified_logits,
             "loss": modified_loss
         }
+
+
+class NoiseMILinkDefender(MIDefender):
+    """
+    MI defense for Link Prediction tasks via edge logit perturbation
+    """
+    name = "NoiseMILinkDefender"
+
+    def __init__(
+            self,
+            noise_type: Literal["reverse_sigmoid", "random", "none"] = "reverse_sigmoid",
+            beta: float = 0.3,
+            gamma: float = 0.8,
+            noise_scale: float = 0.2,
+            temperature: float = 1.0,
+            **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.noise_type = noise_type
+        self.beta = beta
+        self.gamma = gamma
+        self.noise_scale = noise_scale
+        self.temperature = temperature
+
+        if noise_type not in ["reverse_sigmoid", "random", "none"]:
+            raise ValueError(f"Invalid noise_type: {noise_type}")
+
+    def _apply_reverse_sigmoid_binary(
+            self,
+            edge_logits: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        Reverse sigmoid perturbation for binary classification (link prediction)
+        """
+        probs = torch.sigmoid(edge_logits / self.temperature)
+
+        perturbed_temp_logits = self.gamma * edge_logits
+        perturbed_temp_probs = torch.sigmoid(perturbed_temp_logits)
+        r = self.beta * (perturbed_temp_probs - 0.5)
+
+        perturbed_probs = probs - r
+        perturbed_probs = torch.clamp(perturbed_probs, min=1e-7, max=1.0 - 1e-7)
+
+        perturbed_logits = torch.logit(perturbed_probs, eps=1e-7) * self.temperature
+        return perturbed_logits
+
+    def _apply_random_noise(
+            self,
+            edge_logits: torch.Tensor
+    ) -> torch.Tensor:
+        """Add Gaussian noise to edge logits"""
+        noise = torch.randn_like(edge_logits) * self.noise_scale
+        return edge_logits + noise
+
+    def post_batch(
+            self,
+            model_manager: Any,
+            batch: Any,
+            **kwargs
+    ) -> dict:
+        node_emb = model_manager.gnn(batch.x, batch.edge_index)
+        src_emb = node_emb[batch.edge_label_index[0]]
+        dst_emb = node_emb[batch.edge_label_index[1]]
+
+        if hasattr(model_manager.gnn, 'decode'):
+            edge_logits = model_manager.gnn.decode(src_emb, dst_emb).squeeze(-1)
+        else:
+            edge_logits = (src_emb * dst_emb).sum(dim=-1)
+
+        if self.noise_type == "reverse_sigmoid":
+            modified_logits = self._apply_reverse_sigmoid_binary(edge_logits)
+        elif self.noise_type == "random":
+            modified_logits = self._apply_random_noise(edge_logits)
+        else:
+            modified_logits = edge_logits
+
+        edge_labels = batch.edge_label.float()
+        modified_loss = model_manager.loss_function(modified_logits, edge_labels)
+
+        return {
+            "outputs": modified_logits,
+            "loss": modified_loss,
+            "original_logits": edge_logits.detach()
+        }