PaperlayerView exact revision on arXiv

Source-linked AI summary

Calibration in One Go: Jointly Learning Accurate and Calibrated Neural Networks

Nektarios Kalampalikis, Pablo M. Olmos, Isabel Valera

arXiv:2610.13693v1cs.LGcs.AIstat.ML

TL;DR

Neural networks are often miscalibrated, while post-hoc calibration typically requires a separate dataset and optimization stage. COG jointly optimizes the model and calibrator through a bilevel Stackelberg formulation, and achieves strong performance across settings with minimal added complexity, including a parameter-free variant.

  • Problem

    Neural networks are often miscalibrated, while post-hoc calibration typically requires a separate held-out dataset and additional optimization, complicating calibration.

    Source: §1 · pp. 1–2

  • Method

    COG jointly trains the model and calibrator through a bilevel Stackelberg formulation using shared cross-entropy objectives and distinct learning rates.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    Across a broad range of experiments, architectures, and calibration approaches, COG matches or outperforms existing methods in most settings, and COG-FC achieves state-of-the-art calibration without additional trainable parameters.

    Source: §5 · p. 9

  • Takeaways & Limitations

    COG-FC provides a practical end-to-end calibration option without separate calibration data or additional trainable parameters, including in multi-task learning.

    Source: §4.3 · p. 8, §5 · p. 9

    The evaluation could be extended to alternative calibration definitions and emerging metrics, and the paper does not explicitly explore optimistic or pessimistic bilevel optimization.

    Source: §5 · p. 9

Abstract

from arXiv · show

Neural networks (NNs), while highly flexible and versatile, are often poorly calibrated. Post-hoc calibration methods adjust the outputs of an already trained model using a separate calibration dataset and additional calibration parameters. This process adds overhead that requires separate optimization of the calibrator and subsequent model selection. In this paper, we show that standard (post-hoc) calibrators can be trained end-to-end with the backbone via a bilevel Stackelberg formulation. This formulation eliminates the need for a separate calibration set or a dedicated calibration loss by jointly optimizing accuracy and calibration using cross-entropy, introducing only one additional hyperparameter: the calibrator learning rate. Although our framework supports calibrators ranging from temperature and matrix scaling to NN layers, our results show that using the final classification layer itself as the calibrator yields accurate, well-calibrated models with no additional parameters. Our approach outperforms SOTA calibrators with minimal added complexity, while also avoiding the multiplicative overhead of separate calibration components in settings such as multi-task learning.

1 INTRODUCTION

Neural-network probabilities matter when decisions depend on uncertainty, but models are often miscalibrated. The paper presents COG, which jointly trains model and calibrator to avoid a separate calibration stage and data while retaining accuracy and calibration.

1 INTRODUCTION

  • Predicted probabilities are useful for uncertainty-dependent decisions only when they reflect actual predictive uncertainty; neural networks are often miscalibrated.Calibration aligns predicted probabilities with observed outcome frequencies.Source: §1 · p. 1
  • Separate calibration adds pipeline complexity, and multi-task learning may require a calibrator per task while reserving additional calibration data is costly.Post-hoc calibrators can also overfit, motivating regularization or simpler calibrators.Source: §1 · p. 2
  • COG jointly trains the model and calibrator through bilevel optimization, using one forward and backward pass per iteration and only one additional hyperparameter.The calibrator is the leader and the model the follower, with separate learning rates.Source: §1 · p. 2
  • The framework needs neither a calibration-specific loss nor separate calibration data, and it adds no architectural constraints.It supports calibrators from temperature and matrix scaling to neural layers, with at most the trainable parameters of the corresponding post-hoc calibrator.Source: §1 · p. 2
  • COG achieves accuracy and calibration competitive with substantially more complex state-of-the-art methods, while the final classification layer can calibrate without additional parameters.The paper also reports extending the single-stage approach to multi-task learning, where conventional calibration is cumbersome.Source: §1 · p. 2

