Gradient-free variational learning with conditional mixture networks

TL;DR

This paper introduces CAVI-CMN, a fast, gradient-free variational method for training probabilistic conditional mixture networks, offering efficient Bayesian inference with competitive accuracy and scalability.

Contribution

The paper presents a novel gradient-free variational inference approach for conditional mixture networks using conjugacy and Pólya-Gamma augmentation, enabling efficient training without gradients.

Findings

CAVI-CMN achieves competitive or superior accuracy compared to MLE with backpropagation.

The method maintains competitive runtime and scales well with input size and number of experts.

It provides full posterior distributions over model parameters, enhancing uncertainty quantification.

Abstract

Balancing computational efficiency with robust predictive performance is crucial in supervised learning, especially for critical applications. Standard deep learning models, while accurate and scalable, often lack probabilistic features like calibrated predictions and uncertainty quantification. Bayesian methods address these issues but can be computationally expensive as model and data complexity increase. Previous work shows that fast variational methods can reduce the compute requirements of Bayesian methods by eliminating the need for gradient computation or sampling, but are often limited to simple models. We introduce CAVI-CMN, a fast, gradient-free variational method for training conditional mixture networks (CMNs), a probabilistic variant of the mixture-of-experts (MoE) model. CMNs are composed of linear experts and a softmax gating network. By exploiting conditional conjugacy and P\'olya-Gamma augmentation, we furnish Gaussian likelihoods for the weights of both the linear layers and the gating network. This enables efficient variational updates using coordinate ascent variational inference (CAVI), avoiding traditional gradient-based optimization. We validate this approach by training two-layer CMNs on standard classification benchmarks from the UCI repository. CAVI-CMN achieves competitive and often superior predictive accuracy compared to maximum likelihood estimation (MLE) with backpropagation, while maintaining competitive runtime and full posterior distributions over all model parameters. Moreover, as input size or the number of experts increases, computation time scales competitively with MLE and other gradient-based solutions like black-box variational inference (BBVI), making CAVI-CMN a promising tool for deep, fast, and gradient-free Bayesian networks.