Information Theory for Machine Learning

Information theory, introduced by Claude Shannon in 1948, provides the mathematical tools to quantify uncertainty, information, and the fundamental limits of learning. This module develops the theory and its deep connections to machine learning.

1. Entropy and Information

1.1 Shannon Entropy

For a discrete random variable with probability mass function :

where the base of the logarithm determines the units (base 2: bits, base : nats).

Properties:

• with equality iff is deterministic
• with equality iff is uniform
• Concave:

1.2 Differential Entropy

For a continuous random variable with density :

Key difference from discrete entropy: can be negative. For example, if :

which is negative when .

1.3 Joint and Conditional Entropy

Chain Rule:

2. KL Divergence

2.1 Definition

The Kullback-Leibler divergence between distributions and is:

For continuous distributions:

2.2 Fundamental Properties

Theorem (Gibbs' Inequality): with equality iff almost everywhere.

Proof sketch: By Jensen's inequality applied to the convex function :

Non-symmetry: in general. This asymmetry has profound implications.

Data Processing Inequality: If forms a Markov chain, then:

2.3 KL Divergence of Exponential Families

For exponential family :

This connection to the log-partition function reveals the geometry of exponential families.

3. Mutual Information

3.1 Definition

Interpretation: measures how much knowing reduces uncertainty about (and vice versa).

3.2 Properties

• with equality iff
• (symmetric)
• Data processing inequality

3.3 Variational Lower Bounds on MI

For neural estimation of mutual information:

where is any class of functions. This is the Donsker-Varadhan representation.

MINE (Belghazi et al., 2018): Use a neural network to parameterize the critic:

Optimize via the moving average trick for the negative term to reduce variance.

4. Fisher Information

4.1 Definition

For a parametric family , the Fisher information matrix is:

4.2 Properties

Alternative form: Under regularity conditions:

Chain rule: for independent .

Data processing inequality: If is Markov, then .

4.3 Fisher Information and Exponential Families

For exponential family :

The Fisher information equals the Hessian of the log-partition function, connecting information geometry to convex analysis.

5. Cramér-Rao Bound

5.1 Statement

Theorem (Cramér-Rao): For any unbiased estimator of :

That is, the covariance matrix of any unbiased estimator dominates the inverse Fisher information.

5.2 Multivariate Case

For vector-valued parameters:

for each component . Equality holds iff achieves the minimum variance unbiased bound (MVUB).

5.3 Asymptotic Efficiency

Under regularity conditions, the maximum likelihood estimator achieves the Cramér-Rao bound asymptotically:

This is the asymptotic normality of MLE and makes it asymptotically efficient.

6. Information Bottleneck

6.1 Setup

Given joint distribution , find a compressed representation of that retains maximal information about :

where controls the trade-off between compression and prediction.

6.2 Optimal Solution

Theorem (Tishby et al., 1999): The optimal satisfies:

where is a normalization constant.

6.3 Connection to Deep Learning

The information bottleneck principle suggests that deep networks learn to:

1. Extract relevant information about from (compression)
2. Discard irrelevant information in (generalization)

The information plane visualization shows layers organizing along two axes: (complexity) and (relevance).

7. Rate-Distortion Theory

7.1 Definition

Given a distortion measure and rate constraint , the rate-distortion function is:

7.2 Gaussian Source

For and squared distortion:

where .

7.3 Connection to Neural Compression

Neural networks for lossy compression implicitly optimize a Lagrangian version:

where is the Lagrange multiplier trading off distortion for rate.

8. Channel Coding Theorem

8.1 Noisy Channel Coding

A discrete memoryless channel (DMC) has transition probabilities . The channel capacity is:

8.2 Shannon's Theorem

Theorem (Shannon, 1948): For any rate , there exist codes with arbitrarily low error probability. For , the error probability is bounded away from zero.

8.3 Gaussian Channel

For with and power constraint :

9. Sufficient Statistics

9.1 Definition

is a sufficient statistic for if:

That is, the conditional distribution of given does not depend on .

9.2 Factorization Theorem

Theorem (Neyman-Fisher): is sufficient for iff:

for some functions and .

9.3 Minimal Sufficient Statistics

A sufficient statistic is minimal if it is a function of every other sufficient statistic. For exponential families, is minimal sufficient.

10. Connections to Machine Learning

10.1 Cross-Entropy Loss

The cross-entropy loss used in classification:

Minimizing cross-entropy = minimizing KL divergence to the true distribution.

10.2 Variational Autoencoders

The VAE objective (ELBO):

This decomposes into reconstruction (cross-entropy) and regularization (KL divergence).

10.3 Information-Theoretic Regularization

InfoNCE: where is the number of negative samples.

This bounds mutual information from below, enabling contrastive learning.

11. Code: Information-Theoretic Measures

import numpy as np
from scipy.special import digamma, gammaln
from sklearn.neighbors import KernelDensity

def entropy_discrete(p):
    """Compute Shannon entropy of discrete distribution."""
    p = np.asarray(p)
    p = p[p > 0]  # Remove zeros
    return -np.sum(p * np.log(p))


def kl_divergence(p, q):
    """Compute KL divergence from discrete distributions p to q."""
    p = np.asarray(p)
    q = np.asarray(q)
    
    # Filter out zero entries
    mask = (p > 0) & (q > 0)
    p = p[mask]
    q = q[mask]
    
    return np.sum(p * np.log(p / q))


def mutual_information_discrete(p_xy):
    """
    Compute mutual information from joint distribution.
    
    Args:
        p_xy: Joint distribution matrix of shape (num_x, num_y)
    """
    p_x = p_xy.sum(axis=1)
    p_y = p_xy.sum(axis=0)
    
    mi = 0.0
    for i in range(len(p_x)):
        for j in range(len(p_y)):
            if p_xy[i, j] > 0:
                mi += p_xy[i, j] * np.log(p_xy[i, j] / (p_x[i] * p_y[j]))
    
    return mi


def fisher_information_gaussian(mu, sigma2, n_samples=10000):
    """
    Estimate Fisher information for Gaussian N(mu, sigma^2).
    
    For Gaussian: I(mu, sigma^2) = diag(1/sigma^2, 1/(2*sigma^4))
    """
    # Analytical Fisher information
    I_analytical = np.array([
        [1/sigma2, 0],
        [0, 1/(2*sigma2**2)]
    ])
    
    # Monte Carlo estimation
    x = np.random.normal(mu, np.sqrt(sigma2), n_samples)
    
    # Score functions
    score_mu = (x - mu) / sigma2
    score_sigma2 = -0.5/sigma2 + (x - mu)**2 / (2*sigma2**2)
    
    # Fisher information as variance of score
    I_mc = np.array([
        [np.var(score_mu), np.cov(score_mu, score_sigma2)[0, 1]],
        [np.cov(score_mu, score_sigma2)[1, 0], np.var(score_sigma2)]
    ])
    
    return I_analytical, I_mc


def cramers_rao_bound(fisher_info):
    """Compute Cramér-Rao lower bound."""
    return np.linalg.inv(fisher_info)


def information_bottleneck(X, Y, beta, n_clusters=10, max_iter=100):
    """
    Simplified information bottleneck via iterative clustering.
    
    Args:
        X: Input data
        Y: Target data
        beta: Trade-off parameter
        n_clusters: Number of clusters for T
        max_iter: Maximum iterations
    """
    n = len(X)
    
    # Initialize: random clustering
    T = np.random.randint(0, n_clusters, size=n)
    
    for iteration in range(max_iter):
        T_old = T.copy()
        
        # E-step: compute p(y|t)
        p_y_given_t = np.zeros((n_clusters, len(np.unique(Y))))
        for t in range(n_clusters):
            mask = T == t
            if mask.sum() > 0:
                for y_val in np.unique(Y):
                    p_y_given_t[t, int(y_val)] = (mask & (Y == y_val)).sum() / mask.sum()
        
        # M-step: assign each x to t minimizing KL(p(y|x) || p(y|t))
        p_y_given_x = np.zeros((n, len(np.unique(Y))))
        for i, x_val in enumerate(np.unique(X)):
            mask = X == x_val
            for y_val in np.unique(Y):
                p_y_given_x[mask, int(y_val)] = (mask & (Y == y_val)).sum() / mask.sum()
        
        for i in range(n):
            kl_vals = np.zeros(n_clusters)
            for t in range(n_clusters):
                if p_y_given_t[t].sum() > 0:
                    kl_vals[t] = kl_divergence(p_y_given_x[i], p_y_given_t[t])
            
            T[i] = np.argmin(kl_vals)
        
        if np.all(T == T_old):
            break
    
    return T


def mine_mutual_information(x, y, hidden_dim=64, epochs=100, lr=1e-3):
    """
    Estimate mutual information using MINE (Neural Information Estimation).
    
    Args:
        x, y: Data arrays of shape (n, dx) and (n, dy)
        hidden_dim: Hidden dimension of critic network
        epochs: Training epochs
        lr: Learning rate
    """
    import torch
    import torch.nn as nn
    import torch.optim as optim
    
    n = len(x)
    x_tensor = torch.FloatTensor(x)
    y_tensor = torch.FloatTensor(y)
    
    # Critic network
    class Critic(nn.Module):
        def __init__(self):
            super().__init__()
            self.net = nn.Sequential(
                nn.Linear(x.shape[1] + y.shape[1], hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, hidden_dim),
                nn.ReLU(),
                nn.Linear(hidden_dim, 1)
            )
        
        def forward(self, x, y):
            return self.net(torch.cat([x, y], dim=1))
    
    critic = Critic()
    optimizer = optim.Adam(critic.parameters(), lr=lr)
    
    # Running mean for variance reduction
    ma_rate = 0.01
    
    mi_estimates = []
    
    for epoch in range(epochs):
        # Shuffle for negative samples
        perm = torch.randperm(n)
        
        # Positive pairs
        T_xy = critic(x_tensor, y_tensor)
        
        # Negative pairs (marginals)
        T_x_y = critic(x_tensor, y_tensor[perm])
        
        # MINE bound
        mi = torch.mean(T_xy) - torch.log(torch.mean(torch.exp(T_x_y)))
        
        # Optimize
        optimizer.zero_grad()
        (-mi).backward()  # Maximize MI = minimize -MI
        optimizer.step()
        
        mi_estimates.append(mi.item())
    
    return np.mean(mi_estimates[-10:])  # Average over last epochs


def information_plane_analysis(model, X, Y, n_layers=5):
    """
    Compute information plane coordinates for each layer.
    
    Args:
        model: Neural network model
        X, Y: Data
        n_layers: Number of layers to analyze
    """
    import torch
    
    layer_mi_x = []  # I(T; X) for each layer
    layer_mi_y = []  # I(T; Y) for each layer
    
    # Get activations from each layer
    activations = []
    hooks = []
    
    def hook_fn(module, input, output):
        activations.append(output.detach().numpy())
    
    # Register hooks
    for i, layer in enumerate(model.layers[:n_layers]):
        hooks.append(layer.register_forward_hook(hook_fn))
    
    # Forward pass
    with torch.no_grad():
        model(torch.FloatTensor(X))
    
    # Remove hooks
    for hook in hooks:
        hook.remove()
    
    # Estimate MI for each layer
    for act in activations:
        # Simplified: use k-nearest neighbor MI estimation
        mi_x = _estimate_mi_knn(X, act, k=5)
        mi_y = _estimate_mi_knn(Y, act, k=5)
        layer_mi_x.append(mi_x)
        layer_mi_y.append(mi_y)
    
    return layer_mi_x, layer_mi_y


def _estimate_mi_knn(x, y, k=5):
    """Estimate MI using k-nearest neighbors (KSG estimator)."""
    from scipy.spatial import cKDTree
    
    n = len(x)
    d_x = x.shape[1] if x.ndim > 1 else 1
    d_y = y.shape[1] if y.ndim > 1 else 1
    
    # Concatenate
    xy = np.column_stack([x, y]) if x.ndim > 1 else np.column_stack([x, y])
    
    # Build tree
    tree = cKDTree(xy)
    
    # Find k-th nearest neighbor distances
    distances, _ = tree.query(xy, k=k+1)
    eps = distances[:, k]  # k-th neighbor distance
    
    # Count neighbors in marginal spaces
    tree_x = cKDTree(x.reshape(-1, 1) if x.ndim == 1 else x)
    tree_y = cKDTree(y.reshape(-1, 1) if y.ndim == 1 else y)
    
    nx = np.array([len(tree_x.query_ball_point(xi, eps[i])) for i, xi in enumerate(x.reshape(-1, 1) if x.ndim == 1 else x)])
    ny = np.array([len(tree_y.query_ball_point(yi, eps[i])) for i, yi in enumerate(y.reshape(-1, 1) if y.ndim == 1 else y)])
    
    # KSG estimator
    mi = digamma(k) - np.mean(digamma(nx) + digamma(ny)) + digamma(n)
    
    return mi

12. Summary

Information theory provides the language and mathematics for reasoning about learning:

1. Entropy quantifies uncertainty; KL divergence measures distributional distance
2. Mutual information captures statistical dependencies; Fisher information measures parameter sensitivity
3. Cramér-Rao bound sets fundamental limits on estimation accuracy
4. Information bottleneck provides a principled framework for representation learning
5. Channel coding establishes fundamental limits on communication and learning

These concepts unify diverse ML algorithms: VAEs optimize an information-theoretic objective, contrastive learning bounds mutual information, and information bottleneck theory explains deep network generalization.