1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

import numpy as np
import torch
import torch.nn.functional as F
from scipy import stats
 
# Core insight: Loss = -log p(y|x) for some probabilistic model
 
# 1. Cross-Entropy = -log p(y|x) for Categorical distribution
def cross_entropy_loss(logits, targets):
    """
    Cross-entropy for classification.
    Assumes p(y|x) = Categorical(softmax(f(x)))
    Loss = -log p(y=target|x)
    """
    log_probs = F.log_softmax(logits, dim=-1)
    return -log_probs.gather(1, targets.unsqueeze(1)).mean()
 
# 2. MSE = -log p(y|x) for Gaussian distribution (up to constant)
def mse_as_nll(predictions, targets, sigma=1.0):
    """
    MSE for regression.
    Assumes p(y|x) = N(f(x), σ²)
    -log p(y|x) = (y - f(x))² / (2σ²) + log(σ√(2π))
    MSE ignores constant terms.
    """
    mse = ((predictions - targets) ** 2).mean()
    # Full NLL (for proper probabilistic interpretation)
    nll = mse / (2 * sigma**2) + np.log(sigma * np.sqrt(2 * np.pi))
    return mse, nll
 
# 3. MAE = -log p(y|x) for Laplace distribution (up to constant)
def mae_as_nll(predictions, targets, b=1.0):
    """
    MAE (L1 loss) for regression.
    Assumes p(y|x) = Laplace(f(x), b)
    -log p(y|x) = |y - f(x)| / b + log(2b)
    """
    mae = torch.abs(predictions - targets).mean()
    nll = mae / b + np.log(2 * b)
    return mae, nll
 
# Example
print("Loss Functions as Negative Log-Likelihood")
print("=" * 60)
 
# Generate data
torch.manual_seed(42)
preds = torch.tensor([0.5, 1.2, 0.8, 1.5])
targets = torch.tensor([0.6, 1.0, 1.0, 1.3])
 
mse, mse_nll = mse_as_nll(preds, targets)
mae, mae_nll = mae_as_nll(preds, targets)
 
print(f"MSE: {mse:.4f}")
print(f"MSE as Gaussian NLL (σ=1): {mse_nll:.4f}")
print()
print(f"MAE: {mae.item():.4f}")
print(f"MAE as Laplace NLL (b=1): {mae_nll:.4f}")
print()
print("Key insight: Choosing a loss = choosing a noise distribution")
print("MSE → Gaussian noise | MAE → Laplacian noise | Huber → hybrid")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
 
class VAE_InformationTheoretic(nn.Module):
    """
    VAE viewed through information-theoretic lens.
    
    Objective: Maximize I(X; Z) while keeping I(Z; X) close to prior entropy.
    
    ELBO = E_q(z|x)[log p(x|z)] - D_KL(q(z|x) || p(z))
         = Reconstruction          - Rate
         
    Rate-Distortion interpretation:
    - Reconstruction = How well can we decode X from Z?
    - KL term = How many bits does Z use beyond the prior?
    """
    
    def __init__(self, input_dim, latent_dim):
        super().__init__()
        self.latent_dim = latent_dim
        
        # Encoder: q(z|x)
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU()
        )
        self.fc_mu = nn.Linear(128, latent_dim)
        self.fc_logvar = nn.Linear(128, latent_dim)
        
        # Decoder: p(x|z)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, input_dim),
            nn.Sigmoid()
        )
    
    def encode(self, x):
        h = self.encoder(x)
        return self.fc_mu(h), self.fc_logvar(h)
    
    def reparameterize(self, mu, logvar):
        """Reparameterization trick: z = μ + σ * ε, ε ~ N(0,1)"""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z)
        return x_recon, mu, logvar
    
    def loss_function(self, x, x_recon, mu, logvar, beta=1.0):
        """
        ELBO loss with β-VAE weighting.
        
        Returns components for analysis:
        - Reconstruction: -E_q[log p(x|z)] ≈ BCE(x, x_recon)
        - Rate: D_KL(q(z|x) || p(z)) = D_KL(N(μ, σ²) || N(0, 1))
        """
        # Reconstruction loss (negative log-likelihood)
        # For Bernoulli decoder: BCE
        recon = F.binary_cross_entropy(x_recon, x, reduction='sum')
        
        # KL divergence: D_KL(N(μ, σ²) || N(0, 1))
        # = 0.5 * Σ(μ² + σ² - 1 - log(σ²))
        kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        
        # Total loss (negative ELBO)
        total = recon + beta * kl
        
        return total, recon, kl
 
