Scale-invariant Bayesian Neural Networks with Connectivity Tangent Kernel

TL;DR

This paper introduces a scale-invariant Bayesian neural network framework that improves generalization bounds and uncertainty calibration by decomposing parameters into scale and connectivity, addressing issues caused by parameter scaling.

Contribution

It proposes a novel prior and posterior distribution invariant to parameter scaling, enabling more accurate generalization bounds and uncertainty calibration for practical neural network transformations.

Findings

Invariant posterior improves flatness measures

Enhanced uncertainty calibration in Bayesian neural networks

Effective in practical parameter transformation scenarios

Abstract

Explaining generalizations and preventing over-confident predictions are central goals of studies on the loss landscape of neural networks. Flatness, defined as loss invariability on perturbations of a pre-trained solution, is widely accepted as a predictor of generalization in this context. However, the problem that flatness and generalization bounds can be changed arbitrarily according to the scale of a parameter was pointed out, and previous studies partially solved the problem with restrictions: Counter-intuitively, their generalization bounds were still variant for the function-preserving parameter scaling transformation or limited only to an impractical network structure. As a more fundamental solution, we propose new prior and posterior distributions invariant to scaling transformations by \textit{decomposing} the scale and connectivity of parameters, thereby allowing the resulting generalization bound to describe the generalizability of a broad class of networks with the more practical class of transformations such as weight decay with batch normalization. We also show that the above issue adversely affects the uncertainty calibration of Laplace approximation and propose a solution using our invariant posterior. We empirically demonstrate our posterior provides effective flatness and calibration measures with low complexity in such a practical parameter transformation case, supporting its practical effectiveness in line with our rationale.