Neural Bayes: A Generic Parameterization Method for Unsupervised Learning
Authors: Devansh Arpit
Huan Wang
Caiming Xiong
richard socher
Yoshua Bengio
Venue:  (2021)
Keywords: unsupervised learning
clustering
manifold separation
representation learning
Bayes rule

Abstract: We introduce a parameterization method called Neural Bayes which allows computing statistical quantities that are in general difficult to compute and opens avenues for formulating new objectives for unsupervised representation learning. Specifically, given an observed random variable $\mathbf{x}$ and a latent discrete variable $z$, we can express $p(\mathbf{x}|z)$, $p(z|\mathbf{x})$ and $p(z)$ in closed form in terms of a sufficiently expressive function (Eg. neural network) using our parameterization without restricting the class of these distributions. To demonstrate its usefulness, we develop two independent use cases for this parameterization: 

1. Disjoint Manifold Separation: Neural Bayes allows us to formulate an objective which can optimally label samples from disjoint manifolds present in the support of a continuous distribution. This can be seen as a specific form of clustering where each disjoint manifold in the support is a separate cluster. We design clustering tasks that obey this formulation and empirically show that the model optimally labels the disjoint manifolds.

2. Mutual Information Maximization (MIM): MIM has become a popular means for self-supervised representation learning. Neural Bayes allows us to compute mutual information between observed random variables $\mathbf{x}$ and latent discrete random variables $z$ in closed form. We use this for learning image representations and show its usefulness on downstream classification tasks. 

