Calibration and Uncertainty in GNN Predictions

A GNN that produces a confidence score of 95% for a prediction should be correct approximately 95% of the time. When that alignment holds, we say the model is well-calibrated. In practice, deep neural networks — and GNNs in particular — are frequently overconfident: they assign high probability to predictions even when they are wrong. For applications like drug discovery, fraud detection, or medical diagnosis, overconfident predictions are dangerous.

TGraphX provides a calibration module (tgraphx.calibration) that measures calibration quality using the Expected Calibration Error (ECE) metric and applies post-hoc temperature scaling to correct miscalibration. This article explains what calibration means for GNNs, how ECE is computed, how temperature scaling works, and how to use the TGraphX calibration tools.

Why GNNs Are Prone to Miscalibration

Standard GNNs are trained to minimize cross-entropy loss, which encourages the model to push predicted probabilities toward 0 and 1. Over-parameterized models — those with more capacity than the dataset requires — tend to fit training data with extreme confidence, which does not transfer to held-out examples.

GNNs inherit this problem but also face graph-specific amplifiers:

Over-smoothing occurs when stacking many GNN layers causes node embeddings to converge to the same value, making the model's predictions trivially confident on a degraded signal.

Neighborhood aggregation propagates confidence. A node surrounded by confidently predicted neighbors will itself receive confident messages. If those neighbors are incorrectly predicted, the target node's overconfidence is compounded.

Semi-supervised training on sparse labels means the model sees very few examples per class. Small labeled sets amplify overfitting and miscalibration.

Graph homophily assumption mismatch. Many GNNs implicitly assume homophily (connected nodes have similar labels). When a graph is heterophilous, the aggregation aggregates misleading signal and the model compensates by increasing confidence.

The Expected Calibration Error (ECE)

ECE measures the gap between predicted confidence and actual accuracy, averaged over confidence bins. The computation proceeds as follows:

1. Sort all predictions by their confidence score.
2. Group predictions into M equal-width bins based on confidence (e.g., [0.0, 0.1), [0.1, 0.2), ..., [0.9, 1.0]).
3. For each bin, compute the average confidence and the fraction of correct predictions (accuracy).
4. ECE = weighted average of |confidence − accuracy| across bins:

ECE = Σ_m (|B_m| / N) × |acc(B_m) - conf(B_m)|
        

A perfectly calibrated model has ECE = 0. In practice, values below 0.05 are considered good; values above 0.1 indicate significant miscalibration.

Prerequisites

This article assumes familiarity with:

• Softmax outputs and cross-entropy loss
• Basic GNN node classification (see GraphSAGE guide for a practical GNN training example)
• Overfitting and regularization concepts

Install TGraphX:

pip install tgraphx
        

Computing ECE with TGraphX

TGraphX's calibration module provides ECE computation from raw logits:

import torch
        from tgraphx.calibration import compute_ece
        
        # Suppose we have a trained model and a validation set
        # logits: [N, num_classes], labels: [N]
        logits = torch.randn(500, 7)   # raw GNN outputs (before softmax)
        labels = torch.randint(0, 7, (500,))
        
        ece = compute_ece(logits, labels, n_bins=15)
        print(f"ECE: {ece:.4f}")
        

A reliability diagram plots accuracy vs. confidence per bin. Overconfident models have accuracy well below the diagonal; underconfident models sit above it. Computing ECE before and after calibration is the standard way to verify that calibration is working.

Temperature Scaling

Temperature scaling is the simplest and most effective post-hoc calibration method. It introduces a single scalar parameter T > 0 that divides all logits before the softmax:

p_calibrated(y | x) = softmax(logits / T)
        

• T > 1 softens the predictions (reduces confidence), which is the common fix for overconfidence.
• T < 1 sharpens predictions (increases confidence), needed for underconfident models.

Temperature scaling does not change the model's accuracy — it reorders no predictions, it only rescales confidence values. This is a key property: you can tune T on a validation set purely to minimize ECE without touching the trained weights.

T is found by minimizing the negative log-likelihood (NLL) on a held-out validation set with the model weights frozen.

Applying Temperature Scaling with TGraphX

