+1 (415) 366-1364

Models That Know What They Don't Know: Calibration, Conformal Prediction, and Selective Prediction on TensorFlow

Every model your team ships returns a number. The question a risk, ops, or clinical stakeholder actually asks is different: how much should I trust this number? A classifier that says "0.93" is useless if 0.93 means 70% in practice, and a regressor that returns a point estimate gives a downstream automation rule nothing to threshold on.

This tutorial covers the practical uncertainty toolkit we deploy on TensorFlow engagements: calibration (make the probabilities honest), conformal prediction (turn any model into one with finite-sample coverage guarantees), and selective prediction (abstain and route to a human). None of it requires retraining your architecture, and all of it works on a Keras 3 model you already have.

1. Measure the problem before fixing it

Start with a held-out calibration split that the model never saw during training or early stopping. Three splits, not two: train, calibration, test.

import numpy as np, keras, tensorflow as tf

probs_cal = model.predict(x_cal, batch_size=256)      # (n, n_classes)
probs_test = model.predict(x_test, batch_size=256)

The headline metric is Expected Calibration Error (ECE): bin predictions by confidence and compare average confidence to average accuracy in each bin.

def ece(probs, labels, n_bins=15):
    conf = probs.max(axis=1)
    pred = probs.argmax(axis=1)
    correct = (pred == labels).astype(np.float64)
    bins = np.linspace(0.0, 1.0, n_bins + 1)
    total = 0.0
    for lo, hi in zip(bins[:-1], bins[1:]):
        m = (conf > lo) & (conf <= hi)
        if m.sum() == 0:
            continue
        total += (m.mean()) * abs(correct[m].mean() - conf[m].mean())
    return total

print("ECE before:", ece(probs_test, y_test))

Also plot a reliability diagram (confidence on x, observed accuracy on y). Modern deep networks, especially after long training with high capacity, are usually over-confident: the curve sits below the diagonal. Label smoothing, mixup, and heavy augmentation push the other way and can leave you under-confident.

Track ECE alongside accuracy in the same dashboard you use for the TensorBoard profiler and training metrics. A model refresh that improves AUC while doubling ECE is a regression for any thresholded workflow.

2. Temperature scaling: the cheapest real fix

Temperature scaling divides the logits by a single learned scalar T. One parameter, fit on the calibration split, cannot change the ranking of predictions — so accuracy, AUC and every ranking metric are untouched — but it fixes most of the over-confidence.

Export a logits model first (your training model probably ends in softmax):

logits_model = keras.Model(model.inputs, model.layers[-1].input)   # pre-softmax
logits_cal = logits_model.predict(x_cal, batch_size=256)

Fit T by minimising NLL with a few hundred steps of gradient descent:

log_T = tf.Variable(0.0, dtype=tf.float32)          # parameterise log T > 0
opt = keras.optimizers.Adam(0.01)
lg = tf.constant(logits_cal, tf.float32)
yc = tf.constant(y_cal, tf.int32)

for _ in range(500):
    with tf.GradientTape() as tape:
        scaled = lg / tf.exp(log_T)
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=yc, logits=scaled))
    opt.apply_gradients(zip(tape.gradient(loss, [log_T]), [log_T]))

T = float(tf.exp(log_T))
print("temperature:", T)

T > 1 means the model was over-confident. Bake it into the served graph so nothing downstream has to remember it:

scaled_logits = keras.layers.Lambda(lambda z: z / T, name="temperature")(logits_model.output)
calibrated = keras.Model(logits_model.inputs,
                         keras.layers.Activation("softmax")(scaled_logits))
calibrated.export("artifacts/calibrated/1")   # SavedModel for TF Serving

Recompute ECE on the test split. A drop from ~0.08 to ~0.01 on an image or text classifier is typical. For binary problems with skewed classes, isotonic regression on the calibration split sometimes beats temperature scaling, at the cost of being non-parametric and hungrier for data — use it when the calibration split has tens of thousands of rows.

3. Conformal prediction: coverage you can promise

Calibration makes the average probability honest. It does not promise anything about an individual prediction. Split conformal prediction does: given exchangeable data, it produces a prediction set that contains the true label at least 1 - alpha of the time, with no assumptions about the model.

The recipe is three lines of numpy. Score each calibration point by how badly the model ranked the true label, then take a quantile.

alpha = 0.10                      # target 90% coverage
n = len(y_cal)
true_prob = probs_cal[np.arange(n), y_cal]
scores = 1.0 - true_prob                       # non-conformity score
q = np.quantile(scores, np.ceil((n + 1) * (1 - alpha)) / n, method="higher")

At inference, the set is every class whose score falls under the threshold:

def predict_set(p, q):
    return [np.where(1.0 - row <= q)[0] for row in p]

sets = predict_set(probs_test, q)
coverage = np.mean([y in s for y, s in zip(y_test, sets)])
avg_size = np.mean([len(s) for s in sets])
print(f"coverage={coverage:.3f}  avg set size={avg_size:.2f}")

Coverage should land near 0.90 whatever model you plugged in — that is the guarantee. The informative number is average set size: a well-trained, well-calibrated model returns singletons for most inputs and larger sets only where it is genuinely confused. Set size is therefore a far better production health signal than raw accuracy; when it creeps up, the input distribution moved.

