Gradient-free variational learning with conditional mixture networks

Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-08-29 11:23:09

Abstract: Balancing computational efficiency with robust predictive performance is crucial in supervised learning, especially for safety-critical applications. While deep learning models are accurate and scalable, they often lack calibrated predictions and uncertainty quantification. Bayesian methods address these issues but are often computationally expensive. We introduce CAVI-CMN, a gradient-free, fast variational method for training conditional mixture networks (CMNs), a probabilistic variant of the mixture-of-experts (MoE) model. Using conjugate priors and Pólya-Gamma augmentation, we derive efficient updates via coordinate ascent variational inference (CAVI).  We apply this method to train conditional mixture networks on classification tasks from the UCI repository. Our approach, which we name CAVI-CMN, achieves competitive and often superior predictive accuracy compared to backpropagation (i.e., maximum likelihood estimation) while maintaining posterior distributions over model parameters. Moreover, computation time scales in model complexity competitively to both MLE and other gradient-based solutions like black-box variational inference (BBVI), while running overall much faster than BBVI and sampling-based inference and with similar speed to MLE. This combination of probabilistic robustness and computational efficiency positions CAVI-CMN as a building block for constructing discriminative models that are fast, gradient-free, and Bayesian.

Main paper: 
 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., 2007) . Deep neural networks (DNNs) have demonstrated success in this domain, owing to their universal function approximation properties (Park & Sandberg, 1991) and the soft regularization inherited from stochastic gradient descent learning via backpropagation (Amari, 1993) . However, despite its computational efficiency, accuracy, and scalability to increasingly large datasets and models, DNNs trained this way do not provide well-calibrated predictions or uncertainty estimates, and practitioners typically utilize post-hoc calibration methods on validation datasets (Wang et al., 2021; Shao et al., 2020) . This limits the applicability and reliability of using DNNs in safety-critical applications like autonomous driving, medicine, and disaster response (Papamarkou et al., 2024) , 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 θ θ θ as a process of inferencenamely, calculating a posterior distribution over model parameters p θ θ θ | D D D , given observed data D D D (e.g., a data set of input-output pairs D D D = (x x x 1 , y y y 1 ), . . . , (x x x n , y y y n ) ). The resulting posterior distribution captures both expectations about model parameters θ θ θ and their uncertainty. The uncertainty is then incorporated in predictions p(y y y) that are, in principle, well-calibrated to new datapoints coming from the same set. This probabilistic treatment allows methods like Bayesian neural networks (BNNs) (Hernández-Lobato & Adams, 2015) to maintain the expressiveness of deep neural networks while also encoding uncertainty over network weights and thus the network's predictions. However,• CAVI-CMN matches, and sometimes exceeds, the performance of maximum likelihood estimation (MLE) in terms of predictive accuracy, while maintaining the benefits of a full Bayesian approach, yielding well calibrated models that quantify uncertainty. 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. 
