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

import numpy as np
import torch
import torch.nn.functional as F
 
def softmax_naive(z):
    """
    Naive softmax implementation.
    WARNING: Numerically unstable for large values!
    """
    exp_z = np.exp(z)
    return exp_z / np.sum(exp_z)
 
def softmax_stable(z):
    """
    Numerically stable softmax using translation invariance.
    
    Key insight: softmax(z - max(z)) = softmax(z)
    But subtracting max prevents exp overflow.
    """
    z_shifted = z - np.max(z)
    exp_z = np.exp(z_shifted)
    return exp_z / np.sum(exp_z)
 
def softmax_with_temperature(z, temperature=1.0):
    """
    Temperature-scaled softmax.
    
    temperature > 1: Softer distribution (more uniform)
    temperature < 1: Sharper distribution (more peaked)
    temperature → 0:  Approaches argmax (one-hot)
    temperature → ∞:  Approaches uniform distribution
    """
    z_scaled = z / temperature
    return softmax_stable(z_scaled)
 
# Demonstrate instability
print("=== Numerical Stability Demo ===")
z_normal = np.array([1.0, 2.0, 3.0])
z_large = np.array([1000.0, 2000.0, 3000.0])
 
print(f"Normal logits:  {softmax_naive(z_normal)}")
print(f"Large logits (naive): {softmax_naive(z_large)}")  # Will show nan!
print(f"Large logits (stable): {softmax_stable(z_large)}")
 
# Temperature scaling demonstration
print("\n=== Temperature Scaling Demo ===")
logits = np.array([2.0, 1.0, 0.5, 0.0, -1.0])
for temp in [0.1, 0.5, 1.0, 2.0, 10.0]:
    probs = softmax_with_temperature(logits, temp)
    entropy = -np.sum(probs * np.log(probs + 1e-10))
    print(f"T={temp:4.1f}: {probs.round(3)}, entropy={entropy:.2f}")
 
# PyTorch comparison
z_torch = torch.tensor([[-1.0, 0.0, 1.0, 2.0]])
print(f"\nPyTorch softmax: {F.softmax(z_torch, dim=-1)}")
print(f"PyTorch log_softmax: {F.log_softmax(z_torch, dim=-1)}")

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

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
 
def cross_entropy_manual(logits, targets):
    """
    Manual cross-entropy computation.
    Shows the math but is NOT numerically optimal.
    
    Args:
        logits: [batch_size, num_classes]
        targets: [batch_size] (class indices)
    """
    # Stable softmax via log-sum-exp
    log_sum_exp = torch.logsumexp(logits, dim=-1)
    log_probs = logits - log_sum_exp.unsqueeze(-1)
    
    # Select log probability of true class
    batch_size = logits.size(0)
    true_log_probs = log_probs[torch.arange(batch_size), targets]
    
    return -true_log_probs.mean()
 
 
def cross_entropy_with_one_hot(logits, one_hot_targets):
    """
    Cross-entropy using one-hot encoded targets.
    Equivalent to the formula: -sum(y * log(p))
    """
    log_probs = F.log_softmax(logits, dim=-1)
    return -(one_hot_targets * log_probs).sum(dim=-1).mean()
 
 
# Create sample data
batch_size = 4
num_classes = 5
logits = torch.randn(batch_size, num_classes)
targets = torch.randint(0, num_classes, (batch_size,))
 
# One-hot encoding
one_hot = F.one_hot(targets, num_classes).float()
 
# Compare implementations
print("=== Cross-Entropy Implementations ===")
print(f"Manual:           {cross_entropy_manual(logits, targets):.6f}")
print(f"One-hot:          {cross_entropy_with_one_hot(logits, one_hot):.6f}")
print(f"PyTorch CE:       {F.cross_entropy(logits, targets):.6f}")
print(f"PyTorch NLL+LSM:  {F.nll_loss(F.log_softmax(logits, dim=-1), targets):.6f}")
 