import torch
        import torch.nn as nn
        import torch.nn.functional as F
        from tgraphx.calibration import TemperatureScaler
        
        # Assume model is trained and we have validation logits and labels
        val_logits = torch.randn(500, 7)   # outputs from a trained GNN on validation set
        val_labels = torch.randint(0, 7, (500,))
        
        # Fit temperature scaler on validation set
        scaler = TemperatureScaler()
        scaler.fit(val_logits, val_labels)
        
        print(f"Learned temperature: {scaler.temperature.item():.4f}")
        
        # Apply to test logits
        test_logits = torch.randn(300, 7)
        calibrated_probs = scaler.calibrate(test_logits)  # softmax(logits / T)
        print(calibrated_probs.shape)  # [300, 7], probabilities sum to 1
        
        # Measure ECE before and after
        from tgraphx.calibration import compute_ece
        ece_before = compute_ece(test_logits, torch.randint(0, 7, (300,)))
        ece_after = compute_ece(
            calibrated_probs.log(),  # compute_ece expects logits or log-probs
            torch.randint(0, 7, (300,))
        )
        print(f"ECE before: {ece_before:.4f} | ECE after: {ece_after:.4f}")
        

ECE Across GNN Models: What to Expect

GNN architecture and calibration behavior are correlated. The following table shows representative patterns — not claimed benchmark results, since actual ECE values depend heavily on dataset, training setup, and hyperparameters.

Temperature scaling is almost always worth applying after training. It costs one scalar parameter and a short NLL minimization on the validation set.

Bayesian and Ensemble Approaches

Temperature scaling corrects the scale of confidence but does not estimate epistemic uncertainty — the uncertainty arising from limited data. For that, more sophisticated methods exist:

MC Dropout: Run the GNN with dropout active at inference time and average predictions over multiple forward passes. The variance of the averaged probabilities approximates epistemic uncertainty.

Deep Ensembles: Train multiple GNNs with different random seeds and average their predictions. This is more expensive but more reliable than MC Dropout.

Graph-specific uncertainty propagation: Nodes in low-density graph regions (few neighbors) have higher structural uncertainty. This is not captured by standard temperature scaling.

TGraphX's calibration module currently provides ECE computation and temperature scaling. MC Dropout and ensemble methods can be implemented using standard PyTorch patterns alongside TGraphX layers.

Integration with the Reproducibility Module

Calibration results are sensitive to random seeds, train/val splits, and training setup. Before reporting ECE numbers in publications, fix all sources of randomness and use TGraphX's reproducibility context:

from tgraphx.reproducibility import set_reproducible
        
        with set_reproducible(seed=42):
            # Train GNN
            # Compute val_logits
            # Fit temperature scaler
            # Report ECE
            pass
        

For a full treatment of reproducibility in GNN experiments, see the GNN research reproducibility guide.

Limitations and Honest Notes

ECE is sensitive to the number of bins. Fewer bins are less noisy but less precise. The community standard is 15 bins for datasets with at least a few hundred examples. For very small graphs, ECE estimates are high-variance.

Temperature scaling assumes all classes need the same correction. If the model is overconfident for some classes and underconfident for others, temperature scaling cannot correct both simultaneously. Class-conditional calibration or Dirichlet calibration handle this case.

Calibration on graphs violates the i.i.d. assumption. ECE is derived under the assumption that each data point is independent. In a graph, neighboring nodes share information and their predictions are correlated. ECE values computed on a single connected graph component are therefore optimistic estimates of true calibration.

Post-hoc calibration does not fix model errors. A well-calibrated model that is 70% accurate will be 30% wrong, but those wrong predictions will be appropriately less confident. Calibration improves trust in confidence scores but does not improve the underlying predictions.

Frequently Asked Questions

When should I apply calibration?
Before deploying any GNN in a context where confidence scores are used to make decisions — thresholding predictions, prioritizing cases for human review, or propagating predictions to downstream models.

Does calibration change which nodes are predicted to belong to which class?
No. Temperature scaling divides all logits by the same scalar, so the argmax (the predicted class) is unchanged.

Can I use calibration for link prediction or graph classification?
Yes. ECE applies to any binary or multi-class confidence score. For link prediction, it measures whether predicted edge existence probabilities match observed edge frequencies.

Is there a validation set size requirement?
Temperature scaling with as few as 100–200 validation examples can estimate T reliably. Fewer examples lead to noisy ECE estimates, but the temperature fit is robust.

Where is the TGraphX calibration module documented?
See the TGraphX GitHub for source code. The package is available at PyPI.

A Full Calibration Workflow: End to End

Putting everything together, here is a complete workflow from training a GNN to reporting calibrated confidence scores on a node classification task:

import torch
        import torch.nn as nn
        import torch.nn.functional as F
        from tgraphx.layers.sage import TensorGraphSAGELayer
        from tgraphx.calibration import compute_ece, TemperatureScaler
        from tgraphx.reproducibility import set_reproducible
        
        class SimpleGNN(nn.Module):
            def __init__(self, in_dim, hidden_dim, num_classes):
                super().__init__()
                self.conv1 = TensorGraphSAGELayer(in_dim, hidden_dim)
                self.conv2 = TensorGraphSAGELayer(hidden_dim, num_classes)
        
            def forward(self, x, edge_index):
                x = F.relu(self.conv1(x, edge_index))
                return self.conv2(x, edge_index)
        
        with set_reproducible(seed=42):
            # Setup
            N = 1000
            x = torch.randn(N, 64)
            edge_index = torch.randint(0, N, (2, 4000), dtype=torch.long)
            labels = torch.randint(0, 7, (N,))
        
            # Masks: 60% train, 20% val, 20% test (mutually exclusive)
            perm = torch.randperm(N)
            train_mask = torch.zeros(N, dtype=torch.bool)
            val_mask = torch.zeros(N, dtype=torch.bool)
            test_mask = torch.zeros(N, dtype=torch.bool)
            train_mask[perm[:600]] = True
            val_mask[perm[600:800]] = True
            test_mask[perm[800:]] = True
        
            # Train
            model = SimpleGNN(64, 128, 7)
            optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
            for epoch in range(100):
                model.train()
                optimizer.zero_grad()
                logits = model(x, edge_index)
                loss = F.cross_entropy(logits[train_mask], labels[train_mask])
                loss.backward()
                optimizer.step()
        
            # Step 1: Measure ECE before calibration (on val set)
            model.eval()
            with torch.no_grad():
                val_logits = model(x, edge_index)[val_mask]
            ece_before = compute_ece(val_logits, labels[val_mask])
            print(f"ECE before calibration: {ece_before:.4f}")
        
            # Step 2: Fit temperature scaler on val set
            scaler = TemperatureScaler()
            scaler.fit(val_logits, labels[val_mask])
            print(f"Learned temperature T={scaler.temperature.item():.4f}")
        
            # Step 3: Evaluate on test set with calibrated predictions
            with torch.no_grad():
                test_logits = model(x, edge_index)[test_mask]
            calibrated_test_probs = scaler.calibrate(test_logits)
        
            ece_after = compute_ece(
                (calibrated_test_probs + 1e-9).log(),  # log-probs for ECE
                labels[test_mask]
            )
            test_acc = (calibrated_test_probs.argmax(1) == labels[test_mask]).float().mean()
            print(f"ECE after calibration: {ece_after:.4f}")
            print(f"Test accuracy (unchanged by calibration): {test_acc.item():.4f}")
        

This workflow illustrates the key invariant: accuracy is unchanged by calibration, but ECE improves. If calibration substantially changes accuracy, something is wrong — check that you are applying scaler.calibrate() correctly and not accidentally reordering predictions.

When Calibration Is Not Enough: Structural Uncertainty

Temperature scaling corrects miscalibration globally but does not model the graph's structural uncertainty. Some nodes have few neighbors and thus receive less aggregated information — their predictions should carry higher uncertainty than well-connected nodes, but a single global temperature T treats all nodes identically.

A simple structural uncertainty heuristic: nodes with degree below a threshold receive a higher temperature scaling:

def degree_aware_temperature(logits, edge_index, num_nodes, base_T, low_degree_T=3.0, threshold=3):
            """
            Apply higher temperature to low-degree nodes (fewer than `threshold` neighbors).
            """
            degree = torch.zeros(num_nodes)
            degree.scatter_add_(0, edge_index[0], torch.ones(edge_index.shape[1]))
        
            T = torch.full((num_nodes,), base_T)
            T[degree < threshold] = low_degree_T
        
            # Apply per-node temperature
            calibrated = logits / T.unsqueeze(1)
            return F.softmax(calibrated, dim=-1)
        
        base_T = scaler.temperature.item()
        probs_structural = degree_aware_temperature(
            test_logits, edge_index[:, :200],  # test subgraph edges
            num_nodes=test_mask.sum().item(),
            base_T=base_T,
        )
        

This is a heuristic, not a rigorous uncertainty quantification method. For rigorous graph uncertainty, probabilistic GNN methods (GP-GNN, stochastic message passing) are more principled but substantially more complex to implement.