related work: 
 The Mixture-of-Experts (MoE) architecture is a close relative of the CMN model we introduce here. Jacobs et al. (1991) originally introduced MoEs as a way to improve the performance of neural networks by combining the strengths of multiple specialized models (Gormley & Frühwirth-Schnatter, 2019) . 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., 2013) , further bolstered by the MoE's capacity for universal function approximation (Nguyen et al., 2016; Nguyen & Chamroukhi, 2018) . Indeed, the powerful self-attention mechanism employed by transformers has has demonstrated the power and flexibility of MoE models (Movellan & Gabbur, 2020) . Non-Bayesian approaches to MoE typically rely on maximum likelihood estimation (MLE) (Jacobs et al., 1991; Jordan & Jacobs, 1994) , which can suffer from overfitting and poor generalization due to the lack of regularization mechanisms (Bishop & Svenskn, 2003) , 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 & Svenskn, 2003; Mossavat & Amft, 2011) . This Bayesian treatment enables the estimation of model evidence (log marginal-likelihood) and provides a natural framework for model comparison and selection (Svensén, 2003; Zens, 2019) . 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. ( 2013) enabled a range of novel and more computationally efficient algorithms for Bayesian treatment of MoE models (Linderman et al., 2015; He et al., 2019; Sharma et al., 2019; Viroli & McLachlan, 2019; Zens et al., 2023) . 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). 
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.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, 1992; Blundell et al., 2015; Daxberger et al., 2021) .We circumvent these problems by focusing on the Mixture-of-Experts (MoE) models (Jacobs et al., 1991) , 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 & Jordan, 1997; Bishop & Svenskn, 2003) or Pólya-Gamma augmentation (Polson et al., 2013) ).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 formp(Y, Z, Θ) = p(θ 1:K )p(π) N i=1 p(y n |z n , θ 1:K )p(z n |π) ,where y n is an observation, Θ = {θ 1:K , π}, p k (y n |θ 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 |θ k ) = p(y n |z n = k, θ 1:K ). For instance, if each 'expert' likelihood p k (y n |θ k ) is a Gaussian distribution, then the MoE becomes a Gaussian Mixture Model, whereθ k = (µ k , Σ k ).The problem of learning the model's parameters, then becomes one of doing inference over the latent variables Z and parameters Θ 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 Under review as a conference paper at ICLR 2025 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., 2008; Hoffman et al., 2013; Blei et al., 2017) , 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.The conditional mixture network maps from a continuous input vector x x x 0 ∈ R d to its label y ∈ {1, . . . , L}. This is achieved with two layers: a conditional mixture of linear experts, which outputs a joint continuous-discrete latent x x x 1 ∈ R h , z 1 ∈ {1, . . . , K} and a multinomial logistic regression, which maps from the continuous latent x x x 1 to the corresponding label y. The probabilistic mapping can be described in terms of the following operations:z 1 ∼ Mult (z 1 ; x x x 0 , β β β 0 ) x x x 1 = A A A z1 • [x x x 0 ; 1] + u u u z1 , u u u z1 ∼ N (0 0 0, Σ Σ Σ z1 ) y ∼ Mult (y; x x x 1 , β β β 1 )where we pad the input variable x x x 0 with a constant value set to 1, to absorb the bias term within the mapping matrix A A A z1 ∈ R h×d+1 , and where Mult (z; x x x, β β β) denotes a multinomial distribution parameterized with a regressor x x x and logistic regression coefficients β β β. Note that for every pair of regressors and labels (x x x n 0 , y n ), we assume a corresponding pair of latent variables (x x x n 1 , z n 1 ). 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, β β β 0 and β β β 1 , we use Pólya-Gamma augmentation (Polson et al., 2013; Linderman et al., 2015) applied to the stick breaking construction for the multinomial distribution:p z = k|β β β, x x x = π k (β β β, x x x) k-1 j=1 1 -π j (β β β, x x x) π j (β β β, x x x) = 1 1 + exp -β β β j • [x x x; 1] , ∀j < K π K = 1 (1)where for the gating network (input layer) we will have coefficients of dimension β β β 0 ∈ R K-1×d , and for the output likelihood coefficients of dimension β β β 1 ∈ R L-1×h .Given a set of labels Y = y 1 , y 2 , ..., y N , and regressors X X X 0 = x x x 1 0 , x x x 2 0 , ..., x x x N 0 , that define i.i.d input-output pairs x x x n 0 , y n , we write the joint distribution over labels Y , latents X X X 1 , Z 1 , and parameters Θ Θ Θ as:p(Y Y Y , X X X 1 , Z 1 , Θ Θ Θ|X X X 0 ) = p (Θ Θ Θ) N n=1 p β1 β1 β1 y n |x x x 1 p λ λ λ1 x x x n 1 |x x x n 0 , z n 1 p β β β0 z n 1 |x x x n 0 p (Θ Θ Θ) = p (β β β 1 ) p (β β β 0 ) p (λ λ λ 1 ) = L-1 l=1 p(β β β l,1 ) K-1 k=1 p(β β β k,0 ) K j=1 p A A A j , Σ Σ Σ -1 j (2)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., 2007) . However, we use the term Under review as a conference paper at ICLR 2025x x x n 1 z n 1 A A A k Σ Σ Σ -1 k x x x n 0 y n β β β k,0 β β β l,1 a 0 , b 0 M M M 0 , v 0 0 0 0, σ 0 0 0 0, σ 1 K -1 N K L -1Figure 1 : A Bayesian network representation of the two-layer conditional mixture network, with input-output pairs x x x n 0 , y n and latent variables x x x n 1 , z n 1 . Observations are shaded nodes, while latents and parameters are transparent. Prior hyperparameters are shown without boundaries.generative model to emphasize the fact that the model contains priors over latent variables (X X X 1 , Z 1 ), and parametersΘ Θ Θ = β β β 1:L-1,1 , β β β 1:K-1,0 , A A A 1:K , Σ Σ Σ -1 1:K, 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 λ λ λ 1 as notational shorthand for the parameters A A A 1:K , Σ Σ Σ -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:EQUATIONwhere we fixed the prior mean matrix of the linear transformation A A A k to be a matrix of zeros: M M M 0 = 0 0 0. Other hyperparameters of the priors are described in Appendix C.1. In the following section we introduce a mean-field variational inference scheme we use for performing inference and learning in the two-layer CMN.In this section we describe a variational approach for inverting the probabilistic model described in Equation ( 2) and computing an approximate posterior over latents and parameters specified asp X X X 1 , Z 1 , Θ Θ Θ|Y, X X X = p (Y, X X X 1 , Z 1 , Θ Θ Θ, X X X) p Y |X X X ≈ q (Θ Θ Θ) N n=1 q (z n 1 ) q x x x n 1 |z n 1 (4)where q x x x n 1 |z n 1 corresponds to a component specific multivariate normal distribution, and q (z n 1 ) to a multinomial distribution. Importantly, the approximate posterior over parameters q (Θ Θ Θ) further Under review as a conference paper at ICLR 2025 factorizes (Svensén, 2003) asEQUATIONThe 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 lowerbound (ELBO) on the marginal log likelihoodEQUATIONWe maximize the ELBO using an iterative update scheme for the parameters of the approximate posterior, often referred to as variational Bayesian expectation maximization (VBEM) (Beal, 2003) or coordinate ascent variational inference (CAVI) (Bishop & Nasrabadi, 2006; Blei et al., 2017) .The procedure consists of two parts:At a given variational iteration t, we fix the posterior over parameters q (Θ Θ Θ) to their value at the last iteration q t-1 (Θ Θ Θ) (or if t = 1, to randomly-initialized values). With the parameter posterior fixed, we then update the posterior over latent variables by setting them equal to the solution that maximizes L(q), under q t-1 (Θ Θ Θ):q t (X 1 , Z 1 ) ∝ exp E qt-1(Θ Θ Θ) ln p Θ Θ Θ Y, X 1 , Z 1 |X 0 (7)This update of the latents is also known as the 'variational E-step' due to its resemblance to the E-step in expectation maximization (Beal, 2003) .The posterior over latent variables updated in the E-step of Equation (7) q t (X 1 , Z 1 ) is then used to update the posterior over parameters asq t (Θ Θ Θ) ∝ exp    N n=1 E qt(x n 1 ,z n 1 ) ln p Θ Θ Θ y n , x n 1 , z n 1 |x n 0    (8)This update of the parameter posterior is similarly known as the 'variational M-step' due to its resemblance to the step which maximizes the log likelihood with respect to the parameters in the E-M algorithm. More detailed expansions of these equations, including their functional forms, and the PG augmentation scheme needed to turn them into conditionally-conjugate updates, are given in Appendix A and Appendix B. 
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 Θ Θ Θ of the CMN using maximum-likelihood estimation (backpropagation to minimize the negative log likelihood). NUTS-HMC -The No-U-Turn Sampler (NUTS), an extension to Hamiltonian Monte Carlo (HMC) that incorporates adaptive step sizes (Hoffman et al., 2014) . This provides samples from a posterior distribution over Θ Θ Θ.BBVI -Black-Box Variational Inference (BBVI) method (Ranganath et al., 2014) . BBVI maximizes the evidence lower bound (ELBO) using stochastic estimation of its gradients with respect to variational parameters.Under review as a conference paper at ICLR 2025 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., 2014) , and expected calibration error measures how well a model's predictions are calibrated to the class probabilities observed in the data (Guo et al., 2017) . In Figure 2 we show performance for the Pinwheels dataset and in Figure 3 for the Waveform dataset as a function of N . 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 CAVI-CMN's high sample efficiency.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., 2024) . 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., 2017) .The CAVI-CMN approach consistently provided higher WAIC scores in comparison to the MLE algorithm, and with comparable magnitude to those computed using BBVI and NUTS. The results confirm that using fully conjugate priors within the CAVI framework, does not diminish 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.The NUTS algorithm, although considered state-of-the-art in terms of inference robustness and accuracy (for well calibrated models (Gelman et al., 2020) ), is notoriously difficult to apply to largescale problems (Cobb & Jalaian, 2021). Hence, the preferred algorithm of choice for probabilistic machine learning applications have been methods grounded in variational inference, such as black- 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 model parameters increases. To study this, we varied the complexity of the CMN along the following dimensions: the dataset size used for training N , the number of linear experts in the mixture layer K, the dimensionality of the input space d, and the dimensionality of the latent space h. Note that in all runtime experiments, we did not measure NUTS' runtime because its poor scaling and speed is well-documented (Blei et al., 2017; Izmailov et al., 2021; Cobb & Jalaian, 2021) and for all but the smallest dataset sizes and model dimensionalities, it took already many times longer to run than any of the other methods we studied (MLE, BBVI, or CAVI). The runtime performance for varying dataset size N is shown for the Pinwheels dataset in the bottom two subplots of Figure 2 . This shows the total runtime in seconds, and steps until convergence for different algorithms. As expected, all algorithms exhibit an increase in runtime as N 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 runtime scaling of MLE, CAVI, and BBVI (proportional to the runtime of the least complex variant), as we increase the number of parameters along different dimensions of the model structure. Fitting CMNs with CAVI scales competitively with gradient-based methods like BBVI and MLE. However, the rightmost subplot of Figure 4 indicates that as we increase the dimensionality of the latent variable X X 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(x x x n 1 |z n 1 ) = N (x x x n 1 ; µ µ µ n 1 , Σ Σ Σ 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 X X X 1 . Various meth-Under review as a conference paper at ICLR 2025 ods (like low-rank approximations to the covariance structure of q(x x x n 1 |z n 1 ) and further factorization of q(x x x n 1 , z n 1 ) into q(x x x n 1 )q(z n 1 )) could help mitigate the scaling of CAVI-CMN's runtime as we increase the latent dimension h. These additions are fruitful avenue for future research into scaling CAVI-CMN to deeper (more layers) and wider (larger latent dimension) architectures.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. Like BBVI, CAVI-CMN offers a variational Bayesian treatment of latent variables and parameters, while maintaining much faster absolute runtime and quicker convergence, due to the use of coordinate ascent to update the parameters of variational distributions, as opposed to the stochastic gradient updates used in BBVI. 
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 (BBVI) 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., 2013) . Variational methods offer a scalable alternative to sampling-based inference 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 and is slower than CAVI when applied to inference and learning in conditional mixture networks. Hence, we expect CAVI to be a more practical choice for largescale application, especially when further combined with data mini-batching methods (Hoffman et al., 2013) .The UCI 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 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 computed from the data; this means the CAVI algorithm we described is readily amenable to online computation and mini-batching, where sufficient statistics can computed and summed on-the-fly to update model parameters in a streaming fashion (Hoffman et al., 2013; Broderick et al., 2013) . This approach will become necessary when scaling CAVI-CMN to deeper (more than two-layer) models (Viroli & McLachlan, 2019) 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.Under review as a conference paper at ICLR 2025 REPRODUCIBILITY STATEMENT An author-and affiliation-anonymized version of the code for using CAVI and the other 3 methods to fit the CMN model on the pinwheel and 7 UCI datasets is available for download at this link. The main performance figures in the text can be reproduced by running each inference script and changing the --train size parameter, while using the hyperparameters specified in the appendices (e.g., Appendix C.1). 
 
 Appendix: In this section we expand the CAVI (or VBEM) equations provided in Equation ( 7) and Equation ( 8) of the main text to provide their detailed functional form.Under review as a conference paper at ICLR 2025 Recall the ELBO used for CAVI-CMN:EQUATIONCoordinate ascent variational inference entails maximizing L(q) = using an iterative update scheme for the parameters of the approximate posterior, which factorizes into the posterior over parameters q(Θ Θ Θ) and the posterior over latent variables q(X 1 , Z 1 ). The procedure consists of two parts:First, we fix the posterior over the parameters to its value at the previous iteration (or randomlyinitialized values, if at the first iteration). Given the posterior over parameters, we update the posterior over latent variables (variational E-step) asEQUATIONSecond, the posterior over latents that was updated in the E-step, is used to update the posterior over parameters (variational M-step) asEQUATIONIn the variational inference literature, the distinction between latents and parameters is often described in terms of 'local' vs 'global' latent variables (Hoffman et al., 2013) , where local variables are datapoint-specific, and global variables are shared across datapoints. To detail the form of the updates to the parameters of the linear experts in Equation ( 11), i.e. q t (λ λ λ 1 ) = q t (A A A 1:K , Σ Σ Σ -1 1:K ), first we note the form of the approximate posteriors over the latent variables q(X X X 1 , Z 1 ):q X X X 1 |Z 1 = N n=1 K k=1 N (x x x n 1 ; µ µ µ n k,1 , Σ Σ Σ n k,1 ) q (Z 1 ) = N n=1 Cat(z n 1 ; γ γ γ n ) (12)The update to the k th expert's parameters q(A A A k , Σ Σ Σ -1 k ) can written in terms of weighted updates to the Matrix Normal Gamma's canonical parameters M M M k , V V V k , a k and b k , where the weights are provided by the sufficient statistics of {qx x x 1 1 |z 1 1 = k , q x x x 2 1 |z 2 1 = k , . . . , q x x x N 1 |z N 1 = k }:Finally, in the above update equations, we have replaced instances of the multinomial distribution p z|x x x, β β β with its augmented form p ω|z e l(z,x x x,ω ω ω,β β β) where the function l (•) is quadratic with respect to the coefficients β β β and the input variables x x x, leading to tractable update equations.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 Appendix B.4 is applied in practice to both the gating network's parameters β β β 0 as well as those of the final output likelihood for the class label β β β 1 .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|x x x, β β β) = p k (15)The problem of multinomial logistic regression is to identify or estimate the values of regression coefficients β β β that explain the relationship between some dataset of given continuous input regressors X X X = (x x x1 , x x x 2 , . . . , x x x N ) and corresponding categorical labels Y = (y 1 , y 2 , . . . , y N ), y n ∈ 1, 2, . . . , K.We can use a stick-breaking construction to parameterize the likelihood over y using a set of K -1 stick-breaking coefficients: π π π = (π 1 , . . . , π K-1 ). Each coefficient is parameterized with an input regressor x x x, and a corresponding set of regression weights β β β j . Stick-breaking coefficient π j is then given by a sigmoid transform of the product of the regression weights and the input regressors:π j = σ β β β j [x x x; 1] ,where σ β β β j [x x x; 1] = 1 1 + exp -β β β j [x x x; 1] ,andEQUATIONThe outcome likelihood is then obtained via stick breaking transform 1 as followsEQUATIONwhere π K = 1, and β β β K = ⃗ 0.Finally, we can express the likelihood in the form of a Categorical distribution asEQUATIONwhere N k,y = 1 for k ≤ y, and N k,y = 0 otherwise (or N k,y = 1 -k-1 j=1 δ j,y ), and δ k,y = 1 for k = y and is zero otherwise.Under review as a conference paper at ICLR 2025 B.2 P ÓLYA-GAMMA AUGMENTATION The Pólya-Gamma augmentation scheme (Polson et al., 2013; Linderman et al., 2015; Durante & Rigon, 2019) is defined asEQUATIONwhere κ = a -b/2 and p ω|b, 0 is the density of the Pólya-Gamma distribution PG(b, 0) which does not depend on ψ. The useful properties of the Pólya-Gamma are the exponential tilting property expressed asP G(ω; b, ψ) = e -ωψ 2 /2 P G(ω; b, 0) E e -ωψ 2 /2 (20)the expected value of ω, and e -ωψ 2 /2 given asE [ω] = ∞ o ωP G(ω; b, ψ) dω = b 2ψ tanh ψ 2 , E e -ωψ 2 /2 = cosh -b ψ 2 (21)and the Kulback-Leibler divergence between q (ω) = P G(ω; b, ψ) and p (ω) = P G(ω; b, 0) obtained asEQUATIONWe can express the likelihood function in Equation ( 18) using the augmentation as p(y, ω ω ω|ψ ψ ψ) = p y|ψ ψ ψ p ω ω ω|y, ψ ψ ψ =EQUATIONGiven a set of observations D D D = y 1 , . . . , y N the augmented joint distribution can be expressed asp (D D D, Ω Ω Ω, X X X, β β β) = p (β β β) N n=1 p (x x x n ) p ω ω ω n |y n e l(y n ,ψ ψ ψ n ,ω ω ω n )Under review as a conference paper at ICLR 2025We can express the evidence lower-bound (ELBO) asEQUATIONwhere we use the following forms for the approximate posteriorEQUATIONThe mean-field assumption in Equation ( 26) allows the implementation of a simple CAVI algorithm (Wainwright et al., 2008; Beal, 2003; Hoffman et al., 2013; Blei et al., 2017) which sequentially maximizes the evidence lower bound in Equation ( 25) with respect to each factor in q Ω Ω Ω|Y q (X X X) q (β β β), via the following updates:Update to latents ('E-step')q (t,l) (x x x n ) ∝ p (x x x n ) exp E q (t-1) (β β β)q (t,l-1) (ω ω ω n ) l (y n , ψ ψ ψ n , ω ω ω n ) q (t,l) ω n k |y n ∝ p ω n k |y n exp E q (t-1) (β β β)q (t,l) (x x x n ) l k (y n , ψ n k , ω n k )∀n ∈ {1, . . . , N } , and for q (t,0) ω ω ω n |y n = q (t-1,L) ω ω ω n |y n Update to parameters ('M-step')EQUATIONat 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: q (t,l) (x x x n ) ∝ N x xThe 'Banknote Authentication' dataset (Lohweg, 2013) 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 for more information about the dataset.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.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 This blog post has helpful discussion on the stick-breaking form of the multinomial logistic likelihood and provides more intuition behind its functional form. 