# Information quantities
print("VAE Information-Theoretic Analysis")
print("=" * 60)
print()
print("ELBO = E_q[log p(x|z)] - D_KL(q(z|x) || p(z))")
print()
print("Decomposition:")
print("1. E_q[log p(x|z)]: Reconstruction quality")
print("   - How much information about X is preserved in Z?")
print("   - Maximizing this → high I(X; Z)")
print()
print("2. D_KL(q(z|x) || p(z)): Rate / Compression")
print("   - How many 'extra bits' does z use?")
print("   - Minimizing this → low I(X; Z)")
print()
print("β-VAE: Total = Reconstruction + β × KL")
print("  β > 1: More compression, better disentanglement")
print("  β < 1: Better reconstruction, less regularization")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

# Information Bottleneck Framework
# ================================
 
# The IB Objective:
# min I(X; T) - β × I(T; Y)
#
# Equivalently:
# max I(T; Y) - β⁻¹ × I(X; T)
 
# Key insights:
# 1. Optimal T discards information in X not relevant to Y
# 2. β controls compression-prediction tradeoff
# 3. Neural network layers approximately solve IB
 
# The IB Lagrangian:
L = I(X; T) - β × I(T; Y)
 
# Can be rewritten using Markov chain X → T → Y:
# I(T; Y) ≤ I(X; Y)  (data processing inequality)
# I(X; T) ≥ I(T; Y)  (must preserve enough for prediction)
 
# In neural networks:
# - Early layers: High I(X; T), moderate I(T; Y)
# - Later layers: Lower I(X; T), higher I(T; Y)
# - This is "compression" as information flows through network
 
import numpy as np
import matplotlib.pyplot as plt
 
def plot_information_plane():
    """
    Conceptual illustration of the information plane.
    Each point (I(X;T), I(T;Y)) represents a layer's information content.
    """
    # Simulated layer-wise information (conceptual)
    layers = ['Input', 'Hidden1', 'Hidden2', 'Hidden3', 'Output']
    
    # During training, networks compress (reduce I(X;T)) while maintaining I(T;Y)
    I_XT_init = [5.0, 4.8, 4.5, 4.0, 3.5]  # Initially, preserve most input info
    I_TY_init = [0.5, 0.8, 1.0, 1.2, 1.5]  # Gradually improve prediction
    
    I_XT_final = [5.0, 3.5, 2.5, 2.0, 1.8]  # Compress unnecessary info
    I_TY_final = [0.5, 1.2, 1.5, 1.8, 2.0]  # Improve prediction
    
    print("Information Plane Analysis")
    print("=" * 60)
    print(f"{'Layer':<15} {'I(X;T) init':<15} {'I(X;T) final':<15} {'Compression':<15}")
    print("-" * 60)
    for i, layer in enumerate(layers):
        comp = I_XT_init[i] - I_XT_final[i]
        print(f"{layer:<15} {I_XT_init[i]:<15.2f} {I_XT_final[i]:<15.2f} {comp:+.2f}")
    print()
    print("Key observation: Networks compress I(X;T) while improving I(T;Y)")
    print("This is the 'information bottleneck' at work!")
 
plot_information_plane()
 
# Variational Information Bottleneck (VIB)
print()
print("Variational Information Bottleneck (VIB)")
print("-" * 60)
print("Tractable approximation to IB using variational bounds:")
print()
print("max I(T; Y) - β × I(X; T)")
print("≈ max E[log p(Y|T)] - β × D_KL(p(T|X) || r(T))")
print()
print("Where:")
print("  p(T|X): Encoder (stochastic neural network)")
print("  r(T): Variational prior (e.g., N(0, I))")
print("  p(Y|T): Predictor")
print()
print("This is a stochastic neural network with KL regularization!")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

