Replica-exchange Nos\'e-Hoover dynamics for Bayesian learning on large datasets

TL;DR

This paper introduces a novel Bayesian learning method that uses replica-exchange Nosé-Hoover dynamics to efficiently sample from complex, multimodal posterior distributions in large-scale datasets, improving over existing methods.

Contribution

It develops a practical, noise-neutralizing sampling technique with a new acceptance protocol, enabling effective Bayesian inference on large, complex datasets.

Findings

Effective sampling from multimodal posteriors demonstrated on synthetic data.

Significant improvements over strong baselines in deep Bayesian neural networks.

Robustness to mini-batch noise in large-scale Bayesian learning.

Abstract

In this paper, we present a new practical method for Bayesian learning that can rapidly draw representative samples from complex posterior distributions with multiple isolated modes in the presence of mini-batch noise. This is achieved by simulating a collection of replicas in parallel with different temperatures and periodically swapping them. When evolving the replicas' states, the Nos\'e-Hoover dynamics is applied, which adaptively neutralizes the mini-batch noise. To perform proper exchanges, a new protocol is developed with a noise-aware test of acceptance, by which the detailed balance is reserved in an asymptotic way. While its efficacy on complex multimodal posteriors has been illustrated by testing over synthetic distributions, experiments with deep Bayesian neural networks on large-scale datasets have shown its significant improvements over strong baselines.