Stochastic Neural Network with Kronecker Flow
This addresses scalability limitations in variational inference for stochastic neural networks, though it appears incremental as a generalization of existing techniques.
The paper tackles the challenge of scaling variational inference for stochastic neural networks by introducing Kronecker Flow, a method that generalizes the Kronecker product for invertible mappings to capture parameter dependencies. The results show the method is competitive with existing approaches and outperforms baselines in tasks like variational Bayesian neural networks, PAC-Bayes generalization bound estimation, and contextual bandits.
Recent advances in variational inference enable the modelling of highly structured joint distributions, but are limited in their capacity to scale to the high-dimensional setting of stochastic neural networks. This limitation motivates a need for scalable parameterizations of the noise generation process, in a manner that adequately captures the dependencies among the various parameters. In this work, we address this need and present the Kronecker Flow, a generalization of the Kronecker product to invertible mappings designed for stochastic neural networks. We apply our method to variational Bayesian neural networks on predictive tasks, PAC-Bayes generalization bound estimation, and approximate Thompson sampling in contextual bandits. In all setups, our methods prove to be competitive with existing methods and better than the baselines.