Reviews: review: title: title: BDU Review
main_review: main_review: # Overview

Thanks to the authors for the interesting and well-organized paper. The overall quality is good, with clear motivation and claims, as well as significant experimental results to back up those claims. The paper demonstrates a novel application of Coordinate Ascent Variational Inference (CAVI) to fitting fully (approximate) Bayesian models of a particular structure. My general impression is that CAVI has largely fallen out of favor compared to Black Box Variational Inference (BBVI), which allows for more general approximate posteriors. However, the authors show that, for this particular class of models, CAVI yields strong results that are competitive with BBVI (and NUTS) while requiring significantly less wall-clock time. The wall-clock time is in fact similar to MLE, while retaining a full posterior over parameters. The significance of the paper may be reduced by the fact the structure of the model is limited by conjugacy requirements (in order to yield efficient updates): the authors mention that the could be expanded to deeper architectures, but I'm not convinced that this would gain anything if each layer is linear? I suppose it depends on what "deeper architectures" means concretely. However, all research has limitations: The paper's contribution is still meaningful within context.

## Pros
- Clearly written motivation.
- Well-analyzed experimental results that show the strengths of CAVI-CMN compared to other parameter fitting/inference approaches.
  - Experiments include both synthetic and real datasets.
  - Experiments compare to relevant baselines.
  - Experiments illustrate scaling of accuracy and running time based on training set size (and scaling of running time, w.r.t. other parameters, in appendix).
  - Metrics for experiments are well-chosen to illustrate important aspects of classification (e.g., not just accuracy or LPD, but also model calibration) and show advantages of maintaining a full posterior over model parameters.

