Reliability-aware checkpoint selection improves domain generalization results
Reliability-Aware Checkpoint Selection for Domain Generalization
Machine LearningComputer Vision and Pattern Recognition
Summary
Choosing the best saved model during training usually depends on how accurate it is with known data, but this doesn't always guarantee good or reliable predictions on new data that looks different. The authors found a way to pick among top-accuracy models by also considering how trustworthy their predicted probabilities are, without needing any new data or extra training. This method improved prediction reliability on new data while keeping accuracy about the same. Their work shows it's possible to pick models that not only predict well but also provide better probability estimates.
What this means in practice
- •For machine learning engineers: Select more reliable trained models for deployment in new or changing environments without needing data from those environments.
- •For quality assurance teams: Evaluate and improve model confidence measures to support trustworthy AI systems in applications sensitive to prediction reliability.
Authors
Jinshi Liu, Jiahao Li, Pan Liu, Yanfeng Li, Rui Qian, Zhao Tong, Yue Sun, Tao Tan
Abstract
Checkpoint selection in domain generalization often relies on source-validation accuracy, yet the selected checkpoint need not provide reliable probabilities on unseen target domains. Source-target distribution shifts can alter accuracy rankings, while accuracy alone does not measure predictive probability quality. We identify an empirical selection opportunity within fixed training trajectories: reselecting among checkpoints with near-optimal source accuracy can improve mean target probability quality with small observed changes in mean target accuracy. We study accuracy-constrained reliability selection (AC), which retains checkpoints within a tolerance of the best source-validation accuracy and ranks them by source reliability. Our reference rule aggregates within-set normalized negative log-likelihood (NLL) and class-wise calibration error (CwECE) using $D_\infty$. AC uses no target data and requires neither additional training nor weight averaging. We evaluate five domain generalization training algorithms on three benchmarks, using PACS to develop the objectives and a 0.5-percentage-point tolerance. In exploratory aggregation comparisons on 360 OfficeHome and TerraIncognita runs, the reference rule reduces mean target soft-bin squared-gap ECE and CwECE by 0.240% and 0.182%, respectively, and NLL by 0.030 relative to Source-Acc. Mean target accuracy changes by +0.213 percentage points. These results identify opportunities for reliability-aware reselection, while the additional benefit of joint over single-objective ranking remains unresolved.