Stochastic Prototype Embeddings

1 INTRODUCTION . Supervised deep-embedding methods map instances from an input space to a latent embedding space in which same-label pairs are near and different-label pairs are far . The embedding thus captures semantic relationships without discarding inter-class structure . In contrast , consider a standard neural network classifier with a softmax output layer trained with a cross-entropy loss . Although its penultimate layer might be treated as an embedding , the classifier ’ s training objective attempts to orthogonalize all classes and thereby eliminate any information about inter-class structure . Supervised embedding methods are critical for large- and open-set classification tasks , and are popular for few-shot and lifelong learning tasks . Nearly all previous methods for deep embeddings are deterministic : an instance projects to a point in the embedding space . Deterministic embeddings fail to capture uncertainty due either to out-ofdistribution inputs ( e.g. , data corruption ) or label ambiguity ( e.g. , overlapping classes ) . Representing uncertainty is important for many reasons , including robust classification and decision making , informing downstream models , interpreting representations , and detecting out-of-distribution samples . In this article , we propose a method for discovering stochastic embeddings , where each embedded instance is a random variable whose distribution reflects the uncertainty in the embedding space . Our proposed method , the Stochastic Prototype Embedding ( SPE ) , is an extension of the Prototypical Network ( PN ) ( Snell et al. , 2017 ) . As in the PN , our SPE assumes each class can be characterized by a prototype in the embedding space and an instance is classified based on its proximity to a prototype . In the case of the SPE , the embeddings and prototypes are Gaussian random variables , each class instance is assumed to be a Gaussian perturbation of the prototype , and a query instance is classified by marginalizing over the embedding uncertainty . Our main contribution is to show that SPE outperforms the only other fully-formulated method for stochastic supervised embeddings , the Hedged Instance Embedding ( HIB ) ( Oh et al. , 2019 ) , on a superset of the complete battery of experiments used to justify HIB . SPE is also more computation efficient to train than HIB , with complexity comparable to that of the PN , and has no hand-tuned parameters . We also demonstrate that embedding distributions are related to label uncertainty and input ambiguity . Finally , we explore an intriguing emergent property of SPE : that it attains more interpretable representations by disentangling class-discriminative features . 2 RELATED WORK . Supervised embedding methods are popular in the few-shot learning literature ( Koch et al. , 2015 ; Vinyals et al. , 2016 ; Snell et al. , 2017 ; Triantafillou et al. , 2017 ; Finn et al. , 2017 ; Edwards and Storkey , 2017 ; Scott et al. , 2018 ; Ridgeway and Mozer , 2018 ; Mishra et al. , 2018 ) where the goal is to classify query instances based on one or a small number of labeled exemplars of novel classes . These methods operate by embedding the queries and exemplars using a pre-trained network , and classifying each query according to its proximity to the exemplars . Embedding methods are also critical in open-set recognition domains such as face recognition and person re-identification ( Chopra et al. , 2005 ; Li et al. , 2014 ; Yi et al. , 2014 ; Zheng et al. , 2015 ; Schroff et al. , 2015 ; Liu et al. , 2015 ; Ustinova and Lempitsky , 2016 ; Song et al. , 2016 ; Wang et al. , 2017 ) . Loss functions used to obtain embeddings can be characterized according to the number of instances required to specify a loss . To describe these losses , we will use the notation zα for an embedding of class α. Pairwise losses attempt to minimize within-class distances , ||zα − z′α|| , and maximize between-class distances , ||zα−zβ || ( Chopra et al. , 2005 ; Hadsell et al. , 2006 ; Yi et al. , 2014 ) . Triplet losses attempt to ensure within-class instances are closer than between-class instances , ||zα − z′α|| < ||zα − zβ || ( Schroff et al. , 2015 ; Song et al. , 2016 ; Wang et al. , 2017 ) . Quadruplet losses attempt to ensure every within-class pair is closer than every between-class pair , ||zα − z′α|| < ||z′′α − zβ || ( Ustinova and Lempitsky , 2016 ) . Finally , cluster-based losses attempt to use all instances of a class ( Rippel et al. , 2016 ; Fort , 2017 ; Song et al. , 2017 ; Snell et al. , 2017 ; Ridgeway and Mozer , 2018 ) . In particular , the Prototypical Network ( Snell et al. , 2017 ) computes the mean of a set of instances of a class , z̄α , and ensures that additional instances of that class , zα , satisfy a proximity constraint such as ||zα − z̄α|| < ||zα − z̄β || . Cluster-based methods represent state-of-the-art over , in particular , pairwise and triplet losses , as one might expect given the chronology of publications . Probabilistic embedding methods have recently appeared . Allen et al . ( 2019 ) extend PNs via Bayesian nonparametrics to treat each prototype as a mixture distribution , though they do not allow uncertainty in the embedding space , the critical element of our work . Vilnis and McCallum ( 2018 ) propose a method for learning density-based word embeddings , but it is unsupervised . Deep Variational Transfer ( Belhaj et al. , 2018 ) is a generative form of the discriminative model we propose , which requires modeling input distributions ; this work tackled the somewhat different problem of covariate shift . The method closest to ours is an extension of PNs that uses a Mahalanobis distance instead of a Euclidean distance to assess similarity ( Fort , 2017 ) . Although this method lacks probabilistic semantics , it has similarity with SPE ( comparison in Appendix G ) . Two prior methods have been proposed for discovering stochastic embeddings in a supervised setting , i.e. , for few-shot and open-set recognition . The Hedged Instance Embedding ( HIB ) ( Oh et al. , 2019 ) utilizes a probabilistic alternative to the contrastive loss and is trained using a variational approximation to the information bottleneck principle . HIB is critically dependent on a constant , β , that determines characteristics of an information bottleneck ( i.e. , how much of the input entropy is retained in the embedding ) . Choosing this constant is a matter of art . The Oracle-Prioritized Belief Network ( OPBN ) ( Karaletsos et al. , 2016 ) is a generative model that learns a joint distribution over inputs and oracle-provided triplet constraints . The OPBN was not tested on few-shot and open-set recognition because it requires extensions to be applied to classification tasks . In the deterministic setting , Scott et al . ( 2018 ) argue that cluster-based methods outperform pairwise and triplet methods ; thus , we have reason to expect that in a stochastic setting , a cluster-based method like the one we propose in this article , SPE , will outperform pairwise ( HIB ) and triplet ( OPBN ) methods . 3 THE MODEL . The SPE assumes that the latent representation , z , is a Gaussian RV conditioned on the input , x : p ( z|x ) = N ( z ; µx , σ2xI ) ( 1 ) with mean , µx , and variance , σ2x , computed by a deep neural network , similar to a Variational Autoencoder ( Kingma and Welling , 2014 ) . The classification , y , in turn is conditioned on z , with p ( y|z ) taking the same form as in the original PN ( Snell et al. , 2017 ) , to be described shortly . Given an input , a class prediction is made by marginalizing over the embedding uncertainty : p ( y|x ) = ∫ z p ( y|z ) p ( z|x ) dz , ( 2 ) Informally , Figure 1a depicts the relationship between the input , latent , and class representations ; for a formal depiction of SPE , see the chain graph in Appendix F. We train the SPE using the standard paradigm for supervised embedding methods , via a sequence of episodes , each with m instances of n classes . We split the m×n instances into k×n support examples , defining a set S , and ( m− k ) ×n query examples . The support instances for each class c , Sc ∈ S , are used to determine the class prototype , ρc , and the query instances are evaluated to predict class label ( Equation 2 ) . 3.1 FORMING CLASS PROTOTYPES . In the SPE , each class y has an associated prototype in the embedding space , ρy , and each instance i of class y projects to an embedding , zi , in the neighborhood of ρy such that : p ( zi|ρy ) = N ( ρy , σ2 I ) . ( 3 ) With independent support instances and uniform class priors , Bayes rule obtains the posterior of prototype ρy for support set Sy : p ( ρy|Sy ) = ∏ i∈Sy p ( zi|ρy ) ∫ ρ ∏ i∈Sy p ( zi|ρ ) dρ . ( 4 ) This normalized product of Gaussians is itself Gaussian : p ( ρy|Sy ) = N ( µy , σ2yI ) with σ2y = ( ∑ i∈Sy σ̂ −2 xi ) −1 and µy = σ2y ◦ ( ∑ i∈Sy σ̂ −2 xi ◦ µxi ) , ( 5 ) where σ̂2xi = σ 2 xi + σ 2 and ◦ denotes the Hadamard product and { µx , σ2x } are the outputs of the embedding procedure for input x . Essentially , the prototype is a confidence-weighted average of the support instances . This formulation has a clear advantage over the deterministic PN , which is premised on an unweighted average , because it de-emphasizes noisy support instances . 3.2 PREDICTION AND APPROXIMATE INFERENCE . Given a query embedding , z , and uniform class priors , the class posterior is : p ( y|z , S ) ∝ N ( z ; µy , σ̂2yI ) ( 6 ) where σ̂2y = σ 2 y + σ 2 , and µy and σ 2 y are the mean and diagonal covariance of ρy ( Equation 5 ) . Expressing this posterior in terms of the query x ( i.e. , combining Equations 1 , 2 , and 6 ) : p ( y|x , S ) = ∫ z N ( z ; µx , σ2xI ) N ( z ; µy , σ̂2yI ) ∑ cN ( z ; µc , σ̂2cI ) dz . ( 7 ) The class distribution is equivalent to that produced by the deterministic PN as σ2x → 0 when σ2y = σ 2 y′ for all class pairs ( y , y ′ ) . However , in the general case , the integral has no closed form solution ; thus , we must sample to approximate p ( y|x , S ) , both for training and evaluation . We employ two samplers , which we refer to as naïve and intersection . 3.2.1 NAÏVE SAMPLING . A direct approach to approximating the class posterior is to express Equation 2 as an expectation , Ez∼p ( z|x ) [ p ( y|z , S ) ] , and to replace the expectation with the average over a set of samples . We utilize the reparameterization trick of Kingma and Welling ( 2014 ) to train the model . Although this Monte Carlo method is the simplest approach , it is sample-inefficient during training , and when the number of samples is reduced , model performance is impacted .

Summary: By extending prototypical networks, this paper proposes a probabilistic model, i.e., stochastic prototype embedding, that treats embeddings as random variables. The model is very straightforward and easy to understand. The authors make a few assumptions to simplify the problem. For example, the distance between every instance $z_i$ of class $y$ and the class embedding $\rho_y$ follows a Gaussian distribution, and the a softmax prediction for a query embedding also follows a Gaussian distribution. Combining these, the authors give the class posterior for the query. Since the class posterior involves an integral, the paper employs a naive samplying and an intersection sampling. The intersection sampling seems a bit interesting in the sense that the sampler focuses on the intersection of the input distribution and the class distribution and it is more sample-efficient.