# Gradient verification
logits_grad = logits.clone().requires_grad_(True)
loss = F.cross_entropy(logits_grad, targets)
loss.backward()
 
probs = F.softmax(logits, dim=-1)
expected_grad = (probs - one_hot) / batch_size
 
print(f"\n=== Gradient Verification ===")
print(f"Computed gradient (first sample):\n{logits_grad.grad[0]}")
print(f"Expected (p - y):\n{expected_grad[0]}")
print(f"Match: {torch.allclose(logits_grad.grad, expected_grad, atol=1e-5)}")

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

import numpy as np
import torch
 
def logsumexp_naive(z):
    """Naive log-sum-exp. Overflows for large values!"""
    return np.log(np.sum(np.exp(z)))
 
def logsumexp_stable(z):
    """
    Stable log-sum-exp using the max trick.
    
    log(sum(exp(z))) = m + log(sum(exp(z - m)))
    
    This works because:
    - exp(z - m) <= exp(0) = 1 for all z (since m = max(z))
    - At least one term equals 1 (when z_j = m)
    - No overflow possible
    """
    m = np.max(z)
    return m + np.log(np.sum(np.exp(z - m)))
 
def log_softmax_stable(z):
    """Stable log-softmax using LSE."""
    return z - logsumexp_stable(z)
 
def cross_entropy_stable(logits, target_idx):
    """
    Stable cross-entropy from scratch.
    
    CE = -log(p_target) = -log(softmax(z)_target)
       = -(z_target - LSE(z))
       = LSE(z) - z_target
    """
    lse = logsumexp_stable(logits)
    return lse - logits[target_idx]
 
# Demonstrate stability
print("=== Log-Sum-Exp Stability ===")
 
# Small values: both work
z_small = np.array([1.0, 2.0, 3.0])
print(f"Small logits: {z_small}")
print(f"  Naive LSE:  {logsumexp_naive(z_small):.6f}")
print(f"  Stable LSE: {logsumexp_stable(z_small):.6f}")
 
# Large values: naive fails
z_large = np.array([1000.0, 1001.0, 1002.0])
print(f"\nLarge logits: {z_large}")
print(f"  Naive LSE:  {logsumexp_naive(z_large)}")  # inf or overflow
print(f"  Stable LSE: {logsumexp_stable(z_large):.6f}")
 
# Very negative values
z_neg = np.array([-1000.0, -1001.0, -1002.0])
print(f"\nNegative logits: {z_neg}")
print(f"  Stable LSE: {logsumexp_stable(z_neg):.6f}")
 
# PyTorch built-in
z_torch = torch.tensor([1000.0, 1001.0, 1002.0])
print(f"\nPyTorch logsumexp: {torch.logsumexp(z_torch, dim=0):.6f}")
 
# Cross-entropy example
logits = np.array([2.0, 1.0, 0.5, 0.0])
target = 0  # True class
ce = cross_entropy_stable(logits, target)
print(f"\nCross-entropy for target={target}: {ce:.6f}")
print(f"Probability of true class: {np.exp(-ce):.6f}")

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

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from scipy.optimize import minimize_scalar
 
