Gradient-free variational learning with conditional mixture networks

## Abstract

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 demonstrate that conditional mixture networks (CMNs), a probabilistic variant of the mixture-of-experts (MoE) model, are suitable for fast, gradient-free inference and can solve complex classification tasks. CMNs employ linear experts and a softmax gating network. By exploiting conditional conjugacy and Pólya-Gamma augmentation, we furnish Gaussian likelihoods for the weights of both the linear experts 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 benchmarks from the UCI repository. Our method, 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.

## 1 Introduction

Modern machine learning methods attempt to learn functions of complex data (e.g., images, audio, text) to predict information associated with that data, such as discrete labels in the case of classification (Bernardo et al., ). Deep neural networks (DNNs) have demonstrated success in this domain, owing to their universal function approximation properties (Park and Sandberg, ) and the soft regularization inherited from stochastic gradient descent learning via back propagation(Amari, ). However, despite its computational efficiency, accuracy, and scalability to increasingly large datasets and models, DNNs trained this way do not provide well calibrated predictions and uncertainty estimates, and practitioners typically utilize post-hoc calibration methods on validation datasets (Wang et al., ; Shao et al., ). This limits the applicability and reliability of using DNNs in safety-critical applications like autonomous driving, medicine, and disaster response (Papamarkou et al., ), where uncertainty-sensitive decision-making is required.

Bayesian machine learning addresses the issues of poor calibration and uncertainty quantification by offering a probabilistic framework that casts learning model parameters $\boldsymbol{\theta}$ as a process of inference - namely, calculating a posterior distribution over model parameters, given observed data ($\boldsymbol{\mathcal{D}}=\del{(\boldsymbol{x}_{1},\boldsymbol{y}_{1}),\ldots,(
\boldsymbol{x}_{n},\boldsymbol{y}_{n})}$), using Bayes’ rule:

$$ $p\del{\boldsymbol{\theta}\mid\boldsymbol{\mathcal{D}}}=\frac{p\del{\boldsymbol {\theta},\boldsymbol{\mathcal{D}}}}{p\del{\boldsymbol{\theta}}}$ (1) $$

The resulting posterior distribution captures both expectations about model parameters $\boldsymbol{\theta}$ and their uncertainty. The uncertainty is then incorporated in predictions that are, in principle, well-calibrated to novel datapoints coming from the same set. This probabilistic treatment allows methods like Bayesian neural networks (BNNs) (Hernández-Lobato and Adams, ) to maintain the expressiveness of deep neural networks while also encoding uncertainty over network weights and thus the network’s predictions. However, these methods are known to come with a significant increase in computational cost and thus scale poorly when applied to large datasets and high-dimensional models (Izmailov et al., ).

In this paper we introduce a gradient-free variational learning algorithm for a probabilistic variant of a two-layer, feedforward neural network — the conditional mixture network or CMN — and measure its performance on supervised learning benchmarks. This method rests on coordinate ascent variational inference (CAVI) (Wainwright et al., ; Hoffman et al., ) and hence we name it CAVI-CMN. We compare CAVI-CMN to maximum likelihood estimation and two other Bayesian estimation techniques: the No U-Turn Sampler (NUTS) variant of Hamiltonian Monte Carlo (Hoffman et al., ) and black-box variational inference (Ranganath et al., ). We demonstrate that CAVI-CMN maintains the predictive accuracy and scalability of an architecture-matched feedforward neural network fit with maximum likelihood estimation (i.e., gradient descent via backpropagation), while maintaining full distributions over network parameters and generating calibrated predictions, as measured in relationship to state-of-the-art Bayesian methods like NUTS and BBVI.

We summarize the contributions of this work below:

- •
Introduce and derive a variational inference scheme for the conditional mixture network, which we term CAVI-CMN. This relies on the use of conjugate priors for the linear experts and Pólya-Gamma augmentation (Polson et al., ) for the gating network and the final softmax layer.
- •
CAVI-CMN matches, and sometimes exceeds, the performance of maximum likelihood estimation (MLE) in terms of predictive accuracy, while maintaining the probabilistic benefits of being Bayesian, like quantifying uncertainty and have low calibration error. This is shown across a suite of 8 different supervised classification tasks (2 synthetic, 6 real).
- •
CAVI-CMN displays all the benefits explained above while requiring drastically less time to converge and overall runtime than the other state-of-the-art Bayesian methods like NUTS and BBVI.

The rest of this paper is organized as follows: first, we discuss related works include the MoE architecture and existing (Bayesian and non-Bayesian approaches) to fitting these models. We then introduce the probabilistic conditional mixture model and derive a variational inference algorithm for optimizing posterior distributions over its latent variables and parameters. We present experimental results comparing the performance of CAVI-based conditional mixture models with sampling based methods, such as BBVI, NUTS, and traditional MLE based estimation, where gradients of the log likelihood are computed using backpropagation and used to update the network’s parameters. Finally, we discuss the implications of these findings and potential directions for future research.

## 2 Related work

The Mixture-of-Experts (MoE) architecture is a close relative of the CMN model we introduce here. Jacobs et al. () originally introduced MoEs as a way to improve the performance of neural networks by combining the strengths of multiple specialized models (Gormley and Frühwirth-Schnatter, ). MoE models process inputs by averaging the predictions of individual learners or experts, where each expert’s output is weighted by a different mixing coefficient before the averaging. The fundamental idea behind MoE is that the input space can be partitioned in such a way that different experts (models) can be trained to excel in different regions of this space, with a gating network determining the appropriate expert (or combination of experts) for each input. This leads to composable (and sometimes interpretable) latent descriptions of arbitrary input-output relationships (Eigen et al., ), further bolstered by the MoE’s capacity for universal function approximation (Nguyen et al., ; Nguyen and Chamroukhi, ). Indeed, the powerful self-attention mechanism employed by transformers has has demonstrated the power and flexibility of MoE models (Movellan and Gabbur, ). Non-Bayesian approaches to MoE typically rely on maximum likelihood estimation (MLE) (Jacobs et al., ; Jordan and Jacobs, ), which can suffer from overfitting and poor generalization due to the lack of regularization mechanisms (Bishop and Svenskn, ), especially in low data size regimes.

To address these issues, Bayesian approaches to MoE have been developed, which incorporate prior information and yield posterior distributions over model parameters (Bishop and Svenskn, ; Mossavat and Amft, ). This Bayesian treatment enables the estimation of model evidence (log marginal-likelihood) and provides a natural framework for model comparison and selection (Svensén, ; Zens, ). Bayesian MoE models offer significant advantages, such as improved robustness against overfitting and a better understanding of uncertainty in predictions. However, they also introduce computational challenges, particularly when dealing with high-dimensional data and complex model structures.

The introduction of the Pólya-Gamma (PG) augmentation technique in Polson et al. () enabled a range of novel and more computationally efficient algorithms for Bayesian treatment of MoE models (Linderman et al., ; He et al., ; Sharma et al., ; Viroli and McLachlan, ; Zens et al., ). Here we complement these past works, which mostly rest on improving sampling methods with PG augmentation, by introducing a closed-form update rules for MoE’s with linear experts in the form of coordinate ascent variational inference (CAVI).

## 3 Methods

In this section we first motivate the use of conditional mixture models for supervised learning, and then introduce the conditional mixture network (CMN), the probabilistic model whose properties and capabilities we demonstrate in the remainder of the paper.

### 3.1 Conditional mixtures for function approximation

Feedforward neural networks are highly expressive, approximating nonlinear functions through sequences of nonlinear transformations, but the posterior distributions over their weights are intractable, requiring expensive techniques like MCMC or variational inference (MacKay, ; Blundell et al., ; Daxberger et al., ).

We circumvent these problems by focusing on the Mixture-of-Experts (MoE) models (Jacobs et al., ), and particularly a variant of MoE that is amenable to gradient free CAVI parameter updates. MoEs can be made tractable to gradient-free CAVI when the expert likelihoods are constrained to be members of the exponential family (see Section 2 for more details on the MoE architecture), and when the gating network is formulated in such a way to allow exact Bayesian inference (through lower bounds on the log-sigmoid likelihood (Jaakkola and Jordan, ; Bishop and Svenskn, ) or Pólya-Gamma augmentation (Polson et al., )).

The MoE can be reformulated probabilistically as a mixture model by introducing a latent assignment variable, $z^{n}$, leading to a joint probability distribution of the form

$$ $p(Y,Z,\Theta)=p(\theta_{1:K})p(\pi)\prod_{i=1}^{n}p(y^{n}|z^{n},\theta_{1:K})p (z^{n}|\pi)~{},$ $$

where $y^{n}$ is an observation, $\Theta=\{\theta_{1:K},\pi\}$, $p_{k}(y^{n}|\theta_{k})$ is the $k$^th-component’s likelihood and $z^{n}$ is a discrete latent variable that assigns the $n$^th datapoint to one of the $K$ mixture components, i.e. $p_{k}(y^{n}|\theta_{k})=p(y^{n}|z^{n}\!=\!k,\theta_{1:K})$. For instance, if each ‘expert’ likelihood $p_{k}(y^{n}|\theta_{k})$ is a Gaussian distribution, then the MoE becomes a Gaussian Mixture Model, where $\theta_{k}=(\mu_{k},\Sigma_{k})$.

The problem of learning the model’s parameters, then becomes one of doing inference over the latent variables $Z$ and parameters $\Theta$ of the mixture model. However, mixture models are generally not tractable for exact Bayesian inference, so some form of approximation or sampling-based scheme is required to obtain full posteriors over their parameters. However, if each expert (i.e., likelihood distribution) in the MoE belongs to the exponential family, the model becomes conditionally conjugate. This allows for derivation of exact fixed-point updates to an approximate posterior over each expert’s parameters. The approach we propose, CAVI-CMN, does exactly this – we take advantage of the conditional conjugacy of mixture models, along with an augmentation trick for the the gating network, to make all parameters amenable to an approximate Bayesian treatment. The conditionally-conjugate form of the model allows us to use coordinate ascent variational inference to obtain posteriors over the weights of both the individual linear experts and the gating network (Wainwright et al., ; Hoffman et al., ; Blei et al., ), without resorting to costly gradient or sampling computations.

Going forward we use the term conditional mixture networks (CMN) to emphasize (1) the discriminative nature of proposed application of this approach, where the model is designed to predict an output $y$ given an input $x$ and (2) the fact that individual MoE layers can be stacked hierarchically into a feedforward architecture. This makes CMNs particularly suitable for tasks such as supervised classification and regression, where the goal is effectively that of function approximation; predict some output variable $y$ given input regressors $x$.

### 3.2 Conditional mixture network overview

The conditional mixture network maps from a continuous input vector $\boldsymbol{x}_{0}\in\mathbb{R}^{d}$ to its label $y\in\cbr{1,\ldots,L}$. This is achieved with two layers: a conditional mixture of linear experts, which outputs a joint continuous-discrete latent $\del{\boldsymbol{x}_{1}\in\mathbb{R}^{h},z_{1}\in\cbr{1,\ldots,K}}$ and a multinomial logistic regression, which maps from the continuous latent $\boldsymbol{x}_{1}$ to the corresponding label $y$. The probabilistic mapping can be described in terms of the following operations:

$$ $\begin{split}z_{1}&\sim\text{Mult}\del{z_{1};\boldsymbol{x}_{0},\boldsymbol{ \beta}_{0}}\\ \boldsymbol{x}_{1}&=\boldsymbol{A}_{z_{1}}\cdot\sbr{\boldsymbol{x}_{0};1}+ \boldsymbol{u}_{z_{1}},\hskip 28.45274pt\boldsymbol{u}_{z_{1}}\sim N( \boldsymbol{0},\boldsymbol{\Sigma}_{z_{1}})\\ y&\sim\text{Mult}\del{y;\boldsymbol{x}_{1},\boldsymbol{\beta}_{1}}\end{split}$ $$

