aibt-fl 1.0.0__py3-none-any.whl

aibt/__init__.py

@@ -0,0 +1,77 @@
+"""
+AIBT: Adversarial Information Bottleneck Training for Privacy-Preserving Federated Learning
+
+This package provides a complete implementation of the AIBT framework for
+privacy-preserving federated learning, combining:
+- Information Bottleneck (IB) for feature compression
+- Adversarial Training for sensitive attribute suppression
+- Federated Learning for distributed training
+
+Example:
+    >>> from aibt import AIBTFL, create_aibt_model, evaluate_privacy
+    >>> model = create_aibt_model(input_dim=13, num_classes=2)
+    >>> aibt = AIBTFL(model=model, num_clients=10)
+    >>> history = aibt.train(num_rounds=100, local_epochs=5)
+"""
+
+__version__ = "1.0.0"
+__author__ = "AIBT Research Team"
+
+# Core FL implementation
+from aibt.core import AIBTFL
+
+# Neural network models
+from aibt.models import (
+    GradientReversalLayer,
+    GradientReversalFunction,
+    VariationalEncoder,
+    MLPEncoder,
+    Predictor,
+    Adversary,
+    AIBTModel,
+    create_aibt_model,
+)
+
+# Client implementation
+from aibt.client import FLClient, AIBTClient
+
+# Privacy metrics
+from aibt.metrics import (
+    MembershipInferenceAttack,
+    AttributeInferenceAttack,
+    evaluate_membership_inference,
+    evaluate_attribute_inference,
+    evaluate_privacy,
+    evaluate_performance,
+)
+
+# Aggregation utilities
+from aibt.aggregation import fedavg_aggregate
+
+__all__ = [
+    # Version
+    "__version__",
+    # Core
+    "AIBTFL",
+    # Models
+    "GradientReversalLayer",
+    "GradientReversalFunction",
+    "VariationalEncoder",
+    "MLPEncoder",
+    "Predictor",
+    "Adversary",
+    "AIBTModel",
+    "create_aibt_model",
+    # Clients
+    "FLClient",
+    "AIBTClient",
+    # Metrics
+    "MembershipInferenceAttack",
+    "AttributeInferenceAttack",
+    "evaluate_membership_inference",
+    "evaluate_attribute_inference",
+    "evaluate_privacy",
+    "evaluate_performance",
+    # Aggregation
+    "fedavg_aggregate",
+]

aibt/aggregation.py

@@ -0,0 +1,68 @@
+"""
+FedAvg Aggregation utilities for AIBT.
+
+Implements weighted averaging of model parameters for federated learning.
+"""
+
+import torch
+from typing import Dict, List
+
+
+def fedavg_aggregate(
+    client_params: List[Dict[str, torch.Tensor]],
+    client_sizes: List[int]
+) -> Dict[str, torch.Tensor]:
+    """
+    FedAvg aggregation for federated learning.
+    
+    Computes weighted average of client model parameters:
+    w_global = Σ(n_k/n) * w_k
+    
+    Args:
+        client_params: List of model state dicts from clients
+        client_sizes: List of dataset sizes per client
+    
+    Returns:
+        Aggregated model state dict
+    """
+    if not client_params:
+        raise ValueError("No client parameters provided")
+    
+    total_size = sum(client_sizes)
+    aggregated = {}
+    
+    for key in client_params[0].keys():
+        weighted_sum = torch.zeros_like(client_params[0][key], dtype=torch.float32)
+        
+        for i, params in enumerate(client_params):
+            weight = client_sizes[i] / total_size
+            weighted_sum += weight * params[key].float()
+        
+        aggregated[key] = weighted_sum
+    
+    return aggregated
+
+
+def simple_average(
+    client_params: List[Dict[str, torch.Tensor]]
+) -> Dict[str, torch.Tensor]:
+    """
+    Simple averaging of client parameters (unweighted).
+    
+    Args:
+        client_params: List of model state dicts from clients
+    
+    Returns:
+        Aggregated model state dict
+    """
+    if not client_params:
+        raise ValueError("No client parameters provided")
+    
+    num_clients = len(client_params)
+    aggregated = {}
+    
+    for key in client_params[0].keys():
+        stacked = torch.stack([params[key].float() for params in client_params])
+        aggregated[key] = stacked.mean(dim=0)
+    
+    return aggregated

aibt/client.py

