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
 
def elbo_intuition():
    """
    Understand ELBO decomposition for Bayesian neural networks.
    
    ELBO = E_q[log p(D|W)] - KL[q(W) || p(W)]
    
    = Expected log-likelihood under q
    - Penalty for deviating from prior
    
    This is remarkably similar to:
    
    Standard NN loss = -log p(D|W*) + λ||W*||²
    
    = Negative log-likelihood at point estimate
    + L2 regularization (Gaussian prior)
    """
    
    print("ELBO Decomposition")
    print("=" * 50)
    print()
    print("ELBO = E_q[log p(D|W)] - KL[q(W) || p(W)]")
    print("       |_______________|   |______________|")
    print("       Expected data fit   Regularization")
    print()
    print("Interpretation:")
    print("-" * 50)
    print("1. FIRST TERM: Maximize expected log-likelihood")
    print("   → q(W) should assign probability to weights that fit data")
    print()
    print("2. SECOND TERM: Stay close to prior p(W)")
    print("   → Prevents q(W) from overfitting to training data")
    print("   → Acts as regularization")
    print()
    print("Trade-off: Better fit vs. simpler model")
    print()
    
    # Connection to standard training
    print("Connection to Standard Training:")
    print("-" * 50)
    print("If q(W) = δ(W - W*)  (point estimate)")
    print("And p(W) = N(0, σ²I)  (Gaussian prior)")
    print()
    print("Then ELBO ≈ log p(D|W*) - (1/2σ²)||W*||²")
    print("           = log-likelihood - L2 regularization")
    print()
    print("→ Standard training with L2 regularization is")
    print("  a limiting case of variational inference!")
 
 
def kl_divergence_gaussian(mu_q, std_q, mu_p, std_p):
    """
    KL divergence between two 1D Gaussians.
    
    KL[N(μ_q, σ_q²) || N(μ_p, σ_p²)]
    """
    var_q = std_q ** 2
    var_p = std_p ** 2
    
    kl = (
        np.log(std_p / std_q) +
        (var_q + (mu_q - mu_p)**2) / (2 * var_p) -
        0.5
    )
    return kl
 
 
def demonstrate_kl_prior_effect():
    """Show how KL term encourages staying near prior."""
    print()
    print("KL Divergence from Prior")
    print("=" * 50)
    
    # Prior: N(0, 1)
    mu_p, std_p = 0.0, 1.0
    
    # Various posteriors
    posteriors = [
        ("Close to prior: N(0, 1)", 0.0, 1.0),
        ("Shifted mean: N(2, 1)", 2.0, 1.0),
        ("Narrow: N(0, 0.5)", 0.0, 0.5),
        ("Wide: N(0, 2)", 0.0, 2.0),
        ("Far and narrow: N(3, 0.3)", 3.0, 0.3),
    ]
    
    print(f"Prior: N({mu_p}, {std_p}²)")
    print()
    print(f"{'Posterior':<30} {'KL Divergence':>15}")
    print("-" * 45)
    
    for name, mu_q, std_q in posteriors:
        kl = kl_divergence_gaussian(mu_q, std_q, mu_p, std_p)
        print(f"{name:<30} {kl:>15.4f}")
    
    print()
    print("Key insight: Moving far from prior is penalized")
    print("This prevents overfitting by keeping weights conservative")
 
 
elbo_intuition()
demonstrate_kl_prior_effect()

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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207

import numpy as np
from typing import List, Tuple
 
def theoretical_equivalence():
    """
    Demonstrate the theoretical equivalence between dropout training
    and variational inference.
    """
    
    print("Dropout as Variational Inference")
    print("=" * 60)
    
    print("""
    THEOREM (Gal & Ghahramani, 2016):
    
    Training a neural network with:
    - Dropout probability p on each layer
    - L2 weight decay λ
    - Cross-entropy loss
    
    Is equivalent to:
    
    Variational inference with:
    - Gaussian prior p(W) = N(0, I/λ)
    - Approximate posterior q(W) = W * diag(Bernoulli(1-p))
    - KL divergence regularization
    
    The dropout loss approximates the negative ELBO:
    
    L_dropout ≈ -ELBO = -E_q[log p(D|W)] + KL[q(W)||p(W)]
    """)
    
    print("Key implications:")
    print("-" * 60)
    print("1. Dropout trains a distribution over weights, not a point estimate")
    print("2. Each dropout mask samples from the approximate posterior")
    print("3. The posterior captures weight uncertainty")
    print("4. We can use this for uncertainty estimation!")
 
 
