Gradient-free variational learning with conditional mixture networks

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.

• 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).

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.

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).