Information Bottleneck: Understanding Deep Learning Through Information Theory 2025

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize

def compute_ib_curve(px, pxy, beta_values):
    """
    Compute the Information Bottleneck curve for discrete distributions.
    
    For each beta, finds the optimal encoder p(t|x) that minimizes:
    I(X;T) - beta * I(T;Y)
    
    Args:
        px: P(X) marginal distribution
        pxy: P(X,Y) joint distribution
        beta_values: Array of beta values to compute
    
    Returns:
        I_XT_values: I(X;T) for each beta
        I_TY_values: I(T;Y) for each beta
    """
    n_x = len(px)
    n_y = pxy.shape[1]
    
    # Marginal P(Y)
    py = pxy.sum(axis=0)
    
    # Conditional P(Y|X)
    pyx = pxy / px[:, np.newaxis]
    
    I_XT_values = []
    I_TY_values = []
    
    for beta in beta_values:
        # Initialize random encoder
        n_t = min(n_x, 10)  # Number of hidden states
        ptx = np.random.dirichlet(np.ones(n_t), size=n_x)  # P(T|X)
        
        # Blahut-Arimoto style iteration
        for _ in range(100):
            # Compute P(T)
            pt = ptx.T @ px
            
            # Compute P(Y|T)
            pyt = np.zeros((n_y, n_t))
            for t in range(n_t):
                for y in range(n_y):
                    for x in range(n_x):
                        pyt[y, t] += pxy[x, y] * ptx[x, t] / (pt[t] + 1e-10)
            
            # Update P(T|X) using IB equations
            new_ptx = np.zeros_like(ptx)
            for x in range(n_x):
                for t in range(n_t):
                    # IB update rule
                    exp_term = 0
                    for y in range(n_y):
                        if pyx[x, y] > 0 and pyt[y, t] > 0:
                            exp_term += pyx[x, y] * np.log(pyt[y, t] / (py[y] + 1e-10) + 1e-10)
                    
                    new_ptx[x, t] = pt[t] * np.exp(beta * exp_term)
                
                # Normalize
                new_ptx[x] /= (new_ptx[x].sum() + 1e-10)
            
            ptx = new_ptx
        
        # Compute I(X;T)
        pt = ptx.T @ px
        I_XT = 0
        for x in range(n_x):
            for t in range(n_t):
                if ptx[x, t] > 0 and pt[t] > 0:
                    I_XT += px[x] * ptx[x, t] * np.log(ptx[x, t] / pt[t] + 1e-10)
        
        # Compute I(T;Y)
        pty = np.zeros((n_t, n_y))
        for t in range(n_t):
            for y in range(n_y):
                for x in range(n_x):
                    pty[t, y] += px[x] * ptx[x, t] * pyx[x, y]
        
        I_TY = 0
        for t in range(n_t):
            for y in range(n_y):
                if pty[t, y] > 0 and pt[t] > 0 and py[y] > 0:
                    I_TY += pty[t, y] * np.log(pty[t, y] / (pt[t] * py[y] + 1e-10) + 1e-10)
        
        I_XT_values.append(max(0, I_XT))
        I_TY_values.append(max(0, I_TY))
    
    return np.array(I_XT_values), np.array(I_TY_values)

# Example: Compute IB curve for a simple problem
print("=== Information Bottleneck Curve ===")

# Binary classification example
px = np.array([0.5, 0.5])  # P(X)
pxy = np.array([
    [0.45, 0.05],  # X=0: mostly Y=0
    [0.05, 0.45]   # X=1: mostly Y=1
])  # P(X,Y)

betas = np.linspace(0.1, 10, 20)
I_XT, I_TY = compute_ib_curve(px, pxy, betas)

print("Beta values:", betas[:5], "...")
print("I(X;T) values:", I_XT[:5], "...")
print("I(T;Y) values:", I_TY[:5], "...")

import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.neighbors import KernelDensity