Under review as a conference paper at ICLR 2021
NEURAL BAYES: A GENERIC PARAMETERIZATION METHOD FOR UNSUPERVISED LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
We introduce a parameterization method called Neural Bayes which allows comput- ing statistical quantities that in general difficult to compute and opens avenues for formulating new objectives for unsupervised representation learning. Specifi- cally, given an observed random variable x and latent discrete variable 2, we can express p(x 2), p(=[x) and p(=) in closed form in terms of sufficiently expres- sive funetion (Eg. neural network) using our parameterization without restricting the class of these distributions. To demonstrate its usefulness, we develop two independent use cases for this parameterization:
1. Disjoint Manifold Separation: Neural Bayes allows us to formulate an objective which can optimally label samples from disjoint manifolds present in the support of a continuous distribution. This can be seen as a specific form of clustering where each disjoint manifold in the support separate cluster. We design clustering tasks that obey this formulation and empirically show that the model optimally labels the disjoint manifolds.
2. Mutual Information Maximization (MIM): MIM has become a popular means for self-supervised representation learning. Neural Bayes allows us compute mutual information between observed random variables x and latent discrete random variables closed form. We use this for learning image representations and show its usefulness on downstream classification tasks.
INTRODUCTION
We introduce a generic parameterization called Neural Bayes that facilitates unsupervised learning from unlabeled data by categorizing them. Specifically, our parameterization implicitly maps samples from an observed random variable x to a latent discrete space where the distribution p(x) gets segmented into finite number of arbitrary conditional distributions. Imposing different conditions on the latent space z through different objective functions will result in learning qualitatively different representations.
Our parameterization may be used to compute statistical quantities involving observed variables and latent variables that in general difficult to compute (thanks to the discrete latent space), thus providing a flexible framework for unsupervised learning. To illustrate this aspect, we develop two independent cases for this parameterization- disjoint manifold separation (DMS) and mutual information maximization (Linsker, 1988), as described in the abstract. For the manifold separation task, we show experiments on 2D datasets and their tigh-dimensional counter-parts designed as per the problem formulation, and show that the proposed objective can optimally label disjoint manifolds. For the MIM task, we experiment with benchmark image datasets and show that the unsupervised representation learned by the network achieves performance on downstream classification tasks comparable with a closely related MIM method Deep InfoMax (DIM, (Hjelm et al., 2019)). For both objectives we design regularizations necessary to achieve the desired behavior in practice. All the proofs can be found in the appendix.
RELATED WORK
Neural Bayes-DMS: Numerous recent papers have proposed clustering algorithm for unsupervised representation learning such as Deep Clustering (Caron et al., 2018), information based clustering (Ji et al,, 2019), Spectral Clustering (Shaham al., 2018), 8), Assosiative Deep Clustering (Haeusser

---

Under review as a conference paper at ICLR 2021
et al., 2018) etc. Our goal in regards to clustering in Neural Bayes-DMS is in general different from such methods. Our objective is aimed at finding disjoint manifolds in the support of a distribution. It is therefore generalization of traditional subspace clustering methods (where the goal is to find disjoint affine subspaces) (Ma et al., 2008; Liu et al., 2010), to arbitrary manifolds.
Another class of clustering algorithms include mixture models (such kNNs). Our clustering proposal (DMS) is novel compared to this class in two ways: 1. we formulate the clustering problem as that of identifying disjoint manifold in the support of distribution. This is different from assuming K ground truth clusters, where the notion of cluster is ill-defined; 2. the DMS objective in proposition 1 itself novel and we prove its optimality towards labeling disjoint manifolds in theorem 1.
Neural Bayes-MIM: Self-supervised representation learning has attracted lot of attention in recent years. Currently contrastive learning methods and similar variants (such as MoCo (He et al., 2020), SimCLR (Chen et al., 2020), B BYOL Grill et al. (2020)) produce state-of-the-an (SOTA) performance on downstream classification tasks. These methods make use of handcrafted image augmentation methods that exploit priors such as class information is typically associated with object shape and is location invariant. However, since we specifically develop an easier alternative to DIM (which also maximizes the mutual information between the input and the latent representations), for a fair comparison, we compare the performance of our Neural Bayes-MIM, algorithm for representation learning only with DIM. We leave the extension of Neural Bayes MIM algorithm with data augmentation techniques and other advanced regularizations similar to Bachman al. (2019) as future work. Our experiments show that our proposal performs comparably or slightly better compared with DIM. The main advantage of our proposal over DIM is that it offers closed form estimation of MI due discrete latent variables.
We note that the principle of mutual information maximization for representation learning was introduced in Linsker (1988) and Bell & Sejnowski (1995), and since then, number of supervieed methods involving MIM have been proposed. Vincent et al. (2010) showed that auto-encoder based methods achieve this goal implicitly by minimizing the reconstruction error of the input samples under isotropic Gaussian assumption. Deep infomax (DIM, Hjelm et al. (2019)) uses estimators like MINE Belghazi et al. (2018) and poise-contrastive estimation (NCE, Gutmann & Hyvärinen (2010)) to estimate MI and maximize it for both both local and global features in convolutional networks. Contrastive Predictive Coding (Oord et al., 2018) is another approach that maximizes MI by predicting the activations of layer from the layer above using NCE.
We also point out that the estimation of mutual information due to Neural Bayes parameterization in the Neural Bayes-MIM-vl objective (Eq 8) turns out to be identical to the proposed in IMSAT (Hu et al., 2017). However, there are important differences: 1. we introduce regularizations which significantly improve the performance of representations on downstream tasks compared to IMSAT (cf table in (Hu et al., 2017)); 2. we provide theoretical justifications for the parameterization used (lemma 1) and show in theorem why it is feasible to compute high fidelity gradients using this objective in the mini-batch setting even though it contains the term Ex LL (x)]. On the other hand, the justification used in IMSAT is that optimizing using mini-batches equivalent to optimizing an upper bound of the original objective; 3. we perform extensive ablation studies exposing the importance of the introduced regularizations.
3 NEURAL BAYES
Consider a data distribution p(x) from which we have access to i.i.d. samples x E R.. We assume that this marginal distribution is union of K conditionals where the htth conditional's density is denoted by p(x|z = k) E IR+ and the corresponding probability mass denoted by p(= = k) E R+ Here is discrete random variable with K states. We now introduce the parameterization that allows us to implicitly factorize any marginal distribution into conditionals as described above. Aside from the technical details, the key idea behind this parameterization is the Bayes' nule.
Lemma 1.
Let p(x|= = k) and p(z) be any conditional and marginal distribution defined for continuous random variable x and discrete random variable z. If (T # O'Yk E [K]. then there exists non-parametric function T((x): T RR+ K for any given input E R with the
2

---

Under review as a conference paper at ICLR 2021
property Ex= Lx(x) =1 Vx such that,
[x)x)= p(x)
p(x|2 k) Rxxp(x) [x)]' -)==))-Exlak
p(= k|x)
Lx(xx
Thus the function can be seen as a form of soft categorization of input samples. In practice, we use neural network with sufficient capacity and softmax output to realize this function L. We name our parameterization method Neural Bayes and replace L with Lg to denote the parameters of the network. By imposing different conditions the structure of by formulating meaningful objectives, we will get qualitatively different kinds of factorization of the marginal p(x), and therefore the function La will encode the posterior for that factorization. In summary, if one formulates any objective that involves the terms p(x[z), p(2) or p( (x), where x is an observed random variable and is discrete latent random variable, then they be substituted with (x(x))(((x) L& (x) respectively.
E. [x()] Ex (x)) and
On an important note, Neural Bayes parameterization requires using the term Ex [Tx (x)]. through which computing gradient is infeasible general. A general discussion around this can be found in the appendix D. Nonetheless, objectives we propose using our parameterization. In the next two sections, we explore two different ways of factorizing p(x) resulting in qualitatively different goals of unsupervised representation learning.
show that mini-batch gradients can have good fidelity for one of the
4 DISJOINT MANIFOLD SEPARATION (DMS)
In many cases, the support of distribution may be set of disjoint manifolds. In this task, our goal label samples from each disjoint manifold with a distinct value. This formulation can be seen as a generalization of subspace clustering (Ma et al., 2008) where the goal is to identify disjoint affine subspaces. To make the problem concrete, we first formalize the definition of disjoint manifold.
Definition 1. (Connected Set) We say that set S C R" is connected set (disjoint manifold) for any y, E S, there exists a continuous path between x and such that all the points on the path also belong to S.
To identify such disjoint manifolds in distribution, exploit the observation that only partitions that separate one disjoint manifold from others have high divergence between the respective conditional distributions while partitions that cut through disjoint manifold result in conditional distributions with low divergence between them. Therefore, the objective we propose for this task is to partition the unlabeled data distribution p(x) into conditional distributions q: (xx)ss such that a divergence between them is maximized. By doing so we recover the conditional distributions defined over the disjoint manifolds (we prove optimality in theorem 1). We begin with two disjoint manifolds and extend this idea to multiple disjoint manifolds in the appendix B.
Let be a symmetric divergence (Eg. Gensen-Shanme divergence, Wasserstein divergence, etc), and 40 and qi be the disjoint conditional distributions that we want to learn. Then the aforementioned objective can be written formally as follows:
max ](Q((x)lqx(x))
(2)
90-91
(0)0.1)
[10(x)=1 1,
[
s.t.
(1(x) π + go()) (1- π) = p(x).
Since our goal is simply assign labels data samples corresponding to which manifold they belong instead of learning conditional distributions as achieved by Eq. (2), we would like to learn function L(x) which maps samples from disjoint manifolds to distinct labels. To do so, below we derive an objective equivalent to Eq. (2) that learns such function L(x).
Proposition 1. (Neural Bayes-DMS Let L(x): R + [0, 1] be non-parametric functionnfol for any given input R,,", and let J be the Jensen-Shannon Shannon divergence. Define scalars fi fi(x): [_][(x]]
Lix)
3

---

Under review as a conference paper at ICLR 2021
and fu(x):
1-L(x)
Then the objective in Eq. (2) is equivalent to,
[[x)]
1 max Ex [xx)) -log log
fi(x)
)
fo(x)
] log
[i(x) + fo(x)) )
Ex [o(x) log [i(x) fo(x)
L 2
s.l. Ex [x[[()]] & 1).
Optimality: We now prove the optimality of the proposed objective towards discovering disjoint manifolds present in the support of a probability density function p(x).
Theorem 1. (optimality) Let p(x) be probability density function over R" whose support is the union of two non- non-empty connected sets (definition 1) S1 and S2 that are disjoint, i.e. S1 n = 2. Let L(x) [0, 1] belong the class continuous functions which is learned by solving the objective in Eq. (3). Then the objective in Eq. (3) is maximized if and only if one of the following is true:
Vx E Si 0
E
L(x)=
[(x)=
1
VX E S2
0
VX S2
The above theorem proves that optimizing the derived objective over the space of functions L implicitly partitions the data distribution into maximally separated conditionals by assigning distinct label to points in each manifold. Most importantly, the theorem shows that the continuity condition on the function Lxx) plays important role. Without this condition, the network cannot identify disjoint manifolds. Extension to multiple disjoint manifold case can be found in section B in appendix.
4.1 IMPLEMENTATION DETAILS
Prior Collapse: The constraint in proposition is boundary condition required for technical reasons in lemma 1. In practice we do not worry about them because optimization itself avoids situations where Ex [x)] {0, 1}. To see the reason behind this, note that except when initialized in way such that Ex [xx] E {0, 1}, the log terms are negative by definition. Since the denominators of So and fi are Ex [x)] and
Ex [[x)] respectively, the objective is maximized when Ex [x)) moves away from 0 and 1. Thus, for any reasonable initialization, optimization itself pushes Ex [x)]
away from 0 and 1.
Smoothness of (.):: As shown in theorem the proposed objectives can optimally recover disjoint manifolds only when the function L.(.) is continuous. practice we found enforcing the function to be smooth (thus also continuous) helps significantly. Therefore, after experimenting with a handful of heuristics for regularizing Lg. we found the following finite difference Jacobian regularization to be effective (L..) can be scalar vector),
B Σ im1
(o(x)) Lo(x; + C +..)
R₂ = B
(4)
5. is normalized noise vector computed independently for each sample x; in a batch where δ; = la of size as 6, == Xv; Here E gaxB is the matrix containing the batch of samples, and each dimension of v, E B is sampled i.i.d. from a standard Gaussian. This computation ensures that the perturbation lies in the span of data, which we found to be important. Finally Ç is the scale of normalized noise added to all samples in batch. In experiments, since always normalize the datasets to have zero mean and unit variance across all dimensions, we sample Ç ~ (0 0.1²).
Implementation: We implement the binary-partition Neural Bayes-DMS using the Monte-Carlo sampling approximation of the following objective,
[i()) log 1 [o(x)
] + 2 Ex fo(x) log (1
1
f((x))) ) + 3. B. (5) fo(x)
Ex
8
fx(x)
where fi() == Ex [x[o]]]] + and fo(x) := I-Ex][a(x]] + Here = 10 is small scalar used control the continuity of L. The multi-partition case can be implemented in a similar way. Due to the need for computing Ex[Le in the objective, optimizing using gradient descent methods with small batch-sizes is not possible. Therefore we experiment with this method on datasets where gradients can be computed for a very large batch-size needed to approximate the gradient through E (T)) sufficiently well.
[xp(
1-Le(x)
to prevent numerical instability, and 8 is hyper-parameter
4

---

Under review as a conference paper at ICLR 2021
4.2 EXPERIMENTS
Clustering in general is an ill posed problem. However, in our problem setup, the definition is precise, i.e., our goal is to optimally label all the disjoint manifolds present the support of distribution. Since this is unique goal that is not generally considered in literature, as empirical verification,
show qualitative re- sults on 2D synthetic datasets in figure 1. Top 2 sub-figures have 2 clusters and the bottom 2 have 3 clusters. For all experiments we use a 4 layer MLP with 400 hidden units each, batch- norm, ReLU activation, and last layer Softmax activation. In all cases train using Adam optimizer with learning rate of 0.001, batch size of 400 and no weight decay, and trained until convergence. Regularization coefficient / was chosen from [0.5, 6] that resulted optimal clustering. For generality in these experiments, these 2D datasets were projected to high dimen- sions (512) by appending 510 dimensions of O entries to each sample and then randomly ro- tated before performing clustering. The datasets were then projected back the original 2D space for visualizing predictions. Additional experiments can be found in the appendix C.
Figure I: Neural Bayes-DMS network prediction on synthetic high dimensional data (lying on a 2D plane) projected back to dimensions for visual- ization. The different colors denote the label pre- dicted by the network La(.) thresholded at 0.5. The darker shades of colors denote predictions made at training data points while the lighter shades denote predictions the of the space.
MUTUAL INFORMATION MAXIMIZATION (MIM)
Suppose want find discrete latent representation (with K states) for the distribution p(x) such that the mutual information M MI(x. z) is maximized (Linsker, 1988). Such an encoding demands that must be very efficient since has to capture maximum possible information about the continuous distribution p(x) in just K discrete states. Assuming we can learn such an encoding, we are interested in computing p(z[x) since tells us the likelihood of x belonging to each discrete state of z, thereby performing soft categorization which may be useful for downstream tasks. In the proposition below, we show an objective for computing p(z(x) for discrete latent representation that maximizes M MI(x,=).
Proposition 2. (Neural Bayes-MIM-vl] Let L(x) R= + R+ be a non-parametric function for any given imput E I with the property Σ Ei=1 Lx(x) 1 Vx. Consider the following objective,
K
[x(x)
L = =aa max Ex b
Σ
(x(x) (xxx)log
Ex [xx)]x
Then Lī((x) p(=" k(x),, where 2t arg max; MT(x, 2).
The proof essentially involves expressing MI in terms of p(z|x), and p(z), which be substituted using Neural Bayes parameterization. However, the objective proposed in the above theorem poses challenge- the objective contains the term Ex L (x)] for which computing high fidelity gradient in batch setting is problematic (see appendix D). However, we can overcome this problem for the MIM objective because turns out that gradient through certain terms are 0 as shown by the following theorem.
Theorem 2. (Gradient Simplification) Denote,
Lxx(x)
K Σ
Le (x)
J(Ø) -Exx
Lo
Lo log Ex[Lo. (x))
j(() Exx
Lo (x) log
(7)
Ex Lo, (x))
LA=1
where (.) indicates that gradients computed through the argument. Then 0J(0)
OJ(0)
88
5

---

Under review as a conference paper at ICLR 2021
The above theorem implies that as long as we plugin decent estimate of Lx (x)] in the objective, unbiased gradients can be computed using randomly sampled can be re-written as,
Note that the objective
K
K Σ
min
Exx
Σ
Lo (x) (La (x))
+
Ex(La log
(8)
The second term is the negative entropy of the discrete latent representation p(= = = Ex (x))
which acts as a uniform prior. In other words, this term encourages learning a latent code 2 such that all states of activate uniformly over the marginal input distribution x. This is an attribute of distributed representation which is a fundamental goal in deep learning. We can therefore further encourage this behavior by treating the coefficient of this term hyper-parament confirm both the distributed representation behavior of this term as well as the benefit of using hyper-parametee as our coefficient.
In our experiments we
5.1 IMPLEMENTATION DETAILS
Alternative Formulation of Uniform Prior: In practice we found that an alternative formulation of the second term in Eq 8 results in better performance and more interpretable filters. Specifically, we replace it with the following cross-entropy formulation,
K
(0) = Σ
log (x)( + K K
log(1 Ex [x())xx))
(9)
k=1
While both, the second term in Eq 8 well (0) are minimized when Ex[Lx(x]] 1/K, the latter formulation provides much stronger gradients during optimization when Ex [xx()) approaches 1 (see appendix F.1 for details); [x representation. Finally, unbiased gradients can be computed through Eq 9 as long as a good estimate of [L (x) is plugged in. Also note that the condition Ex [xx((x) $ {0, 1} in lemma I is met by the Neural Bayes-MIM objective implicitly during optimization as discussed in the above paragraph in regards to distributed representation.
1 is undesirable since it discourages distributed
Implementation: The final Neural Bayes MIM- objective is,
K
min Ex
Lav (x) log (x) + E) Σ
+ a1+ (2) + B Ro
(10)
where a and 3 are hyper-parametees R. is a smoothness regularization introduced in Eq. e 1 is a small scalar used to prevent numerical instability. Qualitatively, we find that the regularization RR prevents filters from memorizing the input samples. Finally, we apply the first two terms in Eq 10 to all hidden layers of a deep network at different scales (computed by spatially average pooling and applying Softmax). These two regularizations gave a significant performance boost. Thorough implementation details are provided in the appendix E. For brevity, we refer to our final objective as Neural Bayes- MII in the rest of the paper.
On the other hand, to compute a good estimate of gradients, use the following trick. During optimization, we compute gradients using a sufficiently large mini-batch of size MBS (Eg. 500) that fits in memory (so that the estimate of Ex (x)) is reasonable), and accumulate these gradients until BS samples are seen (Eg. 2000), and averaged before updating the parameters to further reduce estimation error.
5.2 EXPERIMENTS
Instead of aiming for state-of-the-at results on self-supervised learning tasks, our goal in this section conduct experiments using Neural Bayes-MIM to understand the behavior of the algorithm, the involved, and compare performance with DIM, a closely related self-supervised learning method. The use of data augmentation, additional architectures and other regularizations similar to Bachman et al. (2019) is left as future work. Therefore, we use the following simple CNN encoder architecture Enc in our experiments: C(200, 3, 1, 0) P(2,2, O, max) C(500, 3, 1,
We use the following shorthand for a) conv layer: C(number of filters, filter size, stride size, padding); b) pooling: P(kernl size, stride, padding, pool mode)

---

Under review as a conference paper at ICLR 2021
Without Scaling
With Scaling
Test Accuracy
80
Accuracy
60
60
8
Figure 2: The effect of Neural Bayes-MIM pyper-parameters on the final test performance for CIFAR-10. A CNN encoder is trained using Neural Bayes-MIM with different configurations- hyper-parameters a, 3 and scaling of states on which the objective is applied. one hidden layer classifier (with 200 units) is then trained using labels on these frozen features to get the final test accuracy. The plots show that performance is significantly worse when a 8 = 0 and no scaling used, showing their important role regularizer. The best performing model reaches 83.59%. Black dotted line is the baseline performance (77%)) when randomly initialized network (with identical architecture) is used as encoder. Green and blue dotted lines are the average of all the green and blue points respectively.
C(700, 3. 1) P(2, 2, O, max) C(1000, 3, 1,). For an input image x of size 32 x 32 x 3, the output of this encoder Enc(x) has size x 2 1000. The encoder is initialized using orthogonal initialization (Saxe et al., 2013), batch normalization (loffe & Szegedy, 2015) is used after each convolution layer and ReLU non-linearities are used. All datasets are normalized to have dimension- wise O mean and unit variance. Early stopping in all experiments is done using the test set (following previous work). We broadly follow the experimental setup of Hjelm et al. (2019). We do not use any data augmentation in our experiments. After training the encoder, we freeze its features and train hidden layer (200 units) classifier to get the final test accuracy.
5.2.1 ABLATION STUDIES
Behavior of Neural Bayes-MIM-rl (Eq 8) vs Neural Bayes-MIM (v2, Eq 9): The experiments and details are discussed in the appendix F.2. The main differences are: 1. majority of the filters learned by the objective are dead, as opposed to the v2 objective which encourages distributed representation; 2. the performance of v2 is better than that of the v1 objective.
We now evaluate the effects of the various components involved in the Neural Bayes-MIM objective- coefficients and B, and applying the objective at different scales of hidden states. We use the CIFAR- 10 dataset for these experiments.
Performance due to Regularizations and State Scaling:
In the first experiment, for each value of the number of different scales considered, we vary a, B and record the final performance, thus capturing the variation in performance due to all these three components, We consider two scaling configurations: 1. no pooling is applied to the hidden layers; 2. for each hidden layer, we spatially average pool the state using a 2 x 2 pooing filter with a stride of 2. For the encoder used in our experiments (which has 4 internal hidden layers post ReLUU, this gives us 4 and 8 states respectively (including the original un-scaled hidden layers) to apply the Neural Bayes-MIl objective. After getting all the states, we apply the Softmax activation each state along the channel dimension so that the Neural Bayes parameterization holds. Thus for states with height and width, the objective is applied to each spatial (x. y) location separately and averaged. Also, for states with height (or width) less than the pooling size, we use the height (or width) as pooling size.
We train Neural Bayes-MIn MIM on the full training set for 100 epochs using Adam with learning rate 0.001 (other Adam hyper-parameet decay. In the first 32 experiments, ar and are sampled uniformly from [o, 10) and [0. 5] respectively. In the next experiments, a is set be 0 while 3 is sampled uniformly. In the next 5 experiments, is set to be while a is sampled uniformly. Thus in total we run 42 experiments for each number of scaling considered.
are standard), mini-batch size 500 and batch size 2000, O weight
Once we get a trained Enc(x), we train a 1 hidden layer (with 200 units) MLP classifier on the frozen features from Enc(x) using the labels in the training set. This training is done for 100 epochs using Adam with learning rate 0.001 (other Adam hyper- parameteess are standard), batch size 128 and weight decay O.
7

---

Under review as a conference paper at ICLR 2021
Enender Dataset
CIFAR-TO
CIFAR-TOD
STT.-10
Random Network
67.97
42.97
53.91
DIM (Hjelm et al., 2019)
80.95
49.74
Neural Bayes-MIM
82.81
55.47
64.84
Table 1: Classification performance of a one hidden layer MLP classifier trained on frozen features from the mentioned encoder models (trained using Neural Bayes-MIM) and datasets (no data augmentation used in experiments we ran). STL-10 was resized to 32 x 32 our runs instead of 64 64 as in DIM due to memory restrictions. Performance reported from DIM paper are their best numbers (omitting STL-10 due to difference in image size).
As baseline for these experiments, tunable hyper-parameters in this case, we perform a grid search on the classifier hyper-papaaae apper-parameter... Specifically, we choose weight decay from {0, lc 5, 5e 5, le 4). batch size from {128, 256). and learning rate from {0.0001, 0.001 ]. This yields a total of 16 configurations. The test accuracy from these runs varied between 58.59% and 67.97%. We consider 67. 67.97% our baseline.
randomly initialed encoder Enc(x). Since there no
The performance of encoders under the aforementioned configurations is shown in figure 2. It is clear that both the hyper-parameters a and especially 8 play an important role in the quality of representations learned. Also, applying Neural Bayes-MIM at different scales of the network states significantly improves the average and best performance.
Filter visualization, convergence experiments and the effect of batch-size are shown in appendix F.3.
5.2.2 FINAL CLASSIFICATION PERFORMANCE
We compare the final test accuracy of Neural Bayes-MIM with baselines- a random encoder (described in ablation studies) and Deep Infomax (Hjelm et al., 2019) on benchmark image datasets- CIFARIO and CIFAR- (Krizhevsky, 2009) and STL-10 (Coates al., 2011). Random Network refers to the use of randomly initialized network. The experimental details for them are identical to those in the ablation above involving a hyper-paramett search over 16 configurations done for each dataset separately.
DIM results are reported from Hjelm et (2019). We omit STL-10 number for DIM because we resize images to much smaller size of 32 x 32 in our runs instead of 64 x 64 as used in DIM.
The following describes the experimental details for Neural Bayes-MIM. We use a and 3 (chosen roughly by examining figure 2), and MBS=500, BS=4000 in all the experiments. Note these values are not tuned for STL-10 and CIFAR-100. For CIFAR-10 10 and STL-10 each, we run learning rate {0.0001, 0.001) } and
configurations of Neural Bayes-MIM over weight decay E {0, 0.00001). ). For each run, we then train a 1 hidden layer (200 units) classifier on top of the frozen features with learning rate E {0.0001, 0.001 J. We report the best performance of all runs. For CIFAR- 100, we take the encoder that produces the best performance on CIFAR-10 10, and train classifier with the 2 learning rates and report the best of the 2 runs.
Table reports the classification performance of all the methods. We note that all experiments were done with CNN architecture without any data augmentation. Neural Bayes-MIM outperforms baseline methods in general.
CONCLUSION
We proposed a parameterization method that can be used to express an arbitrary set of distributions p(x]z), p(z]x) and p(z) in closed form using a neural network with sufficient capacity, which can in turn be used to formulate new objective functions. We formulated two different objectives that use this parameterization which were aimed towards different goals of unsupervised learning- identification of disjoint manifolds in the support of continuous distributions and learning deep network features using the infomax principle. We focused on theoretical analysis for both the objectives and presented preliminary experiments supporting the theoretical claims.
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In Advances in Neural Information Processing Systems, pp. 15509-15519, 2019.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeswar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and R Devon Hjelm. Mine: mutual information neural estimation. arXiv preprint OrXiv:1801.0466. 2018.
Anthony J Bell and Terrence J Sejnowski. An informatiion-maxiiiztioo approach to blind separation and blind deconvolution. Neural computation, 7(6): 129-119999,, 1995.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsuper- vised learning of visual features. Proceedings of the European Conference on Computer Vision (ECCV), pp. 132-149, 2018.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv preprint arXiv:2002.05709 2020.
Adam Coates, Andrew Ng. and Honglak Lee. An analysis single-layer networks in unsupervised feature learning. In Proceedings of the fourteenth international conference artificial intelligence and statistics, pp. 215-223, 2011.
Baan-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07733 2020.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the Thirteenth International Conference Artificial Intelligence and Statistics, pp. 297-304, 2010.
Philip Hacusser, Johannes Plapp. Vladimir Golkov, Elie Aljalbout, and Daniel Cremers. Associative deep clustering: Training a classification network with no labels. In German Conference on Pattern Recognition, pp. 18-32. Springer, 2018.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 9729-9738, 2020.
R Devon Hjelm, Alex Fedorov, Samuel Maroii-Marc,,don, Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. International Conference Learning Representations (ICLR), 2019.
Weihua Hu, Takeru Miyato, Seiya Tokui, Eicchi Matsumoto, and Masashi Sugiyama. Learning discrete representations via information maximizing self- augmeneeedt training. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pp. 1558-1567. JMLR. org, 2017.
Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167, 2015.
Xu Ji, João F Henriques, and Andrea Vedaldi. Invariant information clustering for unsupervised image classification and segmentation. In Proceedings of the IEEE International Conference on Computer Vision, pp. 9865-9874, 2019.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Ralph Linsker. Self-organization in perceptual network. Computer, 21(3):105-117, 1988.
Guangcan Liu, Zhouchen Lin, and Yong Yu. Robust subspace segmentation by low-rank represen- tation. In Proceedings of the 27th international conference on machine learning (ICML-10), pp. 663-670, 2010.

---

Under review as a conference paper at ICLR 2021
Yi Ma, Allen Y Yang. Harm Derksen, and Robert Fossum. Estimation of subspace arrangements with applications in modeling and segmenting mixed data. SIAM review, 50(3):413 458, 2008.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprini 1rXXiv:807..322 2018.
Andrew M Saxe, James L. McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 13/2.6120 2013.
Uri Shaham, Kelly Stanton, Henry Li, Boaz Nadler, Ronen Basri, and Yuval Kluger. Spectralnet: Spectral clustering using deep neural networks. arXiv preprint arXiv:1801.01587 2018.
Pascal Vincent, Hugo Larochelle, Isabelle Lajoie, Yoshua Bengio, and Pierre-Antoine Manzagol. Stacked denoising autoencoders: Learning useful representations in deep network with local denoising criterion. Journal of machine learning research, L!(Dec):3311-3408 2010.
10

---

Under review as a conference paper at ICLR 2021
APPENDIX
A PROOFS FOR NEURAL BAYES-DME (BINARY CASE)
Proposition (Neural Bayes-DMS, proposition 2 in main text) Let L(x) R" + |0, 1] be a non- parametric function for any given input x E R and let J be the Jensen-Shanne divergence. Define
scalars f1(x) Fe [x[](
L(x)
1-L(x) Then the objective in Eq. (2) is equivalent to, and fo(x) = T- [[(x]]
1 max Ex f1(x)- log f1(x) +
fi(x)
)
)
fo(x)
) (((x))
+ Ex fu(x) log
log 2
1, 2
f1(() fo(x))
(11)
s.t. [x[L(x)] & {0,1)
(12)
Proof: Using the Neurul Bayes parameterization from lemma for binary case, we set,
(x(x) :- L(x) p(x)
4x(x) (1 (1-)(x)) p(x)
(13)
Exxp(xx [[x)]
Ex-p(() [[x)]
These parameterizations therefore automatically satisfy the constraints in Eq. (2). Finally, using the definition of JS divergence, the maximization problem in Eq. (2) can be written as,
-. q1( log 0.5 (go(x) + qi (x)) + (x) log 0.5 (ạo(x) + (x)) 90:91 2
41 (x)
(x(x)
max
(14)
Substituting qo and qı with their respective parameterizations and using the definitions of fo(x) and (x) completes the proof. O
main text) Let p(x) be probability density fanction over TR." whose support is the union of two non- non-empty connected sets (definition 1) S₁ and So that are disjoint, i.e. S₁ nS2 = Ø. Let L(x) E [0, 1] belong to the class of continuous functions which is learned by solving the objective in Eq. (3). Then the objective in Eq. (3) is maximized if and only if one of the following true:
Theorem 1. (optimality, Theorem
50 Vx E Si
[1
VX E S
L(x)
L(x)
(15)
Vx S2
VX
S2
Proof: The two cases exist in the theorem due to symmetry. Recall the definition of fo(x) and fi(x) in Eq. (3),
L(x) Ex[[(x)
1- L(x)
fi(x)
fo(x) 1- [xx]
(16)
where L(x) [0, 1] and for a feasible L(x), and therefore т = Ex(L(x) E (0, 1) due to the conditions specified in this theorem. Thus fi(x) [0, [] and fo(x) [0, T- By design, the terms
log
fi(x) 1o(x]
and log
fo(x)
non-positive. Thus, for any x E S U..
[o(x)+
[1(x)
fo(x)
F(x) = fi(x) log
+ fo(x) log
(17)
fi(x) + fo(x)
f1(x) + fo(x)
is maximized only when [(x) = 0 or [(x)= leading to F(x)= = O. Therefore, the objective in Eq. (3) is maximized by setting T.(x) = 0 or T(x) Vx U.2
Finally, since L(x) is continuous function, fx1, S such that L(x1) 0 and L(x2) = 1. We prove this by contradiction. Suppose there exists a pair X2) of this kind. Then along any path connecting x and x2 within Si, there must exist point where L(x) is not continuous since L(x) = or L(x) = 1 Yx E S₁ U S2 to satisfy the maximization condition. This a contradiction. By symmetry, the same argument mentioned in the theorem must be the optimal solution for L(x) in Eq. (3). Thus we have proved the claim. 0
be proved for x1, X2 E Sz. Therefore one the two cases

---

Under review as a conference paper at ICLR 2021
B NEURAL BAYES-DMS: EXTENSION TO MULTIPLE PARTITIONS
In order to extend our proposal to multiple partitions (say K), the idea is to find conditional distribution di (i [K]) corresponding to each of the K partitions such the divergence between conditional distribution of every partition and the conditional distribution of the combined remaining partitions is maximized. Specifically, we propose the following primary objective,
K
1
max * (KK
J (ax (x)) s.t.
(18)
K
9he (x)))= =1 wk [K]]
(19)
9h (x(xx) Th = p(x)
(20)
TT = 1
(21)
1om
where (x) is the conditional distribution corresponding to the full data distribution excluding the partition defined by qie (x). Formally,
Th (x) = p(x) 9k (x) Te
(22)
1
Then the theorem below shows an equivalent way of solving the above objective.
Theorem 2. Let L(x)= R + + be a non-parametric function for any given input x E B with divergence. Define scalars Lx(x) and [x(x) == 1=)x(x) Then the objective in Eq. (18) is equivalent to, fx(x) = [x[
the property Σ [x)x) Vx, and let J be the
1
fx(x)
fa(x) Fr((x) ffx(x)
)
max
2
Ex
Sa(x) llog
fa(x) + fa(x)
[x( log
+ logg 2
k=
(23)
S.l. [x]LL(x)] Th
Here Lx(x) denotes the kth unit of L(x).
Proof: Similar to theorem 1, the main idea is to parameterize Qk and qi. as follows,
Tu. (x(x = Lx(x) p(x)
Exxptx) [Tx (x) = (1 - La (x)) p(x)
(24)
EE Exxp(x) (x)]
To verify that these parameterizations are valid, note that,
[ [x(x)) / (x) p(x)
=1
(25)
(x)epp()]LLa [xx(x)]
Similarly, Sx (x) = 1. To verify that the second constraint is satisfied, we use the above parameter-
ization and substitute Ex [xx) = कy and get.
K Le(x) p(x)
K
Tsk p(x)
Σ Ln(x)
(26)
= p(x)
(27)
where the last equality uses the definition of L(x). Also notice that each Ta E [0, 1] and thus Ex [x)(x)) Ty is feasible for any arbitrary distribution qu (x(x) when L (x) 2
2

---

Under review as a conference paper at ICLR 2021
Finally, using the proposed parameterization have, p(x) q:(x) TT
2k. (xxx
(28)
(x) [x[[x[]]] 7T
= -p(x)
(29)
1 Th
(x)
= =p(x)
(30)
Ex [x((x]]
=f((x) p(x)
(31)
where we have used the fact that Ex [x]]=(x)] Tk- Using the definition of JS divergence, the max problem in Eq. (18) can be written as.
1
Σ k=
gx(x) log (0.5. (x(x) 0.5 (gx(x) + (x)))) (x)) -log 0.5 (xx(x) + qu (x))
(x(x)
max [AVEE
2
(32)
Substituting qs and qs with their respective paraanterrirtsitiionnn and using the definitions of f& fx(xx) and fa (x) completes the proof. O
In terms of implementation, we propose to simply have K output units in the label generating network Le while sharing the rest of the network. Also, we use Softmax activation at the output layer to satisfy the properties of L specified in the above theorem.
C ADDITIONAL NEURAL DAYSS-DMS EXPERIMENTS
We run experiment MNIST. We randomly split the training set into 90% -10% training- validation set In this experiment, we train a CNN with the following architecture: C(100, 3, 1,0) P(2, max) 3, 1, 1,0)-C(200) 3, 122, 2, O, max) 3, 1, P FC(10). Here P(.,.. ........ avg) denotes the entire spatial field is average pooled to result in 1 x height- width, and (10))) denotes a fully connected layer with output dimension 10. Finally, Softmax is applied at the output and the network is trained using the Neural Bayes- Bayes-DMS objective. We optimize the objective using Adam with learning rate 0.001, batch size 5000, weight decay for 100 epochs (other Adam hyper-parameters are kept standard). We use B for the smoothness regularization coefficient. Once this network is trained, we train linear classifier on top of this 10 dimensional output using Adam with identical configurations except batch size of 128 is used. We early stop on the validation set of MNIST and report the test accuracy using that model. The classifier reaches 99.22% test accuracy. This experiment shows that MNIST classes lie on nearly disjoint manifolds and that Neural Bayes-DMr DMS can correctly label them. As baseline, a linear classifier trained on features from randomly initialized identical CNN architecture reaches 42.97%.
avg)
D GRADIENT COMPUTATION PROBLEM FOR THE Ex [x)] TERM
The Neural Bayes parameterization contains the term Fy (x)). Computing unbiased gradient through this term is in general difficult without the use of very large batch-sizes even though the quantity [x[L+(x)) itself may have good estimate using very few samples. For instance, consider the scalar function (t) 1 0.01 sin wt. Consider the scenario when - 0c. The quantity E[)(t)] can be estimated very accurately using even one example. Further, E[}(tt] = 1, hence 8E [A]
O. However, when using finite number of samples, the approximation of @E(P()) of can have a very high variance estimate due to improper cancelling gradient terms from individual samples.
In the case of Neural Bayes-MIl we found that gradients through terms involving [x)] were 0. This allows to estimate gradients for this objective reliably in the mini-batcl setting. But in general it may be challenging to do so and solving objectives using Neural Bayes parameterization may require customized work-around for each objective.
3

---

Under review as a conference paper at ICLR 2021
E IMPLEMENTATION DETAILS OF THE NEURAL BAYES-MIM OBJECTIVE
We apply the Neural Bayes-MIM objective (Eq 10) to all the hidden layers at different scales (using average pooling). We now discuss its implementation details. Consider the CNN architecture used in our experiments- C(200, 3, 1, P(2, O, max) - C(500, 3, 1. 0) - C(700, 3. .) P(2, 0, max) C(1000, 3, 1, 0). Denote h' (i {0, 1, 2, 3)) be the 4 hidden layer ReLU outputs after the 4 convolution layers. For input of size 32 x 32 x 3, all these hidden states have height and width dimension in addition channel dimension. For mini-batch B, these hidden states are therefore 4 dimensional tensors. Let these 4 dimensions for the 1th state be denoted by |B] x C, X H, x W, where the dimensions denote batch-size, number of channels, height and width. Denote S to be the Softmax function applied along the channel dimension, and γ to be P(2, 2, 0, avg). Further, denote h' == PPhh-1) (i E (4,5, 5, 6, 7]) as the scaled version of the original states computed by average pooling, and define numbers C₁, H₁, W, accordingly. Then the total Neural Bayes-MIM objective for this architecture is given by,
H:W: C.
min 6
Σ
A (x)) Log(S((hi (x)) + E)
H.W,
XEB
1=
k 1=1
+ (1+a) (0) + +B- Re
(33)
where,
7
H.W,,
C log Σ
R2( =
Σ
S(h ,(
+
C, 1
log
1
1
Σ
S(h
8 1=0
H.W.
C
C₁
B
h,w=1
xEB
xEB
(34)
and,
IP(h
(( + ( 8111
R2 =
(35)
B
C²
X
where 8 =
8 1 is normalized noise vector computed independently for each sample x in the batch
B as,
8 Xv.
(36)
Here X E RaxB is the matrix containing the batch of samples, and each dimension of v E R sampled i.i.d. from a standard Gaussian. This computation ensures that the perturbation lies in the span of data. Finally ( is the scale of normalized noise added to all samples in a batch. In our experiments, since we always normalize the datasets to have zero mean and unit variance across all dimensions, we sample Ç ~ (0, 0.1²). Note that for the architecture used, P(h results in an output with height and width equal to I, hence the output is effectively a 2D matrix of size [B] x C3. Finally, the gradient form this mini-batch is accumulated and averaged over multiple batches before updating the parameters for more accurate estimate of gradients.
F ADDITIONAL ANALYSIS OF NEURAL BAYES-MIM
F.1 GRADIENT STRENGTH OF UNIFORM PRIOR IN NEURAL BAYES-MIM-VI (Eq 8) vs NEURAL BAYES- (9)
As discussed in the main text, the term,
K
R5 (0) = Σ
(37)
k=1
acts as a uniform prior encouraging the representations to be distributed. However, gradients are much stronger when Ex [Lx (x)) approaches for the alternative cross-entropy formulation,
K
K - K
(0) = log Ex (x))
log(1 Ex)(((x)))
(38)
K
4

---

Under review as a conference paper at ICLR 2021
Figure 3: MNIST filters learned using Neural Bayes- objective (Eq 8) using the configuration 4, 3 Majority of filters dead. For comparison with filters from Neural Bayes-MIM-v2. see figure (bottom).
To see this, note that gradient for (2) is given by,
(R)
Lo, (x)]
OE Lon (x)
log Fx [Lx
(39)
00
00
00
k=1
k=1
K
bl Lo. (x)]
θE 1=k=1 Lo (x)
log Ex|Lo
(40)
30
00
km
01 (x)]
Σ
log Ex Lo, (x))
(41)
0e
where the last equality holds due the linearity of expectation and because Σx Lo (x) by design. On the other hand, gradients for R2²2 (0) is given by,
(R)
K (x) -- Ex Ln(x)]
K
dE Lo. (x)
(42)
00
K
30
When the representation being learned is such that the marginal p(z) peaks along single state k, ie.. Ex [x)] + 1 (making the representation degenerate), the gradient for the ttth term for vl given by.
ØEx Le, (x) 00
log Ex LL, (x)] R 0
(43)
while that for v2 is given by,
K -
01Ex (x)]
DEx[Le, (x)]
lim c-+0
(44)
K [x]]x(x)
1- Ex(L+(x)
00
00
whose magnitude approaches infinity Ex [Ix(x)] gradient strength.
Thus (²) is beneficial in terms of
F.2 EMPIRICAL COMPARISON BETWEEN NEURAL BAYES-MIM-V (Eq 8) AND NEURAL BAYES-MIM-V2 (9)
To empirically understand the difference in behavior of Neural Bayes-MIM objective vl vs v2, we first plot the filters learned by the vl objective and compare it with those learned by the v2 objective. The filters learned by the v1 objective are shown in figure 3 using the configuration = 4, 3 = 4. can be seen that most filters are dead. We tried other configurations as well without any change in the outcome. Since the vI and v2 objective differ only in the formulation of the uniform prior regularization, as explained in the previous section, we believe that vl leads to dead filters because of weak gradients from its regularization term.
In the second set of experiments, we train many models using Neural Bayes-MIM-v2 objectives separately with different hyper-parameter configurations similar to the setting of figure 2. The performance scatter plot is shown figure We find that Neural Bayes- MIM-v2 has better average and best performance compared with Neural Mayes-MIM-v1.
and Neural
5

---

Under review as a conference paper at ICLR 2021
(a)
0,
3
3
(b)
0
(c) a = 4, =4
Figure 4: MNIST filters learned using Neural Bayes-MIM objective Eq 10. Majority of filters are dead when regularization coefficient = 0. Filters memorize input samples when regularization coefficient 0. Using both regularization terms results in filters that mainly capture parts of inputs which are good for distributed representation.
MBS ABS
250
500
2000
3000
50
40.62
42.97
41.41
75
78.91
100
N/A
67.97
66.41
78.12
78.91
250
N/A
76.56
78.91
82.03
84,38
500
N/A
N/A
82.03
78.91
79.69
Table 2: Effect of mini-batch size (MBS) and batch size (BS) final test accuracy of Neural Bayes-MIM on CIFAR-I0. Gradients are computed using batches size MBS and accumulated until BS samples are seen before parameter update. As expected, sufficiently large MBS is needed for computing high fidelity gradients due to [x+(x)) term. Gradient accumulation using BS further helps. All models are trained for the same number of epochs.

---

Under review as a conference paper at ICLR 2021
80
Test Accuracy
60
40
0
20
40
60
80
100
Epochs
Figure 5: Mean (solid line) and standard deviation (error band) of accuracy evolution during MLP classifier training on CIFAR-IO using Neural Bayes-MIM encoder (blue) and random encoder (green). Features learned by Neural Bayes-MIM allow training to start at much higher value (~~ 65% on average) and converge faster as opposed a random encoder (~~ 47% on average).
V2
Test Accuracy
Text 70
70
Figure Performance of Neural Bayes-MIM-vl vs Neural Bayes-MIM-v2. CNN encoder is trained using Neural Bayes-MIM with different configurations- hyper-parameters B and scaling of states on which the objective is applied. A one hidden layer classifier (with 200 units) then trained using labels on these frozen features to get the final test accuracy, Green and blue dotted lines are the average all the green and blue points respectively. Neural Bayes-MIM-v2 has better average and best performance compared with Neural Bayes- -MM--1.
F.3 ADDITIONAL EXPERIMENTS
Visualization of Filters:
We visualize the filters learned by the Neural Bayes-MIM objective on MNIST digits and qualitatively study the effects of the regularizations used. For this train deep fully connected network with 3 hidden layers each of width 500 using Adam with learning rate 0.001, batch size 500, 0 weight decay for 50 epochs (other Adam
are kept standard). We train three configurations: 1. a = 0, 8 = 4; 2. a 4, 3 = 0; 3. a = 4, B = 4. The learned filters are shown in figure 3. We find that the uniform prior regularization (a > prevents dead filters while the smoothness regularization (3 0) prevents input memorization.
Accuracy vs Epochs:
Finally, we plot the evolution of accuracy over epochs for all the models learned in the experiments of figure 2. For Neural Bayes-MIM we use the models with scaling (42 in total), and all 16 models for the random encoder. The convergence plot (figure 5) shows that models pre-trained with Neural Bayes-MIM quickly converge higher test accuracy compared to the bascline.
Effect of Mini-batch size (MBS) and Batch size (BS); During implementation, we proposed to compute gradients using a reasonably large mini-batel of size MBS and accumulate gradients until BS samples are seen. This is done overcome the gradient estimation problem due to the Ey [L (x)] term in Neural Bayes-MIM, Here we evaluate the effect of these two hyper-parameter on the final test performance. We choose MBS from (50, 100, 250, 500) and BS from (50, 250, 500, 2000. 3000).
7

---

Under review as a conference paper at ICLR 2021
For each combination of MBS and BS, train the CNN encoder using Neural Bayes-MIM with a = and 3 = 4 (chosen by examining figure 2); the rest of the training settings are kept identical to those used for figure 2 experiment. Table 2 shows the final test accuracy CIFAR- 10 for each combination of hyper-parameters MBS and BS. We make two observations: 1. using very small MBS (Eg. 50 and 100) typically results in poor (even worse than that of random encoder (67.9 (67.97%)), while larger MBS significantly improves performance; 2. using larger BS further improves performance most cases (even when MBS is small).
G PROOF OF LEMMA 1
Lemma 1. Let p(x|z = k) and p(z) be any conditional and marginal distribution defined for continuous random variable and discrete random variable If [xxx + O. [K],
then there exists non-parametric function T(x) T R+* for any Joy given input E TR" with the
property E [x(x) Vx such that,
(x( p(x)
p(x = k:) Ex-p(x) [x)xx
p(= = k) -0xx[((L())) p(z = k|xx = (x(x)
(45)
and this parameterization consistent.
Proof:
First we show the existence proof. Notice that there exists cane-parrmmtriccuunccii ga (x)
p(x]==k) Vx E supp(p(x)). Denote Gx(xx) = p(2 k)gx(x)) (x). Then,
p(x)
[x]]x(x)] = Exlp(= = k))x(())] = p==== =k)
(46)
and,
Gx(x) p(z = k))g(x) p(x|== = k)) Ex Gxx(x)] p(==
(47)
p(x)
Thus La Gk works. To verify that this parameterization is consistent, note that for any k.
[ p(x|= 1s) = [S (x): p(x)
(48)
Ex-p(x) [x(x)]
where we use the condition E Ex-p(x) [x((x]] + 0 AK: [K]. Secondly, we note that,
Σ p(x = k) p(z [h)= Lx(x) p(x)
K
K
Ex[La [x(L((x))
(49)
= [x( -p(x)
(50)
k=1
= p(x)
(51)
where the last equality is due to the conditions Σ k= [x(x)= = 1V. Thirdly,
K
Σ p(z= = k)
Ex [xx()]]
(52)
k=1
Ex
Lx(x)]
(53)
Finally, we have from Bayes' rule:
p(z = k/x) p(x|z = k) p(z = k.)
(54)
p(x)
[x(x)= -((x)
Excopteel [T (x))
(55)
Exxp(x) [xx(x]]
p(x)
-x(xx)
(56)
where the second equality holds because of the existence and consistency proofs of p(x]2 k) := p((x))p(x) Expep(se) [L)(
and p(z k) := [x]]x(x)] shown above. D
8

---

Under review as a conference paper at ICLR 2021
H PROOFS FOR NEURAL BAYES-MIM
Proposition (Neural Bayes-MIM-vl) (proposition main text) T(x)) : T - -R+** be
non-parametric function for any given input x E IR" with the property E Lx(x) - VX. Consider the following objective.
Lx(x)
L arg max Ex
Σ (x( log Ex [x()
(57)
Then [x((x) p(=2== k|x),, where 2" arg max, MT(x,2). 2).
Proof: Using the Neural Bayes parameterization lemma 1, we have,
p(x,= = = k)) log p(x; 2 k) MI(x, = Σ
K
(58)
p(x)p((= = k))
k=1
p(= = xxx)p(x))l) log p(= k|xx
(59)
p(= k)
k=1
/
(x)x)
(x) p(x)
log
(60)
Ex [1 (x)]
K
La (x)
= pxx)
Σ (x) log
(61)
Ex
[xxxxx
kk=1
Therefore the two objectives are equivalent and we have a closed form estimate of mutual information. Given z" is maximizer of M MI(x, 2), since L is non-parametric function, there exists L such that
p(z" = -k(x) = Li (x) due to lemma O.O
Theorem 3. (Theorem I in main text) Denote,
K
Lo (x)
J(0) = -Ex
Lo (x) log
(62)
Ex Too (x)]
Lo (xx)
((0) Ex
Σ
Len (x) log(
(63)
Ex Lox (x)
where (.) denotes gradients are not computed through the argument. Then 0J(0)
a.(())
09
00
Proof: We note that,
K Lo (x) log
Lou (x)
J(0) = -Ex
(64)
Ex Lox (x)]
-x
Σ
Lup (x) log Log (x)
+ Ex
Lo (x) log Ex Ex (x)]
(65)
Denote the first term by Tí. Then due to chain rule,
0T,
K
LLo (x)
K
Lag (x)
ALa (x)
Ex
log Log (x)
Fx
(66)
00
00
Lop (x)
00
1=11
k
Loo (x)
O (x)
Ex
log Lgg (x)
B
(67)
80
K
ÖLg (x)
k= Lop (x)
Ex
log Lo. (x)
Ex
(68)
00
00
@Lo, (x)
Ex
log La. (x)
(69)
20

---

Under review as a conference paper at ICLR 2021
where the last equality holds due the linearity of expectation and because E Loz (x) design. Now denote the second term by T2. Then due to chain rule.
by
0T2 00
K
LL (x)
K
Lo (x)
2E Loz (x)
Ex
log Ex Len (x))
Ex
(70)
00
Ex Lop (x))
00
k=1
Lop (x)
K
Ex [x)]
3Bx | (x)
Ex
log Ex Lo (x)]
(71)
00
Ex Lop (x)]
00
k=
Lop (x)
Lbp (x)
Ex
log Bxx Lop (x))
(72)
00
k=1
K
ÖLo, (x)
@ x K
Lo, (x)
Ex
log Bx Lo (x)]
Ex
k=1
(73)
00
Le (x)
Ex
log Exx Lo,, (x))
(74)
00
where once again the last equality holds due the linearity of expectation and because Ek Lo, (x) by design. Thus the gradient for is given by.
aJ(0)
K
ØL (x)
LB (x)
-Ex
log
(75)
06
00
Ex Lo, (x))
which is the same 03(8) This concludes the proof.
10


Reviews:
[{"review_id": "vAo_93nMVfL", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary \nThe paper introduces a new function $L(x)$ so that, when optimised under certain objectives defined over continuous observation $x$ and discrete latent $z$, learns the correct clustering probability $p(z|x)$. The loss functions considered are the Jensen-Fisher divergence and muture information. The authors introduces modifications to the principled objectives in practice and demonstrate performance on toy and real image datasets.\n\n# Pros:\n* The functional view on learning the clustering function $p(z|x)$ is interesting\n* This paper explores a few theoretically motivated new unsupervised objectives and practical regularization, all of which seem new and are worth further investigation\n* Thm1 and Thm2 are correct and interesting. \n\n# Cons:\n* The paper is presented in a mysterious language both in the qualititative description and mathematical details. The reader would have to do a lot of guesswork and think a lot to understand what the paper is trying to say (not yet to verify the theoretical statements). I will ask the authors to clarify a lot of them later, and I'm happy to raise the score on satisfactory feedback\n* The heuristics introduced for practical use is not accompanied with enough (though some) justifications other than \"we tried a few\"\n* The experiments seem to require a lot of hand-tuning of hyperparameters and eyeballing for good results, making the methods less useful in practice.\n\n# Recommendation:\nI think this paper is not ready for publication, both in terms of clarity and impact. The idea is worth investigating more.\n\n# Issues and questions:\n1. The notation and statement in the fundamental Lemma 1 is unclear. Are $p(x|z)$ and $p(z)$ from the same joint distribution $p(x,z)$? If so, then they would imply a consistent distribution $p(x)$ which is the data distribution, and the results hold trivially. If they are not from a consistent joint, then the construction of $G_k$ in the proof at Appendix G does not guarantee a normalised distribution such that $\\sum_k G_k(x) = 1$, and the equality of (51) does not hold. Can the authors clarify this, ideally using different notations for each distribution, e.g. $p(z)$, $q(x|z)$ and $pi(x)$ to avoid confusion.\n1. If \"if... then....\" statement is unclear. How can $L$ be involved in the \"if\" part, and also \"exist\" as a consequence in the \"then\" part? \n1. Is the normalisation property $\\sum_k L_k(x) = 1$ a condition or a consequence? \n1. In the definition of connected set, what is a path?\n1. In the proof of Prop 1, why are the constructions of $q_1$ and $q_2$ optimal for the optimisation problem (2)? I think it may be plausible, but from a reader's perspective, I hope to see a clear proof of this.\n1. Can the author provide more justification for the smoothness regularisation (4)? In particular, why is the perturbation in the span of data?\n1. In the experiments at 4.2, I think we would really like to see the effect of regularisation weights on the results, andThe line \"Regularization coefficient β was chosen from [0.5, 6] that resulted in optimal clustering\" seems to suggest that eyeballing on the results is needed, which is only possible for 1- or 2-D problems. \n1. Is there a way to choose the number of clusters?\n1. The authors claims that DMS finds manifolds of subspaces, but in fact it learns a clustering function which does not identify the manifold structure.\n1. Prop 2 raises multiple confusions if not inaccuracies. The objective cannot be written as an argmax. Mutual information is between two random variables rather than symbols (What is $z^*$?).\n1.  In the proof of Prop 2, I do not follow the last sentence. I understand (58)-(61), but why is the posterior the optimal solution to (57)?\n1. What's the connection between the two formualitons? Is it correct that MIM does not require the smoothness assumption needed for DMS?\n\n# Detailed comments:\n1. In the paragraph above Section 4, \"through which computing the gradient is infeasbile\", is this the gradient w.r.t. $\\theta$? I also do not follow the example you listed in the appendix at all. Does $t$ play the same role as $x$?\n1. 7th line in the 2nd to last paragraph of Section 2, \"both both\"\n1. Please use bold for all vectors, such as $\\delta$ in (4).\n1. In the MNIST experiment of 4.2, how do other clustering algorithms perform in downstream linear classification? \n1. Why is (9) called a cross-entropy formulation? Is the minus in (9) sign a typo? Also, $L_k$ should be $L_{\\theta,k}$ here and also in other places. \n1. The comment following (9) is not very obvious at least to me. The reader would have to check it themselves. (There is an extra comma after \"both\")\n1. The comment \"This is an attribute of distributed representation which is a fundamental goal in deep learning.\" is questionable: why is a evenly distributed representation a fundamental goal of deep learning?\n\n\n===== update =====\nNo response is provided so I am maintaining the score. I hope the authors could address all the issues in a future version. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "nTCPqXUjFvH", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work introduces a parameterization called Neural Bayes that facilitates learning representations from unlabeled data by categorizing them, where each data point x is mapped to a latent discrete variable z such that the distribution p(x) is segmented into a finite number of conditional distributions. Imposing different constraints on the latent discrete space will result in learning representations manifesting various properties. Two use cases of the proposed parameterization are studied: disjoint manifold separation and mutual information maximization.\n\nAs argued by the authors, the disjoint manifold separation objective and the optimality as shown in Theorem 1 is very interesting. Such a simple objective can lead to optimal manifold separation but this may also due to the strong restriction of the disjoint manifold assumption. Unfortunately, the authors do not elaborate much on this and only present some preliminary results on some synthetic datasets without much discussion. The results in Figure 1 do not show much benefits of the proposed method without a proper comparison to existing deep subspace clustering methods. Also, it is argued in Section 1 that \"imposing different conditions on the latent space z through different objective functions will result in learning qualitatively different representations\" which is the central theme of this work, it would be nice to cast the DMS objective into this framework as well to show how equation (3) corresponds to a regularization on z.\n\nOn the other hand, the proposed Neural Bayes framework looks identical to a standard mixture model and the analysis follows naturally from the Bayes' rule. Similar ideas of unsupervised representation learning via mixture models have been explored extensively in the literature, DeepCluster, IMSAT, VQ-VAE, to name a few. Almost all methods share the same mixture framework, but use different regularizations on the conditionals or the latent discrete variables. The Neural Bayes-MIM-v1 objective is identical to the IMSAT method as discussed by the authors. This work presents a better theoretical justification, but actual objective used (10) is different from theoretical-backed objective (8) minus the smooth regularization. When $\\alpha \\ne 0$, (10) is different from the optimal result. It would be nice to discuss the tradeoff manifested by controlling the entropy and how it impacts the learned representation and the downstream tasks. It also would be good to include the quantitative comparison of (8) and (9) in the main text to give a complete story.  As in the disjoint manifold separation case, the quantitative results are weak and only selected baselines are compared. In the description of the settings, only the architecture and learning parameters are discussed, but the number of clusters used is never mentioned. What K is used in the experimental settings? What happens if we overcluster the data which is usually the case in practice since we have no idea about the true number of clusters. This might create a serious problem for the disjoint manifold separation case due to the strong disjoint assumptions.\n\nTo conclude, this work seems preliminary and I do not think it passes the bar for acceptance in its current format.", "extra_scores": {}, "extra_text": {}}, {"review_id": "HfykrDJjS_8", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "- Novelty is minimal, given previous work in clustering, latent variable models, etc. \n \n- The structure of the paper (mainly with the two use cases at the core) is rather non-standard. Regardless of that, the ideas proposed in both use cases are very incremental. \n\n- The generalization to arbitrary manifolds is a strength of the paper (++).\n\n- Experimental results are not the most impressive either. In addition, only two baselines have been compared to (e.g. in Table 1), notwithstanding the plethora of recent works on the area, including the works cited in the Related Work. \n\n- p2: Last paragraph in the Related Work: Two things cannot be \"identical\", yet there are \"important differences\" among them.. Please fix the language. \n\n- p2: Last paragraph in the Related Work: The differences referred to as \"important differences\" are very incremental and non-fundamental. Therefore, the Neural Bayes-MIM-v1 objective is not novel. \n\n- The paper needs massive improvements in terms of writing. Ideas need to flow in order and to be much more unequivocal than the current form. Just as an example, the term E_x[L_k(x)] has been thrown in in the related work section prior to the formulation of the notation. \n\n- p4: \"we found the following finite difference Jacobian regularization to be effective.\": Could you please elaborate a bit on that? Did you try other regularization procedures which turned out to be ineffective or less effective?", "extra_scores": {}, "extra_text": {}}, {"review_id": "7sJhlggOniA", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nSummary. \nThe paper proposes Neural Bayes, a special parameterization of the posterior p(z|x) in Bayes' theorem with a neural network when the latent variables z are discrete. It presents two further developments of this idea, that each rely on different assumption on the latent code: 1) Disjoint Manifold Separation (DMS) assumes disjoint support of the conditional distributions p(x|z), and can is shown to be usable for clustering; 2) Mutual Information Maximization (MIM) maximizes the mutual information of x and z, and is shown to (unsupervisedly) generate feature representations, which when combined with an supervised MLP outperform another method on CIFAR-10, CIFAR-100 & STL-10. \n\nReasons for score:\nIn general the applicability of the proposed methods remains unclear - or where it is given (e.g. clustering, or unsupervised feature extraction), insufficient experimental evidence for its usefulness is given. Hence, I vote for rejection of the paper in its current form.\n\nPros:\n- The disjoint manifold separation method is a novel idea to me, and is a very reasonable approach to cluster data\n- MIM is a promising approach to learn unsupervised feature representations \n- the Gradient Simplification is a nice trick to reduce variance of estimates for the MIM\n\n\nCons:\n- Experiments of DMS needs more elaboration:\n\t- how does this behave when the number of possible z is not equal the number of clusters in the data? \n\t- how does the method compare to other clustering methods (e.g. spectral clustering, DBSCAN, or other state-of-the art methods?)\n- Experiments of MIM:\n   - comparison to other self-supervised methods is necessary (e.g. SimCLR or SimCLRv2) to put the method in context to the state-of-the art here, as the only shown application is fitting this context. \n   - the best performing model is chosen; this is not good practice, as it can be seen in Figure 2 that the variance of the performance is very high, hence running more models will also yield better performing models. Please report mean + std-dev (or some other reasonable statistic) for a given parameter set (alpha, beta,..) over a reasonable size of runs.\n- Formal exposition could be clearer, e.g. instead of L_k one could directly call this p(z|k) and mention that it's parametrized by a NN (p_theta(.) then). Or in chapter 4, why is q_i(x) introduced as a \"conditional distribution\" - one could just use p(x|z).\n\nRebuttal:\n- please address my points in Cons.\n\nOther small comments:\n- in the abstract: \"...express p(xjz), p(zjx) and p(z) in closed form\". As written later in the text, Monte-Carlo is needed to calculate p(z) etc. (it's an integral over x), so exact evaluation is not possible in a finite number of operations, and hence it's not closed form.\n- the smoothness of L_theta is reminiscent of recent developments showing that Lipschitz continuity regularization is beneficial for generalization. The relation should be worked out.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Summary: The authors propose a method for representing a posterior
over discrete latent variables in representation learning problems
using a neural network. Two applications are discussed: One are certain
clustering problems, in which clusters are sufficiently
separated. Another is the computation of mutual information of
discrete random variables. This is applied to learning image representations.

Discussion: The authors have not provided a response to the reviews.

Recommendation: Four detailed reviews unanimously recommend rejection. Main
points of criticism are lack of novelty, limited and unconvincing
experimental evaluation, and a poor presentation that also lacks
technical detail. This work is clearly not ready for publication.

Decision: Reject