aibt-fl 1.0.0__py3-none-any.whl

aibt/models.py

@@ -0,0 +1,520 @@
+"""
+Neural Network Components for AIBT Framework
+
+Implements:
+- Gradient Reversal Layer (GRL) for adversarial training
+- Variational Encoder for Information Bottleneck
+- MLP Encoder for tabular data
+- Predictor (Task Head)
+- Adversary (Sensitive Attribute Classifier)
+- Complete AIBTModel
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, List, Optional
+import math
+
+
+# ============================================================================
+# GRADIENT REVERSAL LAYER
+# ============================================================================
+
+class GradientReversalFunction(torch.autograd.Function):
+    """
+    Gradient Reversal Layer for adversarial training.
+    Forward pass: identity function
+    Backward pass: negate gradients and scale by lambda
+    
+    Reference: Ganin & Lempitsky, "Domain-Adversarial Training of Neural Networks"
+    """
+    
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, lambda_grl: float) -> torch.Tensor:
+        ctx.lambda_grl = lambda_grl
+        return x.view_as(x)
+    
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]:
+        return -ctx.lambda_grl * grad_output, None
+
+
+class GradientReversalLayer(nn.Module):
+    """Wrapper module for Gradient Reversal"""
+    
+    def __init__(self, lambda_grl: float = 1.0):
+        super().__init__()
+        self.lambda_grl = lambda_grl
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return GradientReversalFunction.apply(x, self.lambda_grl)
+    
+    def set_lambda(self, lambda_grl: float) -> None:
+        """Update lambda value (useful for scheduling)"""
+        self.lambda_grl = lambda_grl
+
+
+def grad_reverse(x: torch.Tensor, lambda_grl: float = 1.0) -> torch.Tensor:
+    """Functional version of gradient reversal"""
+    return GradientReversalFunction.apply(x, lambda_grl)
+
+
+# ============================================================================
+# VARIATIONAL ENCODER (Information Bottleneck)
+# ============================================================================
+
+class VariationalEncoder(nn.Module):
+    """
+    Variational Encoder for Information Bottleneck.
+    Outputs mean (mu) and log-variance (logvar) for reparameterization trick.
+    
+    z = mu + sigma * epsilon, where epsilon ~ N(0, I)
+    """
+    
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dims: List[int],
+        latent_dim: int,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        
+        self.input_dim = input_dim
+        self.latent_dim = latent_dim
+        
+        # Build encoder layers
+        layers = []
+        prev_dim = input_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.BatchNorm1d(hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Dropout(dropout)
+            ])
+            prev_dim = hidden_dim
+        
+        self.encoder = nn.Sequential(*layers)
+        
+        # Mean and log-variance heads
+        self.fc_mu = nn.Linear(prev_dim, latent_dim)
+        self.fc_logvar = nn.Linear(prev_dim, latent_dim)
+        
+        self._init_weights()
+    
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+    
+    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encode input to mu and logvar"""
+        h = self.encoder(x)
+        mu = self.fc_mu(h)
+        logvar = self.fc_logvar(h)
+        return mu, logvar
+    
+    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """
+        Reparameterization trick: z = mu + sigma * epsilon
+        """
+        if self.training:
+            std = torch.exp(0.5 * logvar)
+            eps = torch.randn_like(std)
+            return mu + eps * std
+        return mu
+    
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward pass.
+        Returns: z (latent), mu, logvar
+        """
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        return z, mu, logvar
+
+
+# ============================================================================
+# MLP ENCODER (For Tabular Data)
+# ============================================================================
+
+class MLPEncoder(nn.Module):
+    """
+    MLP Encoder for tabular datasets.
+    Can be variational or deterministic.
+    """
+    
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dims: List[int],
+        latent_dim: int,
+        variational: bool = True,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        
+        self.input_dim = input_dim
+        self.latent_dim = latent_dim
+        self.variational = variational
+        
+        # Build encoder layers
+        layers = []
+        prev_dim = input_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.BatchNorm1d(hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Dropout(dropout)
+            ])
+            prev_dim = hidden_dim
+        
+        self.encoder = nn.Sequential(*layers)
+        
+        if variational:
+            self.fc_mu = nn.Linear(prev_dim, latent_dim)
+            self.fc_logvar = nn.Linear(prev_dim, latent_dim)
+        else:
+            self.fc_out = nn.Linear(prev_dim, latent_dim)
+        
+        self._init_weights()
+    
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+    
+    def forward(self, x: torch.Tensor):
+        h = self.encoder(x)
+        
+        if self.variational:
+            mu = self.fc_mu(h)
+            logvar = self.fc_logvar(h)
+            
+            if self.training:
+                std = torch.exp(0.5 * logvar)
+                eps = torch.randn_like(std)
+                z = mu + eps * std
+            else:
+                z = mu
+            
+            return z, mu, logvar
+        else:
+            return self.fc_out(h)
+
+
+# ============================================================================
+# PREDICTOR (Task Head)
+# ============================================================================
+
+class Predictor(nn.Module):
+    """
+    Predictor network for task-specific predictions.
+    Takes latent representation z and outputs class probabilities.
+    """
+    
+    def __init__(
+        self,
+        latent_dim: int,
+        hidden_dims: List[int],
+        num_classes: int,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        
+        layers = []
+        prev_dim = latent_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Dropout(dropout)
+            ])
+            prev_dim = hidden_dim
+        
+        layers.append(nn.Linear(prev_dim, num_classes))
+        
+        self.predictor = nn.Sequential(*layers)
+        self._init_weights()
+    
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+    
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        return self.predictor(z)
+
+
+# ============================================================================
+# ADVERSARY (Sensitive Attribute Classifier)
+# ============================================================================
+
+class Adversary(nn.Module):
+    """
+    Adversary network for predicting sensitive attributes.
+    Used with GRL to enforce privacy in latent representation.
+    """
+    
+    def __init__(
+        self,
+        latent_dim: int,
+        hidden_dims: List[int],
+        num_sensitive_classes: int,
+        dropout: float = 0.1
+    ):
+        super().__init__()
+        
+        layers = []
+        prev_dim = latent_dim
+        
+        for hidden_dim in hidden_dims:
+            layers.extend([
+                nn.Linear(prev_dim, hidden_dim),
+                nn.ReLU(inplace=True),
+                nn.Dropout(dropout)
+            ])
+            prev_dim = hidden_dim
+        
+        layers.append(nn.Linear(prev_dim, num_sensitive_classes))
+        
+        self.adversary = nn.Sequential(*layers)
+        self._init_weights()
+    
+    def _init_weights(self) -> None:
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+    
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        return self.adversary(z)
+
+
+# ============================================================================
+# COMPLETE AIBT MODEL
+# ============================================================================
+
+class AIBTModel(nn.Module):
+    """
+    Complete AIBT (Adversarial Information Bottleneck Training) Model.
+    
+    Combines:
+    - Encoder (Variational for IB)
+    - Predictor (Task head)
+    - Adversary (Sensitive attribute classifier)
+    - Gradient Reversal Layer (GRL)
+    
+    Loss: L = L_task + λ₁ L_KL - λ₂ L_adv
+    """
+    
+    def __init__(
+        self,
+        encoder: nn.Module,
+        predictor: nn.Module,
+        adversary: nn.Module,
+        lambda_kl: float = 0.01,
+        lambda_adv: float = 1.0,
+        lambda_grl: float = 1.0
+    ):
+        super().__init__()
+        
+        self.encoder = encoder
+        self.predictor = predictor
+        self.adversary = adversary
+        self.grl = GradientReversalLayer(lambda_grl)
+        
+        self.lambda_kl = lambda_kl
+        self.lambda_adv = lambda_adv
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+        return_latent: bool = False
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward pass.
+        
+        Returns:
+            y_pred: Task predictions
+            s_pred: Sensitive attribute predictions (through GRL)
+            mu: Mean (for KL loss)
+            logvar: Log variance (for KL loss)
+        """
+        # Encode
+        encoder_out = self.encoder(x)
+        
+        if isinstance(encoder_out, tuple) and len(encoder_out) == 3:
+            z, mu, logvar = encoder_out
+        else:
+            z = encoder_out
+            mu = z
+            logvar = torch.zeros_like(z)
+        
+        # Task prediction
+        y_pred = self.predictor(z)
+        
+        # Adversarial prediction (with gradient reversal)
+        z_reversed = self.grl(z)
+        s_pred = self.adversary(z_reversed)
+        
+        if return_latent:
+            return y_pred, s_pred, mu, logvar, z
+        
+        return y_pred, s_pred, mu, logvar
+    
+    def compute_kl_loss(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+        """
+        Compute KL divergence loss for variational bottleneck.
+        KL(q(z|x) || p(z)) where p(z) = N(0, I)
+        
+        KL = -0.5 * sum(1 + log(σ²) - μ² - σ²)
+        """
+        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
+        return kl_loss.mean()
+    
+    def compute_loss(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        s: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, dict]:
+        """
+        Compute complete AIBT loss.
+        
+        L = L_task + λ₁ L_KL - λ₂ L_adv
+        
+        Returns:
+            total_loss: Combined loss for backpropagation
+            loss_dict: Dictionary of individual loss components
+        """
+        y_pred, s_pred, mu, logvar = self.forward(x)
+        
+        # Task loss
+        task_loss = F.cross_entropy(y_pred, y)
+        
+        # KL divergence loss
+        kl_loss = self.compute_kl_loss(mu, logvar)
+        
+        # Adversarial loss (if sensitive attributes provided)
+        if s is not None:
+            adv_loss = F.cross_entropy(s_pred, s)
+        else:
+            adv_loss = torch.tensor(0.0, device=x.device)
+        
+        # Combined loss: L = L_task + λ₁ L_KL - λ₂ L_adv
+        # Note: The minus sign is already handled by GRL
+        total_loss = task_loss + self.lambda_kl * kl_loss + self.lambda_adv * adv_loss
+        
+        loss_dict = {
+            "total_loss": total_loss.item(),
+            "task_loss": task_loss.item(),
+            "kl_loss": kl_loss.item(),
+            "adv_loss": adv_loss.item() if isinstance(adv_loss, torch.Tensor) else adv_loss,
+        }
+        
+        return total_loss, loss_dict
+    
+    def get_latent(self, x: torch.Tensor) -> torch.Tensor:
+        """Get latent representation for privacy analysis"""
+        self.eval()
+        with torch.no_grad():
+            encoder_out = self.encoder(x)
+            if isinstance(encoder_out, tuple):
+                z = encoder_out[0]
+            else:
+                z = encoder_out
+        return z
+    
+    def predict(self, x: torch.Tensor) -> torch.Tensor:
+        """Get task predictions"""
+        self.eval()
+        with torch.no_grad():
+            y_pred, _, _, _ = self.forward(x)
+        return y_pred
+
+
+# ============================================================================
+# MODEL FACTORY
+# ============================================================================
+
+def create_aibt_model(
+    input_dim: int,
+    num_classes: int,
+    latent_dim: int = 64,
+    hidden_dims: Optional[List[int]] = None,
+    num_sensitive_classes: int = 2,
+    lambda_kl: float = 0.01,
+    lambda_adv: float = 1.0,
+    lambda_grl: float = 1.0,
+    dropout: float = 0.1
+) -> AIBTModel:
+    """
+    Factory function to create an AIBT model.
+    
+    Args:
+        input_dim: Input feature dimension
+        num_classes: Number of output classes
+        latent_dim: Dimension of latent representation
+        hidden_dims: Hidden layer dimensions (default: [128, 64])
+        num_sensitive_classes: Number of sensitive attribute classes
+        lambda_kl: KL divergence weight
+        lambda_adv: Adversarial loss weight
+        lambda_grl: Gradient reversal strength
+        dropout: Dropout rate
+    
+    Returns:
+        AIBTModel: Complete AIBT model
+    """
+    if hidden_dims is None:
+        hidden_dims = [128, 64]
+    
+    # Create encoder
+    encoder = MLPEncoder(
+        input_dim=input_dim,
+        hidden_dims=hidden_dims,
+        latent_dim=latent_dim,
+        variational=True,
+        dropout=dropout
+    )
+    
+    # Create predictor
+    predictor = Predictor(
+        latent_dim=latent_dim,
+        hidden_dims=[64],
+        num_classes=num_classes,
+        dropout=dropout
+    )
+    
+    # Create adversary
+    adversary = Adversary(
+        latent_dim=latent_dim,
+        hidden_dims=[64],
+        num_sensitive_classes=num_sensitive_classes,
+        dropout=dropout
+    )
+    
+    # Create complete model
+    model = AIBTModel(
+        encoder=encoder,
+        predictor=predictor,
+        adversary=adversary,
+        lambda_kl=lambda_kl,
+        lambda_adv=lambda_adv,
+        lambda_grl=lambda_grl
+    )
+    
+    return model

aibt/py.typed

@@ -0,0 +1,2 @@
+# Marker file for PEP 561
+# This file indicates that the package supports type hints

aibt/utils.py

@@ -0,0 +1,162 @@
+"""
+Utility functions for AIBT Framework
+"""
+
+import torch
+import numpy as np
+from typing import Tuple, List, Optional
+
+
+def get_device(prefer_gpu: bool = True) -> str:
+    """
+    Get the best available device.
+    
+    Args:
+        prefer_gpu: If True, prefer GPU if available
+    
+    Returns:
+        Device string ('cuda' or 'cpu')
+    """
+    if prefer_gpu and torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+def set_seed(seed: int = 42) -> None:
+    """
+    Set random seed for reproducibility.
+    
+    Args:
+        seed: Random seed
+    """
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def split_data_iid(
+    X: np.ndarray,
+    y: np.ndarray,
+    num_clients: int,
+    seed: int = 42
+) -> List[Tuple[np.ndarray, np.ndarray]]:
+    """
+    Split data IID among clients.
+    
+    Args:
+        X: Features
+        y: Labels
+        num_clients: Number of clients
+        seed: Random seed
+    
+    Returns:
+        List of (X, y) tuples for each client
+    """
+    np.random.seed(seed)
+    n_samples = len(y)
+    indices = np.random.permutation(n_samples)
+    
+    # Split indices among clients
+    splits = np.array_split(indices, num_clients)
+    
+    client_data = []
+    for split in splits:
+        client_data.append((X[split], y[split]))
+    
+    return client_data
+
+
+def split_data_non_iid(
+    X: np.ndarray,
+    y: np.ndarray,
+    num_clients: int,
+    alpha: float = 0.5,
+    seed: int = 42
+) -> List[Tuple[np.ndarray, np.ndarray]]:
+    """
+    Split data non-IID among clients using Dirichlet distribution.
+    
+    Args:
+        X: Features
+        y: Labels
+        num_clients: Number of clients
+        alpha: Dirichlet concentration parameter (smaller = more non-IID)
+        seed: Random seed
+    
+    Returns:
+        List of (X, y) tuples for each client
+    """
+    np.random.seed(seed)
+    n_samples = len(y)
+    num_classes = len(np.unique(y))
+    
+    # Get indices for each class
+    class_indices = {c: np.where(y == c)[0] for c in range(num_classes)}
+    
+    # Sample proportions from Dirichlet distribution
+    proportions = np.random.dirichlet([alpha] * num_clients, num_classes)
+    
+    # Assign samples to clients
+    client_indices = [[] for _ in range(num_clients)]
+    
+    for c in range(num_classes):
+        indices = class_indices[c]
+        np.random.shuffle(indices)
+        
+        # Split class samples according to proportions
+        splits = np.split(
+            indices,
+            (np.cumsum(proportions[c]) * len(indices)).astype(int)[:-1]
+        )
+        
+        for client_id, split in enumerate(splits):
+            client_indices[client_id].extend(split)
+    
+    # Create client datasets
+    client_data = []
+    for indices in client_indices:
+        indices = np.array(indices)
+        if len(indices) > 0:
+            client_data.append((X[indices], y[indices]))
+        else:
+            # Empty client, add minimal data
+            client_data.append((X[:1], y[:1]))
+    
+    return client_data
+
+
+def normalize_data(
+    X_train: np.ndarray,
+    X_test: np.ndarray
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Normalize data using training statistics.
+    
+    Args:
+        X_train: Training features
+        X_test: Test features
+    
+    Returns:
+        Normalized (X_train, X_test)
+    """
+    mean = X_train.mean(axis=0)
+    std = X_train.std(axis=0) + 1e-8
+    
+    X_train_norm = (X_train - mean) / std
+    X_test_norm = (X_test - mean) / std
+    
+    return X_train_norm, X_test_norm
+
+
+def count_parameters(model: torch.nn.Module) -> int:
+    """
+    Count trainable parameters in a model.
+    
+    Args:
+        model: PyTorch model
+    
+    Returns:
+        Number of trainable parameters
+    """
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)