class InformationPlaneAnalyzer:
    """
    Analyze neural network learning dynamics using the Information Plane.
    
    Tracks I(X;T) and I(T;Y) for each layer during training.
    """
    
    def __init__(self, model, layer_indices):
        """
        Args:
            model: PyTorch neural network
            layer_indices: Indices of layers to track
        """
        self.model = model
        self.layer_indices = layer_indices
        self.history = {idx: {'I_XT': [], 'I_TY': []} for idx in layer_indices}
    
    def get_layer_activations(self, x, layer_idx):
        """Extract activations from a specific layer."""
        activations = x
        
        for idx, layer in enumerate(self.model):
            activations = layer(activations)
            if idx == layer_idx:
                return activations
        
        return activations
    
    def estimate_mi_binning(self, x, y, n_bins=30):
        """
        Estimate mutual information using binning.
        
        Args:
            x: First variable (n_samples, n_features)
            y: Second variable (n_samples,) or (n_samples, n_features)
            n_bins: Number of bins
        
        Returns:
            Estimated mutual information
        """
        # Reduce dimensionality if needed
        if x.ndim > 1 and x.shape[1] > 1:
            # Use first principal component
            x_centered = x - x.mean(axis=0)
            _, _, vh = np.linalg.svd(x_centered, full_matrices=False)
            x = x_centered @ vh[0]
        else:
            x = x.flatten()
        
        if y.ndim > 1:
            y = y.flatten()
        
        # Bin the data
        x_bins = np.digitize(x, np.percentile(x, np.linspace(0, 100, n_bins)))
        y_bins = np.digitize(y, np.percentile(y, np.linspace(0, 100, n_bins)))
        
        # Compute joint histogram
        joint_hist = np.histogram2d(x_bins, y_bins, bins=n_bins)[0]
        joint_hist = joint_hist / joint_hist.sum()
        
        # Marginals
        px = joint_hist.sum(axis=1)
        py = joint_hist.sum(axis=0)
        
        # Mutual information
        mi = 0
        for i in range(n_bins):
            for j in range(n_bins):
                if joint_hist[i, j] > 0 and px[i] > 0 and py[j] > 0:
                    mi += joint_hist[i, j] * np.log2(joint_hist[i, j] / (px[i] * py[j]))
        
        return max(0, mi)
    
    def compute_information_plane_coordinates(self, X, Y, epoch=None):
        """
        Compute I(X;T) and I(T;Y) for all tracked layers.
        
        Args:
            X: Input data (torch.Tensor)
            Y: Labels (torch.Tensor)
            epoch: Optional epoch number for logging
        
        Returns:
            Dictionary with coordinates for each layer
        """
        self.model.eval()
        coordinates = {}
        
        X_np = X.detach().cpu().numpy()
        Y_np = Y.detach().cpu().numpy()
        
        with torch.no_grad():
            for layer_idx in self.layer_indices:
                # Get layer activations
                T = self.get_layer_activations(X, layer_idx)
                T_np = T.detach().cpu().numpy()
                
                # Flatten if needed
                if T_np.ndim > 2:
                    T_np = T_np.reshape(T_np.shape[0], -1)
                
                # Estimate I(X;T)
                I_XT = self.estimate_mi_binning(X_np, T_np)
                
                # Estimate I(T;Y)
                I_TY = self.estimate_mi_binning(T_np, Y_np)
                
                coordinates[layer_idx] = {'I_XT': I_XT, 'I_TY': I_TY}
                
                # Store in history
                self.history[layer_idx]['I_XT'].append(I_XT)
                self.history[layer_idx]['I_TY'].append(I_TY)
        
        return coordinates
    
    def get_trajectory(self, layer_idx):
        """Get the full trajectory in information plane for a layer."""
        return {
            'I_XT': np.array(self.history[layer_idx]['I_XT']),
            'I_TY': np.array(self.history[layer_idx]['I_TY'])
        }

# Example network for information plane analysis
class MLPForInfoPlane(nn.Module):
    """Simple MLP for Information Plane experiments."""
    
    def __init__(self, input_dim, hidden_dims, output_dim):
        super().__init__()
        
        layers = []
        prev_dim = input_dim
        
        for h_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, h_dim))
            layers.append(nn.Tanh())  # Tanh often shows clearer IB dynamics
            prev_dim = h_dim
        
        layers.append(nn.Linear(prev_dim, output_dim))
        self.layers = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.layers(x)

print("=== Information Plane Analyzer ===")
print("Ready to track neural network learning dynamics")

# Create example
model = MLPForInfoPlane(10, [64, 32, 16], 2)
analyzer = InformationPlaneAnalyzer(model.layers, layer_indices=[1, 3, 5])

# Simulate training analysis
X = torch.randn(500, 10)
Y = torch.randint(0, 2, (500,))

coords = analyzer.compute_information_plane_coordinates(X, Y, epoch=0)
print("\nLayer coordinates in Information Plane:")
for layer_idx, coord in coords.items():
    print(f"  Layer {layer_idx}: I(X;T)={coord['I_XT']:.3f}, I(T;Y)={coord['I_TY']:.3f}")

