Stochastic Neural Network with Kronecker Flow

TL;DR

This paper introduces the Kronecker Flow, a scalable invertible mapping for stochastic neural networks that captures complex dependencies in high-dimensional variational inference, demonstrating competitive results across various tasks.

Contribution

The paper proposes the Kronecker Flow, a novel scalable parameterization for noise generation in stochastic neural networks, extending the Kronecker product to invertible mappings.

Findings

Competitive performance in variational Bayesian neural networks

Effective in PAC-Bayes bound estimation

Improves approximate Thompson sampling in bandits

Abstract

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.