add dp with options

Author: cyfan11

fedgraph/differential_privacy/__init__.py

@@ -0,0 +1,46 @@
+from .dp_mechanisms import DPMechanism, DPAccountant
+from .server_dp import Server_DP
+from .trainer_dp import Trainer_General_DP
+
+__version__ = "1.0.0"
+__author__ = "FedGraph Team"
+
+__all__ = [
+    "DPMechanism",
+    "DPAccountant", 
+    "Server_DP",
+    "Trainer_General_DP",
+]
+
+# Module-level configuration
+DEFAULT_DP_CONFIG = {
+    "epsilon": 1.0,
+    "delta": 1e-5,
+    "mechanism": "gaussian",
+    "sensitivity": 1.0,
+    "clip_norm": 1.0,
+}
+
+def get_default_config():
+    """Get default DP configuration."""
+    return DEFAULT_DP_CONFIG.copy()
+
+def validate_dp_config(config):
+    """Validate DP configuration parameters."""
+    required_keys = ["epsilon", "delta", "mechanism"]
+    for key in required_keys:
+        if key not in config:
+            raise ValueError(f"Missing required DP parameter: {key}")
+    
+    if config["epsilon"] <= 0:
+        raise ValueError("epsilon must be positive")
+    if config["delta"] <= 0 or config["delta"] >= 1:
+        raise ValueError("delta must be in (0, 1)")
+    
+    valid_mechanisms = ["gaussian", "laplace", "local"]
+    if config["mechanism"] not in valid_mechanisms:
+        raise ValueError(f"mechanism must be one of {valid_mechanisms}")
+    
+    return True
+
+print(f"FedGraph Differential Privacy module loaded (v{__version__})")

fedgraph/differential_privacy/dp_mechanisms.py

@@ -0,0 +1,144 @@
+import torch
+import numpy as np
+import random
+import time
+from typing import Dict, List, Tuple, Optional, Any
+
+class DPMechanism:
+    """
+    Differential Privacy mechanisms for federated learning.
+    
+    Supports multiple DP mechanisms:
+    - Gaussian mechanism
+    - Laplace mechanism  
+    - Local DP with randomized response
+    """
+    
+    def __init__(self, epsilon: float = 1.0, delta: float = 1e-5, 
+                 sensitivity: float = 1.0, mechanism: str = "gaussian"):
+        """
+        Initialize DP mechanism.
+        
+        Parameters
+        ----------
+        epsilon : float
+            Privacy budget (smaller = more private)
+        delta : float  
+            Failure probability for (ε,δ)-DP
+        sensitivity : float
+            L2 sensitivity of the function
+        mechanism : str
+            DP mechanism ("gaussian", "laplace", "local")
+        """
+        self.epsilon = epsilon
+        self.delta = delta
+        self.sensitivity = sensitivity
+        self.mechanism = mechanism
+        
+        # Calculate noise parameters
+        if mechanism == "gaussian":
+            # For (ε,δ)-DP: σ ≥ sqrt(2ln(1.25/δ)) * Δ / ε
+            self.sigma = np.sqrt(2 * np.log(1.25 / delta)) * sensitivity / epsilon
+        elif mechanism == "laplace":
+            # For ε-DP: b = Δ / ε
+            self.scale = sensitivity / epsilon
+        elif mechanism == "local":
+            # For local DP
+            self.p = np.exp(epsilon) / (np.exp(epsilon) + 1)
+        
+        print(f"Initialized {mechanism} DP mechanism:")
+        print(f"  ε={epsilon}, δ={delta}, sensitivity={sensitivity}")
+        if mechanism == "gaussian":
+            print(f"  Gaussian noise σ={self.sigma:.4f}")
+        elif mechanism == "laplace":
+            print(f"  Laplace scale={self.scale:.4f}")
+
+    def add_noise(self, tensor: torch.Tensor) -> torch.Tensor:
+        """
+        Add differential privacy noise to tensor.
+        
+        Parameters
+        ----------
+        tensor : torch.Tensor
+            Input tensor to add noise to
+            
+        Returns
+        -------
+        torch.Tensor
+            Tensor with DP noise added
+        """
+        if self.mechanism == "gaussian":
+            noise = torch.normal(0, self.sigma, size=tensor.shape, device=tensor.device)
+            return tensor + noise
+            
+        elif self.mechanism == "laplace":
+            # Laplace noise using exponential distribution
+            uniform = torch.rand(tensor.shape, device=tensor.device)
+            sign = torch.sign(uniform - 0.5)
+            noise = -sign * self.scale * torch.log(1 - 2 * torch.abs(uniform - 0.5))
+            return tensor + noise
+            
+        elif self.mechanism == "local":
+            # Local DP with randomized response
+            prob_matrix = torch.rand(tensor.shape, device=tensor.device)
+            mask = prob_matrix < self.p
+            # Flip with probability (1-p)
+            noisy_tensor = tensor.clone()
+            noisy_tensor[~mask] = -noisy_tensor[~mask]  # Simple bit flip for demonstration
+            return noisy_tensor
+            
+        else:
+            raise ValueError(f"Unknown mechanism: {self.mechanism}")
+
+    def clip_gradients(self, tensor: torch.Tensor, max_norm: float) -> torch.Tensor:
+        """
+        Clip tensor to bound sensitivity.
+        
+        Parameters
+        ----------
+        tensor : torch.Tensor
+            Input tensor to clip
+        max_norm : float
+            Maximum L2 norm
+            
+        Returns
+        -------
+        torch.Tensor
+            Clipped tensor
+        """
+        current_norm = torch.norm(tensor)
+        if current_norm > max_norm:
+            return tensor * (max_norm / current_norm)
+        return tensor
+
+    def get_privacy_spent(self) -> Tuple[float, float]:
+        """Get privacy budget spent."""
+        return self.epsilon, self.delta
+
+
+class DPAccountant:
+    """
+    Privacy accountant for tracking cumulative privacy loss.
+    """
+    
+    def __init__(self):
+        self.total_epsilon = 0.0
+        self.total_delta = 0.0
+        self.rounds = 0
+        
+    def add_step(self, epsilon: float, delta: float):
+        """Add privacy cost of one step."""
+        # Simple composition (can be improved with advanced composition)
+        self.total_epsilon += epsilon
+        self.total_delta += delta
+        self.rounds += 1
+        
+    def get_total_privacy_spent(self) -> Tuple[float, float]:
+        """Get total privacy spent."""
+        return self.total_epsilon, self.total_delta
+    
+    def print_privacy_budget(self):
+        """Print current privacy budget."""
+        print(f"Privacy Budget Used: ε={self.total_epsilon:.4f}, δ={self.total_delta:.8f}")
+        print(f"Rounds completed: {self.rounds}")
+