import torch
import torch.nn.functional as F
 
class ContrastiveLearner:
    """
    Contrastive learning maximizes I(view_1; view_2).
    
    The InfoNCE bound:
    I(X; Y) >= log(K) - L_NCE
    
    where K is the number of negatives + 1 and L_NCE is the contrastive loss.
    """
    
    def __init__(self, temperature=0.07):
        self.temperature = temperature
    
    def info_nce_loss(self, z_i, z_j, negatives=None):
        """
        Compute InfoNCE loss.
        
        Args:
            z_i: First view representations (batch_size, dim)
            z_j: Second view representations (batch_size, dim)  
            negatives: Optional explicit negatives
        
        If negatives=None, use other samples in batch as negatives.
        """
        batch_size = z_i.size(0)
        
        # Normalize embeddings
        z_i = F.normalize(z_i, dim=1)
        z_j = F.normalize(z_j, dim=1)
        
        # Compute similarity matrix
        representations = torch.cat([z_i, z_j], dim=0)  # (2*batch, dim)
        similarity = representations @ representations.T / self.temperature
        
        # Mask self-similarity (diagonal)
        mask = torch.eye(2 * batch_size, device=z_i.device).bool()
        similarity = similarity.masked_fill(mask, float('-inf'))
        
        # For each z_i, positive is the corresponding z_j (and vice versa)
        labels = torch.cat([
            torch.arange(batch_size, 2 * batch_size),
            torch.arange(batch_size)
        ], dim=0).to(z_i.device)
        
        # Cross-entropy loss
        loss = F.cross_entropy(similarity, labels)
        
        return loss
    
    def estimate_mi_bound(self, loss, num_negatives):
        """
        Estimate MI lower bound from InfoNCE loss.
        
        I(X; Y) >= log(K) - L_NCE
        """
        K = num_negatives + 1  # +1 for the positive
        return torch.log(torch.tensor(K)) - loss
 
# Example
learner = ContrastiveLearner(temperature=0.07)
 
# Simulated representations from two views
batch_size = 256
dim = 128
 
z1 = torch.randn(batch_size, dim)
z2 = z1 + torch.randn(batch_size, dim) * 0.3  # Views are similar
 
loss = learner.info_nce_loss(z1, z2)
mi_bound = learner.estimate_mi_bound(loss, num_negatives=2*batch_size - 2)
 
print("Contrastive Learning Analysis")
print("=" * 60)
print(f"InfoNCE Loss: {loss.item():.4f}")
print(f"Number of negatives: {2*batch_size - 2}")
print(f"MI Lower Bound: {mi_bound.item():.4f} nats")
print()
print("Key insight: Lower loss → higher MI bound → better representations")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

# Information-Theoretic Generalization Bounds
# =============================================
 
# Classic bound (simplified form):
# E[|L_test - L_train|] ≤ √(2 * I(W; S) / n)
#
# Where:
# - W: Learned parameters (output of algorithm)
# - S: Training set
# - n: Number of training samples
# - I(W; S): Mutual information between weights and training data
 
# Key implications:
# 1. Less information extracted → better generalization
# 2. SGD noise limits I(W; S), aiding generalization
# 3. Regularization reduces I(W; S)
# 4. Data augmentation increases effective n
 
# Practical insights:
print("Information-Theoretic View of Regularization")
print("=" * 60)
print()
 
regularizers = {
    "Weight Decay (L2)": "Constrains weights to low-norm region, limiting information capacity",
    "Dropout": "Prevents co-adaptation, reducing I(W; S) by adding noise",
    "Batch Normalization": "Stabilizes representations, limiting individual sample influence",
    "Label Smoothing": "Reduces I(Y_pred; Y_true|X), preventing overfitting to hard labels",
    "Data Augmentation": "Increases effective n without changing I(W; S)",
    "Early Stopping": "Stops before full training data memorization, limiting I(W; S)",
}
 
