Improving Generalization with Flat Hilbert Bayesian Inference

Tuan Truong, Quyen Tran, Ngoc-Quan Pham, Nhat Ho, Dinh Phung, Trung Le

Abstract: We introduce Flat Hilbert Bayesian Inference (FHBI), an algorithm designed to enhance generalization in Bayesian inference. Our approach involves an iterative two-step procedure with an adversarial functional perturbation step and a functional descent step within the reproducing kernel Hilbert spaces. This methodology is supported by a theoretical analysis that extends previous findings on generalization ability from finite-dimensional Euclidean spaces to infinite-dimensional functional spaces. To evaluate the effectiveness of FHBI, we conduct comprehensive comparisons against nine baseline methods on the VTAB-1K benchmark, which encompasses 19 diverse datasets across various domains with diverse semantics. Empirical results demonstrate that FHBI consistently outperforms the baselines by notable margins, highlighting its practical efficacy.

Lay Summary:

Approximate Bayesian Inference techniques have been known to be effective for machine learning models to deal with uncertainty by approximating an unknown target distribution. In our research, we explore ways to improve the generalization ability of the Bayesian Inference models. To achieve this, we incorporate a technique called sharpness-aware minimization (SAM) into Bayesian Inference. SAM focuses on finding model settings that are less sensitive to small changes, leading to more reliable predictions. However, applying SAM directly to each model particle didn't yield significant improvements for the ensemble.

To address this, we developed a novel theoretical framework that formulates the *sharpness on the functional spaces*. Then, we applied this framework to monitor the sharpness of the *transportation functions* which govern the motion of the model particles. By monitoring and controlling this sharpness of the movements, we not only improved each model's ability to generalize but also enhanced the diversity of the model particles, hence yielding a better approximation of the target distribution and enhancing the generalization ability of the final ensemble.

We tested our approach on the VTAB-1K benchmark, which includes 19 diverse datasets from various fields like natural images, medical imaging, and simulated environments. Our method showed both theoretical and practical improvements, indicating its potential for enhancing machine learning predictions across different domains.

Link To Code: https://github.com/tuantruongubc/Flat-Hilbert-Variational-Inference

Primary Area: Probabilistic Methods->Variational Inference

Keywords: Bayesian Inference, Sharpness-aware Minimization, RKHS

Submission Number: 8258