QAC: Data Quality-Aware Uncertainty Calibration for Trustworthy AI

A classifier's predicted probability is its statement of uncertainty about the label, and calibration makes this statement match observed frequencies. Post-hoc calibrators use only the model's output, so two samples with the same output receive the same probability. Consider two patients whose laboratory values are almost the same, but one patient's value is wrong because the blood sample was stored too long before analysis: the model gives both the same prediction and the same confidence, although the second prediction is more likely to be wrong. Laboratories and other data pipelines often record such problems. We call a per-sample value that describes how reliable a data point's input or label is its data quality. When it is not available at training time, it cannot be added to the model's input, but it can be given to the calibrator. A standard result implies that, for the best possible predictor, a data-quality value can improve the prediction only if it carries information about the label that the input does not. We propose Quality-Aware Uncertainty Calibration (QAC), a second calibration stage that can follow any existing calibrator. For each class, it estimates on the calibration set how often the class is correct given the first-stage probability and the quality score; the direction of the effect of quality is learned, not assumed. On seven datasets, nine base calibrators and ten seeds, and compared with the same stage given a constant quality score, QAC lowers the Brier score by $0.001$-$0.016$ on the four datasets where quality is informative (a simulation and three datasets where quality is annotator agreement). It does not help on the two datasets whose quality score is computed from the input, and gives no significant gain on a clinical dataset. The gain is mainly in separating correct from wrong predictions (error-detection AUROC $+0.021$ to $+0.059$), while the separation between classes improves about a quarter as much or less.