class DropoutVI:
    """
    Dropout interpreted as variational inference.
    
    This class demonstrates the equivalence between:
    - Standard dropout training
    - Variational inference with Bernoulli approximate posterior
    """
    
    def __init__(
        self,
        input_dim: int,
        hidden_dim: int,
        output_dim: int,
        dropout_rate: float = 0.5,
        l2_reg: float = 0.01
    ):
        """
        Initialize network.
        
        Args:
            input_dim: Input feature dimension
            hidden_dim: Hidden layer dimension  
            output_dim: Output dimension
            dropout_rate: Dropout probability (corresponds to Bernoulli param)
            l2_reg: L2 regularization (corresponds to prior precision)
        """
        self.p = dropout_rate
        self.l2_reg = l2_reg
        
        # Prior: p(W) = N(0, 1/l2_reg * I)
        self.prior_precision = l2_reg
        
        # Learnable weight matrices
        # These are the "mean" of our approximate posterior
        self.W1 = np.random.randn(input_dim, hidden_dim) * 0.1
        self.b1 = np.zeros(hidden_dim)
        self.W2 = np.random.randn(hidden_dim, output_dim) * 0.1
        self.b2 = np.zeros(output_dim)
    
    def sample_weights(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Sample weights from approximate posterior q(W).
        
        The approximate posterior is:
        q(W) = W * diag(Bernoulli(1-p))
        
        Each call returns a different sample.
        """
        # Sample Bernoulli masks
        mask1 = np.random.binomial(1, 1 - self.p, size=self.W1.shape[1])
        mask2 = np.random.binomial(1, 1 - self.p, size=self.W2.shape[0])
        
        # Apply masks (this is sampling from q)
        W1_sample = self.W1 * mask1 / (1 - self.p)  # Inverted dropout scaling
        W2_sample = (self.W2 * mask2.reshape(-1, 1)) / (1 - self.p)
        
        return W1_sample, self.b1, W2_sample, self.b2
    
    def forward_sample(self, x: np.ndarray) -> np.ndarray:
        """
        Forward pass with a single weight sample from q(W).
        
        This is what happens during dropout training.
        """
        W1, b1, W2, b2 = self.sample_weights()
        
        h = np.maximum(0, x @ W1 + b1)  # ReLU
        out = h @ W2 + b2
        
        return out
    
    def compute_kl_term(self) -> float:
        """
        Compute KL divergence from prior.
        
        For Bernoulli dropout with rate p and prior N(0, 1/λ):
        
        KL ≈ (λ/2) * (1-p) * ||W||² + const
        
        This is (approximately) the L2 regularization term!
        """
        l2_norm_sq = np.sum(self.W1 ** 2) + np.sum(self.W2 ** 2)
        
        # KL divergence approximation
        kl = 0.5 * self.l2_reg * (1 - self.p) * l2_norm_sq
        
        return kl
    
    def elbo_loss(
        self, 
        x: np.ndarray, 
        y: np.ndarray, 
        num_samples: int = 1
    ) -> float:
        """
        Compute (negative) ELBO loss.
        
        This is what we're actually minimizing during dropout training:
        
        -ELBO = E_q[-log p(y|x,W)] + KL[q(W)||p(W)]
              = Expected loss + KL regularization
        """
        # Monte Carlo estimate of expected log-likelihood
        total_nll = 0.0
        for _ in range(num_samples):
            pred = self.forward_sample(x)
            # Negative log-likelihood (cross-entropy for classification)
            nll = np.mean((pred - y) ** 2)  # MSE for simplicity
            total_nll += nll
        
        expected_nll = total_nll / num_samples
        
        # KL divergence term
        kl = self.compute_kl_term()
        
        # Negative ELBO (what we minimize)
        neg_elbo = expected_nll + kl
        
        return neg_elbo
 
 
def demonstrate_vi_interpretation():
    """Show that dropout loss approximates negative ELBO."""
    np.random.seed(42)
    
    print()
    print("Dropout Loss ≈ Negative ELBO")
    print("=" * 60)
    
    # Create network
    model = DropoutVI(
        input_dim=10,
        hidden_dim=64,
        output_dim=2,
        dropout_rate=0.5,
        l2_reg=0.01
    )
    
    # Sample data
    x = np.random.randn(32, 10)
    y = np.random.randn(32, 2)
    
    # Compute ELBO loss
    elbo_loss = model.elbo_loss(x, y, num_samples=10)
    
    # Standard dropout training loss (1 sample)
    pred = model.forward_sample(x)
    mse = np.mean((pred - y) ** 2)
    l2_term = 0.5 * model.l2_reg * (
        np.sum(model.W1 ** 2) + np.sum(model.W2 ** 2)
    )
    dropout_loss = mse + l2_term
    
    print(f"ELBO loss (10 samples):      {elbo_loss:.4f}")
    print(f"Standard dropout loss:       {dropout_loss:.4f}")
    print()
    print("Components:")
    print(f"  Expected NLL (data fit):   {mse:.4f}")
    print(f"  KL / L2 (regularization):  {l2_term:.4f}")
    print()
    print("✓ Dropout training minimizes an approximation of the negative ELBO")
    print("  This means we're implicitly doing variational inference!")
 
 
theoretical_equivalence()
demonstrate_vi_interpretation()

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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187

import numpy as np
from typing import Tuple
 
class MCDropoutModel:
    """
    Monte Carlo Dropout for uncertainty quantification.
    
    Uses dropout at inference time to approximate the predictive distribution
    of a Bayesian neural network.
    """
    
    def __init__(
        self,
        input_dim: int,
        hidden_dims: list,
        output_dim: int,
        dropout_rate: float = 0.5
    ):
        """Initialize network with multiple hidden layers."""
        self.dropout_rate = dropout_rate
        self.layers = []
        
        dims = [input_dim] + hidden_dims + [output_dim]
        for i in range(len(dims) - 1):
            W = np.random.randn(dims[i], dims[i+1]) * np.sqrt(2.0 / dims[i])
            b = np.zeros(dims[i+1])
            self.layers.append((W, b))
    
    def forward(self, x: np.ndarray, dropout: bool = True) -> np.ndarray:
        """
        Forward pass with optional dropout.
        
        Args:
            x: Input tensor
            dropout: If True, apply dropout (for training or MC inference)
        """
        h = x
        for i, (W, b) in enumerate(self.layers[:-1]):
            h = h @ W + b
            h = np.maximum(0, h)  # ReLU
            
            if dropout:
                mask = np.random.binomial(1, 1 - self.dropout_rate, h.shape)
                h = h * mask / (1 - self.dropout_rate)
        
        # Output layer (no dropout, no activation)
        W, b = self.layers[-1]
        out = h @ W + b
        return out
    
    def mc_predict(
        self, 
        x: np.ndarray, 
        num_samples: int = 100
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Monte Carlo Dropout prediction.
        
        Runs multiple forward passes with dropout to approximate
        the predictive distribution.
        
        Args:
            x: Input tensor of shape (batch, features)
            num_samples: Number of MC samples
        
        Returns:
            mean: Predictive mean (batch, output_dim)
            variance: Predictive variance (batch, output_dim)
        """
        # Collect predictions from multiple forward passes
        predictions = []
        for _ in range(num_samples):
            pred = self.forward(x, dropout=True)
            predictions.append(pred)
        
        predictions = np.stack(predictions, axis=0)  # (samples, batch, output_dim)
        
        # Compute mean and variance
        mean = predictions.mean(axis=0)
        variance = predictions.var(axis=0)
        
        return mean, variance
    
    def standard_predict(self, x: np.ndarray) -> np.ndarray:
        """Standard prediction without dropout (point estimate only)."""
        return self.forward(x, dropout=False)
 
 
def softmax(x: np.ndarray) -> np.ndarray:
    """Compute softmax along last axis."""
    exp_x = np.exp(x - x.max(axis=-1, keepdims=True))
    return exp_x / exp_x.sum(axis=-1, keepdims=True)
 
 
def entropy(probs: np.ndarray) -> np.ndarray:
    """Compute entropy of probability distribution."""
    return -np.sum(probs * np.log(probs + 1e-10), axis=-1)
 
 
def demonstrate_mc_dropout():
    """Demonstrate MC Dropout for uncertainty estimation."""
    np.random.seed(42)
    
    print("Monte Carlo Dropout Demonstration")
    print("=" * 60)
    
    # Create model
    model = MCDropoutModel(
        input_dim=10,
        hidden_dims=[64, 64],
        output_dim=3,  # 3-class classification
        dropout_rate=0.5
    )
    
    # Generate test samples
    # Sample 1: Random data (high uncertainty expected)
    x_random = np.random.randn(1, 10)
    
    # Sample 2: Similar to training-like distribution
    x_training = np.random.randn(1, 10) * 0.5  # Lower variance
    
    # Sample 3: Out-of-distribution (very high uncertainty expected)
    x_ood = np.random.randn(1, 10) * 5.0  # High variance
    
    print("\nComparing predictions with uncertainty:")
    print("-" * 60)
    
    for name, x in [("Random", x_random), 
                     ("Training-like", x_training), 
                     ("Out-of-distribution", x_ood)]:
        # Standard prediction (no uncertainty)
        standard_pred = model.standard_predict(x)
        standard_probs = softmax(standard_pred)
        
        # MC Dropout prediction (with uncertainty)
        mc_mean, mc_var = model.mc_predict(x, num_samples=100)
        mc_probs = softmax(mc_mean)
        
        # Collect all probability predictions for entropy calculation
        mc_preds = []
        for _ in range(100):
            pred = model.forward(x, dropout=True)
            mc_preds.append(softmax(pred))
        mc_preds = np.stack(mc_preds, axis=0)
        
        # Uncertainty measures
        mean_probs = mc_preds.mean(axis=0)
        pred_entropy = entropy(mean_probs)[0]
        
        # Variance of softmax outputs (another uncertainty measure)
        prob_variance = mc_preds.var(axis=0).mean()
        
        print(f"\n{name} sample:")
        print(f"  Standard prediction: {standard_probs[0].round(3)}")
        print(f"  MC mean prediction:  {mc_probs[0].round(3)}")
        print(f"  Predictive entropy:  {pred_entropy:.4f}")
        print(f"  Probability variance: {prob_variance:.4f}")
        print(f"  Interpretation: {'High' if pred_entropy > 0.8 else 'Low'} uncertainty")
    
    # Show how MC sample count affects estimate quality
    print("\n" + "=" * 60)
    print("Effect of MC Sample Count")
    print("-" * 60)
    
    x = np.random.randn(1, 10)
    sample_counts = [1, 5, 10, 50, 100, 500]
    
    print(f"{'Samples':<10} {'Mean variance':<15} {'Var of mean':<15}")
    print("-" * 40)
    
    for n in sample_counts:
        means = []
        for _ in range(20):  # Repeat to estimate variance of estimate
            mean, _ = model.mc_predict(x, num_samples=n)
            means.append(mean[0])
        
        mean_of_means = np.mean(means, axis=0)
        var_of_means = np.var(means, axis=0).mean()
        _, variance = model.mc_predict(x, num_samples=n)
        
        print(f"{n:<10} {variance.mean():<15.4f} {var_of_means:<15.4f}")
    
    print("\n✓ More samples → more stable estimates")
    print("  But diminishing returns after ~50-100 samples")
 
 
demonstrate_mc_dropout()

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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple, Optional
 
class MCDropoutNet(nn.Module):
    """
    Neural network with MC Dropout support.
    
    Key design decisions:
    1. Dropout layers can be enabled during evaluation
    2. Efficient batched MC sampling
    3. Multiple uncertainty metrics computed together
    """
    
    def __init__(
        self,
        input_dim: int,
        hidden_dims: list,
        output_dim: int,
        dropout_rate: float = 0.5
    ):
        super().__init__()
        
        self.dropout_rate = dropout_rate
        
        # Build layers
        layers = []
        dims = [input_dim] + hidden_dims
        for i in range(len(dims) - 1):
            layers.append(nn.Linear(dims[i], dims[i+1]))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout_rate))
        
        # Output layer (no dropout)
        layers.append(nn.Linear(dims[-1], output_dim))
        
        self.network = nn.Sequential(*layers)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.network(x)
    
    def enable_mc_dropout(self):
        """Enable dropout during evaluation for MC sampling."""
        for module in self.modules():
            if isinstance(module, nn.Dropout):
                module.train()
    
    def disable_mc_dropout(self):
        """Disable dropout for standard evaluation."""
        for module in self.modules():
            if isinstance(module, nn.Dropout):
                module.eval()
    
    def mc_forward(
        self, 
        x: torch.Tensor, 
        num_samples: int = 100
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Monte Carlo forward pass.
        
        Args:
            x: Input tensor (batch, features)
            num_samples: Number of MC samples
        
        Returns:
            mean: Predictive mean (batch, output_dim)
            variance: Predictive variance (batch, output_dim)
            all_preds: All predictions (samples, batch, output_dim)
        """
        self.enable_mc_dropout()
        
        # Efficient: expand input and run in one forward pass
        batch_size = x.size(0)
        
        # Stack input num_samples times
        x_expanded = x.unsqueeze(0).expand(num_samples, -1, -1)
        x_expanded = x_expanded.reshape(-1, x.size(-1))
        
        # Single forward pass (more efficient than loop)
        with torch.no_grad():
            all_preds = self(x_expanded)
        
        # Reshape back
        all_preds = all_preds.reshape(num_samples, batch_size, -1)
        
        # Compute statistics
        mean = all_preds.mean(dim=0)
        variance = all_preds.var(dim=0)
        
        return mean, variance, all_preds
 
 
class UncertaintyMetrics:
    """Compute various uncertainty metrics from MC samples."""
    
    @staticmethod
    def predictive_entropy(mc_probs: torch.Tensor) -> torch.Tensor:
        """
        Entropy of the mean predictive distribution.
        
        Args:
            mc_probs: Softmax probabilities (samples, batch, classes)
        
        Returns:
            Entropy per sample (batch,)
        """
        mean_probs = mc_probs.mean(dim=0)  # Average over MC samples
        entropy = -torch.sum(mean_probs * torch.log(mean_probs + 1e-10), dim=-1)
        return entropy
    
    @staticmethod
    def expected_entropy(mc_probs: torch.Tensor) -> torch.Tensor:
        """
        Expected entropy (average entropy of individual predictions).
        
        Args:
            mc_probs: Softmax probabilities (samples, batch, classes)
        
        Returns:
            Expected entropy per sample (batch,)
        """
        # Entropy of each MC sample
        entropies = -torch.sum(mc_probs * torch.log(mc_probs + 1e-10), dim=-1)
        # Average over MC samples
        return entropies.mean(dim=0)
    
    @staticmethod
    def mutual_information(mc_probs: torch.Tensor) -> torch.Tensor:
        """
        Mutual information: I(y; W | x, D)
        
        This captures epistemic uncertainty (model uncertainty).
        High MI = model is uncertain about which weights to use.
        
        MI = H(E[p]) - E[H(p)] = predictive_entropy - expected_entropy
        """
        pred_entropy = UncertaintyMetrics.predictive_entropy(mc_probs)
        exp_entropy = UncertaintyMetrics.expected_entropy(mc_probs)
        return pred_entropy - exp_entropy
    
    @staticmethod
    def variation_ratio(mc_probs: torch.Tensor) -> torch.Tensor:
        """
        Variation ratio: fraction of MC samples that don't predict the mode.
        
        Simple, interpretable uncertainty measure.
        0 = all samples agree, 1 = maximum disagreement.
        """
        # Get predicted class for each MC sample
        predicted_classes = mc_probs.argmax(dim=-1)  # (samples, batch)
        
        # Find mode (most common prediction)
        modes, _ = predicted_classes.mode(dim=0)  # (batch,)
        
        # Fraction that disagree with mode
        disagreement = (predicted_classes != modes.unsqueeze(0)).float()
        variation_ratio = disagreement.mean(dim=0)
        
        return variation_ratio
 
 
def demonstrate_pytorch_mc_dropout():
    """Demonstrate MC Dropout in PyTorch."""
    torch.manual_seed(42)
    
    print("PyTorch MC Dropout Implementation")
    print("=" * 60)
    
    # Create model
    model = MCDropoutNet(
        input_dim=10,
        hidden_dims=[64, 32],
        output_dim=5,  # 5-class classification
        dropout_rate=0.5
    )
    
    # Test input
    x = torch.randn(8, 10)
    
    # Standard prediction
    model.eval()
    with torch.no_grad():
        standard_pred = model(x)
        standard_probs = F.softmax(standard_pred, dim=-1)
    
    # MC Dropout prediction
    mean, variance, all_preds = model.mc_forward(x, num_samples=100)
    
    # Compute softmax for all predictions
    all_probs = F.softmax(all_preds, dim=-1)
    
    # Compute uncertainty metrics
    pred_entropy = UncertaintyMetrics.predictive_entropy(all_probs)
    mi = UncertaintyMetrics.mutual_information(all_probs)
    var_ratio = UncertaintyMetrics.variation_ratio(all_probs)
    
    print(f"\nBatch of {x.size(0)} samples, {model.network[-1].out_features} classes")
    print("-" * 60)
    print(f"\nPredictive entropy (uncertainty from output distribution):")
    print(f"  {pred_entropy.numpy().round(3)}")
    print(f"\nMutual information (epistemic uncertainty):")
    print(f"  {mi.numpy().round(3)}")
    print(f"\nVariation ratio (fraction of disagreeing samples):")
    print(f"  {var_ratio.numpy().round(3)}")
    
    # Flag high-uncertainty predictions
    print(f"\nHigh-uncertainty samples (MI > 0.5):")
    high_uncertainty = (mi > 0.5).nonzero().squeeze().tolist()
    if isinstance(high_uncertainty, int):
        high_uncertainty = [high_uncertainty]
    print(f"  Indices: {high_uncertainty if high_uncertainty else 'None'}")
 
 
# Run demonstration
if __name__ == "__main__":
    demonstrate_pytorch_mc_dropout()