Two refinements worth the effort:

  • Adaptive sets (APS/RAPS). Instead of 1 - p_true, accumulate sorted probabilities until the true label is reached. Sets get slightly larger on average but coverage becomes much more even across classes instead of being concentrated on easy ones.
  • Mondrian / class-conditional conformal. Compute a separate quantile per class or per segment (region, device type, customer tier). This buys per-group coverage, which is what a fairness or regulatory review will actually ask for.

For regression, the same machinery gives intervals. Train a quantile regressor with pinball loss, then conformalise the residuals — conformalized quantile regression:

def pinball(tau):
    def loss(y_true, y_pred):
        e = y_true - y_pred
        return tf.reduce_mean(tf.maximum(tau * e, (tau - 1) * e))
    return loss

# model outputs [lo, hi] for tau = alpha/2 and 1 - alpha/2
lo_cal, hi_cal = qmodel.predict(x_cal).T
e = np.maximum(lo_cal - y_cal, y_cal - hi_cal)                 # conformity score
q = np.quantile(e, np.ceil((len(y_cal) + 1) * (1 - alpha)) / len(y_cal), method="higher")
# served interval: [lo - q, hi + q]

This composes nicely with the probabilistic forecasting setup in our time-series tutorial: train the quantile heads once, conformalise per horizon.

4. Selective prediction: knowing when to abstain

Most business value from uncertainty comes from routing. Keep the predictions you trust, send the rest to a human or a slower, more expensive model.

conf = probs_test.max(axis=1)
for cov_target in [1.0, 0.9, 0.8, 0.7]:
    thr = np.quantile(conf, 1 - cov_target)
    m = conf >= thr
    acc = (probs_test[m].argmax(1) == y_test[m]).mean()
    print(f"coverage {m.mean():.2f} -> accuracy {acc:.4f} (threshold {thr:.3f})")

That risk–coverage table is the artifact to put in front of a stakeholder. "At 80% automation we are 99.1% accurate, and the 20% we defer is where the errors live" is a decision they can make. Pick the threshold from the calibration split, not the test split, and re-fit it on every model refresh.

Conformal set size works as an abstention rule too, and it is usually better behaved: abstain whenever len(set) != 1.

5. Model uncertainty, not just data uncertainty

Temperature scaling and conformal prediction both treat the model as fixed. They quantify aleatoric uncertainty — noise in the data. They will happily give you a confident, well-calibrated, wrong answer on an input unlike anything in training. For that epistemic piece:

  • MC dropout. Keep dropout active at inference and average several stochastic passes. Cheap to implement (model(x, training=True) in a loop), but the variance estimate is crude.
  • Deep ensembles. Train 3–5 models with different seeds and average their probabilities. Still the strongest baseline in the literature, and it parallelises perfectly on the training setup from our multi-GPU tutorial. The cost is N× training and N× serving — often worth it for a small tabular or tree-adjacent model, rarely worth it for a 7B LLM.
  • TensorFlow Probability layers. tfp.layers.DenseVariational and friends give you a Bayesian last layer with a fraction of ensemble cost; a common compromise is a deterministic backbone plus a variational head.

Pair any of these with an out-of-distribution check on inputs, and wire the drift detector from our TensorFlow Data Validation tutorial to the same alert channel. Epistemic uncertainty spikes and feature drift usually fire together; when only one fires, that is informative.

6. What breaks in production

Exchangeability. Conformal's guarantee needs calibration and production data to be exchangeable. Under temporal drift they are not, and coverage quietly decays. Mitigations: re-calibrate on a rolling recent window, or use adaptive conformal inference, which nudges alpha up or down based on observed coverage. Monitor realised coverage on the labels that do come back.

Calibration drift. T fitted in March is wrong in September. Re-fit on every retrain, and log ECE and conformal set size as first-class metrics in your TFX evaluation step so a badly calibrated model fails blessing instead of reaching users.

Threshold leakage. Fitting T, q, or the abstention threshold on the test split invalidates everything you report. Three splits, always.

Post-processing in the wrong place. If temperature lives in an application-layer Python helper, the batch job, the mobile build, and the canary will each drift out of sync. Put it in the SavedModel graph, and keep q in the model's metadata next to the weights so canary rollouts like the ones in our KServe tutorial ship the model and its calibration together.

LLM outputs. Token-level log-probabilities from a fine-tuned Gemma are not a trustworthy confidence score for a whole generated answer. For generation, conformalise a task-level score instead — a retrieval relevance score, a verifier model's judgement, or agreement across sampled completions — and treat it with the same calibration/eval discipline as our generative AI evaluation tutorial.

Checklist

  1. Three splits: train / calibration / test.
  2. Report ECE and a reliability diagram alongside accuracy.
  3. Fit temperature on calibration; bake it into the SavedModel.
  4. Add split conformal; track coverage and average set size in monitoring.
  5. Publish a risk–coverage table and set the abstention threshold with the business.
  6. Use ensembles or a variational head where epistemic uncertainty matters.
  7. Re-calibrate on every retrain; alert on coverage decay.

Shipping a model that knows what it does not know is usually a two-week change, not a re-architecture — and it is often what moves a pilot into production. If your team is weighing automation thresholds, regulatory coverage requirements, or a human-in-the-loop design on the TensorFlow stack, get in touch and we will walk through your current evaluation setup.