fedgraph/differential_privacy/server_dp.py

@@ -0,0 +1,120 @@
+import torch
+import time
+from typing import Dict, List, Tuple, Optional, Any
+
+from ..server_class import Server
+from .dp_mechanisms import DPMechanism, DPAccountant
+
+
+class Server_DP(Server):
+    """
+    Enhanced server class with Differential Privacy support for FedGCN.
+    Extends the original Server class to support DP in pre-training aggregation.
+    """
+    
+    def __init__(self, feature_dim: int, args_hidden: int, class_num: int, 
+                 device: torch.device, trainers: list, args: Any):
+        super().__init__(feature_dim, args_hidden, class_num, device, trainers, args)
+        
+        # DP configuration
+        self.use_dp = getattr(args, 'use_dp', False)
+        
+        if self.use_dp:
+            self.dp_epsilon = getattr(args, 'dp_epsilon', 1.0)
+            self.dp_delta = getattr(args, 'dp_delta', 1e-5)
+            self.dp_sensitivity = getattr(args, 'dp_sensitivity', 1.0)
+            self.dp_mechanism = getattr(args, 'dp_mechanism', 'gaussian')
+            self.dp_clip_norm = getattr(args, 'dp_clip_norm', 1.0)
+            
+            # Initialize DP mechanism
+            self.dp_mechanism_obj = DPMechanism(
+                epsilon=self.dp_epsilon,
+                delta=self.dp_delta,
+                sensitivity=self.dp_sensitivity,
+                mechanism=self.dp_mechanism
+            )
+            
+            # Privacy accountant
+            self.privacy_accountant = DPAccountant()
+            
+            print(f"Server initialized with Differential Privacy:")
+            print(f"  Mechanism: {self.dp_mechanism}")
+            print(f"  Privacy parameters: ε={self.dp_epsilon}, δ={self.dp_delta}")
+            print(f"  Sensitivity: {self.dp_sensitivity}")
+            print(f"  Clipping norm: {self.dp_clip_norm}")
+
+    def aggregate_dp_feature_sums(self, local_feature_sums: List[torch.Tensor]) -> Tuple[torch.Tensor, Dict]:
+        """
+        Aggregate feature sums with differential privacy.
+        
+        Parameters
+        ----------
+        local_feature_sums : List[torch.Tensor]
+            List of local feature sums from trainers
+            
+        Returns
+        -------
+        Tuple[torch.Tensor, Dict]
+            Aggregated feature sum with DP noise and statistics
+        """
+        aggregation_start = time.time()
+        
+        # Step 1: Clip individual contributions
+        clipped_sums = []
+        clipping_stats = []
+        
+        for i, local_sum in enumerate(local_feature_sums):
+            original_norm = torch.norm(local_sum).item()
+            clipped_sum = self.dp_mechanism_obj.clip_gradients(local_sum, self.dp_clip_norm)
+            clipped_norm = torch.norm(clipped_sum).item()
+            
+            clipped_sums.append(clipped_sum)
+            clipping_stats.append({
+                'trainer_id': i,
+                'original_norm': original_norm,
+                'clipped_norm': clipped_norm,
+                'was_clipped': original_norm > self.dp_clip_norm
+            })
+        
+        # Step 2: Aggregate clipped sums
+        aggregated_sum = torch.stack(clipped_sums).sum(dim=0)
+        
+        # Step 3: Add DP noise
+        noisy_aggregated_sum = self.dp_mechanism_obj.add_noise(aggregated_sum)
+        
+        aggregation_time = time.time() - aggregation_start
+        
+        # Step 4: Update privacy accountant
+        self.privacy_accountant.add_step(self.dp_epsilon, self.dp_delta)
+        
+        # Statistics
+        dp_stats = {
+            'aggregation_time': aggregation_time,
+            'clipping_stats': clipping_stats,
+            'num_clipped': sum(1 for stat in clipping_stats if stat['was_clipped']),
+            'pre_noise_norm': torch.norm(aggregated_sum).item(),
+            'post_noise_norm': torch.norm(noisy_aggregated_sum).item(),
+            'noise_magnitude': torch.norm(noisy_aggregated_sum - aggregated_sum).item(),
+            'privacy_spent': self.privacy_accountant.get_total_privacy_spent()
+        }
+        
+        return noisy_aggregated_sum, dp_stats
+
+    def print_dp_stats(self, dp_stats: Dict):
+        """Print differential privacy statistics."""
+        print("\n=== Differential Privacy Statistics ===")
+        print(f"Aggregation time: {dp_stats['aggregation_time']:.4f}s")
+        print(f"Trainers clipped: {dp_stats['num_clipped']}/{len(dp_stats['clipping_stats'])}")
+        print(f"Pre-noise norm: {dp_stats['pre_noise_norm']:.4f}")
+        print(f"Post-noise norm: {dp_stats['post_noise_norm']:.4f}")
+        print(f"Noise magnitude: {dp_stats['noise_magnitude']:.4f}")
+        
+        total_eps, total_delta = dp_stats['privacy_spent']
+        print(f"Total privacy spent: ε={total_eps:.4f}, δ={total_delta:.8f}")
+        
+        # Per-trainer clipping details
+        clipped_trainers = [stat for stat in dp_stats['clipping_stats'] if stat['was_clipped']]
+        if clipped_trainers:
+            print("Clipped trainers:")
+            for stat in clipped_trainers:
+                print(f"  Trainer {stat['trainer_id']}: {stat['original_norm']:.4f} -> {stat['clipped_norm']:.4f}")