def simulate_information_plane_dynamics():
    """
    Simulate the two-phase learning dynamic in the information plane.
    
    Phase 1: Fitting - increase both I(X;T) and I(T;Y)
    Phase 2: Compression - decrease I(X;T), maintain I(T;Y)
    """
    # Simulate epochs
    epochs = 100
    
    # Phase 1: Fitting (epochs 0-40)
    # I(X;T) increases, I(T;Y) increases
    I_XT_fit = np.linspace(0.5, 4.0, 40) + 0.2 * np.random.randn(40)
    I_TY_fit = np.linspace(0.2, 3.0, 40) + 0.1 * np.random.randn(40)
    
    # Phase 2: Compression (epochs 40-100)
    # I(X;T) decreases, I(T;Y) stable
    I_XT_comp = np.linspace(4.0, 2.0, 60) + 0.1 * np.random.randn(60)
    I_TY_comp = 3.0 + 0.05 * np.random.randn(60)  # Stable around 3.0
    
    I_XT = np.concatenate([I_XT_fit, I_XT_comp])
    I_TY = np.concatenate([I_TY_fit, I_TY_comp])
    
    # Ensure non-negative
    I_XT = np.maximum(I_XT, 0)
    I_TY = np.maximum(I_TY, 0)
    
    return I_XT, I_TY

# Simulate dynamics for multiple layers
print("=== Two-Phase Learning Dynamics ===\n")

layers_dynamics = {}
for layer in ['Layer 1 (near input)', 'Layer 2 (middle)', 'Layer 3 (near output)']:
    I_XT, I_TY = simulate_information_plane_dynamics()
    layers_dynamics[layer] = {'I_XT': I_XT, 'I_TY': I_TY}

# Show key epochs
for layer, dynamics in layers_dynamics.items():
    print(f"{layer}:")
    print(f"  Epoch 0:   I(X;T)={dynamics['I_XT'][0]:.2f}, I(T;Y)={dynamics['I_TY'][0]:.2f}")
    print(f"  Epoch 40:  I(X;T)={dynamics['I_XT'][40]:.2f}, I(T;Y)={dynamics['I_TY'][40]:.2f}")
    print(f"  Epoch 100: I(X;T)={dynamics['I_XT'][-1]:.2f}, I(T;Y)={dynamics['I_TY'][-1]:.2f}")
    print()

def analyze_compression_mechanisms():
    """
    Analyze different compression mechanisms in neural networks.
    """
    print("=== Compression Mechanisms in DNNs ===\n")
    
    # 1. ReLU Sparsity Analysis
    def relu_sparsity(activations):
        """Measure compression via ReLU sparsity."""
        return (activations <= 0).mean()
    
    # Simulate activations before and after ReLU
    pre_relu = np.random.randn(1000, 256)
    post_relu = np.maximum(pre_relu, 0)
    
    sparsity = relu_sparsity(pre_relu)
    print(f"1. ReLU Compression:")
    print(f"   Pre-ReLU zeros: {(pre_relu == 0).mean():.2%}")
    print(f"   Post-ReLU zeros: {(post_relu == 0).mean():.2%}")
    print(f"   Information reduction: ~{(post_relu == 0).mean() * 100:.0f}% of activations zeroed")
    print()
    
    # 2. Dropout Compression
    def dropout_compression(x, p=0.5):
        """Analyze dropout as compression."""
        mask = np.random.binomial(1, 1-p, size=x.shape)
        return x * mask / (1-p)  # Scale to maintain expected value
    
    original = np.random.randn(1000, 256)
    dropped = dropout_compression(original, p=0.5)
    
    print(f"2. Dropout Compression (p=0.5):")
    print(f"   Active neurons: {(dropped != 0).mean():.2%}")
    print(f"   Effective capacity reduction: 50%")
    print()
    
    # 3. Weight Decay Effect
    def analyze_weight_decay(initial_weights, decay_rate=0.01, epochs=100):
        """Simulate weight decay compression."""
        weights = initial_weights.copy()
        weight_norms = [np.linalg.norm(weights)]
        
        for _ in range(epochs):
            # Simulate gradient + weight decay
            gradient = np.random.randn(*weights.shape) * 0.1
            weights = weights - 0.01 * gradient - decay_rate * weights
            weight_norms.append(np.linalg.norm(weights))
        
        return weight_norms
    
    initial_w = np.random.randn(256, 128)
    norms = analyze_weight_decay(initial_w)
    
    print(f"3. Weight Decay Compression:")
    print(f"   Initial weight norm: {norms[0]:.2f}")
    print(f"   Final weight norm: {norms[-1]:.2f}")
    print(f"   Reduction: {(1 - norms[-1]/norms[0]):.1%}")
    print()
    
    # 4. Batch Normalization Effect
    def batch_norm_compression(x):
        """Analyze BatchNorm compression."""
        # Normalize
        mean = x.mean(axis=0)
        std = x.std(axis=0) + 1e-5
        normalized = (x - mean) / std
        
        # Entropy before and after
        def estimate_entropy(data):
            hist, _ = np.histogram(data.flatten(), bins=50, density=True)
            hist = hist[hist > 0]
            return -np.sum(hist * np.log2(hist + 1e-10))
        
        return estimate_entropy(x), estimate_entropy(normalized)
    
    data = np.random.randn(1000, 64) * 3 + 5  # Non-standard distribution
    entropy_before, entropy_after = batch_norm_compression(data)
    
    print(f"4. Batch Normalization Compression:")
    print(f"   Entropy before BN: {entropy_before:.2f} bits")
    print(f"   Entropy after BN: {entropy_after:.2f} bits")
    print(f"   Compression: {entropy_before - entropy_after:.2f} bits reduced")