class TemperatureScaledModel(nn.Module):
    """
    Wrapper that applies learned temperature scaling to a trained model.
    """
    def __init__(self, base_model, initial_temperature=1.0):
        super().__init__()
        self.base_model = base_model
        # Temperature as a learnable parameter (log space for positivity)
        self.log_temperature = nn.Parameter(
            torch.log(torch.tensor(initial_temperature))
        )
    
    @property
    def temperature(self):
        return torch.exp(self.log_temperature)
    
    def forward(self, x):
        logits = self.base_model(x)
        return logits / self.temperature
    
    def calibrate(self, val_loader, max_iters=100, lr=0.01):
        """
        Learn optimal temperature on validation data.
        Only temperature is updated; base model is frozen.
        """
        # Freeze base model
        for param in self.base_model.parameters():
            param.requires_grad = False
        
        optimizer = torch.optim.LBFGS([self.log_temperature], lr=lr, max_iter=max_iters)
        criterion = nn.CrossEntropyLoss()
        
        # Collect all validation data
        all_logits = []
        all_labels = []
        
        with torch.no_grad():
            for x, y in val_loader:
                logits = self.base_model(x)
                all_logits.append(logits)
                all_labels.append(y)
        
        all_logits = torch.cat(all_logits)
        all_labels = torch.cat(all_labels)
        
        def closure():
            optimizer.zero_grad()
            scaled_logits = all_logits / self.temperature
            loss = criterion(scaled_logits, all_labels)
            loss.backward()
            return loss
        
        optimizer.step(closure)
        
        print(f"Learned temperature: {self.temperature.item():.4f}")
        return self.temperature.item()
 
 
def expected_calibration_error(probs, labels, n_bins=15):
    """
    Compute Expected Calibration Error (ECE).
    
    ECE measures the average difference between confidence and
    accuracy across confidence bins. Lower is better.
    """
    confidences = probs.max(dim=-1).values
    predictions = probs.argmax(dim=-1)
    accuracies = (predictions == labels).float()
    
    bin_boundaries = torch.linspace(0, 1, n_bins + 1)
    ece = torch.tensor(0.0)
    
    for i in range(n_bins):
        in_bin = (confidences > bin_boundaries[i]) &                  (confidences <= bin_boundaries[i + 1])
        
        prop_in_bin = in_bin.float().mean()
        if prop_in_bin > 0:
            avg_confidence = confidences[in_bin].mean()
            avg_accuracy = accuracies[in_bin].mean()
            ece += torch.abs(avg_confidence - avg_accuracy) * prop_in_bin
    
    return ece.item()
 
 
# Demonstration with synthetic data
print("=== Temperature Scaling Demo ===")
 
# Simulate overconfident predictions
# Logits that give ~98% confidence predictions
np.random.seed(42)
n_samples = 1000
n_classes = 10
 
# True labels
true_labels = torch.randint(0, n_classes, (n_samples,))
 
# Overconfident logits: correct class gets high logit
logits = torch.randn(n_samples, n_classes) * 0.5
logits[torch.arange(n_samples), true_labels] += 4.0  # Make correct class dominant
 
# But some are wrong
n_wrong = 100
wrong_idx = torch.randperm(n_samples)[:n_wrong]
wrong_labels = (true_labels[wrong_idx] + torch.randint(1, n_classes, (n_wrong,))) % n_classes
logits[wrong_idx, wrong_labels] = logits[wrong_idx, true_labels[wrong_idx]] + 1
 
# Before temperature scaling
probs_before = F.softmax(logits, dim=-1)
ece_before = expected_calibration_error(probs_before, true_labels)
accuracy = (probs_before.argmax(dim=-1) == true_labels).float().mean()
 
print(f"Accuracy: {accuracy:.4f}")
print(f"ECE before scaling: {ece_before:.4f}")
print(f"Mean confidence: {probs_before.max(dim=-1).values.mean():.4f}")
 
# Find optimal temperature
def nll_at_temp(temp):
    scaled_probs = F.softmax(logits / temp, dim=-1)
    log_probs = torch.log(scaled_probs + 1e-10)
    nll = -log_probs[torch.arange(n_samples), true_labels].mean()
    return nll.item()
 
result = minimize_scalar(nll_at_temp, bounds=(0.1, 10.0), method='bounded')
optimal_temp = result.x
 
# After temperature scaling
probs_after = F.softmax(logits / optimal_temp, dim=-1)
ece_after = expected_calibration_error(probs_after, true_labels)
 
