Minimal Random Code Learning with Mean-KL Parameterization

TL;DR

This paper introduces a novel Mean-KL parameterization for variational Bayesian neural network compression, leading to faster convergence, improved robustness, and more meaningful distributions compared to traditional methods.

Contribution

It proposes a Mean-KL parameterization for MIRACLE that simplifies training and enhances robustness and interpretability of compressed neural network weights.

Findings

Faster convergence in variational training.

More robust and meaningful weight distributions.

Improved compression performance with heavier tails.

Abstract

This paper studies the qualitative behavior and robustness of two variants of Minimal Random Code Learning (MIRACLE) used to compress variational Bayesian neural networks. MIRACLE implements a powerful, conditionally Gaussian variational approximation for the weight posterior $Q_{\mathbf{w}}$ and uses relative entropy coding to compress a weight sample from the posterior using a Gaussian coding distribution $P_{\mathbf{w}}$. To achieve the desired compression rate, $D_{\mathrm{KL}}[Q_{\mathbf{w}} \Vert P_{\mathbf{w}}]$ must be constrained, which requires a computationally expensive annealing procedure under the conventional mean-variance (Mean-Var) parameterization for $Q_{\mathbf{w}}$. Instead, we parameterize $Q_{\mathbf{w}}$ by its mean and KL divergence from $P_{\mathbf{w}}$ to constrain the compression cost to the desired value by construction. We demonstrate that variational training with Mean-KL parameterization converges twice as fast and maintains predictive performance after compression. Furthermore, we show that Mean-KL leads to more meaningful variational distributions with heavier tails and compressed weight samples which are more robust to pruning.