analyze_compression_mechanisms()

def generalization_vs_compression_analysis():
    """
    Demonstrate the relationship between compression and generalization.
    """
    print("=== Compression-Generalization Relationship ===\n")
    
    # Simulate different model complexities
    n_samples = 1000
    
    # Model A: High compression (low I(X;T))
    I_XT_A = 2.0  # Low information about input
    train_error_A = 0.10
    
    # Model B: Medium compression
    I_XT_B = 5.0
    train_error_B = 0.05
    
    # Model C: Low compression (high I(X;T))
    I_XT_C = 10.0
    train_error_C = 0.01  # Better training fit
    
    # Generalization bound: sqrt(I(X;T) / n)
    gen_gap_A = np.sqrt(I_XT_A / n_samples)
    gen_gap_B = np.sqrt(I_XT_B / n_samples)
    gen_gap_C = np.sqrt(I_XT_C / n_samples)
    
    # Estimated test error
    test_error_A = train_error_A + gen_gap_A
    test_error_B = train_error_B + gen_gap_B
    test_error_C = train_error_C + gen_gap_C
    
    print("Compression vs Generalization Analysis:")
    print("-" * 60)
    print(f"{'Model':<10} {'I(X;T)':<10} {'Train Err':<12} {'Gen Gap':<12} {'Test Err':<10}")
    print("-" * 60)
    print(f"{'A (High)':<10} {I_XT_A:<10.1f} {train_error_A:<12.2%} {gen_gap_A:<12.3f} {test_error_A:<10.2%}")
    print(f"{'B (Med)':<10} {I_XT_B:<10.1f} {train_error_B:<12.2%} {gen_gap_B:<12.3f} {test_error_B:<10.2%}")
    print(f"{'C (Low)':<10} {I_XT_C:<10.1f} {train_error_C:<12.2%} {gen_gap_C:<12.3f} {test_error_C:<10.2%}")
    print("-" * 60)
    print(f"\nModel B achieves best test error through balanced compression!")
    print(f"Model C overfits despite lowest training error.")

generalization_vs_compression_analysis()

import torch
import torch.nn as nn
import torch.nn.functional as F

