PyTorch Calibration

Calibration is the match between a model's stated confidence and how often it is right. A reliability table groups predictions by confidence and compares each group's confidence with its accuracy.

Confidence is the top probability

Softmax turns the two scores into probabilities. The confidence is the larger one, and it belongs to the predicted class.

print(f"First row probabilities: {float(probabilities[0][0]):.3f} and {float(probabilities[0][1]):.3f}")
print(f"Confidence: {float(confidence[0]):.3f}, predicted: {int(predicted[0])}, true: {int(test_y[0])}")

One bin, counted by hand

A mask picks the rows inside one bin. The bin's accuracy is the share of those rows the model got right.

top_bin = confidence >= 0.9
print(f"Rows in the top bin: {int(top_bin.sum())}")
print(f"Right in the top bin: {int(correct[top_bin].sum())}")

What perfect calibration would predict

A calibrated model's confidences add up to its number of right answers. Summing the top bin's confidences predicts how many rows it should get right.

print(f"Predicted right in top bin: {float(confidence[top_bin].sum()):.1f}")
print(f"Actually right in top bin: {int(correct[top_bin].sum())}")

Example

import random

import torch
from torch import nn

POSITIVE = ["great", "lovely", "bright", "fun", "warm"]
NEGATIVE = ["awful", "dull", "cold", "boring", "grim"]
FILLER = ["the", "a", "movie", "story", "scene", "really", "very"]
VOCAB = POSITIVE + NEGATIVE + FILLER + ["not"]
WORD_ID = {word: i for i, word in enumerate(VOCAB)}
LENGTH = 5


def make_data(rng, n_plain, n_negated):
    data = []
    for _ in range(n_plain):
        label = rng.randrange(2)
        side = POSITIVE if label == 1 else NEGATIVE
        words = [rng.choice(side) for _ in range(3)]
        words += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(words)
        data.append((words, label))
    for _ in range(n_negated):
        side_label = rng.randrange(2)
        side = POSITIVE if side_label == 1 else NEGATIVE
        rest = [rng.choice(side) for _ in range(2)]
        rest += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(rest)
        data.append((["not"] + rest, 1 - side_label))
    rng.shuffle(data)
    return data


def encode(data):
    x = torch.tensor([[WORD_ID[w] for w in words] for words, _ in data])
    y = torch.tensor([label for _, label in data])
    return x, y


class ReviewClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(len(VOCAB), 16)
        self.position = nn.Embedding(LENGTH, 16)
        layer = nn.TransformerEncoderLayer(
            d_model=16, nhead=2, dim_feedforward=32,
            dropout=0.0, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.head = nn.Linear(16, 2)

    def forward(self, ids):
        h = self.encoder(self.embedding(ids) + self.position.weight)
        return self.head(h.mean(dim=1))


rng = random.Random(42)
train_x, train_y = encode(make_data(rng, 96, 64))

torch.manual_seed(0)
model = ReviewClassifier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loss_fn = nn.CrossEntropyLoss()

for step in range(25):
    optimizer.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    optimizer.step()

test_rng = random.Random(123)
test_x, test_y = encode(make_data(test_rng, 240, 160))

model.eval()
with torch.no_grad():
    probabilities = torch.softmax(model(test_x), dim=1)
confidence, predicted = probabilities.max(dim=1)
correct = predicted == test_y

print(f"Test accuracy: {float(correct.float().mean()):.3f}")
print("Confidence bin  Count  Mean confidence  Accuracy")
edges = [0.5, 0.6, 0.7, 0.8, 0.9, 1.01]
for low, high in zip(edges[:-1], edges[1:]):
    in_bin = (confidence >= low) & (confidence < high)
    count = int(in_bin.sum())
    mean_confidence = float(confidence[in_bin].mean())
    bin_accuracy = float(correct[in_bin].float().mean())
    print(f"{low:.1f}-{min(high, 1.0):.1f}          {count:>5}  "
          f"{mean_confidence:.3f}            {bin_accuracy:.3f}")

Key points

Confidence: the largest softmax probability in one prediction row
Calibration: stated confidence and measured accuracy agree across the whole range
Bin: a slice of the confidence range, like 0.9 to 1.0, scored as one group
Reliability table: one row per bin, holding the count, the mean confidence, and the accuracy
A bin with mean confidence above its accuracy is overconfident
A bin with mean confidence below its accuracy is underconfident

In this example, the program trains the transformer review classifier from the Advanced transformer entry at a reduced size. One encoder layer and 25 training steps keep the model honest but imperfect. The test set holds 400 seeded reviews the training never saw. Softmax turns each score row into probabilities, and the top probability becomes the confidence. Overall test accuracy prints first at 0.877. The table then scores five bins. The crowded top bin holds 262 rows. Its mean confidence is 0.978 and its accuracy is 0.954, so the model is slightly overconfident there. The middle bins track their confidences closely. The two lowest bins hold few rows, and their accuracy swings show it.

Note

Small bins swing on a few rows, so the table reads best where the counts are large.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

import random

import torch
from torch import nn

POSITIVE = ["great", "lovely", "bright", "fun", "warm"]
NEGATIVE = ["awful", "dull", "cold", "boring", "grim"]
FILLER = ["the", "a", "movie", "story", "scene", "really", "very"]
VOCAB = POSITIVE + NEGATIVE + FILLER + ["not"]
WORD_ID = {word: i for i, word in enumerate(VOCAB)}
LENGTH = 5


def make_data(rng, n_plain, n_negated):
    data = []
    for _ in range(n_plain):
        label = rng.randrange(2)
        side = POSITIVE if label == 1 else NEGATIVE
        words = [rng.choice(side) for _ in range(3)]
        words += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(words)
        data.append((words, label))
    for _ in range(n_negated):
        side_label = rng.randrange(2)
        side = POSITIVE if side_label == 1 else NEGATIVE
        rest = [rng.choice(side) for _ in range(2)]
        rest += [rng.choice(FILLER) for _ in range(LENGTH - 3)]
        rng.shuffle(rest)
        data.append((["not"] + rest, 1 - side_label))
    rng.shuffle(data)
    return data


def encode(data):
    x = torch.tensor([[WORD_ID[w] for w in words] for words, _ in data])
    y = torch.tensor([label for _, label in data])
    return x, y


class ReviewClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(len(VOCAB), 16)
        self.position = nn.Embedding(LENGTH, 16)
        layer = nn.TransformerEncoderLayer(
            d_model=16, nhead=2, dim_feedforward=32,
            dropout=0.0, batch_first=True)
        self.encoder = nn.TransformerEncoder(layer, num_layers=1)
        self.head = nn.Linear(16, 2)

    def forward(self, ids):
        h = self.encoder(self.embedding(ids) + self.position.weight)
        return self.head(h.mean(dim=1))


rng = random.Random(42)
train_x, train_y = encode(make_data(rng, 96, 64))

torch.manual_seed(0)
model = ReviewClassifier()
optimizer = torch.optim.Adam(model.parameters(), lr=0.02)
loss_fn = nn.CrossEntropyLoss()

for step in range(25):
    optimizer.zero_grad()
    loss = loss_fn(model(train_x), train_y)
    loss.backward()
    optimizer.step()

test_rng = random.Random(123)
test_x, test_y = encode(make_data(test_rng, 240, 160))

model.eval()
with torch.no_grad():
    probabilities = torch.softmax(model(test_x), dim=1)
confidence, predicted = probabilities.max(dim=1)
correct = predicted == test_y

print(f"Test accuracy: {float(correct.float().mean()):.3f}")
print("Confidence bin  Count  Mean confidence  Accuracy")
edges = [0.5, 0.6, 0.7, 0.8, 0.9, 1.01]
for low, high in zip(edges[:-1], edges[1:]):
    in_bin = (confidence >= low) & (confidence < high)
    count = int(in_bin.sum())
    mean_confidence = float(confidence[in_bin].mean())
    bin_accuracy = float(correct[in_bin].float().mean())
    print(f"{low:.1f}-{min(high, 1.0):.1f}          {count:>5}  "
          f"{mean_confidence:.3f}            {bin_accuracy:.3f}")

Exercise

What does the top bin's row (mean confidence 0.978, accuracy 0.954) show?

Slight overconfidence in that bin Perfect calibration in that bin Strong underconfidence in that bin A bug in the binning code