print(f"\nOptimal temperature: {optimal_temp:.4f}")
print(f"ECE after scaling: {ece_after:.4f}")
print(f"Mean confidence after: {probs_after.max(dim=-1).values.mean():.4f}")

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

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class LabelSmoothingCrossEntropy(nn.Module):
    """
    Cross-entropy loss with label smoothing.
    
    Instead of one-hot targets [0, 0, 1, 0, 0], uses
    soft targets [ε/K, ε/K, 1-ε+ε/K, ε/K, ε/K]
    """
    def __init__(self, smoothing=0.1, reduction='mean'):
        super().__init__()
        self.smoothing = smoothing
        self.reduction = reduction
    
    def forward(self, logits, targets):
        n_classes = logits.size(-1)
        
        # Compute log-softmax
        log_probs = F.log_softmax(logits, dim=-1)
        
        # One-hot encoding of targets
        targets_one_hot = F.one_hot(targets, n_classes).float()
        
        # Smooth the labels
        smooth_targets = targets_one_hot * (1 - self.smoothing) +                          self.smoothing / n_classes
        
        # Cross-entropy with smooth targets
        loss = -(smooth_targets * log_probs).sum(dim=-1)
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss
 
 
class LabelSmoothingCrossEntropyEfficient(nn.Module):
    """
    Efficient version that avoids explicit one-hot construction.
    
    The smoothed loss decomposes into two terms:
    1. (1 - ε) * CE with true label
    2. ε * mean of all log-probs (KL with uniform)
    """
    def __init__(self, smoothing=0.1, reduction='mean'):
        super().__init__()
        self.smoothing = smoothing
        self.reduction = reduction
    
    def forward(self, logits, targets):
        log_probs = F.log_softmax(logits, dim=-1)
        
        # Term 1: CE with true class
        nll_loss = F.nll_loss(log_probs, targets, reduction='none')
        
        # Term 2: Mean log-prob (negative entropy encouragement)
        smooth_loss = -log_probs.mean(dim=-1)
        
        # Combine
        loss = (1 - self.smoothing) * nll_loss + self.smoothing * smooth_loss
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss
 
 
# Compare with PyTorch built-in
print("=== Label Smoothing Comparison ===")
 
batch_size = 32
n_classes = 10
logits = torch.randn(batch_size, n_classes)
targets = torch.randint(0, n_classes, (batch_size,))
 
# Our implementations
ce_smooth_v1 = LabelSmoothingCrossEntropy(smoothing=0.1)
ce_smooth_v2 = LabelSmoothingCrossEntropyEfficient(smoothing=0.1)
 
# PyTorch built-in (PyTorch 1.10+)
ce_pytorch = nn.CrossEntropyLoss(label_smoothing=0.1)
 
# Standard CE for comparison
ce_standard = nn.CrossEntropyLoss()
 
print(f"Standard CE:               {ce_standard(logits, targets):.6f}")
print(f"Label Smoothing (v1):      {ce_smooth_v1(logits, targets):.6f}")
print(f"Label Smoothing (v2):      {ce_smooth_v2(logits, targets):.6f}")
print(f"PyTorch Label Smoothing:   {ce_pytorch(logits, targets):.6f}")
 
# Effect on training dynamics
print("\n=== Effect on Confidence ===")
# Model with high confidence
confident_logits = torch.tensor([[10.0, 0.0, 0.0, 0.0, 0.0]])
target = torch.tensor([0])  # Correct
 
print(f"Predicted probs: {F.softmax(confident_logits, dim=-1)}")
print(f"Standard CE loss: {ce_standard(confident_logits, target):.6f}")
print(f"Smoothed CE loss: {ce_pytorch(confident_logits, target):.6f}")
print("Note: Smoothed loss is higher because it penalizes overconfidence")

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

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple
 