class DeepVIB(nn.Module):
    """
    Deep Variational Information Bottleneck implementation.
    
    Combines:
    - Stochastic encoder for compression
    - KL regularization for controlling I(X;Z)
    - Classifier for maintaining I(Z;Y)
    """
    
    def __init__(self, input_dim, hidden_dims, z_dim, num_classes, beta=1e-3):
        """
        Args:
            input_dim: Input feature dimension
            hidden_dims: List of hidden layer dimensions
            z_dim: Bottleneck dimension
            num_classes: Number of output classes
            beta: Information constraint weight
        """
        super().__init__()
        
        self.z_dim = z_dim
        self.beta = beta
        
        # Encoder network
        encoder_layers = []
        prev_dim = input_dim
        for h_dim in hidden_dims:
            encoder_layers.extend([
                nn.Linear(prev_dim, h_dim),
                nn.BatchNorm1d(h_dim),
                nn.ReLU()
            ])
            prev_dim = h_dim
        
        self.encoder = nn.Sequential(*encoder_layers)
        
        # Bottleneck parameters (mean and log-variance)
        self.fc_mu = nn.Linear(prev_dim, z_dim)
        self.fc_logvar = nn.Linear(prev_dim, z_dim)
        
        # Classifier from bottleneck
        self.classifier = nn.Sequential(
            nn.Linear(z_dim, z_dim),
            nn.ReLU(),
            nn.Linear(z_dim, num_classes)
        )
    
    def encode(self, x):
        """
        Encode input to bottleneck distribution parameters.
        
        Returns:
            mu: Mean of q(z|x)
            logvar: Log variance of q(z|x)
        """
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        """
        Reparameterization trick for backpropagation.
        
        z = mu + std * epsilon, where epsilon ~ N(0, I)
        """
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            return mu + eps * std
        else:
            return mu  # Use mean for evaluation
    
    def forward(self, x):
        """
        Forward pass through VIB.
        
        Returns:
            logits: Class predictions
            mu: Bottleneck mean
            logvar: Bottleneck log-variance
        """
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        logits = self.classifier(z)
        return logits, mu, logvar
    
    def compute_loss(self, x, y):
        """
        Compute VIB loss: Classification + beta * KL Divergence.
        
        Args:
            x: Input batch
            y: Target labels
        
        Returns:
            total_loss: VIB objective
            ce_loss: Cross-entropy loss
            kl_loss: KL divergence from prior
        """
        logits, mu, logvar = self(x)
        
        # Classification loss (negative log-likelihood)
        ce_loss = F.cross_entropy(logits, y)
        
        # KL divergence: KL(q(z|x) || p(z)) where p(z) = N(0, I)
        # Closed form: 0.5 * sum(mu^2 + var - 1 - log(var))
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1)
        kl_loss = kl_loss.mean()
        
        # Total VIB loss
        total_loss = ce_loss + self.beta * kl_loss
        
        return total_loss, ce_loss.item(), kl_loss.item()
    
    def get_bottleneck_info(self, x):
        """
        Analyze bottleneck representation.
        
        Returns information about the learned representation.
        """
        with torch.no_grad():
            mu, logvar = self.encode(x)
            z = self.reparameterize(mu, logvar)
            
            # Compute statistics
            info = {
                'mean_mu': mu.mean().item(),
                'std_mu': mu.std().item(),
                'mean_var': logvar.exp().mean().item(),
                'z_norm': z.norm(dim=1).mean().item(),
                'active_dims': (mu.std(dim=0) > 0.1).sum().item()
            }
        
        return info

# Example usage
print("=== Deep Variational Information Bottleneck ===\n")

# Create model
input_dim = 784  # e.g., MNIST
hidden_dims = [512, 256]
z_dim = 32
num_classes = 10

model = DeepVIB(input_dim, hidden_dims, z_dim, num_classes, beta=1e-3)
print(f"Model architecture:")
print(f"  Input: {input_dim}")
print(f"  Hidden: {hidden_dims}")
print(f"  Bottleneck: {z_dim}")
print(f"  Output: {num_classes}")
print(f"  Beta: {model.beta}")

# Simulate training step
x = torch.randn(32, input_dim)
y = torch.randint(0, num_classes, (32,))

total_loss, ce_loss, kl_loss = model.compute_loss(x, y)
print(f"\nTraining step:")
print(f"  CE Loss: {ce_loss:.4f}")
print(f"  KL Loss: {kl_loss:.4f}")
print(f"  Total Loss: {total_loss.item():.4f}")

bottleneck_info = model.get_bottleneck_info(x)
print(f"\nBottleneck analysis:")
for key, value in bottleneck_info.items():
    print(f"  {key}: {value:.4f}")

class BetaScheduler:
    """
    Schedule beta during training for optimal IB learning.
    """
    
    def __init__(self, initial_beta, final_beta, warmup_epochs, total_epochs):
        """
        Args:
            initial_beta: Starting beta (usually 0 or small)
            final_beta: Target beta value
            warmup_epochs: Epochs to reach final beta
            total_epochs: Total training epochs
        """
        self.initial_beta = initial_beta
        self.final_beta = final_beta
        self.warmup_epochs = warmup_epochs
        self.total_epochs = total_epochs
    
    def get_beta(self, epoch):
        """Get beta value for current epoch."""
        if epoch < self.warmup_epochs:
            # Linear warmup
            progress = epoch / self.warmup_epochs
            return self.initial_beta + progress * (self.final_beta - self.initial_beta)
        else:
            return self.final_beta
    
    def __repr__(self):
        return f"BetaScheduler({self.initial_beta} -> {self.final_beta}, warmup={self.warmup_epochs})"