where we pad the input variable $\boldsymbol{x}_{0}$ with a constant value set to $1$, to absorb the bias term within the mapping matrix $\boldsymbol{A}_{z_{1}}\in\mathbb{R}^{h\times d+1}$, and where $\text{Mult}\del{z;\boldsymbol{x},\boldsymbol{\beta}}$ denotes a multinomial distribution parameterized with a regressor $\boldsymbol{x}$ and logistic regression coefficients $\boldsymbol{\beta}$. Note that for every pair of regressors and labels $\del{\boldsymbol{x}^{n}_{0},y^{n}}$, we assume a corresponding pair of latent variables $\del{\boldsymbol{x}_{1}^{n},z_{1}^{n}}$. Written in this way, it becomes clear than CMN is a mixture of linear transforms that is capable of modeling non-linear transfer functions via a piecewise linear approximation.

In order to obtain a normally distributed posterior over the multinomial logistic regression weights, $\boldsymbol{\beta}_{0}$ and $\boldsymbol{\beta}_{1}$, we use Pólya-Gamma augmentation (Polson et al., ; Linderman et al., ) applied to the stick breaking construction for the multinomial distribution:

$$ $\begin{split}p\del{z=k|\boldsymbol{\beta},\boldsymbol{x}}&=\pi_{k}\del{ \boldsymbol{\beta},\boldsymbol{x}}\prod_{j=1}^{k-1}\del{1-\pi_{j}\del{ \boldsymbol{\beta},\boldsymbol{x}}}\\ \pi_{j}\del{\boldsymbol{\beta},\boldsymbol{x}}&=\frac{1}{1+\exp\cbr{- \boldsymbol{\beta}_{j}\cdot\sbr{\boldsymbol{x};1}}},\forall j<K\\ \pi_{K}&=1\end{split}$ (2) $$

where for the gating network (input layer) we will have coefficients of dimension $\boldsymbol{\beta}_{0}\in\mathbb{R}^{K-1\times d}$, and for the output likelihood coefficients of dimension $\boldsymbol{\beta}_{1}\in\mathbb{R}^{L-1\times h}$.

### 3.3 Generative model for the conditional mixture network

Figure: Figure 1: A Bayesian network representation of the two-layer conditional mixture network, with input-output pairs $\boldsymbol{x}^{n}_{0},y^{n}$ and latent variables $\boldsymbol{x}^{n}_{1},z^{n}_{1}$. Observations are shaded nodes, while latents and parameters are transparent.

Given a set of labels $Y=\cbr{y^{1},y^{2},...,y^{N}}$, and regressors $\boldsymbol{X}_{0}=\cbr{\boldsymbol{x}^{1}_{0},\boldsymbol{x}^{2}_{0},...,
\boldsymbol{x}^{N}_{0}}$, that define i.i.d input-output pairs $\boldsymbol{x}^{n}_{0},y^{n}$, we write the joint distribution over labels $Y$, latents $\boldsymbol{X}_{1},Z_{1}$, and parameters $\boldsymbol{\Theta}$ as:

$$ $\begin{split}p(\boldsymbol{Y},\boldsymbol{X}_{1},\boldsymbol{Z}_{1}, \boldsymbol{\Theta}|\boldsymbol{X}_{0})&=p\del{\boldsymbol{\Theta}}\prod_{n=1} ^{N}p_{\boldsymbol{\beta_{1}}}\del{y^{n}|\boldsymbol{x}_{1}}p_{\boldsymbol{ \lambda}_{1}}\del{\boldsymbol{x}_{1}^{n}|\boldsymbol{x}^{n}_{0},z_{1}^{n}}p_{ \boldsymbol{\beta}_{0}}\del{z_{1}^{n}|\boldsymbol{x}^{n}_{0}}\\ p\del{\boldsymbol{\Theta}}&=p\del{\boldsymbol{\beta}_{1}}p\del{\boldsymbol{ \beta}_{0}}p\del{\boldsymbol{\lambda}_{1}}\\ &=\prod_{l=1}^{L-1}p(\boldsymbol{\beta}_{l,1})\prod_{k=1}^{K-1}p(\boldsymbol{ \beta}_{k,0})\prod_{j=1}^{K}p\del{\boldsymbol{A}_{j},\boldsymbol{\Sigma}_{j}^{ -1}}\end{split}$ (3) $$

Note that this model structure, with input and target variables, is often referred to as a discriminative model, as opposed to a generative model (Bernardo et al., ). However, we use the term generative model to emphasize the fact that the model contains priors over latent variables ($\boldsymbol{X}_{1},Z_{1}$), and parameters $\del{\boldsymbol{\Theta}=\del{\boldsymbol{\beta}_{1:L-1,1},\boldsymbol{\beta}_
{1:K-1,0},\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}_{1:K}^{-1}}}$, and that we are estimating posteriors over these quantities, by maximizing a lower bound on marginal likelihood of the observed target variables $Y$. Note that going forward, we will sometimes use $\boldsymbol{\lambda}_{1}$ as notational shorthand for the parameters $\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K}$ of the first layer’s linear experts.

We specify the following conditionally conjugate priors for the parameters of the two-layer CMN:

$$ $\begin{split}p\del{\boldsymbol{A}_{k}|\boldsymbol{\Sigma}^{-1}_{k}}&=\mathcal{ MN}\del{\boldsymbol{A}_{k};\boldsymbol{0},\boldsymbol{\Sigma}_{k},v_{0} \boldsymbol{I}_{d+1}}\\ p\del{\boldsymbol{\Sigma}^{-1}_{k}\equiv\text{diag}\del{\boldsymbol{\sigma}_{k }^{-2}}}&=\prod_{i=1}^{h}\Gamma\del{\sigma_{k,i}^{-2};a_{0},b_{0}}\\ p\del{\boldsymbol{\beta}_{k,0}}&=\mathcal{N}\del{\boldsymbol{\beta}_{k,0}; \boldsymbol{0},\sigma^{2}_{0}\boldsymbol{I}_{d+1}}\\ p\del{\boldsymbol{\beta}_{l,1}}&=\mathcal{N}\del{\boldsymbol{\beta}_{l,1}; \boldsymbol{0},\sigma^{2}_{1}\boldsymbol{I}_{h+1}}\end{split}$ (4) $$

In the following section we introduce a mean-field variational inference scheme we use for performing inference and learning in the two-layer CMN.

### 3.4 Coordinate ascent variational inference with conjugate priors

In this section we detail a variational approach for inverting the probabilistic model described in Equation 3 and computing an approximate posterior over latents and parameters specified as

$$ $p\del{\boldsymbol{X}_{1},\boldsymbol{Z}_{1},\boldsymbol{\Theta}|Y,\boldsymbol{ X}}=\frac{p\del{Y,\boldsymbol{X}_{1},\boldsymbol{Z}_{1},\boldsymbol{\Theta}, \boldsymbol{X}}}{p\del{Y|\boldsymbol{X}}}\approx q\del{\boldsymbol{\Theta}} \prod_{n=1}^{N}q\del{z_{1}^{n}}q\del{\boldsymbol{x}_{1}^{n}|z_{1}^{n}}$ (5) $$

where $q\del{\boldsymbol{x}_{1}^{n}|z_{1}^{n}}$ corresponds to a component specific multivariate normal distribution, and $q\del{z_{1}^{n}}$ to a multinomial distribution. Importantly, the approximate posterior over parameters $q\del{\boldsymbol{\Theta}}$ further factorizes (Svensén, ) as

$$ $\begin{split}q\del{\boldsymbol{\Theta}}&=\prod_{l=1}^{L-1}q\del{\boldsymbol{ \beta}_{l,1}}\prod_{k=1}^{K-1}q\del{\boldsymbol{\beta}_{k,0}}\underbrace{\prod _{j=1}^{K}q\del{\boldsymbol{A}_{j},\boldsymbol{\Sigma}_{j}^{-1}}}_{=q( \boldsymbol{\lambda}_{1})}\\ q\del{\boldsymbol{\beta}_{l,1}}&=\mathcal{N}\del{\boldsymbol{\beta}_{l,1}; \boldsymbol{\mu}_{l,1},\boldsymbol{\Sigma}_{l,1}}\\ q\del{\boldsymbol{\beta}_{l,0}}&=\mathcal{N}\del{\boldsymbol{\beta}_{l,0}; \boldsymbol{\mu}_{k,0},\boldsymbol{\Sigma}_{k,0}}\\ q\del{\boldsymbol{A}_{j}|\boldsymbol{\Sigma}_{j}^{-1}}&=\mathcal{MN}\del{ \boldsymbol{A}_{j};\boldsymbol{M}_{j},\boldsymbol{\Sigma}_{j},\boldsymbol{V}_{ j}}\\ q\del{\boldsymbol{\Sigma}_{j}^{-1}}&=\prod_{i=1}^{h}\Gamma\del{\sigma_{i,j}^{- 2};a_{j},b_{i,j}}\end{split}$ (6) $$

The above form of the approximate posterior allows us to define tractable conditionally conjugate updates for each factor. This becomes evident from the following expression for the evidence lower-bound (ELBO) on the marginal log likelihood

$$ $\mathcal{L}(q)=\mathbb{E}_{q\del{\boldsymbol{X}_{1},\boldsymbol{Z}_{1}}q\del{ \boldsymbol{\Theta}}}\sbr{\sum_{n=1}^{N}\ln\frac{p_{\boldsymbol{\Theta}}\del{y ^{n},\boldsymbol{x}^{n}_{1},z_{1}^{n}|\boldsymbol{x}^{n}_{0}}}{q\del{z_{1}^{n} }q\del{\boldsymbol{x}_{1}^{n}|z_{1}^{n}}}}+\mathbb{E}_{q\del{\boldsymbol{ \Theta}}}\sbr{\ln\frac{p\del{\boldsymbol{\beta}_{1}}p\del{\boldsymbol{\beta}_{ 0}}p\del{\boldsymbol{\lambda}_{1}}}{q\del{\boldsymbol{\beta}_{1}}q\del{ \boldsymbol{\beta}_{0}}q\del{\boldsymbol{\lambda}_{1}}}}$ (7) $$

We maximize the ELBO using an iterative update scheme for the parameters of the approximate posterior, often referred to as variational Bayesian expectation maximisation (VBEM) (Beal, ) or coordinate ascent variational inference (CAVI) (Bishop and Nasrabadi, ; Blei et al., ). The procedure consists of two parts:

First, we fix the posterior over the parameters (to randomly initialized values). Given the posterior over parameters, we update the posterior over latent variables (variational E-step) as

$$ $\begin{split}q_{t}\del{\boldsymbol{x}^{n}_{1}|z^{n}_{1}}&\propto\exp\cbr{ \mathbb{E}_{q_{t-1}\del{\boldsymbol{\beta}_{1}}q_{t-1}\del{\boldsymbol{\lambda }_{1}}}\sbr{\ln p_{\boldsymbol{\beta}_{1}}\del{y^{n}|\boldsymbol{x}_{1}^{n}}+ \ln p_{\boldsymbol{\lambda}_{1}}\del{\boldsymbol{x}_{1}^{n}|\boldsymbol{x}^{n} _{0},z_{1}^{n}}}}\\ q_{t}\del{z^{n}_{1}}&\propto\exp\cbr{\mathbb{E}_{q_{t-1}\del{\boldsymbol{ \Theta}}}\sbr{\left\langle\ln p_{\boldsymbol{\beta}_{1},\boldsymbol{\lambda}_{ 1}}\del{y^{n},\boldsymbol{x}_{1}^{n}|\boldsymbol{x}^{n}_{0},z_{1}^{n}}\right \rangle_{q_{t}\del{\boldsymbol{x}^{n}_{1}|z^{n}_{1}}}+\ln p_{\boldsymbol{\beta }_{0}}\del{z^{n}_{1}|\boldsymbol{x}^{n}_{0}}}}\end{split}$ (8) $$

Second, the posterior over latents that was updated in the E-step, is used to update the posterior over parameters (variational M-step) as