class MultiClassClassifier(nn.Module):
    """
    Production-ready multi-class classification model.
    
    Design principles:
    1. Output K logits (no softmax during training)
    2. Use CrossEntropyLoss for training
    3. Temperature scaling for calibration (optional)
    4. Label smoothing for regularization (optional)
    """
    def __init__(
        self,
        input_dim: int,
        num_classes: int,
        hidden_dims: list = [256, 128],
        dropout: float = 0.1,
        use_batch_norm: bool = True,
    ):
        super().__init__()
        
        self.num_classes = num_classes
        
        # Build feature extractor
        layers = []
        prev_dim = input_dim
        
        for hidden_dim in hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            if use_batch_norm:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(dropout))
            prev_dim = hidden_dim
        
        self.features = nn.Sequential(*layers)
        
        # Classification head: K output units, NO activation
        self.classifier = nn.Linear(prev_dim, num_classes)
        
        # Temperature for calibration (learned post-training)
        self.register_buffer('temperature', torch.tensor(1.0))
    
    def forward(self, x) -> torch.Tensor:
        """
        Returns raw logits (NOT probabilities).
        Shape: [batch_size, num_classes]
        """
        features = self.features(x)
        logits = self.classifier(features)
        return logits
    
    def predict_proba(
        self, 
        x, 
        use_temperature: bool = True
    ) -> torch.Tensor:
        """
        Returns class probabilities.
        """
        with torch.no_grad():
            logits = self.forward(x)
            if use_temperature:
                logits = logits / self.temperature
            return F.softmax(logits, dim=-1)
    
    def predict(self, x) -> torch.Tensor:
        """Returns class predictions."""
        probs = self.predict_proba(x)
        return probs.argmax(dim=-1)
    
    def set_temperature(self, temperature: float):
        """Set calibration temperature."""
        self.temperature.fill_(temperature)
 
 
class MultiClassTrainer:
    """
    Trainer with full multi-class best practices.
    """
    def __init__(
        self,
        model: MultiClassClassifier,
        learning_rate: float = 1e-3,
        label_smoothing: float = 0.0,
        class_weights: Optional[torch.Tensor] = None,
    ):
        self.model = model
        
        # Loss function with optional smoothing and weighting
        self.criterion = nn.CrossEntropyLoss(
            weight=class_weights,
            label_smoothing=label_smoothing,
        )
        
        self.optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=learning_rate,
            weight_decay=1e-4,
        )
        
        # Learning rate scheduler
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer, T_max=100
        )
    
    def train_step(self, x, y) -> Tuple[float, float]:
        """Single training step."""
        self.model.train()
        self.optimizer.zero_grad()
        
        logits = self.model(x)
        loss = self.criterion(logits, y)
        
        loss.backward()
        
        # Gradient clipping for stability
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
        
        self.optimizer.step()
        
        # Compute accuracy
        with torch.no_grad():
            preds = logits.argmax(dim=-1)
            accuracy = (preds == y).float().mean().item()
        
        return loss.item(), accuracy
    
    @torch.no_grad()
    def evaluate(self, x, y) -> Tuple[float, float]:
        """Evaluate on data."""
        self.model.eval()
        
        logits = self.model(x)
        loss = self.criterion(logits, y)
        
        preds = logits.argmax(dim=-1)
        accuracy = (preds == y).float().mean().item()
        
        return loss.item(), accuracy
 
 
# Usage example
model = MultiClassClassifier(
    input_dim=784,       # e.g., MNIST
    num_classes=10,
    hidden_dims=[512, 256],
    dropout=0.2,
)
 
# For imbalanced data, compute class weights
class_counts = torch.tensor([1000, 500, 100, 50, 20, 800, 600, 400, 200, 100])
class_weights = 1.0 / class_counts
class_weights = class_weights / class_weights.sum() * len(class_weights)
 
trainer = MultiClassTrainer(
    model,
    learning_rate=1e-3,
    label_smoothing=0.1,  # Use label smoothing
    class_weights=class_weights,
)
 
# Training iteration
x_batch = torch.randn(64, 784)
y_batch = torch.randint(0, 10, (64,))
 
loss, acc = trainer.train_step(x_batch, y_batch)
print(f"Loss: {loss:.4f}, Accuracy: {acc:.4f}")