Turning Uncertainty into Control: Bi-level Training with Editable Bayesian Layers

Xinyue Hu, Bowen Xiang, Zhibin Duan, Bo Chen, Yuxin Li, Xuefei Cao, Hongwei Liu, Mingyuan Zhou

Keywords: Uncertainty Estimation, Interpretability

Abstract:

As deep learning systems are increasingly deployed in high-stakes settings, it is
essential not only to quantify predictive uncertainty but also to use it to steer train-
ing and improve downstream decision. Yet dense parameter updates often cause
collateral interference due to the entangled nature of parameter. In this paper,
We propose Bayesian Feature Reweighting (BFR), a framework that turns cali-
brated, instance-level uncertainty into training-time control signals while enforc-
ing sparse, localized parameter updates. Concretely, BFR employs a Bayesian
Last Layer to obtain well-calibrated predictive uncertainty with negligible over-
head and imposes sparsity on the decision layer to concentrate changes on task-
relevant parameters, thereby limiting unintended perturbations. It then tinkers
the model via uncertainty-guide bi-level optimization. Across vision and lan-
guage benchmarks, BFR consistently improves group robustness while maintain-
ing competitive accuracy and yields smaller parameter changes. Code is provided
in the supplementary material.

Primary Area: interpretability and explainable AI

Submission Number: 9203