## Cons
- Section **2.1 The conditional mixture network** was not very clear.
- A more thorough analysis might also show the scaling of accuracy with respect to number of experts in the mixture, perhaps for some subset of problems.
- Although it may be a little out-of-scope, I found myself wondering how this particular model (i.e., a CMN, fit by whatever means) compares to a more general model, such as a fully Bayesian neural network with some reasonable architecture. In other words, the paper does a good job of arguing that CAVI-CMN produces the best CMN, is such a model particularly _good_ at classification, compared to other potentially more flexible models?

# Suggestions
## Major
- Section **2.1 The conditional mixture network** could be much clearer. It's not so much that the structure of the model itself is hard to grasp once you've digested it, but the presentation could use some work.
  - Instead of "a condtional mixture of linear experts, which outputs a joint continuous-discrete latent $\left( \mathbf{x}\_{1} \in \mathbb{R}^{h}, z_{1} \in \\{ 1, \ldots, K \\} \right)$", I would write "a condtional mixture of linear experts, **each of** which outputs a joint continuous-discrete latent $\left( \mathbf{x}\_{1} \in \mathbb{R}^{h}, z_{1} \in \\{ 1, \ldots, K \\} \right)$".
  - A conceptual diagram would help immensely.
  - Although I recognize that labeling can often be a matter of taste, I would re-label as follows:
    - $\mathbf{x}_{0} \rightarrow \mathbf{x}$
    - $\mathbf{x}_{1} \rightarrow \mathbf{z}\_{1}$
    - $z_{1} \rightarrow z\_{2}$

  and similarly for capital-letter collections. This would make the inputs $x$, outputs $y$, and all latents $z$. Even if you decide to keep the labeling as-is, I don't see any reason that $z_{1}$/$Z_{1}$ needs the subscript.