2 RELATED WORKS

Prior work addresses calibration through post-hoc transformations, multi-stage training, altered training objectives, and bilevel optimization. COG differs by turning post-hoc calibration into a train-time procedure within standard backpropagation, without held-out calibration data or expensive optimization.

Post-hoc calibration

  • Post-hoc methods fit a transformation on held-out calibration data; parametric options include Platt, vector, matrix, and temperature scaling, while nonparametric options include histogram binning and isotonic regression.Flexible calibrators can overfit, especially with limited data, and multiclass nonparametric methods often use one-vs-rest schemes.Source: §2 · pp. 2–3

Progressive training approaches

  • Progressive training methods retrain final layers or progressively freeze earlier layers, sometimes adding a Bayesian final layer, and can be viewed as bilevel or multi-level optimization.These methods connect calibration issues to over-compression in later layers.Source: §2 · p. 3

Train-time calibration

  • Train-time methods add calibration penalties or specialized losses, whereas COG changes the optimization structure and can use standard objectives such as cross-entropy or combine with calibration-aware losses.Examples include regularization, differentiable bin-membership terms, label smoothing, and focal loss.Source: §2 · p. 3

Bilevel Optimization and Calibration

  • Earlier bilevel approaches use meta-learning or validation-sample weighting to enhance existing calibration methods or optimize a different objective, often requiring multiple optimization steps.COG instead incorporates calibration into standard backpropagation and avoids held-out calibration data.Source: §2 · p. 3

3 TOWARDS CALIBRATION IN ONE GO

COG casts calibration as a bilevel problem: the calibrator leads and the predictive model follows. It approximates this hierarchy through joint gradient updates, with the calibrator learning more slowly, and supports several calibrator designs including the final classification layer itself.

(Post-hoc) Calibration as a Bilevel Optimization Problem

  • The framework treats predictive accuracy as the follower-level objective and calibration as the leader-level objective, with the calibrator anticipating the classifier’s response.The calibrated classifier is the calibrator applied to the network output.Source: §3 · p. 3
  • Post-hoc calibration trains the model first, freezes its parameters, and then fits calibration parameters on separate data, so the calibration objective cannot influence the learned model solution.The paper motivates joint optimization because neural-network training can have multiple equivalent local optima.Source: §3 · p. 4

Calibration in One Go (COG)

  • COG uses a first-order Stackelberg-inspired scheme that updates calibrator and model parameters together using gradients of their respective objectives and distinct learning rates.The calibrator is the leader; the classifier is the follower.Source: §3 · p. 4
  • Using the same loss and training minibatch at both levels lets COG compute both gradients in one forward and one backpropagation pass; the calibrator learning rate is lower than the model’s.The faster follower timescale preserves the Stackelberg hierarchy, but these updates do not guarantee an optimal Stackelberg solution.Source: §3 · p. 4

From Post-hoc Calibration to COG

  • COG supports temperature scaling, matrix scaling, multi-layer TST, and a parameter-free variant that uses the final classification layer as calibrator.The framework also supports replacing both level objectives with label smoothing; its variants are lightweight end-to-end alternatives, including in multi-task settings.Source: §3 · pp. 4–5, Table 1 · p. 5

4 EXPERIMENTS

The experiments compare jointly trained COG calibrators with post-hoc methods across image-classification datasets, calibration families, distribution shifts, architectures, and multi-task learning. COG variants generally match or outperform corresponding post-hoc methods, while COG-FC achieves strong results without extra calibration parameters.

4 EXPERIMENTS

  • The evaluation uses ResNet101 models trained from scratch on SVHN, CIFAR-10, and CIFAR-100, and compares matched-parameter post-hoc and COG pairs.COG variants learn calibration jointly on the training set, whereas post-hoc calibrators are fitted on the validation set.Source: §4 · p. 5, §4.1 · p. 5