fedgraph/differential_privacy/trainer_dp.py

@@ -0,0 +1,61 @@
+import torch
+import time
+from typing import Dict, List, Tuple, Optional, Any
+
+from ..trainer_class import Trainer_General
+from ..utils_nc import get_1hop_feature_sum
+
+class Trainer_General_DP(Trainer_General):
+    """
+    Enhanced trainer class with Differential Privacy support.
+    """
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.use_dp = getattr(self.args, 'use_dp', False)
+        
+        if self.use_dp:
+            print(f"Trainer {self.rank} initialized with DP support")
+
+    def get_dp_local_feature_sum(self) -> Tuple[torch.Tensor, Dict]:
+        """
+        Get local feature sum with optional client-side DP preprocessing.
+        
+        Returns
+        -------
+        Tuple[torch.Tensor, Dict]
+            Local feature sum and computation statistics
+        """
+        computation_start = time.time()
+        
+        # Compute feature sum (same as original)
+        new_feature_for_trainer = torch.zeros(
+            self.global_node_num, self.features.shape[1]
+        ).to(self.device)
+        new_feature_for_trainer[self.local_node_index] = self.features
+        
+        one_hop_neighbor_feature_sum = get_1hop_feature_sum(
+            new_feature_for_trainer, self.adj, self.device
+        )
+        
+        computation_time = time.time() - computation_start
+        
+        # Compute statistics for DP
+        feature_sum_norm = torch.norm(one_hop_neighbor_feature_sum).item()
+        data_size = one_hop_neighbor_feature_sum.element_size() * one_hop_neighbor_feature_sum.nelement()
+        
+        stats = {
+            'trainer_id': self.rank,
+            'computation_time': computation_time,
+            'feature_sum_norm': feature_sum_norm,
+            'data_size': data_size,
+            'shape': one_hop_neighbor_feature_sum.shape
+        }
+        
+        print(f"Trainer {self.rank} - DP feature sum computed:")
+        print(f"  Norm: {feature_sum_norm:.4f}")
+        print(f"  Shape: {one_hop_neighbor_feature_sum.shape}")
+        print(f"  Computation time: {computation_time:.4f}s")
+        
+        return one_hop_neighbor_feature_sum, stats
+