- You should cite the origin of WAIC (I believe this is S. Watanabe 2010 or 2013) not just the Vehtari paper.
- I would like to see error bars on Table 1: you've already run each experiment 16 times, and these are presumably the means? So why not include a standard deviation?
- Appendix Section **C.3 No U-Turn Sampler** should either report some statistic of convergence (split R-hat, IAC, or ESS) for NUTS, or at least report that some minimum threshold of stationarity was reached during each run.
- I don't think there's any use in reporting the "steps to convergence", since a step is so inherently different for each algorithm (at least, it's very different between CAVI/BBVI and CAVI/MLE). You could argue that it's useful to show the scaling behavior with training set size, but this is already better quantified by the wall-clock runtime.

## Minor
- Even though CAVI-CMN is "only" competitive with BBVI and NUTS on WAIC, it would still be nice to bold the best statistic in each column. Then the reader can compare CAVI-CMN at a glance by looking from the top row to the bolded item. This would, in fact, make it clearer that CAVI is competitive, even if it is not consistently the "best".
- I'm sure it's implicitly in the code provided, but at the bottom of page 13 in Appendix Section **D Dataset Descriptions** it might be nice to describe how the random parameters were initialized (i.e. as samples from the priors, or uniformly in some range, etc.).
title: title: The proposed method is well-developed, the paper is well-presented, and it provides a solid empirical comparison to other standard methods.
main_review: main_review: The authors introduce CAVI-CMN, a gradient-free variational method for training conditional mixture networks. The paper's key contribution is that CAVI-CMN outperforms other variational methods, such as BBVI, in terms of speed, while offering comparable efficiency to MLE and maintaining strong predictive accuracy.

The presentation of the paper is quite good. The empirical results clearly show the advantage of using CAVI-CMN compared to the baseline method. 

Some minor comments:
1. The empirical results are impressive, but I wonder about their scalability on larger models or datasets than those currently presented.
2. Although the PG augmentation scheme is explained in the appendix, it is not clearly described in the main paper. For instance, it is mentioned at the beginning of Section 2 and in Section 2.2, but a clearer one-sentence explanation might be helpful.
3. In Figure 1 and all other performance comparison figures, the green line labeled "CAVI" should be "CAVI-CMN," correct?
4  For advanced readers, it is clear why NUTS are not present in Figure 1, as it is expected to be the slowest among the methods. However, it might be better to include a clarification somewhere.
5. Unless I am missing something, I believe a more detailed description of the neural network used, such as the number of parameters, would help the reader better understand the scale of the model.