# Example beta schedules
print("=== Beta Scheduling Strategies ===\n")

# Strategy 1: Gradual warmup
scheduler_warmup = BetaScheduler(
    initial_beta=0,
    final_beta=1e-3,
    warmup_epochs=50,
    total_epochs=200
)

# Strategy 2: Start with some compression
scheduler_immediate = BetaScheduler(
    initial_beta=1e-4,
    final_beta=1e-3,
    warmup_epochs=20,
    total_epochs=200
)

print("Strategy 1: Gradual warmup from 0")
for epoch in [0, 25, 50, 100]:
    print(f"  Epoch {epoch}: beta = {scheduler_warmup.get_beta(epoch):.2e}")

print("\nStrategy 2: Quick warmup with initial compression")
for epoch in [0, 10, 20, 100]:
    print(f"  Epoch {epoch}: beta = {scheduler_immediate.get_beta(epoch):.2e}")

class ConvVIB(nn.Module):
    """
    Convolutional VIB for image classification.
    
    Applies information bottleneck principle to CNN features.
    """
    
    def __init__(self, z_dim=64, num_classes=10, beta=1e-3):
        super().__init__()
        
        self.beta = beta
        self.z_dim = z_dim
        
        # Convolutional encoder
        self.conv_encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        
        # Bottleneck
        self.fc_mu = nn.Linear(128, z_dim)
        self.fc_logvar = nn.Linear(128, z_dim)
        
        # Classifier
        self.classifier = nn.Linear(z_dim, num_classes)
    
    def encode(self, x):
        """Encode image to bottleneck distribution."""
        h = self.conv_encoder(x)
        h = h.view(h.size(0), -1)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar
    
    def reparameterize(self, mu, logvar):
        """Reparameterization trick."""
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            return mu + eps * std
        return mu
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        logits = self.classifier(z)
        return logits, mu, logvar

print("=== Convolutional VIB for Images ===")
model = ConvVIB(z_dim=64, num_classes=10, beta=1e-3)
print(f"Bottleneck dimension: {model.z_dim}")
print(f"Beta: {model.beta}")

# Test with dummy image
x = torch.randn(4, 1, 28, 28)  # MNIST-like
logits, mu, logvar = model(x)
print(f"\nOutput shapes:")
print(f"  Logits: {logits.shape}")
print(f"  Mu: {mu.shape}")
print(f"  Logvar: {logvar.shape}")

class AttentionIBAnalysis:
    """
    Analyze information flow in self-attention through IB framework.
    """
    
    @staticmethod
    def attention_entropy(attention_weights):
        """
        Compute entropy of attention weights.
        
        Higher entropy = more distributed attention = less compression
        Lower entropy = focused attention = more compression
        
        Args:
            attention_weights: (batch, heads, seq_len, seq_len)
        
        Returns:
            Mean entropy per head
        """
        # Avoid log(0)
        eps = 1e-10
        
        # Entropy along last dimension (attention distribution)
        entropy = -torch.sum(
            attention_weights * torch.log(attention_weights + eps),
            dim=-1
        )
        
        # Mean over batch and sequence
        return entropy.mean(dim=(0, 2))  # Per head
    
    @staticmethod
    def analyze_layer_compression(attention_weights_list):
        """
        Analyze compression across transformer layers.
        
        Args:
            attention_weights_list: List of attention weights per layer
        
        Returns:
            Dictionary with compression analysis
        """
        results = {
            'layer_entropies': [],
            'compression_ratio': []
        }
        
        max_entropy = None
        
        for i, attn in enumerate(attention_weights_list):
            entropy = AttentionIBAnalysis.attention_entropy(attn)
            mean_entropy = entropy.mean().item()
            results['layer_entropies'].append(mean_entropy)
            
            if max_entropy is None:
                max_entropy = mean_entropy
            else:
                results['compression_ratio'].append(1 - mean_entropy/max_entropy)
        
        return results

# Demonstrate attention analysis
print("=== Transformer Attention Analysis ===\n")

# Simulate attention weights for different layers
seq_len = 128
n_heads = 8
batch_size = 4

# Early layer: distributed attention (high entropy)
attn_early = F.softmax(torch.randn(batch_size, n_heads, seq_len, seq_len) * 0.5, dim=-1)

