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: The paper proposes stochastic prototype embeddings (SPE) for few-shot learning. The method is an extension of Prototypical Networks (PN, [1]) with Gaussian embeddings. The idea is to take representation uncertainty into account when classifying objects which makes the model more robust to input and label noise. The authors propose an efficient sampling algorithm to train SPE which outperforms naive Monte Carlo sampling. They conduct a range of experiments on few-shot learning tasks on a synthetic dataset, Omniglot and N-digit MNIST and compare to Prototypical Networks and previous stochastic embedding state-of-the-art method HIB [3]; SPE was shown to outperform prior work in most settings. The authors also plot interpretable disentangled representations learned in 2D embedding space.