$$ $\begin{split}q_{t}\del{\boldsymbol{\beta}_{1}}&\propto\exp\cbr{\sum_{n=1}^{N} \mathbb{E}_{q_{t}\del{\boldsymbol{x}^{n}_{1},z_{1}^{n}}}\sbr{\ln p_{ \boldsymbol{\beta}_{1}}\del{y^{n}|\boldsymbol{x}_{1}^{n}}}}\\ q_{t}\del{\boldsymbol{\beta}_{0}}&\propto\exp\cbr{\sum_{n=1}^{N}\mathbb{E}_{q_ {t}\del{z^{n}_{1}}}\sbr{\ln p_{\boldsymbol{\beta}_{1}}\del{z_{1}^{n}| \boldsymbol{x}^{n}_{0}}}}\\ q_{t}\del{\boldsymbol{\lambda}_{1}}&\propto\exp\cbr{\sum_{n=1}^{N}\mathbb{E}_{ q_{t}\del{\boldsymbol{x}^{n}_{1},z^{n}_{1}}}\left[\ln p_{\boldsymbol{\lambda}_ {1}}\del{\boldsymbol{x}^{n}_{1}|z^{n}_{1},\boldsymbol{x}^{n}_{0}}\right]}\end{split}$ (9) $$

In the variational inference literature, the distinction between latents and parameters is often described in terms of ‘local’ vs ‘global’ latent variables (Hoffman et al., ), where local variables are datapoint-specific, and global variables are shared across datapoints. The exact form of the updates to the parameters of the linear experts in Equation 9, i.e. $q_{t}(\boldsymbol{\lambda}_{1})=q_{t}(\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}
^{-1}_{1:K})$ are found in Appendix B.

Importantly, the update equations described in Equation 8 and in the first two lines of Equation 9 are not computationally tractable without an additional approximation, known as Pólya-Gamma augmentation of the multinomial distribution. The full details of the augmentation procedure are described in Section A.1. Here we will briefly sketch the main steps and describe the high level, augmented update equations. The Pólya-Gamma augmentation introduces datapoint-specific auxiliary variables $\del{\boldsymbol{\omega}_{1}^{n},\boldsymbol{\omega}_{0}^{n}}$, that help us transform the log-probability of the multinomial distribution into a quadratic function (Polson et al., ; Linderman et al., ) over coefficients ($\boldsymbol{\beta}_{1},\boldsymbol{\beta}_{0}$), and latents $\boldsymbol{x}_{1}^{n}$. This quadratic form enables tractable update of $q\del{\boldsymbol{x}_{1}^{n}|z_{1}^{n}}$ in the form of a multivariate normal distribution, and a tractable updating of posteriors over coefficients $q\del{\boldsymbol{\beta}_{1}}$ and $q\del{\boldsymbol{\beta}_{0}}$.

With the introduction of the auxiliary variables the variational expectation and maximisation steps are expressed as

$$ $\begin{split}\text{Update latents (`E-step')}\\ q_{t}\del{\boldsymbol{x}^{n}_{1}|z^{n}_{1}}&\propto\exp\cbr{\mathbb{E}_{q_{t-1 }\del{\boldsymbol{\beta}_{1}}q_{t-1}\del{\boldsymbol{\lambda}_{1}}}\sbr{\left \langle l\del{y^{n},\boldsymbol{x}_{1}^{n},\boldsymbol{\omega}_{1}^{n}, \boldsymbol{\beta}_{1}}\right\rangle_{q_{t-1}\del{\boldsymbol{\omega}_{1}^{n}| y^{n}}}+\ln p_{\boldsymbol{\lambda}_{1}}\del{\boldsymbol{x}_{1}^{n}| \boldsymbol{x}^{n}_{0},z_{1}^{n}}}}\\ q_{t}\del{\boldsymbol{\omega}_{1}|y^{n}}&\propto p\del{\boldsymbol{\omega}_{1} ^{n}|y_{n}}\exp\cbr{\mathbb{E}_{q_{t-1}\del{\boldsymbol{\beta}_{1}}q_{t}\del{ \boldsymbol{x}_{1}^{n}|z_{1}^{n}}}\sbr{l\del{y^{n},\boldsymbol{x}_{1}^{n}, \boldsymbol{\omega}_{1}^{n},\boldsymbol{\beta}_{1}}}}\\ q_{t}\del{\boldsymbol{\omega}_{0}|z_{1}^{n}}&\propto p\del{\boldsymbol{\omega} _{0}^{n}|z_{1}^{n}}\exp\cbr{\mathbb{E}_{q_{t-1}\del{\boldsymbol{\beta}_{0}}} \sbr{l\del{z_{1}^{n},\boldsymbol{x}^{n}_{0},\boldsymbol{\omega}_{0}^{n}, \boldsymbol{\beta}_{0}}}}\\ q_{t}\del{z^{n}_{1}}&\propto\exp\cbr{\mathbb{E}_{q_{t-1}\del{\boldsymbol{ \Theta}}}\sbr{\bar{l}_{z_{1}^{n},t}\del{y^{n},\boldsymbol{\beta}_{1}}+R_{z_{1} ^{n},t}\del{\boldsymbol{x}^{n}_{0},\boldsymbol{\lambda}_{1}}+\bar{l}_{t}\del{z _{1}^{n},\boldsymbol{x}^{n}_{0},\boldsymbol{\beta}_{0}}}}\\ \text{Update parameters (`M-step')}\\ q_{t}\del{\boldsymbol{\beta}_{1}}&\propto\exp\cbr{\sum_{n=1}^{N}\mathbb{E}_{q_ {t}\del{\boldsymbol{x}^{n}_{1},z_{1}^{n}}q_{t}\del{\boldsymbol{\omega}^{n}_{1} |y^{n}}}\sbr{l(y^{n},\boldsymbol{x}^{n}_{1},\boldsymbol{\beta}_{1},\boldsymbol {\omega}^{n}_{1})}}\\ q_{t}\del{\boldsymbol{\beta}_{0}}&\propto\exp\cbr{\sum_{n=1}^{N}\mathbb{E}_{q_ {t}\del{z^{n}_{1}}q_{t}\del{\boldsymbol{\omega}^{n}_{0}|z^{n}_{1}}}\sbr{l(z^{n }_{1},\boldsymbol{x}^{n}_{0},\boldsymbol{\beta}_{0},\boldsymbol{\omega}^{n}_{0 })}}\end{split}$ (10) $$

where we omitted terms whose form did not change. $R_{z_{1}^{n},t}\del{\boldsymbol{x}^{n}_{0},\boldsymbol{\lambda}_{1}}$ reflects a contribution to $q(z^{n}_{1})$ that depends on the expected log partition of the linear (Matrix Normal Gamma) likelihood $p_{\boldsymbol{\lambda}_{1}}(\boldsymbol{x}^{n}_{1}|\boldsymbol{x}^{n}_{0},z^{
n}_{1})$. Note that the updates to each subset of posteriors (latents or parameters) have an analytic form due to the conditional conjugacy of the model. Importantly, both priors and posterior of the auxiliary variables are Pólya-Gamma distributed (Polson et al., ).

Finally, in the above update equations, we have replaced instances of the multinomial distribution $p\del{z|\boldsymbol{x},\boldsymbol{\beta}}$ with its augmented form $p\del{\omega|z}e^{l\del{z,\boldsymbol{x},\boldsymbol{\omega},\boldsymbol{\beta
}}}$ where the function $l\del{\cdot}$ is quadratic with respect to the coefficients $\boldsymbol{\beta}$ and the input variables $\boldsymbol{x}$, leading to tractable update equations.

## 4 Results

To evaluate the effectiveness of the CAVI-based approach, we compared it to other approximate inference algorithms, using several real and synthetic datasets. We compared CMNs fit with CAVI to the following three approaches:

MLE — We obtained point estimates for the parameters $\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K},\boldsymbol{\beta}_{0},
\boldsymbol{\beta}_{1}$ of the CMN using maximum-likelihood estimation. For gradient-based optimization of the loss function (the negative log likelihood), we used the AdaBelief optimizer with parameters set to its default values as introduced in Zhuang et al. () ($\alpha=1e-3$, $\beta_{1}=0.9$, $\beta_{2}=0.999$), and run the optimization for $20,000$ steps. This implements deterministic gradient descent, not stochastic gradient descent, because we fit the model in ‘full-batch’ mode, i.e., without splitting the data into mini-batches and updating model parameters using noisy gradient estimates.

NUTS-HMC — We use the No-U-Turn Sampler (NUTS), an extension to the Hamiltonian Monte Carlo (HMC) samplers, that incorporates adaptive step sizes (Hoffman et al., ). Markov Chain Monte Carlo converges in distribution to samples from a target distribution, so for this method we obtain samples from a joint distribution $p(\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K},\boldsymbol{\beta}_{0},
\boldsymbol{\beta}_{1}|Y,\boldsymbol{X}_{0})$ that approximate the true posterior. We used 800 warm-up steps, 16 independent chains, and 64 samples for each chain.

BBVI — Black-Box Variational Inference (BBVI) method (Ranganath et al., ). In constrast, to CAVI, the BBVI maximizes evidence lower bound (ELBO) using stochastic estimation of the gradients of variational parameters. While BBVI does not require conjugate relationships in the generative model, we use the same CMN model and variational distributions as we use for CAVI-CMN, in order to ensure fair comparison. For stochastic optimization, we used the AdaBelief optimizer with learning rate $\alpha=5e-3$ (other hyperparameters same as for MLE), used 8 samples to estimate the ELBO gradient (the num_particles argument of the Trace_ELBO() class), and ran the optimizer for $20,000$ steps).

For the Bayesian methods (CAVI, NUTS, and BBVI), we used the same form for the CMN priors (see Equation 4 for their parameterization) and fixed the prior parameters to the following values, used for all datasets: $v_{0}=10$, $a_{0}=2$, $b_{0}=1$, $\sigma_{0},\sigma_{1}=5$. For all datasets, we fixed the dimension of the continuous latent $\boldsymbol{x}_{1}$ to be $h=L-1$, where $L$ is the number of classes. For the Pinwheels dataset (see Section 4.1 below), we set the number of linear experts (and hence the dimension of the discrete latent $\boldsymbol{z}_{1}$) to be $K=10$, while for all other datasets we used $K=20$.

Figure: Figure 2: Performance and runtime results of the different inference algorithms on the ‘Pinwheel’ dataset. The standard deviation (vertical lines) of the performance metric is depicted together with the mean estimate (circles) over different runs. The top row of subplots show performance metrics across training set sizes: test accuracy (top left); log predictive density (top center), and expected calibration error (top right). The bottom row shows runtime metrics as a function of increasing training set size: the number of iterations required to achieve convergence (lower left); and the total runtime, estimated using the product of the number of iterations to convergence and the average cost (in seconds) for running one iteration (lower right). The number of iterations required for convergence was calculated by determining the number of gradient steps (or M steps, for CAVI) taken before the ELBO (or negative log likelihood, for MLE) reached 95% of its maximum value (see Appendix E for details on how these metrics were computed).
Refer to caption: extracted/5820796/figures/pinwheel_performance_and_runtime_plot_noOURS.png

