Neural Bayes: A Generic Parameterization Method for Unsupervised Learning

1 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 z where the distribution p ( x ) gets segmented into a 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 are 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 use 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 high-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 . 2 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 et al. , 2018 ) , Assosiative Deep Clustering ( Haeusser 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 a 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 as 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 a 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 is itself novel and we prove its optimality towards labeling disjoint manifolds in theorem 1 . Neural Bayes-MIM : Self-supervised representation learning has attracted a 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 ) , BYOL Grill et al . ( 2020 ) ) produce state-of-the-art ( 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 et 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 a closed form estimation of MI due to 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 , a number of self-supervised 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 noise-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 a 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-v1 objective ( Eq 8 ) turns out to be identical to the one 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 3 in ( Hu et al. , 2017 ) ) ; 2. we provide theoretical justifications for the parameterization used ( lemma 1 ) and show in theorem 2 why it is feasible to compute high fidelity gradients using this objective in the mini-batch setting even though it contains the term Ex [ Lk ( x ) ] . On the other hand , the justification used in IMSAT is that optimizing using mini-batches is 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 ∈ Rn . We assume that this marginal distribution is a union of K conditionals where the kth conditional ’ s density is denoted by p ( x|z = k ) ∈ R+ and the corresponding probability mass denoted by p ( z = k ) ∈ R+ . Here z is a 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 ’ rule . Lemma 1 . Let p ( x|z = k ) and p ( z ) be any conditional and marginal distribution defined for continuous random variable x and discrete random variable z . If Ex∼p ( x ) [ Lk ( x ) ] 6= 0 ∀k ∈ [ K ] , then there exists a non-parametric function L ( x ) : Rn → R+K for any given input x ∈ Rn with the property ∑K k=1 Lk ( x ) = 1 ∀x such that , p ( x|z = k ) = Lk ( x ) · p ( x ) Ex∼p ( x ) [ Lk ( x ) ] , p ( z = k ) = Ex [ Lk ( x ) ] , p ( z = k|x ) = Lk ( x ) ( 1 ) Thus the function L can be seen as a form of soft categorization of input samples . In practice , we use a neural network with sufficient capacity and softmax output to realize this function L. We name our parameterization method Neural Bayes and replace L with Lθ to denote the parameters of the network . By imposing different conditions on the structure of z by formulating meaningful objectives , we will get qualitatively different kinds of factorization of the marginal p ( x ) , and therefore the function Lθ will encode the posterior for that factorization . In summary , if one formulates any objective that involves the terms p ( x|z ) , p ( z ) or p ( z|x ) , where x is an observed random variable and z is a discrete latent random variable , then they can be substituted with Lk ( x ) ·p ( x ) Ex [ Lk ( x ) ] , Ex [ Lk ( x ) ] and Lk ( x ) respectively . On an important note , Neural Bayes parameterization requires using the term Ex [ Lk ( x ) ] , through which computing gradient is infeasible in general . A general discussion around this can be found in the appendix D. Nonetheless , we show that mini-batch gradients can have good fidelity for one of the 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 . 4 DISJOINT MANIFOLD SEPARATION ( DMS ) . In many cases , the support of a distribution may be a set of disjoint manifolds . In this task , our goal is to 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 a disjoint manifold . Definition 1 . ( Connected Set ) We say that a set S ⊂ Rn is a connected set ( disjoint manifold ) if for any x , y ∈ S , there exists a continuous path between x and y such that all the points on the path also belong to S. To identify such disjoint manifolds in a distribution , we 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 a 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 qi ( x ) ’ s such that a divergence between them is maximized . By doing so we recover the conditional distributions defined over the disjoint manifolds ( we prove its optimality in theorem 1 ) . We begin with two disjoint manifolds and extend this idea to multiple disjoint manifolds in the appendix B . Let J be a symmetric divergence ( Eg . Jensen-Shannon divergence , Wasserstein divergence , etc ) , and q0 and q1 be the disjoint conditional distributions that we want to learn . Then the aforementioned objective can be written formally as follows : max q0 , q1 π∈ ( 0,1 ) J ( q0 ( x ) ||q1 ( x ) ) ( 2 ) s.t . ∫ x q0 ( x ) = 1 , ∫ x q1 ( x ) = 1 , q1 ( x ) · π + q0 ( x ) · ( 1− π ) = p ( x ) . Since our goal is to simply assign labels to data samples x corresponding to which manifold they belong instead of learning conditional distributions as achieved by Eq . ( 2 ) , we would like to learn a 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 a function L ( x ) . Proposition 1 . ( Neural Bayes-DMS ) Let L ( x ) : Rn → [ 0 , 1 ] be a non-parametric function for any given input x ∈ Rn , and let J be the Jensen-Shannon divergence . Define scalars f1 ( x ) : = L ( x ) Ex [ L ( x ) ] and f0 ( x ) : = 1−L ( x ) 1−Ex [ L ( x ) ] . Then the objective in Eq . ( 2 ) is equivalent to , max L 1 2 · Ex [ f1 ( x ) · log ( f1 ( x ) f1 ( x ) + f0 ( x ) ) ] + 1 2 · Ex [ f0 ( x ) · log ( f0 ( x ) f1 ( x ) + f0 ( x ) ) ] + log 2 s.t . Ex [ L ( x ) ] /∈ { 0 , 1 } . ( 3 ) 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 a probability density function over Rn whose support is the union of two non-empty connected sets ( definition 1 ) S1 and S2 that are disjoint , i.e . S1 ∩ S2 = ∅ . Let L ( x ) ∈ [ 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 is true : L ( x ) = { 0 ∀x ∈ S1 1 ∀x ∈ S2 or L ( x ) = { 1 ∀x ∈ S1 0 ∀x ∈ 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 a distinct label to points in each manifold . Most importantly , the theorem shows that the continuity condition on the function L ( x ) plays an important role . Without this condition , the network can not identify disjoint manifolds . Extension to multiple disjoint manifold case can be found in section B in appendix .

Summary: The 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.