# Middle layer: more focused
attn_middle = F.softmax(torch.randn(batch_size, n_heads, seq_len, seq_len) * 1.0, dim=-1)

# Late layer: very focused (low entropy)
attn_late = F.softmax(torch.randn(batch_size, n_heads, seq_len, seq_len) * 2.0, dim=-1)

attention_weights = [attn_early, attn_middle, attn_late]
analysis = AttentionIBAnalysis.analyze_layer_compression(attention_weights)

print("Layer-wise attention entropy (lower = more compression):")
for i, entropy in enumerate(analysis['layer_entropies']):
    print(f"  Layer {i+1}: {entropy:.4f} nats")

print("\nCompression increases through layers, consistent with IB theory")

class IBPruning:
    """
    Use Information Bottleneck principle for neural network pruning.
    
    Identifies and removes dimensions that contribute little to I(T;Y).
    """
    
    def __init__(self, model, target_compression=0.5):
        """
        Args:
            model: DeepVIB model
            target_compression: Fraction of dimensions to keep
        """
        self.model = model
        self.target_compression = target_compression
    
    def compute_dimension_importance(self, data_loader, n_batches=50):
        """
        Estimate importance of each bottleneck dimension.
        
        Uses variance of the posterior as a proxy for importance:
        High variance → model uses this dimension → important
        Low variance → dimension not used → can be pruned
        """
        self.model.eval()
        
        variances = []
        
        with torch.no_grad():
            for i, (x, y) in enumerate(data_loader):
                if i >= n_batches:
                    break
                
                mu, logvar = self.model.encode(x)
                var = logvar.exp()
                variances.append(var)
        
        # Stack and compute mean variance per dimension
        all_var = torch.cat(variances, dim=0)
        dim_importance = all_var.mean(dim=0)
        
        return dim_importance
    
    def get_pruning_mask(self, dim_importance):
        """
        Generate mask for keeping important dimensions.
        
        Args:
            dim_importance: Importance score per dimension
        
        Returns:
            Binary mask (1 = keep, 0 = prune)
        """
        n_dims = len(dim_importance)
        n_keep = int(n_dims * self.target_compression)
        
        # Keep dimensions with highest importance
        _, indices = torch.topk(dim_importance, n_keep)
        
        mask = torch.zeros(n_dims)
        mask[indices] = 1
        
        return mask
    
    def apply_pruning(self, mask):
        """
        Apply pruning mask to model.
        
        Creates a new model with reduced bottleneck dimension.
        """
        n_keep = int(mask.sum().item())
        keep_indices = torch.where(mask == 1)[0]
        
        print(f"Pruning bottleneck: {len(mask)} -> {n_keep} dimensions")
        print(f"Compression ratio: {n_keep/len(mask):.1%}")
        
        return keep_indices

print("=== IB-based Model Pruning ===")

# Create example model and pruner
model = DeepVIB(784, [256], 64, 10, beta=1e-3)
pruner = IBPruning(model, target_compression=0.5)

# Simulate importance scores (in practice, computed from data)
dim_importance = torch.rand(64)  # Random for demonstration
dim_importance[0:16] *= 0.1  # Make some dimensions unimportant

mask = pruner.get_pruning_mask(dim_importance)
keep_indices = pruner.apply_pruning(mask)
print(f"Kept dimensions: {keep_indices.tolist()[:10]}...")

class InterpretableVIB(nn.Module):
    """
    VIB with interpretable bottleneck dimensions.
    
    Each bottleneck dimension corresponds to an interpretable concept.
    """
    
    def __init__(self, input_dim, z_dim, num_classes, concept_names=None, beta=1e-3):
        super().__init__()
        
        self.z_dim = z_dim
        self.beta = beta
        self.concept_names = concept_names or [f"concept_{i}" for i in range(z_dim)]
        
        # Encoder
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        
        self.fc_mu = nn.Linear(128, z_dim)
        self.fc_logvar = nn.Linear(128, z_dim)
        
        # Concept-specific classifiers
        self.concept_classifiers = nn.ModuleList([
            nn.Linear(1, 1) for _ in range(z_dim)
        ])
        
        # Final classifier
        self.classifier = nn.Linear(z_dim, num_classes)
    
    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        
        if self.training:
            std = torch.exp(0.5 * logvar)
            z = mu + std * torch.randn_like(std)
        else:
            z = mu
        
        logits = self.classifier(z)
        return logits, mu, logvar
    
    def explain_prediction(self, x):
        """
        Explain prediction using bottleneck concepts.
        
        Returns contribution of each concept to the prediction.
        """
        with torch.no_grad():
            mu, _ = self.encode(x)
            logits = self.classifier(mu)
            
            # Get classifier weights
            weights = self.classifier.weight  # (num_classes, z_dim)
            
            # Contribution of each concept
            contributions = mu.unsqueeze(1) * weights.unsqueeze(0)  # (batch, classes, z_dim)
            
            # Predicted class
            pred_class = logits.argmax(dim=1)
            
            explanations = []
            for i in range(x.size(0)):
                sample_contributions = contributions[i, pred_class[i]]
                
                explanation = {
                    'prediction': pred_class[i].item(),
                    'concepts': {}
                }
                
                for j, name in enumerate(self.concept_names):
                    explanation['concepts'][name] = sample_contributions[j].item()
                
                explanations.append(explanation)
        
        return explanations