@@ -0,0 +1,259 @@
+"""
+Federated Learning Client implementations for AIBT.
+
+Implements:
+- FLClient: Base federated learning client
+- AIBTClient: AIBT-specific client with IB and adversarial training
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from typing import Dict, Tuple, Optional
+import copy
+import numpy as np
+
+
+class FLClient:
+    """
+    Federated Learning Client.
+    Each client has local data and performs local training.
+    """
+    
+    def __init__(
+        self,
+        client_id: int,
+        model: nn.Module,
+        train_data: Tuple[np.ndarray, np.ndarray],
+        sensitive_data: Optional[np.ndarray] = None,
+        batch_size: int = 32,
+        learning_rate: float = 0.001,
+        device: str = "cpu"
+    ):
+        self.client_id = client_id
+        self.device = device
+        self.batch_size = batch_size
+        self.learning_rate = learning_rate
+        
+        # Create local model (deep copy)
+        self.model = copy.deepcopy(model).to(device)
+        
+        # Prepare data
+        X, y = train_data
+        self.n_samples = len(y)
+        
+        # Check for sensitive attributes
+        self.has_sensitive = sensitive_data is not None
+        
+        if self.has_sensitive:
+            dataset = TensorDataset(
+                torch.FloatTensor(X),
+                torch.LongTensor(y),
+                torch.LongTensor(sensitive_data)
+            )
+        else:
+            dataset = TensorDataset(
+                torch.FloatTensor(X),
+                torch.LongTensor(y)
+            )
+        
+        self.dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            drop_last=len(dataset) > batch_size
+        )
+        
+        # Optimizer
+        self.optimizer = torch.optim.Adam(
+            self.model.parameters(),
+            lr=learning_rate
+        )
+    
+    def set_model_params(self, params: Dict[str, torch.Tensor]) -> None:
+        """Load global model parameters"""
+        self.model.load_state_dict(params)
+    
+    def get_model_params(self) -> Dict[str, torch.Tensor]:
+        """Get local model parameters"""
+        return copy.deepcopy(self.model.state_dict())
+    
+    def train_epoch(self) -> Dict[str, float]:
+        """Train for one epoch"""
+        self.model.train()
+        total_loss = 0.0
+        correct = 0
+        total = 0
+        
+        for batch in self.dataloader:
+            if self.has_sensitive:
+                X, y, s = batch
+                s = s.to(self.device)
+            else:
+                X, y = batch
+                s = None
+            
+            X, y = X.to(self.device), y.to(self.device)
+            
+            self.optimizer.zero_grad()
+            
+            # Forward pass
+            output = self.model(X)
+            if isinstance(output, tuple):
+                y_pred = output[0]
+            else:
+                y_pred = output
+            
+            loss = F.cross_entropy(y_pred, y)
+            loss.backward()
+            self.optimizer.step()
+            
+            total_loss += loss.item() * len(y)
+            pred = y_pred.argmax(dim=1)
+            correct += (pred == y).sum().item()
+            total += len(y)
+        
+        return {
+            "loss": total_loss / total,
+            "accuracy": correct / total
+        }
+    
+    def train(self, local_epochs: int) -> Dict[str, float]:
+        """Train for multiple epochs"""
+        metrics = {}
+        for epoch in range(local_epochs):
+            metrics = self.train_epoch()
+        return metrics
+    
+    def evaluate(self, test_data: Tuple[np.ndarray, np.ndarray]) -> Dict[str, float]:
+        """Evaluate model on test data"""
+        self.model.eval()
+        X, y = test_data
+        X_tensor = torch.FloatTensor(X).to(self.device)
+        y_tensor = torch.LongTensor(y).to(self.device)
+        
+        with torch.no_grad():
+            output = self.model(X_tensor)
+            if isinstance(output, tuple):
+                logits = output[0]
+            else:
+                logits = output
+            
+            loss = F.cross_entropy(logits, y_tensor).item()
+            pred = logits.argmax(dim=1)
+            accuracy = (pred == y_tensor).float().mean().item()
+        
+        return {"test_loss": loss, "test_accuracy": accuracy}
+
+
+class AIBTClient(FLClient):
+    """
+    AIBT Client with Information Bottleneck and Adversarial Training.
+    
+    Loss = L_task + λ₁ L_KL - λ₂ L_adv
+    
+    Training procedure:
+    1. Forward pass through encoder → z (with reparameterization)
+    2. Task prediction via predictor
+    3. Adversarial prediction via GRL + adversary
+    4. Compute combined loss
+    5. Update encoder and predictor
+    """
+    
+    def __init__(
+        self,
+        *args,
+        lambda_kl: float = 0.01,
+        lambda_adv: float = 1.0,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.lambda_kl = lambda_kl
+        self.lambda_adv = lambda_adv
+    
+    def train_epoch(self) -> Dict[str, float]:
+        """
+        Train for one epoch using AIBT algorithm.
+        
+        Algorithm 1/2 from the paper:
+        For each batch:
+            1. z = Encoder(x) with reparameterization
+            2. y_hat = Predictor(z)
+            3. s_hat = Adversary(GRL(z))
+            4. L = L_task + λ₁ L_KL - λ₂ L_adv
+            5. Update encoder and predictor
+        """
+        self.model.train()
+        total_loss = 0.0
+        total_task_loss = 0.0
+        total_kl_loss = 0.0
+        total_adv_loss = 0.0
+        correct = 0
+        total = 0
+        
+        for batch in self.dataloader:
+            if self.has_sensitive:
+                X, y, s = batch
+                s = s.to(self.device)
+            else:
+                X, y = batch
+                s = None
+            
+            X, y = X.to(self.device), y.to(self.device)
+            
+            self.optimizer.zero_grad()
+            
+            # Forward pass through AIBT model
+            if hasattr(self.model, 'compute_loss'):
+                # Using AIBTModel with built-in loss computation
+                loss, loss_dict = self.model.compute_loss(X, y, s)
+                
+                # Get predictions for accuracy
+                with torch.no_grad():
+                    y_pred, _, _, _ = self.model(X)
+                    pred = y_pred.argmax(dim=1)
+                
+                total_task_loss += loss_dict["task_loss"] * len(y)
+                total_kl_loss += loss_dict["kl_loss"] * len(y)
+                total_adv_loss += loss_dict["adv_loss"] * len(y)
+            else:
+                # Manual computation for non-AIBT models
+                output = self.model(X)
+                if isinstance(output, tuple):
+                    y_pred = output[0]
+                else:
+                    y_pred = output
+                
+                loss = F.cross_entropy(y_pred, y)
+                pred = y_pred.argmax(dim=1)
+            
+            loss.backward()
+            self.optimizer.step()
+            
+            total_loss += loss.item() * len(y)
+            correct += (pred == y).sum().item()
+            total += len(y)
+        
+        return {
+            "loss": total_loss / total,
+            "task_loss": total_task_loss / total if total_task_loss > 0 else 0,
+            "kl_loss": total_kl_loss / total if total_kl_loss > 0 else 0,
+            "adv_loss": total_adv_loss / total if total_adv_loss > 0 else 0,
+            "accuracy": correct / total
+        }
+    
+    def get_latent_representations(self, X: np.ndarray) -> np.ndarray:
+        """Get latent representations for privacy analysis"""
+        self.model.eval()
+        X_tensor = torch.FloatTensor(X).to(self.device)
+        
+        with torch.no_grad():
+            if hasattr(self.model, 'get_latent'):
+                z = self.model.get_latent(X_tensor)
+            else:
+                z = self.model.encoder(X_tensor)
+                if isinstance(z, tuple):
+                    z = z[0]
+        
+        return z.cpu().numpy()

aibt/core.py

@@ -0,0 +1,287 @@
+"""
+AIBT FL: Adversarial Information Bottleneck Training for Federated Learning
+
+Main federated learning implementation combining:
+- Information Bottleneck (IB) for feature compression
+- Adversarial Training for sensitive attribute suppression
+- Federated Averaging for model aggregation
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Tuple, Optional, Any
+import copy
+import numpy as np
+
+from aibt.client import AIBTClient, FLClient
+from aibt.aggregation import fedavg_aggregate
+
+
+class AIBTFL:
+    """
+    AIBT FL: Adversarial Information Bottleneck Training for Federated Learning.
+    
+    Combines:
+    - Information Bottleneck (IB) for feature compression
+    - Adversarial Training for sensitive attribute suppression
+    - Federated Learning for distributed training
+    
+    Loss: L = L_task + λ₁ L_KL - λ₂ L_adv
+    
+    Args:
+        model: Neural network model (preferably AIBTModel)
+        num_clients: Number of federated learning clients
+        device: Device for computation ('cpu' or 'cuda')
+        learning_rate: Learning rate for local training
+        batch_size: Batch size for local training
+        lambda_kl: Weight for KL divergence loss (Information Bottleneck)
+        lambda_adv: Weight for adversarial loss
+        lambda_grl: Strength of gradient reversal
+        latent_dim: Dimension of latent representation
+        num_sensitive_classes: Number of sensitive attribute classes
+    
+    Example:
+        >>> from aibt import AIBTFL, create_aibt_model
+        >>> model = create_aibt_model(input_dim=13, num_classes=2)
+        >>> aibt = AIBTFL(model=model, num_clients=10)
+        >>> aibt.setup_clients(client_datasets)
+        >>> history = aibt.train(num_rounds=100, local_epochs=5)
+    """
+    
+    def __init__(
+        self,
+        model: nn.Module,
+        num_clients: int,
+        device: str = "cpu",
+        learning_rate: float = 0.001,
+        batch_size: int = 32,
+        lambda_kl: float = 0.01,
+        lambda_adv: float = 1.0,
+        lambda_grl: float = 1.0,
+        latent_dim: int = 128,
+        num_sensitive_classes: int = 2
+    ):
+        self.device = device
+        self.num_clients = num_clients
+        self.learning_rate = learning_rate
+        self.batch_size = batch_size
+        
+        self.lambda_kl = lambda_kl
+        self.lambda_adv = lambda_adv
+        self.lambda_grl = lambda_grl
+        self.latent_dim = latent_dim
+        self.num_sensitive_classes = num_sensitive_classes
+        self.method_name = "AIBT"
+        
+        # Global model
+        self.global_model = copy.deepcopy(model).to(device)
+        
+        # Update model parameters if it's an AIBT model
+        if hasattr(self.global_model, 'lambda_kl'):
+            self.global_model.lambda_kl = lambda_kl
+            self.global_model.lambda_adv = lambda_adv
+        if hasattr(self.global_model, 'grl'):
+            self.global_model.grl.set_lambda(lambda_grl)
+        
+        # Clients (initialized later)
+        self.clients: List[AIBTClient] = []
+        
+        # Metrics tracking
+        self.current_round = 0
+        self.round_metrics: List[Dict[str, float]] = []
+    
+    def setup_clients(
+        self,
+        client_datasets: List[Tuple[np.ndarray, np.ndarray]],
+        sensitive_data: Optional[List[np.ndarray]] = None
+    ) -> None:
+        """
+        Initialize clients with their local data.
+        
+        Args:
+            client_datasets: List of (X, y) tuples, one per client
+            sensitive_data: Optional list of sensitive attribute arrays
+        """
+        self.clients = []
+        
+        for i, data in enumerate(client_datasets):
+            s_data = sensitive_data[i] if sensitive_data else None
+            client = self._create_client(i, data, s_data)
+            self.clients.append(client)
+    
+    def _create_client(
+        self,
+        client_id: int,
+        train_data: Tuple[np.ndarray, np.ndarray],
+        sensitive_data: Optional[np.ndarray] = None
+    ) -> AIBTClient:
+        """Create an AIBT client"""
+        return AIBTClient(
+            client_id=client_id,
+            model=self.global_model,
+            train_data=train_data,
+            sensitive_data=sensitive_data,
+            batch_size=self.batch_size,
+            learning_rate=self.learning_rate,
+            device=self.device,
+            lambda_kl=self.lambda_kl,
+            lambda_adv=self.lambda_adv
+        )
+    
+    def distribute_model(self) -> None:
+        """Send global model to all clients"""
+        global_params = self.global_model.state_dict()
+        for client in self.clients:
+            client.set_model_params(copy.deepcopy(global_params))
+    
+    def train_round(self, local_epochs: int) -> Dict[str, float]:
+        """
+        Execute one round of AIBT federated training.
+        
+        Following Algorithm 3 from the paper:
+        1. Distribute global model to clients
+        2. Each client trains using Algorithm 1/2
+        3. Aggregate updates using FedAvg
+        4. Update global model
+        """
+        self.current_round += 1
+        
+        # Distribute global model
+        self.distribute_model()
+        
+        # Local training
+        client_params = []
+        client_sizes = []
+        client_metrics = []
+        
+        for client in self.clients:
+            # Local AIBT training
+            metrics = client.train(local_epochs)
+            client_metrics.append(metrics)
+            
+            # Collect parameters
+            client_params.append(client.get_model_params())
+            client_sizes.append(client.n_samples)
+        
+        # Aggregate using FedAvg
+        aggregated_params = fedavg_aggregate(client_params, client_sizes)
+        self.global_model.load_state_dict(aggregated_params)
+        
+        # Compute round metrics
+        avg_loss = np.mean([m["loss"] for m in client_metrics])
+        avg_acc = np.mean([m["accuracy"] for m in client_metrics])
+        avg_task_loss = np.mean([m.get("task_loss", 0) for m in client_metrics])
+        avg_kl_loss = np.mean([m.get("kl_loss", 0) for m in client_metrics])
+        avg_adv_loss = np.mean([m.get("adv_loss", 0) for m in client_metrics])
+        
+        round_metrics = {
+            "round": self.current_round,
+            "train_loss": avg_loss,
+            "train_accuracy": avg_acc,
+            "task_loss": avg_task_loss,
+            "kl_loss": avg_kl_loss,
+            "adv_loss": avg_adv_loss,
+        }
+        
+        self.round_metrics.append(round_metrics)
+        return round_metrics
+    
+    def evaluate(self, test_data: Tuple[np.ndarray, np.ndarray]) -> Dict[str, float]:
+        """Evaluate global model on test data"""
+        self.global_model.eval()
+        
+        X, y = test_data
+        X_tensor = torch.FloatTensor(X).to(self.device)
+        y_tensor = torch.LongTensor(y).to(self.device)
+        
+        with torch.no_grad():
+            output = self.global_model(X_tensor)
+            if isinstance(output, tuple):
+                logits = output[0]
+            else:
+                logits = output
+            
+            loss = F.cross_entropy(logits, y_tensor).item()
+            pred = logits.argmax(dim=1)
+            accuracy = (pred == y_tensor).float().mean().item()
+        
+        return {"test_loss": loss, "test_accuracy": accuracy}
+    
+    def get_latent_representations(self, X: np.ndarray) -> np.ndarray:
+        """Get latent representations from global model"""
+        self.global_model.eval()
+        X_tensor = torch.FloatTensor(X).to(self.device)
+        
+        with torch.no_grad():
+            if hasattr(self.global_model, 'get_latent'):
+                z = self.global_model.get_latent(X_tensor)
+            elif hasattr(self.global_model, 'encoder'):
+                z = self.global_model.encoder(X_tensor)
+                if isinstance(z, tuple):
+                    z = z[0]
+            else:
+                raise AttributeError("Model does not have encoder")
+        
+        return z.cpu().numpy()
+    
+    def train(
+        self,
+        num_rounds: int,
+        local_epochs: int,
+        test_data: Optional[Tuple[np.ndarray, np.ndarray]] = None,
+        verbose: bool = True
+    ) -> Dict[str, Any]:
+        """
+        Complete AIBT federated training loop.
+        
+        Args:
+            num_rounds: Number of federated learning rounds
+            local_epochs: Number of local training epochs per round
+            test_data: Optional (X_test, y_test) for evaluation
+            verbose: Print progress
+        
+        Returns:
+            Dictionary with training history and final metrics
+        """
+        history = {
+            "round_metrics": [],
+            "test_metrics": []
+        }
+        
+        for round_idx in range(num_rounds):
+            # Training round
+            round_metrics = self.train_round(local_epochs)
+            
+            # Evaluation
+            if test_data is not None:
+                test_metrics = self.evaluate(test_data)
+                round_metrics.update(test_metrics)
+            
+            history["round_metrics"].append(round_metrics)
+            
+            if verbose:
+                print(f"Round {round_idx + 1}/{num_rounds} | "
+                      f"Loss: {round_metrics['train_loss']:.4f} | "
+                      f"Acc: {round_metrics['train_accuracy']:.4f}", end="")
+                if test_data is not None:
+                    print(f" | Test Acc: {round_metrics.get('test_accuracy', 0):.4f}")
+                else:
+                    print()
+        
+        # Final evaluation
+        if test_data is not None:
+            history["final_test_metrics"] = self.evaluate(test_data)
+        
+        return history
+    
+    def get_model(self) -> nn.Module:
+        """Get the global model"""
+        return self.global_model
+    
+    def get_communication_cost(self) -> int:
+        """Calculate communication cost in bytes"""
+        total_bytes = 0
+        for param in self.global_model.parameters():
+            total_bytes += param.numel() * param.element_size()
+        return total_bytes