Why Softmax Probabilities Lie: Temperature Scaling & Calibrating Medical AI Models

The Overconfidence Crisis in Modern Deep Networks

In 2017, Guo et al. published "On Calibration of Modern Neural Networks", revealing a startling paradox: while modern deep neural networks (ResNets, EfficientNets, DenseNets) achieved far superior classification accuracy than older models like LeNet, their probability calibration was dramatically worse.

A perfectly calibrated model is one where, out of 100 predictions assigned a confidence score of 0.80, exactly 80 predictions are correct. In modern deep networks, however, the combination of depth, batch normalization, weight decay, and cross-entropy optimization pushes uncalibrated softmax values toward extreme overconfidence (0.95 to 0.999), even on misclassified instances.

Measuring Expected Calibration Error (ECE)

To quantify miscalibration mathematically, we partition all predictions into \(M\) equally-spaced confidence bins \(B_m\). The Expected Calibration Error is the weighted average difference between bin accuracy and bin confidence:

ext{ECE} = \sum_{m=1}^M rac{|B_m|}{N} \left| ext{acc}(B_m) - ext{conf}(B_m) ight|

Temperature Scaling Implementation in PyTorch

Temperature scaling is the simplest and most effective post-processing calibration technique. It introduces a single scalar parameter \(T > 0\) to scale logits prior to the softmax activation without altering classification ranking (preserving top-1 accuracy):

import torch
from torch import nn, optim

class ModelWithTemperature(nn.Module):
    '''Wraps a trained model to calibrate probabilities via learned Temperature parameter.'''
    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model
        # Initialize temperature parameter to 1.0 (identity)
        self.temperature = nn.Parameter(torch.ones(1) * 1.5)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        logits = self.model(x)
        return self.temperature_scale(logits)

    def temperature_scale(self, logits: torch.Tensor) -> torch.Tensor:
        # Expand temperature to match logit dimensionality
        temperature = self.temperature.unsqueeze(1).expand(logits.size(0), logits.size(1))
        return logits / temperature

    def calibrate(self, valid_loader: torch.utils.data.DataLoader):
        '''Optimizes temperature on validation set via Negative Log Likelihood (NLL).'''
        self.eval()
        nll_criterion = nn.CrossEntropyLoss()
        
        # 1. Collect all validation logits and true labels
        logits_list = []
        labels_list = []
        with torch.no_grad():
            for inputs, targets in valid_loader:
                logits = self.model(inputs)
                logits_list.append(logits)
                labels_list.append(targets)
        logits = torch.cat(logits_list)
        labels = torch.cat(labels_list)

        # 2. Optimize temperature parameter using L-BFGS
        optimizer = optim.LBFGS([self.temperature], lr=0.01, max_iter=50)

        def eval_step():
            optimizer.zero_grad()
            loss = nll_criterion(self.temperature_scale(logits), labels)
            loss.backward()
            return loss

        optimizer.step(eval_step)
        print(f"Optimal Learned Temperature: {self.temperature.item():.3f}")