print("=== Interpretable VIB ===")

# Create model with named concepts
concept_names = ['shape', 'size', 'texture', 'color', 'position']
model = InterpretableVIB(
    input_dim=100,
    z_dim=5,
    num_classes=10,
    concept_names=concept_names,
    beta=1e-3
)

# Test explanation
x = torch.randn(2, 100)
explanations = model.explain_prediction(x)

print("\nPrediction explanations:")
for i, exp in enumerate(explanations):
    print(f"\nSample {i+1}: Predicted class {exp['prediction']}")
    print("Concept contributions:")
    for concept, contribution in exp['concepts'].items():
        print(f"  {concept}: {contribution:+.3f}")

class RobustVIB(nn.Module):
    """
    VIB for learning adversarially robust representations.
    
    The information constraint naturally limits the information
    an adversary can exploit.
    """
    
    def __init__(self, input_dim, z_dim, num_classes, beta=1e-2):
        super().__init__()
        
        self.z_dim = z_dim
        self.beta = beta
        
        # Higher beta = more compression = more robust
        # but potentially less accurate
        
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 512),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Dropout(0.2)
        )
        
        self.fc_mu = nn.Linear(256, z_dim)
        self.fc_logvar = nn.Linear(256, z_dim)
        
        self.classifier = nn.Sequential(
            nn.Linear(z_dim, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x, num_samples=1):
        """
        Forward with multiple samples for robustness.
        
        Args:
            x: Input
            num_samples: Number of stochastic samples (more = more robust)
        """
        mu, logvar = self.encode(x)
        
        if num_samples == 1:
            z = self.reparameterize(mu, logvar)
            return self.classifier(z), mu, logvar
        else:
            # Multiple samples for robust prediction
            std = torch.exp(0.5 * logvar)
            logits_list = []
            
            for _ in range(num_samples):
                z = mu + std * torch.randn_like(std)
                logits_list.append(self.classifier(z))
            
            # Average predictions
            avg_logits = torch.stack(logits_list).mean(dim=0)
            return avg_logits, mu, logvar
    
    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)
    
    def reparameterize(self, mu, logvar):
        if self.training:
            std = torch.exp(0.5 * logvar)
            return mu + std * torch.randn_like(std)
        return mu

def test_robustness(model, x, epsilon=0.1, num_attacks=10):
    """
    Test robustness against random perturbations.
    
    Returns accuracy under perturbation.
    """
    model.eval()
    
    with torch.no_grad():
        # Clean prediction
        clean_logits, _, _ = model(x, num_samples=5)
        clean_pred = clean_logits.argmax(dim=1)
        
        # Perturbed predictions
        correct = 0
        for _ in range(num_attacks):
            noise = torch.randn_like(x) * epsilon
            perturbed_x = x + noise
            
            perturbed_logits, _, _ = model(perturbed_x, num_samples=5)
            perturbed_pred = perturbed_logits.argmax(dim=1)
            
            correct += (perturbed_pred == clean_pred).float().sum().item()
        
        consistency = correct / (num_attacks * x.size(0))
    
    return consistency

print("=== Robust VIB ===")

# Compare robustness at different beta values
for beta in [1e-4, 1e-3, 1e-2]:
    model = RobustVIB(784, 32, 10, beta=beta)
    x = torch.randn(100, 784)
    
    consistency = test_robustness(model, x, epsilon=0.1)
    print(f"Beta={beta:.0e}: Prediction consistency under noise = {consistency:.1%}")