4.1 COMPARING POST-HOC METHODS AND THEIR COG COUNTERPARTS

  • Across datasets and metrics, COG variants generally match or outperform their post-hoc counterparts; COG-FC performs on par with more expressive COG variants without extra trainable parameters.Post-hoc Matrix Scaling fails to converge on CIFAR-100, while COG-MS trains successfully.Source: §4.1 · pp. 5–6, Figure 2 · p. 6

4.2 COMPARISON ACROSS CALIBRATION FAMILIES

  • A COG variant is top-ranked in two of three in-distribution cases, and COG-FC ranks first on CIFAR-100 while closely matching top-ranked models in reliability diagrams.The reported reliability-diagram difference includes a deviation for CIFAR-10 around 25% confidence.Source: §4.2 · p. 6, Figure 3 · p. 7
  • Across three datasets, a COG variant achieves the best OOD-detection performance in each, while covariate-shift results vary by dataset and shift type.COG+LS generally improves over COG-FC under covariate shift; the paper distinguishes shift robustness from in-distribution calibration.Source: §4.2 · p. 7, Table 3 · p. 8, Table 4 · p. 8
  • With a Vision Transformer backbone, COG-FC again achieves the best overall performance, suggesting the approach is not specific to convolutional architectures.The reported ViT comparison uses CIFAR-10.Source: §4.2 · p. 7, Table 5 · p. 7

4.3 USE CASE: MULTI-TASK LEARNING

  • On CelebA’s 40 binary tasks, all COG variants outperform their corresponding baselines, and COG-FC has the best overall performance.COG-FC achieves higher Accuracy and lower Brier score than Temperature Scaling on 32 of 40 tasks, and than Label Smoothing on all 40.Source: §4.3 · pp. 8–9, Table 6 · p. 9
  • Across MTL methods, task difficulty ordering is similar, with Big Lips consistently among the most challenging tasks, suggesting difficulty is primarily task-dependent.The evaluation reports unweighted averages across tasks.Source: §4.3 · p. 8

5 CONCLUSION

COG reframes neural-network calibration as bilevel, Stackelberg-inspired joint optimization. The authors report strong performance across diverse settings, including multi-task learning, while COG-FC adds no trainable parameters.

5 CONCLUSION

  • COG matches or outperforms existing methods in most settings, with COG-FC achieving state-of-the-art calibration performance without additional trainable parameters.The conclusion also reports stable optimization, good OOD-detection performance, and strong overall and task-specific results in multi-task learning without separate calibration procedures.Source: §5 · p. 9
  • The evaluation could extend to alternative calibration definitions and metrics; the experiments focus on computer vision, while other differentiable models and domains remain possible extensions.The authors also leave optimistic and pessimistic bilevel optimization unexplored and identify more heterogeneous multi-task settings as future work.Source: §5 · p. 9

AI USE STATEMENT

  • The authors used generative AI for code assistance and text polishing, reviewed the assisted work, and take responsibility for the final content.They state that AI was not used for the other required-disclosure tasks.Source: AI USE STATEMENT · pp. 9–10

REPRODUCIBILITY STATEMENT

  • Experiments used fixed random seeds, and supplementary materials provide code and configuration files for reproducing the reported results.Appendix B contains training, checkpoint-selection, calibration, experimental-setting, and hardware details.Source: REPRODUCIBILITY STATEMENT · p. 10

A STACKELBERG AND GRADIENT-DESCENT ASCENT

  • Bilevel optimization has applications including hyperparameter optimization, multi-task learning, architecture search, robustness, generalization, and calibration.Gradient-based methods typically approximate the best-response Jacobian through unrolled updates or implicit differentiation involving the Hessian.Source: A STACKELBERG AND GRADIENT-DESCENT… · p. 14