for reg, explanation in regularizers.items():
    print(f"{reg}:")
    print(f"  {explanation}")
    print()
 
# PAC-Bayes bounds (information-theoretic flavor)
print("PAC-Bayes Bounds")
print("-" * 60)
print()
print("Generalization Error ≤ Train Error + √(D_KL(posterior || prior) / n)")
print()
print("This connects:")
print("- KL divergence between learned distribution and prior")
print("- Sample size n")
print("- Generalization gap")
print()
print("Insight: Learning 'close to prior' generalizes better")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

import numpy as np
from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
from sklearn.metrics import log_loss
from scipy.stats import entropy
from scipy.special import rel_entr
 
class InformationTheoreticML:
    """
    Toolkit for IT-informed machine learning practice.
    """
    
    @staticmethod
    def feature_importance_mi(X, y, task='classification'):
        """
        Compute MI-based feature importance.
        Superior to correlation for nonlinear relationships.
        """
        if task == 'classification':
            mi_scores = mutual_info_classif(X, y, random_state=42)
        else:
            mi_scores = mutual_info_regression(X, y, random_state=42)
        
        return mi_scores
    
    @staticmethod
    def prediction_entropy(probs):
        """
        Compute entropy of predicted probabilities.
        High entropy = uncertain model.
        """
        return entropy(probs, axis=1) / np.log(probs.shape[1])  # Normalized
    
    @staticmethod
    def calibration_analysis(y_true, y_pred_proba):
        """
        Information-theoretic calibration analysis.
        
        Well-calibrated: predicted entropy ≈ actual uncertainty
        """
        # Cross-entropy (bits needed using model's predictions)
        ce = log_loss(y_true, y_pred_proba)
        
        # True entropy (if we knew perfect probabilities)
        # Approximate from empirical frequencies
        class_freq = np.bincount(y_true) / len(y_true)
        true_entropy = entropy(class_freq)
        
        return {
            'cross_entropy': ce,
            'true_entropy': true_entropy,
            'calibration_gap': ce - true_entropy,  # Ideally small
        }
    
    @staticmethod
    def kl_divergence(p, q):
        """
        Compute D_KL(P || Q) for discrete distributions.
        """
        return np.sum(rel_entr(p, q))
    
    @staticmethod
    def vae_kl_monitor(mu, log_var):
        """
        Monitor KL term during VAE training.
        Useful for detecting posterior collapse.
        """
        kl = -0.5 * np.sum(1 + log_var - mu**2 - np.exp(log_var), axis=1)
        return {
            'mean_kl': np.mean(kl),
            'std_kl': np.std(kl),
            'kl_per_dim': np.mean(-0.5 * (1 + log_var - mu**2 - np.exp(log_var)), axis=0),
            'collapsed_dims': np.sum(np.mean(-0.5 * (1 + log_var - mu**2 - np.exp(log_var)), axis=0) < 0.01),
        }
 
# Example usage
toolkit = InformationTheoreticML()
 
# Feature importance
np.random.seed(42)
X = np.random.randn(1000, 10)
y = (X[:, 0]**2 + X[:, 1] > 0.5).astype(int)  # Nonlinear in X0, linear in X1
 
mi_scores = toolkit.feature_importance_mi(X, y)
print("MI-based Feature Importance (nonlinear data)")
print("=" * 50)
for i, score in enumerate(mi_scores):
    importance = "HIGH" if i < 2 else "low"
    print(f"Feature {i}: MI = {score:.4f} ({importance})")
 
print()
 
# Prediction uncertainty
probs = np.array([
    [0.1, 0.1, 0.8],  # Confident
    [0.33, 0.33, 0.34],  # Uncertain
    [0.05, 0.05, 0.9],  # Very confident
])
entropies = toolkit.prediction_entropy(probs)
print("Prediction Entropy Analysis")
print("-" * 50)
for i, ent in enumerate(entropies):
    print(f"Sample {i+1}: Entropy = {ent:.4f} (normalized)")