Figure: Figure 3: Performance and runtime results of the different models on the ‘Waveform Domains’ dataset. The waveforms dataset consists of synthetic data generated to classify three different waveform patterns. Each instance is described by 21 continuous attributes. See [here](https://archive.ics.uci.edu/dataset/107/waveform+database+generator+version+1) for more information about the dataset. Descriptions of each subplot are same as in the Figure 2 legend.
Refer to caption: extracted/5820796/figures/waveform_performance_and_runtime_plot_noOURS.png

### 4.1 Comparison on synthetic datasets

For synthetic datasets we selected the Pinwheels and the Waveform Domains (Breiman and Stone, ) datasets. The pinwheels dataset is a synthetic dataset designed to test the model’s ability to handle nonlinear decision boundaries and data with non-Gaussian densities. The dataset consists of multiple clusters arranged in a pinwheel pattern, posing a challenging task for mixture models (Johnson et al., ) due to the curved and elongated spatial distributions of the data. See Appendix C for the parameters we used to simulate the pinwheels dataset. Similarly, the Waveform Domains dataset consists of synthetic data generated to classify three different waveform patterns, where each class is described by 21 continuous attributes (Breiman and Stone, ).

We fit all inference methods using different training set sizes, where each next training set was twice as large as the previous (for pinwheels: we trained using train sizes 50 to 1600; for waveform domains: train sizes 60 to 3840). This was done in order to study the robustness of performance in the low data regime. For each training size, we used the same test-set to evaluate performance (for pinwheels, 500 examples; for waveform domains: 1160 datapoints). For each inference method and examples set size, we fit using the same batch of training data, but with 16 randomly-initialized models (different initial posterior samples or parameters).

We assess the performance of the different inference methods using three main metrics: predictive accuracy (Test Accuracy), log-predictive density (LPD), and expected calibration error (ECE). Log predictive density is a common measure of predictive accuracy for methods that output probabilities (Gelman et al., ), and expected calibration error measures how well a model’s predictions are calibrated to the class probabilities observed in the data (Guo et al., ). In Figure 2 we visualize each of these metrics for the Pinwheels dataset and in Figure 3 for the Waveform dataset as a function of training set size. The CAVI-based approach achieves comparable log predictive density and calibration error to the other two Bayesian methods, which all outperform maximum likelihood estimation in LPD and ECE. This holds across training set sizes, indicating better sample efficiency.

### 4.2 Comparison on real-world datasets

To further validate the performance of CAVI-CMN, we conducted experiments using 6 real-world classification datasets from the UCI Machine Learning Repository (Kelly et al., ). Table 1 summarizes the performance of the different algorithms on all 7 different UCI datasets (the Waveform domains dataset and the 6 real datasets), using the widely-applicable information criterion (WAIC) as a measure of performance. WAIC is an approximate estimate of leave-one-out cross-validation (Vehtari et al., ).

The CAVI-CMN approach consistently provided higher WAIC scores in comparison to the MLE algorithm, and WAIC scores that were on par with BBVI and NUTS. The results confirm that using fully conjugate priors within the CAVI framework, does not diminishes the inference and the predictive performance of the algorithm, when compared to the state-of-the-art Bayesian methods like NUTS and BBVI. Importantly, CAVI-CMN offers substantial advantages in terms of computational efficiency as explored in the next section.

**Table 1: Comparison of widely-applicable information criterion (WAIC) for different methods evaluated on 7 different UCI datasets.**
|  | Rice | Breast Cancer | Waveform | Vehicle Silh. | Banknote | Sonar | Iris |
| --- | --- | --- | --- | --- | --- | --- | --- |
| CAVI | -0.1820 | -0.0504 | -0.2921 | -0.3281 | -0.0206 | -0.1544 | -0.0747 |
| MLE | -0.3599 | -0.3133 | -0.5759 | -0.7437 | -0.3133 | -0.3133 | -0.5514 |
| NUTS | -0.1278 | -0.0324 | -0.3753 | -0.3767 | -0.0110 | -0.0306 | -0.0413 |
| BBVI | -0.1739 | -0.0763 | -0.3618 | -0.4154 | -0.0382 | -0.0583 | -0.1544 |

### 4.3 Runtime comparison

The NUTS algorithm, although considered state-of-the-art in terms of inference robustness and accuracy (for well calibrated models (Gelman et al., )), is notoriously difficult to apply to large-scale problems (Cobb and Jalaian, ). Hence, the preferred algorithm of choice for probabilistic machine learning applications have been methods grounded in variational inference, such as black-box variational inference (BBVI) (Ranganath et al., ) and stochastic variational inference (SVI) (Hoffman et al., ).

In this subsection, we analyze the runtime efficiency of the MLE and BBVI algorithms for CMN models, in comparison to a CAVI-based approach. The focus is on comparing the computation time as the number of parameters increases along different components of the model.

To ensure comprehensive comparison, we varied the complexity of the models by adjusting the number of components, the dimensionality of the input space, and the number of datapoints. These modifications effectively increase the number of parameters allowing us to observe how each algorithm scales with model complexity.

Figure: Figure 4: Relative scaling of fitting time in seconds for Maximum Likelihood, BBVI, and CAVI, as a function of the number of parameters. The number of parameters itself was manipulated in three illustrative ways: changing the input dimension $d$, changing the number of linear experts $K$ in the conditional mixture layer, and changing the dimensionality of the continuous latent variable $h$.
Refer to caption: extracted/5820796/figures/relative_runtime_plot_normal_scale.png

The runtime performance for varying datasize on Pinwheel dataset is summarized in the bottom two subplots of Figure 2 which shows the runtime in seconds, and steps until convergence for different algorithms. We used the steps until convergence to assess the runtime for each algorithm. As expected, all algorithms exhibit an increase in runtime as the number of training data increases (which also scales the number of parameters for BBVI and CAVI). However the rate of increase varies significantly across different algorithms, with CAVI-CMN approach showing the best scaling behavior.

Similarly, in Figure 4 we plot the relative runtimes of Maximum Likelihood, CAVI, and BBVI (proportional to the runtime of the least complex variant), as we increase the number of parameters along different elements of CMN. This shows how fitting CMNs with CAVI scales competitively with gradient-based methods like BBVI and Maximum Likelihood Estimation. However, the rightmost subplot indicates that as we increase the dimensionality of the latent variable $\boldsymbol{X}_{1}$, CAVI-CMN scales more dramatically than the other two methods. This inherits from the computational overhead of matrix operations required by storing multivariate Gaussians posteriors over each continuous latent, i.e., $q(\boldsymbol{x}^{n}_{1}|z^{n}_{1})=\mathcal{N}(\boldsymbol{x}^{n}_{1};
\boldsymbol{\mu}^{n}_{1},\boldsymbol{\Sigma}^{n}_{1})$. Running the CAVI algorithm involves operations (like matrix inversions and matrix-vector products) whose (naive) complexity is quadratic in matrix size. This explains the nonlinear scaling of runtime as a function of $h$, the dimension of $\boldsymbol{X}_{1}$. There are several ways to address this issue:

- •
Low-Rank Approximations: One might use low-rank approximations to the covariance matrix $\boldsymbol{\Sigma}^{n}_{1}$ (e.g., Cholesky or eigendecompositions).
- •
Diagonal Covariance Structure: Further constrain the covariance structure of $q(\boldsymbol{x}^{n}_{1})$ by forcing the latent dimensions to be independent in the posterior, i.e., $q(\boldsymbol{x}^{n}_{1}|z^{n}_{1})=\prod_{i=1}^{h}\mathcal{N}(x^{n}_{i,1};\mu
^{n}_{i,1},(\sigma^{n}_{i,1})^{2})$. This would then mean that the number of parameters to store would only grow as $K(2h)$ in the size of the dataset, rather than as $K(h+O(h^{2}))$.
- •
Full Mean-Field Approximation: Enforce a full mean-field approximation between $\boldsymbol{X}_{1}$ and $Z_{1}$, so that one only needs to store $q(\boldsymbol{x}^{n}_{1})q(z^{n}_{1})$ rather than $q(\boldsymbol{x}^{n}_{1}|z^{n}_{1})q(z^{n}_{1})$. This would reduce the number of multivariate normal parameters that would have to be stored and operated upon by a factor of $K$.
- •
Shared Conditional Covariance Structure: Assume that the conditional covariance structure is shared across all training data points, i.e., $\boldsymbol{\Sigma}^{n}_{1}=\boldsymbol{\Sigma}_{1}$, for all $n\in\{1,2,\ldots,n\}$.

All of these adjustments would help mitigate the quadratic runtime scaling of CAVI-CMN as the latent dimension $h$ increases.

In summary, both sets of runtime analyses (both absolute and relative) suggest CAVI-CMN may be an attractive alternative to BBVI suitable for large-scale and time-sensitive applications, which similarly offers a fully Bayesian treatment of latent variables and parameters, while maintaining fast absolute runtime and time-to-convergence.

## 5 Conclusion

We demonstrate that the CAVI-based approach for conditional mixture networks (CMN) significantly outperforms the traditional maximum likelihood estimation (MLE) based approach, in terms of predictive performance and calibration. The improvement in probabilistic performance over the MLE based approaches can be attributed to implicit regularisation via prior information, and proper handling of posterior uncertainty over latent states and parameters, leading to a better representation of the underlying data, reflected in improved calibration error and log predictive density, even in low data regimes.

One of the key advantages of the CAVI-based approach is its computational efficiency compared to the other Bayesian inference methods such as Black-Box variational inference and the No-U-turn sampler (NUTS). While NUTS can sample from the full joint posterior distribution, which maximizes performance in terms of inference quality, this comes at the expense of substantial computational resources, especially for high dimensional and complex models (Hoffman et al., ). The variational methods offer a scalable alternative to this answer, in the form of methods like black-box variational inference (BBVI). Although BBVI is highly efficient in comparison to NUTS, it takes longer to converge than CAVI when applied to CMN. Hence, we expect CAVI to be a more practical choice for large-scale application, especially when further combined with data mini-batching methods (Hoffman et al., ).

The benchmark results show that CAVI-CMN algorithm achieves comparable performance to BBVI and NUTS in terms of predictive accuracy, log-predictive density and expected calibration error, while being significantly faster. This balance between predictive likelihood and calibration (jointly viewed as indicators of sample efficiency) is particularly important in real-world applications where robust prediction, reflective of underlying uncertainty, are crucial.

Furthermore, a straightforward mixture of linear components present in CMN, offers additional interoperability benefits. By using the conditionally conjugate priors, and a corresponding mean-field approximation over latent variables and model parameters, we facilitate easier interpretation of the model parameters and their uncertainties. This is particularly valuable in domains where understanding the underlying data-generating process is as important as the predictive performance, such as in healthcare, finance, and scientific research. Another important point is that the conjugate form of the CMN means that variational updates end up resembling sums of sufficient statistics collected from the data; this means the CAVI algorithm we described is readily amenable to online computation and minibatching, where sufficient statistics can computed and summed on-the-fly to update model parameters in a streaming fashion (Hoffman et al., ). This approach will become necessary when scaling CAVI-CMN to deeper (more than two-layer) models (Viroli and McLachlan, ) and larger datasets, where storing all the sufficient statistics of the data in memory becomes prohibitive.

Overall, these findings underscore the practical advantages of CAVI-CMN and highlight its promise as a new tool for fast probabilistic machine learning.

## Code availability

The code for using CAVI and the other 3 methods to fit the CMN model on the pinwheel and UCI datasets, is available from the cavi-cmn repository, which can be found at the following link: [https://github.com/VersesTech/cavi-cmn](https://github.com/VersesTech/cavi-cmn).

## Acknowledgments and Disclosure of Funding

The authors would like to thank the members of the VERSES Machine Learning Foundations group for critical discussions and feedback that improved the quality of this work, with special thanks to Tommaso Salvatori, Tim Verbelen, Magnus Koudahl, Toon van der Maele, Hampus Linander, and Karl Friston.

## Appendix A Variational Bayesian Multinomial Logistic Regression

In this section, we focus on a single multinomial logistic regression model (not in the context of the CMN), but the ensuing variational update scheme derived in Section A.4 is applied in practice to both the gating network’s parameters $\boldsymbol{\beta}_{0}$ as well as those of the final output likelihood for the class label $\boldsymbol{\beta}_{1}$.

### A.1 Stick-breaking reparameterization of a multinomial distribution

Multinomial logistic regression considers the probability that an outcome variable $y$ belongs to one of $K$ mutually-exclusive classes or categories. The probability of $y$ belonging to the $k$^th class is given by the categorical likelihood:

$$ $p(y=k|\boldsymbol{x},\boldsymbol{\beta})=p_{k}$ (11) $$

The problem of multinomial logistic regression is to identify or estimate the values of regression coefficients $\boldsymbol{\beta}$ that explain the relationship between some dataset of given continuous input regressors $\boldsymbol{X}=(\boldsymbol{x}^{1},\boldsymbol{x}^{2},\ldots,\boldsymbol{x}^{N})$ and corresponding categorical labels $Y=(y^{1},y^{2},\ldots,y^{N}),y^{n}\in{1,2,\ldots,K}$.

We can use a stick-breaking construction to parameterize the likelihood over $y$ using a set of $K-1$ stick-breaking coefficients: $\boldsymbol{\pi}=\del{\pi_{1},\ldots,\pi_{K-1}}$. Each coefficient is parameterized with an input regressor $\boldsymbol{x}$, and a corresponding set of regression weights $\boldsymbol{\beta}_{j}$. Stick-breaking coefficient $\pi_{j}$ is then given by a sigmoid transform of the product of the regression weights and the input regressors:

$$ $\begin{split}\pi_{j}&=\sigma\del{\boldsymbol{\beta}_{j}\sbr{\boldsymbol{x};1}} ~{},\\ \text{where}\hskip 5.69054pt\sigma\del{\boldsymbol{\beta}_{j}\sbr{\boldsymbol{ x};1}}&=\frac{1}{1+\exp\cbr{-\boldsymbol{\beta}_{j}\sbr{\boldsymbol{x};1}}}~{} ,\\ \text{and}\hskip 5.69054pt\boldsymbol{\beta}_{j}\sbr{\boldsymbol{x};1}&=\sum_{ i=1}^{d}w_{j,i}x_{i}+a_{j}.\end{split}$ (12) $$

The outcome likelihood is then obtained via stick breaking transform(^1^11[This blog post](https://gregorygundersen.com/blog/2020/07/01/multinomial-binomial/) has helpful discussion on the stick-breaking form of the multinomial logistic likelihood and provides more intuition behind its functional form.) as follows

$$ $p_{k}=\pi_{K}\prod_{j=1}^{K-1}(1-\pi_{j})=\sigma\del{\boldsymbol{\beta}_{K} \sbr{\boldsymbol{x};1}}\prod_{j=1}^{K-1}\del{1-\sigma\del{\boldsymbol{\beta}_{ j}\sbr{\boldsymbol{x};1}}}=\prod_{j=1}^{K-1}\frac{\exp\cbr{\boldsymbol{\beta}_ {j}\sbr{\boldsymbol{x};1}}}{1+\exp\cbr{\boldsymbol{\beta}_{j}\sbr{\boldsymbol{ x};1}}}$ (13) $$

where $\pi_{K}=1$, and $\boldsymbol{\beta}_{K}=\vec{0}$.

Finally, we can express the likelihood in the form of a Categorical distribution as

$$ $\text{Cat}\del{y;\boldsymbol{x},\boldsymbol{\beta}}=\prod_{k=1}^{K-1}\frac{ \del{\exp\cbr{\boldsymbol{\beta}_{k}\sbr{\boldsymbol{x};1}}}^{\delta_{k,y}}}{ \del{1+\exp\cbr{\boldsymbol{\beta}_{k}\sbr{\boldsymbol{x};1}}}^{N_{k,y}}}~{}.$ (14) $$

where $N_{k,y}=1$ for $k\leq y$, and $N_{k,y}=0$ otherwise (or $N_{k,y}=1-\sum_{j=1}^{k-1}\delta_{j,y}$), and $\delta_{k,y}=1$ for $k=y$ and is zero otherwise.

### A.2 Pólya-Gamma augmentation

The *Pólya-Gamma augmentation* scheme [Polson et al., , Linderman et al., , Durante and Rigon, ] is defined as

$$ $\frac{\del{e^{\psi}}^{a}}{\del{1+e^{\psi}}^{b}}=2^{-b}e^{\kappa\psi}\int_{0}^{ \infty}e^{-\omega\psi^{2}/2}p(\omega)\dif\omega$ (15) $$

where $\kappa=a-b/2$ and $p\del{\omega|b,0}$ is the density of the Pólya-Gamma distribution $\text{PG}(b,0)$ which does not depend on $\psi$. The useful properties of the Pólya-Gamma are the exponential tilting property expressed as

$$ $PG(\omega;b,\psi)=\frac{e^{{-\omega\psi^{2}/2}}PG(\omega;b,0)}{\mathbb{E}\sbr{ e^{-\omega\psi^{2}/2}}}$ (16) $$

the expected value of $\omega$, and $e^{-\omega\psi^{2}/2}$ given as

$$ $\begin{split}\mathbb{E}\sbr{\omega}&=\int_{o}^{\infty}\omega PG(\omega;b,\psi) \dif\omega=\frac{b}{2\psi}\tanh\del{\frac{\psi}{2}},\\ \mathbb{E}\sbr{e^{-\omega\psi^{2}/2}}&=\cosh^{-b}\del{\frac{\psi}{2}}\end{split}$ (17) $$

and the Kulback-Leibler divergence between $q\del{\omega}=PG(\omega;b,\psi)$ and $p\del{\omega}=PG(\omega;b,0)$ obtained as

$$ $D_{KL}\sbr{q\del{\omega}||p\del{\omega}}=-\mathbb{E}\sbr{\omega}\frac{\psi^{2} }{2}+b\ln\cosh\del{\frac{\psi}{2}}=-\frac{b\psi}{4}\tanh\del{\frac{\psi}{2}}+b \ln\cosh\del{\frac{\psi}{2}}~{}.$ (18) $$

We can express the likelihood function in Equation 14 using the augmentation as

$$ $\begin{split}p(y,\boldsymbol{\omega}|\boldsymbol{\psi})&=p\del{y|\boldsymbol{ \psi}}p\del{\boldsymbol{\omega}|y,\boldsymbol{\psi}}=\prod_{k=1}^{K-1}2^{-b_{k ,y}}e^{\kappa_{k,y}\psi_{k}-\omega_{k}\psi_{k}^{2}/2}\text{PG}(\omega_{k};b_{k ,y},0)\\ p\del{y|\boldsymbol{\psi}}&=\prod_{k=1}^{K-1}2^{-b_{k,y}}e^{\kappa_{k,y}\psi_{ k}}\int_{0}^{\infty}e^{-\omega_{k}\psi_{k}^{2}/2}\text{PG}(\omega_{k};b_{k,y}, 0)\dif\omega_{k}\\ p\del{\boldsymbol{\omega}|y,\boldsymbol{\psi}}&=\prod_{k=1}^{K-1}\text{PG}\del {\omega_{k};b_{k,y},\psi_{k}}\end{split}$ (19) $$

where $b_{k,y}\equiv N_{k,y}$, $\kappa_{k,y}=\delta_{k,y}-N_{k,y}/2$, and $\psi_{k}=\boldsymbol{\beta_{k}}\sbr{\boldsymbol{x};1}$. Given a prior distribution $p\del{\boldsymbol{\psi}}=p\del{\boldsymbol{\beta}}p\del{\boldsymbol{x}}$, we can write the joint $p\del{y,\boldsymbol{\omega},\boldsymbol{\psi}}$ as

$$ $\begin{split}p\del{y,\boldsymbol{\omega},\boldsymbol{\psi}}&=p\del{\boldsymbol {\omega}|y}p\del{\boldsymbol{\psi}}e^{l\del{y,\boldsymbol{\psi},\boldsymbol{ \omega}}}~{},\\ l\del{y,\boldsymbol{\psi},\boldsymbol{\omega}}&=\sum_{k=1}^{K-1}l_{k}\del{y, \psi_{k},\omega_{k}}~{},\\ l_{k}\del{y,\psi_{k},\omega_{k}}&=\kappa_{y,k}\psi_{k}-b_{y,k}\ln 2-\omega_{k} \psi_{k}^{2}/2~{}.\end{split}$ (20) $$

### A.3 Evidence lower-bound

Given a set of observations $\boldsymbol{\mathcal{D}}=\del{y^{1},\ldots,y^{N}}$ the augmented joint distribution can be expressed as

$$ $p\del{\boldsymbol{\mathcal{D}},\boldsymbol{\Omega},\boldsymbol{X},\boldsymbol{ \beta}}=p\del{\boldsymbol{\beta}}\prod_{n=1}^{N}p\del{\boldsymbol{x}^{n}}p\del {\boldsymbol{\omega}^{n}|y^{n}}e^{l\del{y^{n},\boldsymbol{\psi}^{n}, \boldsymbol{\omega}^{n}}}$ $$

We can express the evidence lower-bound (ELBO) as

$$ $\begin{split}\mathcal{L}(q)&=E_{q\del{\boldsymbol{\Omega}}q\del{\boldsymbol{X} }q\del{\boldsymbol{\beta}}}\sbr{-\ln q\del{\boldsymbol{\beta}}+\sum_{n=1}^{N} \ln\frac{p\del{y^{n},\boldsymbol{\psi}^{n},\boldsymbol{\omega}^{n}}}{q\del{ \boldsymbol{\omega}^{n}}q\del{\boldsymbol{x}^{n}}}}\\ &=E_{q\del{\boldsymbol{\Omega}}q\del{\boldsymbol{X}}q\del{\boldsymbol{\beta}}} \sbr{\ln\frac{p\del{\boldsymbol{\beta}}}{q\del{\boldsymbol{\beta}}}+\sum_{n=1} ^{N}l\del{y^{n},\boldsymbol{\psi}^{n},\boldsymbol{\omega}^{n}}+\ln\frac{p\del{ \boldsymbol{\omega}^{n}|y^{n}}}{q\del{\boldsymbol{\omega}^{n}}}+\ln\frac{p\del {\boldsymbol{x}^{n}}}{q\del{\boldsymbol{x}^{n}}}}\\ &\geq\ln p\del{\boldsymbol{\mathcal{D}}}\end{split}$ (21) $$

where we use the following forms for the approximate posterior

$$ $\begin{split}q\del{\boldsymbol{\Omega}|Y}&=\prod_{n=1}^{N}q\del{\boldsymbol{ \omega}^{n}|y^{n}}=\prod_{n=1}^{N}\prod_{k=1}^{K-1}PG\del{b_{k,y^{n}},\xi_{k,n }}~{},\\ q\del{\boldsymbol{X}}&=\prod_{n=1}^{N}q\del{\boldsymbol{x}^{n}}=\prod_{n=1}^{N }\mathcal{N}\del{\boldsymbol{x}^{n};\boldsymbol{\mu}^{n},\boldsymbol{\Sigma}^{ n}}~{},\\ q\del{\boldsymbol{\beta}}&=\prod_{k=1}^{K-1}\mathcal{N}\del{\boldsymbol{\beta} _{k};\boldsymbol{\mu}_{k},\boldsymbol{\Sigma}_{k}}~{}.\end{split}$ (22) $$

### A.4 Coordinate ascent variational inference

The mean-field assumption in Equation 22 allows the implementation of a simple CAVI algorithm [Wainwright et al., , Beal, , Hoffman et al., , Blei et al., ] which sequentially maximizes the evidence lower bound in Equation 21 with respect to each factor in $q\del{\boldsymbol{\Omega}|Y}q\del{\boldsymbol{X}}q\del{\boldsymbol{\beta}}$, via the following updates:

$$ $\begin{split}\text{Update to latents (`E-step')}&\\ q^{\del{t,l}}\del{\boldsymbol{x}^{n}}&\propto p\del{\boldsymbol{x}^{n}}\exp \cbr{\mathbb{E}_{q^{\del{t-1}}\del{\boldsymbol{\beta}}q^{\del{t,l-1}}\del{ \boldsymbol{\omega}^{n}}}\sbr{l\del{y^{n},\boldsymbol{\psi}^{n},\boldsymbol{ \omega}^{n}}}}\\ q^{\del{t,l}}\del{\omega^{n}_{k}|y^{n}}&\propto p\del{\omega^{n}_{k}|y^{n}} \exp\cbr{\mathbb{E}_{q^{\del{t-1}}\del{\boldsymbol{\beta}}q^{\del{t,l}}\del{ \boldsymbol{x}^{n}}}\sbr{l_{k}\del{y^{n},\psi^{n}_{k},\omega^{n}_{k}}}}\\ \forall n&\in\cbr{1,\ldots,N},\text{ and for }q^{\del{t,0}}\del{\boldsymbol{ \omega}^{n}|y^{n}}=q^{\del{t-1,L}}\del{\boldsymbol{\omega}^{n}|y^{n}}\\ \text{Update to parameters (`M-step')}&\\ q^{\del{t}}\del{\boldsymbol{\beta}_{k}}&\propto\exp\cbr{\sum_{n=1}^{N}\mathbb{ E}_{q^{\del{t}}\del{\boldsymbol{x}^{n}}q^{\del{t}}\del{\boldsymbol{\omega}^{n} |y^{n}}}\sbr{l\del{y^{n},\boldsymbol{\psi}^{n},\boldsymbol{\omega}^{n}}}}\\ \end{split}$ (23) $$

at each iteration $t$, and multiple local iteration $l$ during the variational expectation step—until the convergence of the ELBO.

Specifically, the update equations for the parameters of the latents (the ‘E-step’) are:

$$ $\begin{split}q^{\del{t,l}}\del{\boldsymbol{x}^{n}}&\propto\mathcal{N}\del{ \boldsymbol{x}^{n};0,-2\boldsymbol{\lambda}_{2,0}}\exp\cbr{\sum_{k=1}^{K} \kappa_{k,y^{n}}\text{Tr}\del{\boldsymbol{\mu}_{k}^{\del{t-1}}\sbr{\boldsymbol {x}^{n};1}^{T}}-\frac{\langle\omega_{k}\rangle}{2}\text{Tr}\del{\boldsymbol{M} _{k}^{\del{t-1}}\sbr{\boldsymbol{x}^{n};1}\sbr{\boldsymbol{x}^{n};1}^{T}}}\\ \boldsymbol{\lambda}^{\del{n,t,l}}_{1}&=\sum_{k=1}^{K-1}\cbr{\kappa_{k,y^{n}} \sbr{\boldsymbol{\mu}_{k}^{\del{t-1}}}_{1:D}-\langle\omega^{n}_{k}\rangle_{t,l -1}\sbr{\boldsymbol{M}^{\del{t-1}}_{k}}_{D+1,1:D}}\\ \boldsymbol{\lambda}^{\del{n,t,l}}_{2}&=\boldsymbol{\lambda}_{2,0}-\frac{1}{2} \sum_{k=1}^{K-1}\langle\omega^{n}_{k}\rangle_{t,l-1}\sbr{\boldsymbol{M}_{k}}_{ 1:D,1:D}\\ \boldsymbol{M}_{k}^{(t-1)}&=\boldsymbol{\Sigma}_{k}^{\del{t-1}}+\boldsymbol{ \mu}_{k}^{\del{t-1}}\sbr{\boldsymbol{\mu}_{k}^{\del{t-1}}}^{T}\end{split}$ (24) $$

and

$$ $\begin{split}q^{\del{t,l}}\del{\omega^{n}_{k}|y^{n}}&\propto e^{-\omega^{n}_{k }\langle\psi_{k}^{2}\rangle/2}PG\del{\omega^{n}_{k};b_{k,y^{n}},0}\\ \xi^{n}_{k}&=\sqrt{\mathbb{E}_{q^{\del{t-1}}\del{\boldsymbol{\beta}}q^{\del{t, l}}\del{\boldsymbol{x}^{n}}}\sbr{\psi_{k}^{2}}}\\ \xi^{n}_{k}&=\sqrt{\text{Tr}\del{\boldsymbol{M}_{k}^{\del{t-1}}\hat{ \boldsymbol{M}}^{\del{n,t,l}}}}\\ \text{where }\hat{\boldsymbol{M}}^{\del{n,t,l}}&=\begin{pmatrix}\boldsymbol{M} ^{\del{n,t,l}}&\boldsymbol{\mu}^{\del{n,t,l}}\\ \sbr{\boldsymbol{\mu}^{\del{n,t,l}}}^{T}&1\end{pmatrix}~{},\text{ and } \boldsymbol{M}^{\del{n,t,l}}=\boldsymbol{\Sigma}^{\del{n,t,l}}+\boldsymbol{\mu }^{\del{n,t,l}}\sbr{\boldsymbol{\mu}^{\del{n,t,l}}}^{T}.\end{split}$ (25) $$

Similarly, for the parameter updates (‘M-step’) we get

$$ $\begin{split}q^{\del{t}}\del{\boldsymbol{\beta}_{k}}&\propto\mathcal{N}\del{ \boldsymbol{\beta}_{k};0,-2\boldsymbol{\lambda}^{\prime}_{2,0}}\exp\cbr{\sum_{ n=1}^{N}\kappa_{k,y^{n}}\text{Tr}\del{\hat{\boldsymbol{\mu}}^{\del{n,t}} \boldsymbol{\beta}_{k}^{T}}-\frac{\langle\omega_{k}\rangle^{n}_{t}}{2}\text{Tr }\del{\hat{\boldsymbol{M}}_{i}^{\del{t}}\boldsymbol{\beta}_{k}\boldsymbol{ \beta}_{k}^{T}}}\\ \boldsymbol{\lambda}^{\del{t}}_{k,1}&=\sum_{i}\kappa_{k,y^{n}}\hat{\boldsymbol {\mu}}^{\del{n,t}}\\ \boldsymbol{\lambda}^{\del{t}}_{k,2}&=\boldsymbol{\lambda}^{\prime}_{2,0}- \frac{1}{4}\sum_{n=1}^{N}\frac{b_{k,y^{n}}}{\xi_{k}^{\del{n,t}}}\tanh\del{ \frac{\xi_{k}^{\del{n,t}}}{2}}\hat{\boldsymbol{M}}^{\del{n,t}}\end{split}$ (26) $$

where $\hat{\boldsymbol{\mu}}^{\del{n,t}}=\sbr{\boldsymbol{\mu}^{\del{n,t}};1}$.

## Appendix B Variational Bayesian Mixture of Linear Transforms

The variational ‘M-step’ in Equation 9 to update the parameters of the linear experts reduces to a straightforward form when each expert parameterizes a multivariate Gaussian likelihood over the latent variables $\boldsymbol{X}_{1}$ with Matrix Normal Gamma priors over the parameters of the linear function that maps $\boldsymbol{X}_{0}$ to a distribution over $\boldsymbol{X}_{1}$.

Recall the form of the update to the posterior parameters of the linear experts:

$$ $\begin{split}q\del{\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K}}& \propto\exp\cbr{\sum_{n=1}^{N}\mathbb{E}_{q(\boldsymbol{x}^{n}_{1},z^{n}_{1})} \left[\ln p(\boldsymbol{x}^{n}_{1}|z^{n}_{1},\boldsymbol{x}^{n}_{0}, \boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K})\right]}\end{split}$ (27) $$

The approximate posteriors $q\del{\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K}}$ and $q(\boldsymbol{X}_{1},Z_{1})$ have the following form:

$$ $\begin{split}q\del{\boldsymbol{A}_{1:K},\boldsymbol{\Sigma}^{-1}_{1:K}}&=\prod _{k=1}^{K}\left(\mathcal{MN}(\boldsymbol{A}_{k};\boldsymbol{M}_{k},\boldsymbol {\sigma}_{k}^{-2}I,\boldsymbol{V}_{k})\prod_{i=1}^{h}\Gamma(\sigma_{i,k}^{-2}; a_{k},b_{i,k})\right)\\ q\del{\boldsymbol{X}_{1}|Z_{1}}&=\prod_{n=1}^{N}\prod_{k=1}^{K}\mathcal{N}( \boldsymbol{x}^{n}_{1};\boldsymbol{\mu}^{n}_{k,1},\boldsymbol{\Sigma}^{n}_{k,1 })\\ q\del{Z_{1}}&=\prod_{n=1}^{N}\text{Cat}(z^{n}_{1};\boldsymbol{\gamma}^{n})\end {split}$ $$

The parameters of the $k$^th expert $q(\boldsymbol{A}_{k},\boldsymbol{\Sigma}^{-1}_{k})$ can written in terms of weighted updates to the Matrix Normal Gamma’s canonical parameters $\boldsymbol{M}_{k},\boldsymbol{V}_{k},a_{k}$ and $b_{k}$:

$$ $\begin{split}\boldsymbol{V}^{-1}_{k}&=\boldsymbol{V}^{-1}_{k,0}+\sum_{n=1}^{N} \gamma^{n}_{k}\boldsymbol{x}^{n}_{0}\left(\boldsymbol{x}^{n}_{0}\right)^{\top} \\ \boldsymbol{M}_{k}&=\left(\boldsymbol{M}_{k,0}\boldsymbol{V}^{-1}_{k,0}+\sum_{ n=1}^{N}\gamma^{n}_{k}\boldsymbol{\mu}^{n}_{k,1}\left(\boldsymbol{x}^{n}_{0} \right)^{\top}\right)\boldsymbol{V}_{k}\\ a_{k}&=a_{k,0}+\frac{\sum_{n=1}^{N}\gamma^{n}_{k}}{2}\\ b_{i,k}&=b_{i,k,0}+\frac{1}{2}\left(\sum_{n=1}^{N}\gamma^{n}_{k}\left[ \boldsymbol{\Sigma}^{n}_{k,1}+\boldsymbol{\mu}^{n}_{k,1}(\boldsymbol{\mu}^{n}_ {k,1})^{\top}\right]_{ii}-\left[\boldsymbol{M}_{k}\boldsymbol{V}_{k}^{-1} \boldsymbol{M}^{T}_{k}\right]_{ii}+\left[\boldsymbol{M}_{k,0}\boldsymbol{V}_{k ,0}^{-1}\boldsymbol{M}_{k,0}^{T}\right]_{ii}\right)\end{split}$ (28) $$

where the notation $\left[\cdot\right]_{ii}$ selects the $i$^th element of the diagonal of the matrix in the brackets.

## Appendix C Dataset Descriptions

We fit all inference methods using different training set sizes, where each next training set was twice as large as the previous. For each training size, we used the same test-set to evaluate performance. The test set was ensured to have the same relative class frequencies as in the training set(s). For each inference method and examples set size, we fit using the same batch of training data, but with 16 randomly-initialized models (different initial posterior samples or parameters).

### C.1 Pinwheels Dataset

The pinwheels dataset is a synthetic dataset designed to test a model’s ability to handle nonlinear decision boundaries and data with non-Gaussian densities [Johnson et al., ]. The structure of the pinwheels dataset is determined by 4 parameters: the number of clusters or distinct spirals; the angular deviation, which defines how far the spiralling clusters deviate from the origin; the tangential deviation, which defines the noise variance of 2-D points within each cluster; and the angular rate, which determines the curvature of each spiral. For evaluating the four methods (CAVI-CMN, MLE, BBVI, and NUTS) on the synthetic pinwheels dataset, we generated a dataset with 5 clusters, with an angular deviation of 0.7, tangential deviation of 0.3 and angular rate of 0.2. We selected these values by looking at the maximum achieved test accuracy across all the methods for different parameter combinations and tried to upper-bound it 80%, which provides a low enough signal-to-noise ratio to be able to meaningfully show differences in probabilistic metrics like calibration and WAIC. For pinwheels, we trained using train sizes 50 to 1600, doubling the number of training examples at each successive training set size. We tested using 500 held-out test examples generated using the same parameters as used for the training set(s).

### C.2 Waveform Domains Dataset

The Waveform Domains dataset consists of synthetic data generated to classify three different waveform patterns, where each class is described by 21 continuous attributes [Breiman and Stone, ]. For waveform domains, we fit each model on train sizes ranging from 60 to 3840 examples, and tested on a held-out size of 1160 examples. See [here](https://archive.ics.uci.edu/dataset/107/waveform+database+generator+version+1) for more information about the dataset.

### C.3 Vehicle Silhouettes Dataset

This dataset involves classifying vehicle silhouettes into one of four types (bus, van, or two car models) based on features extracted from 2D images captured at various angles []. We fit each model on train sizes ranging from 20 to 650 examples, and tested on a held-out size of 205 examples. See [here](https://archive.ics.uci.edu/dataset/149/statlog+vehicle+silhouettes) for more information about the dataset.

### C.4 Rice Dataset

The Rice dataset contains measurements related to the classification of rice varieties, specifically Cammeo and Osmancik [mis, ]. We fit each model on train sizes ranging from 40 to 2560 examples, and tested on a held-out size of 1250. See [here](https://archive.ics.uci.edu/dataset/545/rice+cammeo+and+osmancik) for more information about the dataset.

### C.5 Breast Cancer Dataset

The ‘Breast Cancer Diagnosis’ dataset [Wolberg et al., ] contains features extracted from breast mass images, which are then used to classify tumors as malignant or benign. See [here](https://archive.ics.uci.edu/dataset/17/breast+cancer+wisconsin+diagnostic) for more information about the dataset. We fit each model on train sizes ranging from 25 to 400 examples, and tested on a held-out size of 169.

### C.6 Sonar (Mines vs Rocks) Dataset

The Sonar (Mines vs Rocks) dataset consists of sonar signals bounced off metal cylinders and rocks under various conditions. The dataset includes 111 patterns from metal cylinders (mines) and 97 patterns from rocks. Each pattern is represented by 60 continuous attributes corresponding to the energy within specific frequency bands []. The task is to classify each pattern as either a mine (M) or a rock (R). For this dataset, we fit each model on train sizes ranging from 8 to 128 examples and tested on a held-out size of 80 examples. See [here](https://archive.ics.uci.edu/dataset/151/connectionist+bench+sonar+mines+vs+rocks) for more information about the dataset.

### C.7 Banknote Authentication Dataset

The ‘Banknote Authentication’ dataset [Lohweg, ] contains features extracted from images of genuine and forged banknotes. It is primarily used for binary classification tasks to distinguish between authentic and counterfeit banknotes. See [here](https://archive.ics.uci.edu/dataset/267/banknote+authentication) for more information about the dataset.

## Appendix D UCI Performance Results

In Figures 5 to 9 we report the same performance and runtime metrics as in Figure 2 for 7 UCI datasets, and find that with the exception of the Sonar dataset, CAVI performs competitively with or better than MLE on all datasets, and always outperforms MLE in terms of LPD and ECE. Runtime scaling is similar as reported for the Pinwheels dataset in the main text; CAVI-CMN always converges in fewer steps and is faster than BBVI, and either outperforms or is competitive with MLE in terms of runtime.

Figure: Figure 5: Performance and runtime results of the different models on the ‘Vehicle Silhouettes’ dataset. Descriptions of each subplot are same as in the Figure 2 legend.
Refer to caption: extracted/5820796/figures/statlog_performance_and_runtime_plot_noOURS.png

Figure: Figure 6: Performance and runtime results of the different models on the ‘Rice’ dataset. Descriptions of each subplot are same as in the Figure 2 legend.
Refer to caption: extracted/5820796/figures/rice_performance_and_runtime_plot_noOURS.png

Figure: Figure 7: Performance and runtime results of the different models on the ‘Breast Cancer’ dataset. Descriptions of each subplot are same as in the Figure 2 legend.
Refer to caption: extracted/5820796/figures/breast_cancer_performance_and_runtime_plot_noOURS.png

Figure: Figure 8: Performance and runtime results of the different models on the ‘Connectionist Bench (Sonar, Mines vs. Rocks)’ dataset. Descriptions of each subplot are same as in the Figure 2 legend..
Refer to caption: extracted/5820796/figures/connectionist_bench_performance_and_runtime_plot_noOURS.png

Figure: Figure 9: Performance and runtime results of the different models on the ‘Banknote Authentication’ dataset. Descriptions of each subplot are same as in the Figure 2 legend.
Refer to caption: extracted/5820796/figures/banknote_performance_and_runtime_plot_noOURS.png

## Appendix E Model Convergence Determination

For each inference algorithm, the number of iterations taken to converge was determined by running each algorithm for a sufficiently high number of gradient (respectively, CAVI update) steps such that the ELBO (or log likelihood - LL - for MLE) stopped significantly changing. This was determined (through anecdotal inspection over many different initializations and runs across the different UCI datasets) to be 20,000 gradient steps for BBVI, 20,000 gradient steps for MLE, and 500 combined CAVI update steps for CAVI-CMN. To determine the time taken to sufficiently converge, we recorded the value of the ELBO or LL at each iteration, and fit an exponential decay function to the negative of each curve. The parameters of the estimated exponential decay were then used to determine the time at which the curve decayed to 95% decay of its value. This time was reported as the number of steps taken to converge.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 2
questions: - **Broader Range of Gradient-Based Optimization Techniques**: To strengthen the comparative analysis, it would be beneficial to extend the comparison to include both [1] and the AdaBelief optimizer [2]. This broader range of gradient-based optimization techniques could provide deeper insights into the performance and robustness of the proposed method, offering a more comprehensive evaluation than focusing solely on AdaBelief [2].

- **Weak Theoretical Justification:** While CAVI-CMN effectively employs conditional conjugacy and Pólya-Gamma augmentation, the paper lacks theoretical justification or clear proofs regarding convergence rates and optimality, as seen in related works [5,6,7]. This absence may raise concerns about the reliability and consistency of CAVI-CMNs, particularly in complex models. Including such theoretical justification, alongside the presented empirical evidence, would substantially enhance the paper's contribution and support its novel approach.

- **Overly Simplistic Conditional Mixture of Experts**: Lines 179-183 point out that the current architecture is relatively simple, utilizing only a linear expert with discrete output and a pair of latent variables. It remains unclear how the proposed method would scale to more complex models, such as those with nonlinear experts, continuous outputs, or deeper architectural layers such as [6,7]. Can this approach be adapted to effectively handle these additional complexities?

- **Model Selection Architecture**: The paper lacks clarity on the criteria used for model selection, specifically regarding the number of experts $K$ and the dimensionality of the latent space $h$ for real datasets. Providing a detailed explanation or methodology for selecting these parameters would enhance understanding of the architecture’s adaptability and effectiveness, especially in practical applications.

- **Scalability Challenges and Potential Extensions for CAVI-CMN in High-Dimensional Architectures:** Figure 4 suggests that as the dimensionality of the latent variable $X_1$ increases, CAVI-CMN exhibits more pronounced scaling challenges than the other two methods as the latent dimension $h$ grows. This raises questions about the current method's adaptability to deeper (additional layers) and wider (larger latent dimensions) architectures. The authors should provide additional details on how incorporating techniques like low-rank approximations or further factorization could enhance the model’s scalability. Further investigation and simulations demonstrating the impact of these approaches would offer valuable insights into the model's ability to handle increasingly complex, high-dimensional architectures efficiently.

References:

[1] Xie, X., Zhou, P., Li, H., Lin, Z., \& Yan, S. (2024). Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models. IEEE Transactions on Pattern Analysis and Machine Intelligence.

[2] Zhuang, J., Tang, T., Ding, Y., Tatikonda, S. C., Dvornek, N., Papademetris, X., \& Duncan, J. (2020). Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33, 18795-18806.

[3] Bhattacharya, A., Pati, D., \& Yang, Y. (2023). On the convergence of coordinate ascent variational inference. arXiv preprint arXiv:2306.01122.

[4] Zhang, F., \& Gao, C. (2020). Convergence rates of variational posterior distributions. The Annals of Statistics, 48(4), 2180-2207.

[5] Wang, B., \& Titterington, D. M. (2006). Convergence properties of a general algorithm for calculating variational Bayesian estimates for a normal mixture model. Bayesian Analysis, 1(3), 625-650.

[6] Viroli, C., \& McLachlan, G. J. (2019). Deep Gaussian mixture models. Statistics and Computing, 29, 43-51.

[7] Kock, L., Klein, N., \& Nott, D. J. (2022). Variational inference and sparsity in high-dimensional deep Gaussian mixture models. Statistics and Computing, 32(5), 70.
soundness: 3
strengths: - The CAVI-CMNs approach introduces a highly efficient gradient-free Bayesian inference method specifically designed for CMNs. Its reliance CAVI paired with conjugate priors and Pólya-Gamma augmentation, demonstrates innovation through a novel combination of existing ideas, especially for probabilistic models that typically depend on gradient-based methods. This efficiency could mark an important step forward in Bayesian neural network inference.

- The paper ensures reproducibility by including comprehensive instructions and corresponding Python codes for nearly all numerical experiments.

-  The paper proposes future directions for extending CAVI-CMNs to support deeper architectures and online learning contexts, demonstrating foresight in adapting the method for broader applicability. Such adaptability may foster additional research and application of CAVI-CMNs, underscoring its potential for significant impact.
summary: The paper introduces Coordinate Ascent Variational Inference for Conditional Mixture Networks (CAVI-CMNs), a framework designed to achieve efficient, scalable Bayesian inference for probabilistic models, specifically focusing on CMNs. By leveraging conditional conjugacy and Pólya-Gamma augmentation, this approach enables effective variational updates without relying on gradient-based optimization, positioning it as a computationally viable alternative to traditional Bayesian inference methods such as Maximum Likelihood Estimation, Black-Box Variational Inference, and the No-U-Turn Sampler.  The model's performance is validated on two synthetic datasets (Pinwheels and Waveform Domains) and six real-world classification datasets from the UCI Machine Learning Repository, demonstrating that it maintains competitive predictive accuracy and runtime efficiency even with increasing model complexity. This efficiency is particularly advantageous for applications requiring rapid, uncertainty-aware predictions, including in fields such as autonomous systems, healthcare, and finance.
weaknesses: **Insufficient Experimental Validation:** The reliance on smaller UCI and synthetic datasets might not adequately demonstrate the robustness or scalability of the approach in real-world, high-dimensional applications. Without validation on challenging or large-scale datasets, such as the CIFAR-10 dataset[1] and the ImageNet[2] or real-time data streams, the method’s claims of scalability and efficiency could be viewed as unconvincing, limiting the impact and generalizability of the results.

**Enhance Consistency and Clarity in Mathematical Notation:**

+ The sample size should be consistently denoted as $N$ throughout the paper. For example, Line 049 uses $n$, whereas Line 147 switches to $N$.

+ Lines 147-153 contain an indexing typo where the index should be written as $n=1$ instead of $i=1$.

+ Line 168: The notation for multivariate input and scalar output is inconsistent, with input variables in bold and outputs in regular font. Adopting a consistent notation style, as outlined in the ICLR guidelines, would improve clarity.

+ **Line 246-256**: Could the authors clarify the specific distributions, such as $\mathcal{MN}$ and $\Gamma$? It would enhance readability if the full names or definitions of these distributions were provided upon first mention. Additionally, would a more flexible prior for $\mathbf{\Sigma}_k$ improve the model's performance?

+ **Line 246-256**: $p(Y|\mathbf{X})$ should be $p(Y,\mathbf{X})?$

- **Baseline Comparison**: **Line 322-323**: For parameter point estimates, in addition to comparing with Maximum Likelihood Estimation using backpropagation, the authors should consider comparing their AVCI (variational EM algorithm) with a standard EM algorithm. This comparison would help elucidate the advantages and limitations of the proposed method relative to traditional approaches, providing a clearer understanding of its effectiveness.

- **Lines 114 and 669-671**: The citations are incorrect. The correct citation should be: Bishop, Christopher M., and Markus Svenskn. "Bayesian Hierarchical Mixtures of Experts." UAI'03 Proceedings of the Nineteenth conference on Uncertainty in Artificial Intelligence. Morgan Kaufmann Publishers Inc., 2003.


References:

[1] Krizhevsky, A., \& Hinton, G. (2009). Learning multiple layers of features from tiny images.

[2] Deng, J., Dong, W., Socher, R., Li, L. J., Li, K., \& Fei-Fei, L. (2009, June). Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition (pp. 248-255). Ieee.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: It is not clear to me whether the lower bounds for CAVI higher than the BBVI-ELBO in the experiments? If not, why does it yield better results in terms of calibration etc? Is this due to the choice of priors, or due to the biases of the variational distribution with the Poly-Gamma augmentation?

Why is there a Gaussian prior for the bias term u in line 181, which seems to be absent in the subsequent model description?
soundness: 3
strengths: While Poly-Gamma augmentations have been used for models with other likelihoods or for Gibbs samplers, their usage for multinomial logistic regression models is new as far as I am aware. The paper aims to address an important problem in that discriminative classifiers can yield not well-calibrated predictions when trained to maximize the log-likelihood, while Bayesian approaches can be computationally expensive. In their arguably simple examples, the authors demonstrate that their approach does indeed improve BBVI or NUTS in terms of calibration or WAIC.
summary: The submission introduces a coordinate ascent variational inference (CAVI) approach for learning a mixture of linear transformations of an input vector to predict a label, termed conditional mixture network (CMN). A Poly-Gamma augmentation scheme is considered to arrive at a conditionally conjugate model as commonly required for CAVI. The suggested approach appears competitive to black-box VI or NUTS in terms of model calibration, while requiring less compute time.
weaknesses: The empirical evaluation seems a bit limited. In particular, it would useful if the authors would consider additional baselines. In particular \
(i) variational approaches beyond BBVI for the same model, or at least with some standard control variates [1,2,3], \
(ii) stochastic-gradient MCMC algorithms [4,5] that can yield better uncertainties [6,7] and are more scalable than NUTS, \
(iii) different models (e.g. deep ensembles [8], BNNs) to see better assess how the reported results (accuracies, expected calibration error, etc) compare to previous models.

The authors claim that ‘CAVI-CMT achieves SOTA-like performance with absolute runtime comparable to a backpropagation-based MLE approach’. However, the used MLE baseline does not use stochastic gradients (i.e. mini-batches), which is arguably a more common approach, and I expect to converge faster? 

The authors claim that ‘using fully conjugate priors within the CAVI framework does not diminish the inference and predictive performance’. However, it is not clear to me why different choices of priors would perform. Does the BBVI or MCMC baseline use different priors?



[1] Tucker, George, et al. "Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models." Advances in Neural Information Processing Systems 30 (2017).\
[2] Grathwohl, Will, et al. "Backpropagation through the Void: Optimizing control variates for black-box gradient estimation." International Conference on Learning Representations. 2018.\
[3] Richter, Lorenz, et al. "VarGrad: a low-variance gradient estimator for variational inference." Advances in Neural Information Processing Systems 33 (2020): 13481-13492.\
[4] Welling, Max, and Yee W. Teh. "Bayesian learning via stochastic gradient Langevin dynamics." Proceedings of the 28th international conference on machine learning (ICML-11). 2011.\
[5] Chen, Tianqi, Emily Fox, and Carlos Guestrin. "Stochastic gradient hamiltonian monte carlo." International conference on machine learning. PMLR, 2014.\
[6] Zhang, Ruqi, et al. "Cyclical Stochastic Gradient MCMC for Bayesian Deep Learning." International Conference on Learning Representations, 2019.\
[7] Maddox, Wesley J., et al. "A simple baseline for bayesian uncertainty in deep learning." Advances in neural information processing systems 32 (2019).\
[8] Lakshminarayanan, Balaji, Alexander Pritzel, and Charles Blundell. "Simple and scalable predictive uncertainty estimation using deep ensembles." Advances in neural information processing systems 30 (2017).

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 1
questions: - Why were the experiments limited to such small datasets? Would CAVI-CMN still be competitive in terms of performance and runtime on larger, modern datasets?

- How does the method compare with more recent advancements in variational inference or Bayesian deep learning, particularly those leveraging deep architectures and large-scale datasets?

- Could the authors elaborate on the practical use cases where CAVI-CMN would be preferred over simpler, gradient-based methods, given its specialized nature?
soundness: 2
strengths: The paper introduces a method that eliminates the need for gradient-based optimization, which is a significant advantage for reducing computational costs.
summary: The paper presents a method called CAVI-CMN, a gradient-free learning algorithm using conditional mixture networks (CMNs) for Bayesian inference.
The method applies coordinate ascent variational inference (CAVI) to a probabilistic mixture-of-experts model, aiming to balance expressiveness, scalability, and uncertainty quantification in classification tasks.
The authors propose the use of conjugate priors and Polya-Gamma augmentation to allow for efficient inference without gradient-based optimization.
Experiments on synthetic and small UCI datasets demonstrate some competitiveness compared to alternatives like maximum likelihood estimation (MLE) and black-box variational inference (BBVI).
weaknesses: - Heavy notation: The notation is often overly complex, and impacts readability and understanding. It lacks consistency across sections, making it difficult to follow key points. For instance, the sadly non-numbered equation on page 3 is very unclear:

     1. on line 150, authors define $p_k$, yet this does not appear (at least explicitly) in the equation above

     2. there is no explanation nor definition about $\pi$. What is $p(\pi)$? and $p(z^n|\pi)$?

     3. on line 150, authors say that y^n is an observation, yet on line 49 and 168, y is intended to be a label.

     4. x does not appear in the LHS of the equation, nor in the explanation below it, but only on the RHS.

Overall, this confusion makes the problem setting very unclear, and impacts next derivations.

- The goal of the paper is somewhat puzzling. It is unclear whether the primary focus is computational efficiency, uncertainty quantification, or scalability, as the paper lacks a clear articulation of its main contributions.

- Small datasets: The experiments are performed on small datasets from the UCI repository, spanning from 4 up to only 60 features.

- I'm not an expert in this field, but I have the feeling that relevant literature is missing, as most of the references in the paper are quite old. In particular, the two main competitors for the experiments, NUTS and BBVI, both date back to 2014.


minors:
- Table 1 does not report somehow the best results, and it is hard to get a message from it.
- it is unclear to me how the first layer of the network is able to output a joint continuous-discrete output (L177).

---

confidence: 5
contribution: 3
initial_rating: 3
presentation: 2
questions: 1.  Is this the first paper to propose the conditional mixture network?  If not, please give the citation.  If so, please explain why no experimental comparisons to Bayesian and non-Bayesian ReLU NNs are given.


2.  Normalization layers (BatchNorm, LayerNorm) are often necessary for modern deep learning architectures.  Is there a way to incorporate them within the CMN?  And even better yet, define them within the same generative framework?


3.  Why is gradient-based learning spoken of so critically in the paper?  In the past, it was largely believed that coordinate ascent algorithms would be better than noisy, gradient-based ones, but with our modern knowledge about how many critical points are not minima and the benefits of noise during optimization, I don’t think many of those past critiques still hold today.
soundness: 1
strengths: **Generative model of piecewise linear representations**:  Forming representations from piecewise linear functions is a core component of modern deep learning architectures—with the best example being ReLU activations.  This work presents an interesting generative perspective of these representations.  While not explored in this paper, I could see this leading to a better way to specify informative priors for feature learning (i.e. draw the linear models from distributions over particular basis function or groups).
summary: This paper proposes a model called the “conditional mixture network” (CMNs).  The idea is that a mixture of linear models is defined such that a given feature vector is transformed by one of the components.  This creates a hidden representation ($X_1$) that is then used as the features for a final linear model, which makes the classification decision.  This model deviates from predecessors such as mixture density networks and mixtures of experts as they usually define the mixture on the output distribution over the response / label.  CMNs, on the other hand, define the mixture to be on the (linear) transformation between the input ($X_0$) and first hidden layers ($X_1$).  The paper claims that the major benefit of this model over traditional Bayesian neural networks is that parameter estimation can be done without gradient-based stochastic optimization.  Rather, as the two-layer CMN admits conditional conjugacy, a mean-field coordinate ascent algorithm can be devised to estimate the posterior distribution (Section 3.4).  The experiments focus on validating that this coordinate ascent algorithm is superior to maximum likelihood estimation (MLE), black-box variational inference (BBVI), and Hamiltonian Monte Carlo (as implemented by NUTS).  Results are reported on synthetic data (Pinwheels and Waveform Domains) and on seven datasets from the UCI repository (Rice, Breast Cancer, Waveform, Vehicle Silh, Banknote, Sonar, and Iris).  In general, the coordinate ascent algorithm is reported to match the predictive performance of BBVI and NUTS while, like BBVI and NUTS, being superior to MLE.  Moreover, the coordinate ascent VI algorithm has a far superior runtime to NUTS and faster convergence compared to BBVI.
weaknesses: 1.  **Conflict between method’s presentation and experimental evaluation**:  The Introduction and beginning of Section 3.1 clearly motivates the work from the perspective of improving the uncertainty quantification abilities of deep learning architectures.  Thus, when arriving at the experiments section (4), I was expecting the result to compare the performance of the CMN with natural baselines, such as a two-layer NN with ReLU activations, since this would be a non-Bayesian piecewise linear model, and a two-layer Bayesian NN.  However, the experimental section’s main hypothesis is to examine if the coordinate ascent VI algorithm improves upon NUTS and BBVI.  While this is certainly a sensible and necessary experimental comparison to make, I don’t see how this line of experimentation alone is sufficient.  As far as I understand (see corresponding question #1 below), this is the first paper to describe the CMN model.  Hence, the generative model itself needs to be justified as well as any particular choice of inference algorithm.  I believe that adding comparisons to small Bayesian and non-Bayesian neural architectures would go a long way to justifying the CMN.  


2.  **Black-box variational inference is not representative of the state of the art**:  In line 66, the paper claims that the coordinate ascent algorithm is benchmarked against “state-of-the-art Bayesian methods like NUTS and BBVI”.  I agree that NUTS is representative of the SOTA for Hamiltonian Monte Carlo implementations.  However, BBVI is not, as using REINFORCE to estimate the gradients usually has much higher variance than more modern alternatives.  The paper should also compare with differentiable non-centered parameterizations for the weights [Kingma & Welling, ICML 2014] (a.k.a. “Bayes by Backdrop” [Blundell et al., ICML 2015]).  The mixture gating variable will need additional tricks for differentiating through the discrete variable [Graves, 2016; Maddison et al., 2017].  There are also non-stochastic alternatives [Wu et al., ICLR 2019].  For a method that blends the benefits of variational inference and MCMC, I would also look at particle methods such as Stein Variational Gradient Descent [Liu & Wang, 2016].


3.  **Weak justification for why a mixture is a good choice for representation learning**:  As stated above, the related works section draws parallels to mixture density networks and mixtures of experts.  However, these models define mixtures over the output distribution, and I don’t think the same intuition can be extended to motivate why mixtures of linear models is the best construction for the hidden representation.  For example, using $Z_1,d$ to denote the gating variable for the $d$th hidden dimension, one could also place the prior $Z_{1,d} \sim \text{Bernoulli}(p)$.  This would have the same representational capacity as the propose CMN (linear transformation) but allow model capacity to be directly controlled by $p$.  This construction would also be similar to ReLU NNs since one can think of a subset of the linear models activating to define the transformation.  There might also be some interesting connection to dropout regularization.  I’m not saying that this dimension-wise Bernoulli construction is obviously better; I’m just saying that the paper should address why the mixture is better motivated than some obvious alternatives.  This relates back to weakness #1.


4.  **No discussion of extending the architecture to increased depths**:  While the paper argues that CMNs could serve as a building block for deeper architectures (see paragraph before the start of 3.2), there is no discussion of how this might be done from the generative perspective or for the inference algorithms.  


Blundell, Charles, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. "Weight uncertainty in neural networks." In Proceedings of the 32nd International Conference on International Conference on Machine Learning-Volume 37, pp. 1613-1622. 2015.

Graves, Alex. "Stochastic backpropagation through mixture density distributions." arXiv preprint arXiv:1607.05690 (2016).

Kingma, Diederik, and Max Welling. "Efficient gradient-based inference through transformations between bayes nets and neural nets." International Conference on Machine Learning. PMLR, 2014

Liu, Qiang, and Dilin Wang. "Stein variational gradient descent: A general purpose bayesian inference algorithm." Advances in neural information processing systems 29 (2016).

Maddison, Chris J., Andriy Mnih, and Yee Whye Teh. "The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables." International Conference on Learning Representations. 2017

Wu, Anqi, Sebastian Nowozin, Edward Meeds, Richard E. Turner, José Miguel Hernández-Lobato, and Alexander L. Gaunt. "Deterministic Variational Inference for Robust Bayesian Neural Networks." In International Conference on Learning Representations. 2019.