A.1 DEFINITION OF EQUILIBRIA

  • A Stackelberg equilibrium accounts for the follower’s optimal response, unlike Nash equilibrium, where each player’s strategy is an optimal response to the others.The leader’s total derivative captures the follower’s implicit response; the solution set can include multiple follower optima.Source: §A.1 · p. 15
  • Gradient descent–ascent is computationally simple for approximating Stackelberg problems, but targets Nash equilibria and may converge to spurious stable points.It can yield valid but restrictive Stackelberg solutions and may fail to converge reliably to Nash equilibria because of rotational dynamics.Source: §A.1 · p. 15

B.1 TRAINING DETAILS

The appendix details training, calibration, checkpoint selection, and hyperparameter tuning, including configurations for different datasets, backbones, and COG variants.

B.1 TRAINING DETAILS

  • ResNet-101 experiments ran for 300 epochs in single-task and 100 in multitask settings, using Adam with batch size 256 and, unless stated otherwise, five random seeds.The experiments used PyTorch 2.11.Source: §B.1 · p. 16
  • COG used cross-entropy in both settings, with learning-rate schedules varying by experiment; the ViT ablation used up to 400 epochs and calibrator and backbone learning rates of 1e-5 and 1e-4.For COG, the learning rate was reduced tenfold every ten epochs, or every three epochs for CelebA; multitask losses used equal task weights. Table 8 reports selected configurations.Source: §B.1 · p. 16, Table 8 · p. 17
  • Post-hoc calibrators were fitted on the validation set, and model checkpoints were selected by validation Brier score; label smoothing used a value of 0.1.Temperature Scaling used its original implementation, while other post-hoc calibrators and evaluation metrics used Probmetrics 1.2.0.Source: §B.1 · p. 16

B.2 MODEL SELECTION

  • COG hyperparameters were tuned primarily over the two learning rates using validation cross-entropy, while baseline learning rate was 1e-4 and TST used its reported default settings.Subsequent experiments generally used the best configurations from the tuning procedure.Source: §B.2 · p. 17

B.3 DATASET DETAILS

The appendix specifies dataset splits, preprocessing, validation use, and hardware requirements for the experiments.

B.3 DATASET DETAILS

  • CIFAR-10, CIFAR-100, and SVHN used validation fractions of 20%, 5%, and 20%, respectively, with training-only random crops and horizontal flips.Validation and test images were evaluated without stochastic augmentation; CelebA used predefined splits and 32 × 32 resized, center-cropped images.Source: §B.3 · p. 17
  • The held-out validation set was used for both checkpoint selection and fitting post-hoc calibrators, and all 40 tasks were included in multitask experiments.Datasets were normalized using dataset-specific channel-wise means and standard deviations.Source: §B.3 · p. 17

B.4 HARDWARE REQUIREMENTS

  • Experiments used six CPU cores, 16 GB of RAM, and a GPU with at least 8 GB of VRAM.They were conducted on a high-performance computing cluster with automatic job scheduling.Source: §B.4 · p. 17

C EXTENDED RESULTS

The extended results compare COG variants with post-hoc methods across calibration metrics and binning choices. They report generally favorable COG comparisons while showing that metric rankings can depend on binning.

C EXTENDED RESULTS

  • The extended experiments additionally report ECE, RMSCE, and MCE as percentages, and include NLL where it is absent from the main text.They provide calibration errors under different norms and bin counts to reduce dependence on any single metric configuration.Source: C EXTENDED RESULTS · pp. 17–18, §C.1 · p. 18
  • ECE has documented limitations: its positive bias increases with model accuracy, its value depends on binning, and it evaluates calibration only for the top-1 predicted class.The passage also notes that standard calibration metrics alone can make comparisons between in-hoc and post-hoc methods problematic.Source: C EXTENDED RESULTS · p. 18

C.1 COMPARISON OF COG AND POST-HOC APPROACHES

  • Across almost all pairwise comparisons, COG variants outperform corresponding post-hoc methods on ECE; COG-TST performs best overall, with Matrix Scaling a notable exception on CIFAR-100.The same general COG advantage appears for RMSCE, while Matrix Scaling substantially worsens the Baseline on CIFAR-100.Source: §C.1 · p. 19, Table 9 · p. 19, Table 10 · p. 19
  • Method rankings can change with bin count across calibration metrics, with the effect most dramatic for MCE, whose values can rise substantially as bins increase.The tables report ECE, RMSCE, and MCE across varying bin counts.Source: §C.1 · p. 19, Table 11 · p. 19, Table 9 · p. 19

C.2 ALL METHODS COMPARED

Across the compared methods, COG-FC achieves the highest overall rank across both backbones without additional trainable parameters. Extended calibration and shift results show competitive COG-FC performance, alongside bin-sensitive rankings and task-level variation.

C.2 ALL METHODS COMPARED

  • COG-FC achieves the highest overall rank across both backbones without additional trainable parameters, end-to-end and without a held-out calibration set.Table 12 ranks methods by mean rank across accuracy, NLL, and Brier for each dataset.Source: §C.2 · p. 20, Table 12 · p. 20
  • Calibration-method rankings can switch across bin counts for ECE, RMSCE, and MCE.Tables 13–15 report these errors for multiple bin choices.Source: §C.2 · p. 20, Table 13 · p. 20, Table 14 · p. 21

C.3 PERFORMANCE UNDER COVARIATE SHIFTS AND OOD

  • Under covariate shifts and across calibration-error metrics and binning configurations, COG-FC is competitive, while V-TST performs particularly strongly on the considered calibration errors.Tables 18–20 report calibration errors averaged across considered covariate shifts.Source: §C.3 · p. 22, Table 18 · p. 22, Table 19 · p. 23
  • On CelebA, COG-FC has higher accuracy than Temperature Scaling and Label Smoothing on all 40 tasks, and better Brier scores on 32 tasks versus Temperature Scaling and all 40 versus Label Smoothing.Figure 6 presents the per-task metrics, with tasks ordered by mean Brier.Source: Figure 6 · p. 24

C.4 VIT ABLATION

This subsection reports calibration metrics for a ViT trained from scratch on CIFAR-10.

C.4 VIT ABLATION

  • The ViT is trained from scratch on CIFAR-10 using the architecture described by Jordahn & Olmos (2024).The results are reported in Table 21.Source: §C.4 · p. 24
  • Table 21 compares ViT calibration errors on CIFAR-10 across multiple bin counts.Source: Table 21 · p. 24

C.5 MULTI-TASK LEARNING

On CelebA multi-task learning, COG-FC and COG+LS improve on their counterparts, whereas COG-TS performs worse than Temperature Scaling. Per-task ECE results show shared task difficulty patterns and variation across methods.

C.5 MULTI-TASK LEARNING

  • On CelebA, COG-FC and COG+LS consistently improve on their counterparts in averaged calibration errors, while COG-TS performs worse than Temperature Scaling.The reported overall calibration error is the unweighted average of per-task errors, computed independently for each task across norms.Source: §C.5 · p. 24, Table 22 · p. 24
  • Across CelebA tasks, methods identify the same most challenging tasks, including Big Lips, while some tasks show substantial ECE variance, particularly for Label Smoothing.The Label Smoothing variability is consistent with the corresponding per-task Brier-score variability; Figure 7 plots ECE-100 for the best-performing models.Source: §C.5 · p. 24, Figure 7 · p. 25

D RESULTS WITH A WRN BACKBONE

With a WRN-28-10 backbone, calibration performance differs by dataset: COG variants dominate on SVHN, (V)-TST models perform particularly well on CIFAR-10, and COG variants are generally strong on CIFAR-100. The section also reports out-of-distribution detection and covariate-shift results, with settings otherwise identical to those used with the ResNet backbone.

D RESULTS WITH A WRN BACKBONE

Loading 2610.13693v1…