Title: No title found

Abstract: Supervised deep-embedding methods project inputs of a domain to a representational space in which same-class instances lie near one another and different-class instances lie far apart. We propose a probabilistic method that treats embeddings as random variables. Extending a state-of-the-art deterministic method, Prototypical Networks (Snell et al., 2017) , our approach supposes the existence of a class prototype around which class instances are Gaussian distributed. The prototype posterior is a product distribution over labeled instances, and query instances are classified by marginalizing relative prototype proximity over embedding uncertainty. We describe an efficient sampler for approximate inference that allows us to train the model at roughly the same space and time cost as its deterministic sibling. Incorporating uncertainty improves performance on few-shot learning and gracefully handles label noise and out-of-distribution inputs. Compared to the state-of-the-art stochastic method, Hedged Instance Embeddings (Oh et al., 2019) , we achieve superior large-and open-set classification accuracy. Our method also aligns class-discriminating features with the axes of the embedding space, yielding an interpretable, disentangled representation.

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 , σ 2 x I)

(1) with mean, µ x , and variance, σ 2

x , 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: 2).

EQUATION

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, z i , in the neighborhood of ρ y such that:

EQUATION

With independent support instances and uniform class priors, Bayes rule obtains the posterior of prototype ρ y for support set S y :

EQUATION

)

This normalized product of Gaussians is itself Gaussian:

EQUATION

where σ2 xi = σ 2 xi + σ 2 and • denotes the Hadamard product and {µ x , σ 2 x } 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:

EQUATION

where σ2 y = σ 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):

EQUATION

The class distribution is equivalent to that produced by the deterministic PN as σ 2 x → 0 when σ 2 y = σ 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, E z∼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.

3.2.2. INTERSECTION SAMPLING:

In Equation 7, the product of Gaussian densities in the numerator can be rewritten:

EQUATION

where

σ 2 xy = (σ -2 x + σ-2 y ) -1 and µ xy = σ 2 xy • (σ -2 x • µ x + σ-2 y • µ y ).

Substituting Equation 8 into Equation 7,

EQUATION

)

By approximating the expectation with samples from N µ xy , σ 2 xy I , we obtain an elegant importance sampler that focuses on the intersection of the input distribution and a given class distribution, as illustrated in Figure 1b . During training with a cross-entropy loss, we need only sample for the known (target) class y. As we will demonstrate, this method is more robust and significantly more sample efficient than the naïve sampler, requiring only a single sample to train effectively.

4. EXPERIMENTAL RESULTS:

We report on three sets of experiments. In Section 4.1, we demonstrate, using a synthetic data set, that SPE infers the generative structure of a domain, disentangles class-discriminating features, and provides meaningful estimates of label uncertainty and input noise. In Section 4.2, we show that SPE obtains state-of-the-art results on few-shot learning via a comparison to its deterministic sibling, PN, the previous state-of-the-art method. We evaluate on a standard data set used to compare methods in the few-shot learning literature, Omniglot (Lake et al., 2015) . In Section 4.3, we show that SPE obtains state-of-the-art results on large-set classification via a comparison to the only other fully developed stochastic method for supervised embeddings, HIB (Oh et al., 2019) . We evaluate on the only data set that Oh et al. (2019) used to explore HIB, a multi-digit variant of MNIST. For details regarding network architectures and hyperparameters, see Appendix A, and for simulation details, including the choice of initialization for σ 2 , see Appendix B.

4.1. SYNTHETIC COLOR-ORIENTATION DATA SET:

The data set consists of 64 × 64 pixel images of 'L' shapes, with four classes that are distinguished by orientation, color, or both (Figure 2a ). Instances are sampled from a class-conditional isotropic Gaussian distribution in the generative space. (The isotropy of these qualitatively different dimensions comes from the fact that both can be mapped as directional quantities.) Because classes overlap on both color and orientation dimensions, elicited embeddings should indicate increased uncertainty near class boundaries. Full details of the synthetic data set can be found in Appendix A.2.

We trained a two-dimensional, intersection-sampling SPE on samples from this domain, using two instances per class to form prototypes. Classification accuracy of held-out samples is approximately 86%. Accounting for class overlap, a Bayes optimal classifier has an accuracy of approximately 87%. For visualization, Figure 2b presents a 5 × 5 array of examples with the class centroids in the corners and the other examples obtained by linear interpolation in the generative space. The resulting embeddings are presented in Figure 2c . Although the correspondence between Figures 2b and 2c seems trivial (mirror one set along the horizontal axis to obtain the other set), remember that the input space is 64 × 64 dimensional and the latent space is 2 dimensional. The network has captured the structure of the domain by disentangling the two factors of variation. Further, the embedding variance encodes label ambiguity; instances halfway between two classes on one dimension have maximal variance along that dimension. Label ambiguity is one type of uncertainty. An equally important source of uncertainty comes from noisy or out-of-distribution (OOD) inputs. We examined OOD inputs generated in two different ways. In the left panel of Figure 3 , we show the consequence of adding pixel hue noise to the four class centroids. Only one of these centroids is shown along the abscissa, but all four are used to make the graph, with many samples per noise level. The grey and black bars in the graph indicate variance on the horizontal and vertical dimensions of the embedding space, respectively. As pixel hue noise increases, uncertainty in color grows but uncertainty in orientation does not. In the right panel of Figure 3 , we show the consequence of shortening the leg-length of the shape. Shortening the legs removes cues that can be used both for determining color and orientation. As a result, the uncertainty grows on both dimensions.

4.2. OMNIGLOT:

The Omniglot data set contains images of labeled, handwritten characters from diverse alphabets.

Omniglot is one of the standard data sets for comparing methods in the few-shot learning literature. The data set contains 1623 unique characters, each with 20 instances. Following Snell et al. (2017) , each grayscale image is resized from 32 × 32 to 28 × 28, and we augment the original classes with all 90 • rotations, resulting in 6492 total classes. We train PNs and SPEs episodically, where a training episode contains 60 randomly sampled classes and 5 query instances per class.

To compare the relative effectiveness of naïve and intersection samplers, we train the SPE on Omniglot varying both the sampler and the number of samples drawn per training query, denoted by s. We evaluate in a 1-shot 20-class setting, where shot refers to the number of support examples used to compute each prototype. Figure 4 shows test classification accuracy as the number of samples drawn per training trial (s) increases. As we previously stated, the intersection-sampling SPE is far more sample efficient, to the point that the intersection sampler with s = 1 outperforms the naïve sampler with s = 81. We have verified that the pattern in Figure 4 is consistent across simulations; consequently, we present only intersection-sampling SPE results in the remainder of the article, and all SPEs are trained with a single sample (s = 1) per query. This choice causes the SPE to be on par with the PN in time and space requirements, even though using more samples may boost classification accuracy, as suggested by the trend in Figure 4 .

Figure 5 is a visualization of a 2D embedding learned by the intersection-sampling SPE on Omniglot. All classes shown in the figure were held-out during training. Omniglot characters clearly vary along more than two dimensions, so a 2D SPE cannot learn a fully-disentangled representation as it did with the synthetic data set. However, we can still interpret the axes of the embedding. The horizontal axis appears to represent character complexity, with single-stroke characters on the left and many-stroke characters on the right. The vertical axis appears to encode the aspect ratio of the characters, with horizontally extended characters on the bottom and vertically extended characters on the top. 5) discount noisier support examples. Although the SPE is still superior when only the query is corrupted, the benefit is small. We also compared PN and SPE using a 64D embedding, but with high dimensional embeddings, both methods are near ceiling on this data set, resulting in comparable performance between the two methods. (See Appendix D for additional results, broken down by condition.)

To emphasize, SPE outperforms the PN, arguably the leading few-shot learning method, especially when inputs are corrupted, at essentially the same computational cost for training. And by providing an estimate of uncertainty associated with embedded instances, the SPE offers the possibility of detecting OOD samples and informing downstream systems that operate on the embedding. (Masson and Loftus, 2003) . 

Character Complexity Aspect Ratio of Character:



4.3. N-DIGIT MNIST:

The N -digit MNIST data set was proposed to evaluate HIB (Oh et al., 2019) ; it is formed by horizontal concatenation of N MNIST digit images. The resulting images are 28 × 28N . To compare with HIB, we study 2and 3-digit MNIST, and use a network architecture identical to that in Oh et al. (2019) . Oh et al. (2019) split the data into a training set (with 70% of the total classes), a seen test set, and an unseen test set. For 2-digit MNIST, the seen test set has the same 70 of 100 classes as the training set and the unseen test set has the remaining 30 classes. For 3-digit MNIST, the training set has 700 classes, the seen and unseen test sets each have a sample of 100 of the 700 seen or 300 unseen classes, respectively. We use the same train and test data splits as Oh et al. ( 2019), but we further divide the training split to include a validation set for early stopping.

Figure 7 shows two views of the 2D embedding learned by the SPE on the 2-digit MNIST test set. Each number is a class label; for example, 71, located in the lower left of the embedding, is the class in which the first of the two MNIST digits is a 7 and the second is a 1. The location of a label in the space corresponds to the mean of its prototype. In the left plot, each class is colored according to the first digit. The right plot is the same embedding, but each prototype is colored according to the second digit. The SPE learns an incredibly robust factorial representation in which the horizontal dimension represents the first digit of a class and the vertical dimension represents the second digit. A black bounding box indicates the unseen test classes, classes not presented during training. Impressively, the unseen test classes are embedded in exactly the positions where they belong, indicating that the SPE can discover relationships among classes that allow it to generalize to classes it has never seen during training. Furthermore, the embedding has captured inter-class similarity structure by placing visually similar digits close to one another. For example, on both the vertical and horizontal bands, nines (teal) and fours (purple) are adjacent, and fives (brown) and threes (red) are adjacent. HIB discovers a clean decomposition along one dimension (Oh et al., 2019) , but the second dimension is somewhat more entangled, suggesting that the SPE learns a more robust representation. Additionally, embeddings for the unseen class are not presented for HIB. The ability to sensibly embed novel classes is essential for any model that will be used for open-set recognition or few-shot learning. As we show in Appendix F, PNs do not obtain a clean compositional structure.

Figure 6b ,c compare N -digit MNIST test accuracy on seen and unseen classes, respectively.1 Each bar is the mean test accuracy across the Cartesian product of conditions specified by the number of MNIST digits in each image, N ∈ {2, 3}, and the dimensionality of the embedding, D ∈ {2, 3}. As in the Omniglot simulation, we varied whether support and query instances were clean or corrupted. The SPE outperforms HIB in all six comparisons. In the 24 individual conditions, SPE is worse on only 7. As in the Omniglot simulation, SPE shines best when support instances may be corrupted. (Appendix A.3 provides tabular results by condition, not only for HIB and SPE, but also their deterministic counterparts, contrastive loss and PN. Because the deterministic methods perform consistently worse than the stochastic methods, we omit the deterministic methods from the figure.)

Whereas SPE is a discriminative model with a specified classification procedure, Oh et al. (2019) had the freedom to design one. They use all available data-roughly 140 examples per class-and perform leave-one-out 5-nearest-neighbor classification. To be consistent with our episodic test procedure, the SPE uses only 50 support instances per class to form prototypes. It is particularly impressive that the SPE, based on a single stored prototype and approximately 1/3 the labeled data, outperforms a nonparametric method that is able to model arbitrary distributions in latent space.

5. DISCUSSION AND CONCLUSIONS:

We proposed the Stochastic Prototype Embedding (SPE) as a method for obtaining supervised embeddings which encode class-label uncertainty in their distributions. Such methods are useful for large-and open-set classification as well as few-shot learning, particularly for domains with ambiguous inputs or label noise. We compared SPE to the only fully-developed alternative method, the Hedged Instance Embedding (HIB), on the complete battery of tasks used to evaluate HIB. On these large-set classification tasks, SPE consistently outperforms HIB. Beyond its performance gains, SPE has no hand tuned parameters, whereas HIB has constant β that determines characteristics of an information bottleneck (i.e., how much of the input entropy is retained in the embedding). Although one could simply set β = 0, doing so would encourage the net to perform like a softmax classifier and discard all information about inter-class similarity. Such similarities are essential in order to generalize to unseen classes (e.g., Figure 7 ). SPE, which is a nondeterministic extension of the Prototypical Net (PN), matches or outperforms the PN on few-shot learning. Because the SPE extends the PN, it seems unlikely to fare worse; but because it can handle uncertainty in both the query and support set, it can fare better, particularly when the embedding space is low dimensional and the support instances may be corrupted. Extensions have been proposed to the PN compatible with ours (e.g., Allen et al., 2019) ; combining methods may potentially attain even stronger few-shot learning performance under uncertainty.

We proposed and evaluated an intersection sampler to train the SPE, which makes the SPE as time and space efficient for training as the deterministic PN, and more efficient for training than HIB, which relies on about 8 samples per item. (Our evaluation method for SPE presently involves drawing 200 samples from the naive sampler, though this conservative decision was arbitrary and not tuned.) An unanticipated virtue of SPE is its ability to obtain interpretable, disentangled representations (Figures 2, 5, 7 ). Because uncertainty is encoded in a diagonal covariance matrix, any classification ambiguity maps to uncertainty in the value of individual features of the embedding. Thus, classdiscriminating feature dimensions must align with the principle axes of the embedding space. In contrast to traditional unsupervised disentangling methods, which aim to discover the underlying generative factors of a domain, the SPE obtains a supervised analog in which the underlying classdiscriminative factors are represented explicitly. This representation facilitates generalization to novel unseen classes and is therefore valuable for few-shot and lifelong learning paradigms.

A NETWORK ARCHITECTURES AND HYPERPARAMETERS:

A.1 OMNIGLOT For all Omniglot experiments, the network consisted of four convolutional blocks. The first three blocks had a convolutional layer with 64 filters, a 3 × 3 kernel, zero-padding of length 1, and a stride of 1, followed by a batch normalization layer, ReLU activation, and 2 × 2 max-pooling. The fourth and final block had a convolutional layer with 2d filters, a 3 × 3 kernel, zero-padding of length 1, and a stride of 1, followed by 2 × 2 max-pooling, where d represents the dimensionality of the embedding space. The flattened output of the network is a vector of length 2d, where the first d elements were considered the mean of the Gaussian distribution and the remaining d elements were the diagonal covariance entries. The weights were initialized using He initialization and the biases with the following uniform distribution: U(-

1 √ fan in , 1 √ fan in ).

All Omniglot models were trained with an initial learning rate of 0.001 which was cut in half every 50 epochs. The models were stopped early using a patience parameter when performance on the validation set no longer increased.

A.2 SYNTHETIC DATA:

The images in the synthetic data set are 64 × 64 pixels in size. For orientation, we chose class centers at 90 • and 180 • , with a standard deviation of 30 • . For color, we manipulated the hue and kept value and saturation constant. Like orientation, hue is a circular quantity. If hue ranges from 0 to 360 degrees, we chose color class centers and standard deviation in the same way as orientation. Additionally, we add noise to a minority (15%) of the images used to train the model. For these, we add Gaussian noise to the hue of each pixel inside the shape. The standard deviation of the hue noise was chosen uniformly between 18 • and 54 • . We also added noise to the leg lengths of the L shapes. The leg length was chosen uniformly between 10% and 98% of its original length. See Figure 3 for some examples.

The network followed an architecture similar to the one we used for Omniglot, except that we added two additional blocks of convolution, batch normalization, ReLU, and max-pooling because the images are larger. We used 2 instances per class to form prototypes and 8 samples per query instance during training. We used a learning rate of 0.0001 and the models were stopped early using a patience parameter when performance on the validation set no longer increased.

A.3 N-DIGIT MNIST:

For all N -digit MNIST experiments, we constructed an architecture which we believe to be identical to that used for HIB MNIST experiments, based on code provided by the authors (Oh et al., 2019) . The network consisted of two convolutional blocks followed by two fully-connected layers. The convolutional blocks each contained a convolutional layer, followed by an ReLU activation, and 2 × 2 max-pooling. The first convolutional layer had 6 filters, a 5 × 5 kernel, zero-padding of length 2, and a stride of 1. The second convolutional layer was identical to the first, but had 16 filters instead of 6. The output of the second convolutional block was flattened, passed through a fully-connected layer with 120 units, an ReLU activation, and a final fully-connected layer with 2d units, where d represents the dimensionality of the embedding space. Like the Omniglot architectures, the first d entries in the output vector are treated as the mean and the remaining d elements as the diagonal covariance entries. The weights were initialized using a Xavier uniform initialization and biases were initialized to zero.

The PN and SPE are trained episodically with all performance results in the main article measured as the mean over 1000 random test episodes. All N -digit MNIST models were trained with an initial learning rate of 0.001 which was cut in half every 50 epochs. The models were stopped early using a patience parameter when performance on the validation set no longer increased. 

B SIMULATION DETAILS:

For all SPE models, σ 2 = softplus (γ) ,

where γ is a trainable parameter. We initialize γ using the following prescription:

γ = |S|γ 2/d

0 , where |S| is the number of support examples per episode during training and d is the dimensionality of the embedding. We chose this prescription for two reasons: (1) as the number of support examples increases, the variance of the prototype distribution approaches zero, so scaling linearly by |S| tends to provide a stronger training signal early on, and (2) the amount of noise in the projection of an embedding should scale with the dimensionality of the embedding space as to maintain unit-volume. All models used γ 0 = 0.01.

The variance of each dimension i, σ 2:

xi , is guaranteed to be non-negative by using a softplus transfer function.

Whether trained with the naïve or intersection sampler, we evaluate model performance using the naïve sampler with 200 samples. This approach ensures that we are comparing the quality of models based only on the method by which they were trained.

C SPE VARIANTS:

We assumed only diagonal covariance matrices in this work. Switching to a full covariance matrix would require matrix inversion, which is ordinarily infeasible, but because one purpose of deep embeddings is visualization, there may be interesting cases involving 2D embeddings where the cost of inversion is trivial. However, using a diagonal covariance matrix causes class-discriminating features to be aligned with the axes of the latent space, as we argued in the main article, and this alignment is a virtue for interpretation.

D TABULAR RESULTS:

D.1 OMNIGLOT Table 1 : Test classification accuracy (%) on Omniglot with both a 2D and 64D embedding for clean-support/clean-query, corrupt-support/clean-query, and clean-support/corrupt-query. PN is our implementation of Prototypical Networks (Snell et al., 2017) . SPE is our model. SPE is trained with intersection sampling (1 sample per trial). Reported accuracy for each experimental configuration is the mean over 1000 random test episodes. The algorithm for applying corruption was identical to the scheme used in Oh et al. (2019) . A random rectangular-sized occlusion of black pixels was determined by first sampling a patch width, L x , and patch height, L y , from a uniform distribution, L x , L y ∼ U(0, 28), and then sampling the top-left corner coordinates, T L x ∼ U(0, 28 -L x ), T L y ∼ U(0, 28 -L y ). This resulted in an occlusion of area L x × L y . Note that if L x = 0 or L y = 0, the image was left unoccluded. Figure 8 shows examples of occluded 2-digit images.

For Omniglot, we only trained/validated on corrupted imagery if the test set contained a corrupted support or corrupted query set. When testing on clean support and clean query, the training and validation sets were left unoccluded. When testing on corrupted imagery, the training and validation sets corrupted each character independently with a probability of 0.2.

The training and validation sets for N -digit MNIST corrupted each digit of each image independently with a probability of 0.2, regardless of test imagery. This matched Oh et al. (2019) .

During testing on both data sets, we considered both clean and corrupt support sets, as well as clean and corrupt query sets. A clean set was one in which all digits/characters were unoccluded. A corrupt set occluded each digit/character in each image according to the procedure described above.

F VISUALIZATIONS OF EMBEDDINGS DISCOVERED BY PROTOTYPICAL NETWORKS The model proposed by Fort (2017) shares with SPE the notion that prototypes are a confidenceweighted average of support instances (Equation 5), and the embedding procedure produces a scaling matrix for computing a Mahalanobis distance, similar to the Gaussian covariance matrix of SPE. However, Fort (2017) does not treat the embedding as stochastic, for to do so, it would need to marginalize over the uncertainty in the embedding to predict a class label. This marginalization is the core of a probabilistic model and is the critical component of SPE. Fort (2017) also obtain best results with a spherical scaling matrix for the Mahalanobis distance, whereas a critical property of SPE is that the uncertainty varies on each dimension of the latent space; our disentangling and uncertainty results all hinge on using a more flexible diagonal covariance matrix. And Fort (2017) omits the σ noise term which turns out to be critical both to obtain a well-formed probabilistic model (Figure 11 ) and for the model to work in practice.

:

HIB results are fromOh et al. (2019). We thank the authors for providing us results on unseen classes, which were not included in their publication.

Figures and Tables:

Figure 2: (a) Samples from the four classes in our synthetic data set. In each plot, class centroids are circled, along with samples spanning ±2 standard deviations in both orientation and color. A sample's transparency is set according to its class-conditional likelihood. Both dimensions can be coded as directional variables. The class centroids on each dimension are 90 • apart with standard deviation of 30 • . (b) A set of examples, with the four class centroids located in the corners and other examples obtained by linear interpolation in the generative space. (c) The 2D stochastic prototype embedding for the examples in (b). The shape is plotted at the mean of p(z|x), and the outlines of the ovals represent equiprobability contours at 0.4 standard deviations.

Figure 3: Synthetic data set: uncertainty on the two embedding dimensions as it becomes more difficult to discern the hue (left) and orientation (right).

Figure 4: Test classification accuracy as a function of number of training samples per query instance for a naïve-sampling and intersection-sampling 2D SPE on a 1-shot, 20-class Omniglot task. Performance is a mean over 5 replications of running the model, showing ±1 standard error of the mean.

Figure 6a compares the PN and SPE with 2D embeddings on Omniglot test classes. Each bar is the mean accuracy across four conditions: 1-shot/5-class, 5-shot/5-class, 1-shot/20-class, and 5-shot/20class. The first pair of bars perform the standard comparison in which the (1 or 5 instance) support set is used to obtain an embedding for each class, prototypes are formed, and query instances are classified. SPE is reliably better than the PN. Because the Omniglot data are carefully curated, the instances have little noise and therefore offer little opportunity to leverage SPE's assessment of uncertainty. Consequently, we corrupted instances by masking out rectangular regions of the input, as proposed by Oh et al. (2019). (See Appendix E for details.) The second and third sets of bars in Figure 6a correspond to the situations where the support and query instances are corrupted, respectively. SPE's advantage over PN increases significantly when the support instances are corrupted due to the fact that SPE's confidence-weighted prototypes (Equation5) discount noisier support examples. Although the SPE is still superior when only the query is corrupted, the benefit is small. We also compared PN and SPE using a 64D embedding, but with high dimensional embeddings, both methods are near ceiling on this data set, resulting in comparable performance between the two methods. (See Appendix D for additional results, broken down by condition.)

Figure 5: Two-dimensional embedding learned by the SPE on the Omniglot test set. Each square thumbnail image in the figure is a randomly-sampled instance from one of 125 randomly-sampled test classes and the location of the image represents the location of the class prototype. The images have a gray bounding box for visualization purposes only.

Figure 7: Two-dimensional embedding learned by the SPE on the 2-digit MNIST test set. A class is specified by a two-digit number. In both figures, the location of the class corresponds to the mean of the prototype in the test set using 140 support instances. The digits surrounded by a black border are classes that were not seen during training. In the left and right figures, the prototypes are colored according to the first and second digit of the class, respectively.

For 2-digit MNIST, each episode in training, validation, and seen-class testing contained all 70 classes and 50 support instances per class. For testing of unseen classes, each episode contained all 30 classes. For 3-digit MNIST, each episode contained 100 classes and either 20 support instances per class for training/validation or 50 support instances per class for seen-and unseen-class testing.

Figure 8: Examples of occluded 2-digit sequences. Occlusion is based on random rectangles that black out portions of each digit.

Figure 9: Two-dimensional embedding learned by the PN on the Omniglot test set. Each square thumbnail image in the figure is a randomly-sampled instance from one of 125 randomly-sampled test classes and the location of the image represents the location of the class prototype. The images have a gray bounding box for visualization purposes only.

Figure 11: Chain graph representing SPE

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, S c ∈ S, are used to determine the class prototype, ρ c , and the query instances are evaluated to predict class label (Equation

Table content: <table><tr><td>𝒙</td><td>𝒛</td><td>𝒚</td><td>class 2</td></tr><tr><td/><td/><td/><td/><td>𝝁 𝒚</td><td>𝝁 𝒙 𝝁 𝒙𝒚</td></tr><tr><td>(a)</td><td/><td/><td>class 1</td><td>(b)</td></tr><tr><td/><td/><td colspan="2">class 3</td></tr><tr><td colspan="5">Figure 1: (a) Illustration of the stochastic prototype embedding. The model learns a mapping from</td></tr><tr><td colspan="5">input space, x, to embedding space, z, in which same-class instances are near and different-class</td></tr><tr><td colspan="5">instances are far. Embeddings are represented as Gaussian random variables. Prototypes, noted as</td></tr><tr><td colspan="5">+ symbols in the embedding, are formed via a confidence-weighted average of the embeddings of</td></tr><tr><td colspan="5">instances known to belong to a class (support instances). Prototype uncertainty is depicted with</td></tr><tr><td colspan="5">the dotted ovals. Given the prototypes, a prediction of class y is made for a query instance by</td></tr><tr><td colspan="5">marginalizing a softmax prediction over the embedding space. (b) Depiction of intersection sampler.</td></tr><tr><td>Informally,</td><td/><td/><td/></tr></table>

Test classification accuracy (%) on 2and 3-digit MNIST for clean-support/clean-query, corrupt-support/clean-query, and clean-support/corrupt-query. N : number of digits in each image; D: dimensionality of the embedding. Contrastive and HIB results fromOh et al. (2019). PN is our implementation of Prototypical Networks(Snell et al., 2017). SPE is our model. SPE is trained with intersection sampling (1 sample per trial). Reported accuracy for PN and SPE for each experimental configuration is the mean over 1000 random test episodes.

Table content: <table><tr><td colspan="3">2D CLEAN SUPPORT, CLEAN QUERY</td><td/><td/><td/></tr><tr><td/><td colspan="4">1-SHOT, 5-CLASS 5-SHOT, 5-CLASS 1-SHOT, 20-CLASS 5-SHOT, 20-CLASS</td><td>MEAN</td></tr><tr><td>PN</td><td>75.7</td><td>82.6</td><td>45.0</td><td>55.9</td><td>64.8</td></tr><tr><td>SPE</td><td>76.9</td><td>82.3</td><td>49.7</td><td>55.3</td><td>66.1</td></tr><tr><td colspan="3">2D CORRUPT SUPPORT, CLEAN QUERY</td><td/><td/><td/></tr><tr><td/><td colspan="4">1-SHOT, 5-CLASS 5-SHOT, 5-CLASS 1-SHOT, 20-CLASS 5-SHOT, 20-CLASS</td><td>MEAN</td></tr><tr><td>PN</td><td>50.0</td><td>65.9</td><td>23.6</td><td>31.7</td><td>42.8</td></tr><tr><td>SPE</td><td>50.7</td><td>73.9</td><td>25.6</td><td>41.6</td><td>48.0</td></tr><tr><td colspan="3">2D CLEAN SUPPORT, CORRUPT QUERY</td><td/><td/><td/></tr><tr><td/><td colspan="4">1-SHOT, 5-CLASS 5-SHOT, 5-CLASS 1-SHOT, 20-CLASS 5-SHOT, 20-CLASS</td><td>MEAN</td></tr><tr><td>PN</td><td>48.9</td><td>52.3</td><td>21.7</td><td>25.6</td><td>37.1</td></tr><tr><td>SPE</td><td>47.8</td><td>52.3</td><td>22.8</td><td>26.8</td><td>37.4</td></tr><tr><td colspan="3">64D CLEAN SUPPORT, CLEAN QUERY</td><td/><td/><td/></tr><tr><td/><td colspan="4">1-SHOT, 5-CLASS 5-SHOT, 5-CLASS 1-SHOT, 20-CLASS 5-SHOT, 20-CLASS</td><td>MEAN</td></tr><tr><td>PN</td><td>98.5</td><td>99.6</td><td>94.9</td><td>98.6</td><td>97.9</td></tr><tr><td>SPE</td><td>98.5</td><td>99.5</td><td>94.9</td><td>98.6</td><td>97.9</td></tr><tr><td colspan="3">64D CORRUPT SUPPORT, CLEAN QUERY</td><td/><td/><td/></tr><tr><td/><td colspan="4">1-SHOT, 5-CLASS 5-SHOT, 5-CLASS 1-SHOT, 20-CLASS 5-SHOT, 20-CLASS</td><td>MEAN</td></tr><tr><td>PN</td><td>85.0</td><td>98.7</td><td>68.5</td><td>95.6</td><td>87.0</td></tr><tr><td>SPE</td><td>85.7</td><td>98.8</td><td>69.3</td><td>95.7</td><td>87.4</td></tr><tr><td colspan="3">64D CLEAN SUPPORT, CORRUPT QUERY</td><td/><td/><td/></tr><tr><td/><td colspan="4">1-SHOT, 5-CLASS 5-SHOT, 5-CLASS 1-SHOT, 20-CLASS 5-SHOT, 20-CLASS</td><td>MEAN</td></tr><tr><td>PN</td><td>80.9</td><td>84.9</td><td>66.7</td><td>74.7</td><td>76.8</td></tr><tr><td>SPE</td><td>80.3</td><td>84.8</td><td>66.3</td><td>74.7</td><td>76.5</td></tr></table>

Peer review (iclr2020):
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. However, in my opinion there are a few weaknesses in the paper in its current state: (1) the clarity of sections 3.1-3.2 could be significantly improved as currently the proposed probabilistic model framework is confusing and not well-defined; (2) the experimental results are provided only for embedding spaces of dimensionality 2-3 which significantly hinders performance of Prototypical Networks compared to having a much higher dimensional embedding space, so it would help to see comparisons of SPE and PN using high-dimensional embeddings; (3) the central idea and proposed model seem to be very close to those of [2] which is mentioned in the related work, so, please, list differences with this prior work in the updated version. For these reasons, I recommend a weak reject for the paper. I explain these issues in more detail below. (1) SPE model assumes that each embedding z is Gaussian-distributed p(z|x) = N(z; mu_x, sigma^2_x I) where parameters of the distribution mu_x and sigma_x are outputted by an embedding neural network. In equation (3) the author define each class prototype as rho_y = z_i + eps where z_i are instances of class y and eps is noise ~ N(0, sigma^2_eps I), so p(rho_y | z_i) = N(rho_y | z_i, sigma^2_eps I) for every z_i from class y. This definition is confusing to me since rho_y is redefined for each z_i, so it is not clear what the generative process for rho_y is. In the Introduction section it is mentioned “...each class instance is assumed to be a Gaussian perturbation of the prototype” which suggests z_i = rho_y + eps for every z_i, so graphical model would be rho_y -> z_i instead of z_i -> rho_y implied in section 3.1. Please, clarify in the rebuttal which the graphical model is implied in SPE. Is it x_i -> z_i -> rho_y? How is then the distribution of rho_y defined given all z_i from class y? Further, in the equation (4) the likelihood of rho_y given all x_i from class y is proportional to product of p(rho_y | x_i) so the expression was factorized over the condition which is the reflection of the authors’ assumption of consistency of the prototype. p(rho_y | x_1, … x_n) could be also written as proportional to the product of p(x_i | rho) times p(rho) using Bayes rule. It would help if authors could explain the differences between the two possible expressions for p(rho_y | x_1 … x_n) and why the former is chosen. Also, in equation (6) the right-hand-side is p(z | y, S) = p(z | rho_y, S) = N(z | mu_y, hat{sigma}^2_y I): why is the variance is hat{sigma}^2_y (so extra sigma^2_eps added to sigma^2_y) instead of just sigma^2_y defined in equation (5)? The noise term with sigma^2_eps seems to be already included in sigma^2_y. (2) In all presented experiments, the dimensionality of the embedding is 2-3 instead of high dimensional embedding spaces in the original Prototypical Network paper which causes a significant drop in performance: e.g., on Omniglot 1-shot 5-class classification, the accuracy dropped from 98.8% to 75.7%. SPE outperformed PN, but an important aspect of the PN model, embedding space dimensionality, was changed from the original in a way that hindered the performance very much, so SPE’s advantage is unclear. It is mentioned in section 4.2 that “We also compared PN and SPE using a 64D embedding, but with high dimensional embeddings, both methods are near ceiling on this data set, resulting in comparable performance between the two methods. (See Appendix D for additional results, broken down by condition.)” but the results on 64D embeddings are not added and discussed in the appendix. If the performance on high dimensional embeddings in the standard setting is similar for SPE and PC, it would help to see the 64D embedding experiment with corrupted support and query data, because if SPE outperformed PN on the corrupted data setting, it would be a fair advantage. It is interesting to see what the 2D embedding space looks like in the visualizations (Figures 2, 5, 7), and SPE is learning interpretable disentangles representations. However, for better classification performance it makes sense to use higher dimensional embedding space to allow learning more expressive representations. For Figures 2c, 5 and 7, it would be helpful to see how the embedding space of PN looks like and compare it visually to the one of SPE, e.g. whether PN has disentangled representations. It seems like we could get a similar embedding space to, for example, Figure 2c and the label uncertainty would come from the roughly equal distance of object representation to different class prototypes. Please, add respective embedding visualizations for PN and explain what the advantages of SPE learned representations over PN’s representations are. (3) The prior work of [2] referenced in the paper seems to have a very similar method (judging from Figure 2 and equations 5-6 which compute parameters of the prototype distribution, [2]). Please, list the differences with this prior work and provide experimental comparison if possible. Other questions: 1. In Figure 2c, please plot and highlight the learned distribution of the class prototypes p(rho | S). 2. Section 4.1: did each object in the training set have only one label or multiple labels (since classes are overlapping)? 3. Experiments: please add HIB method to synthetic data and Omniglot experiments if possible. It would also be interesting to compare to PN on miniImageNet (section 3.2 of [1]). Another interesting experiment which could show advantage of SPE could be out-of-distribution data detection (e.g., comparing likelihoods p(z | y) for in-distribution representations z and out-of-distribution z). [1] Snell, Jake, Kevin Swersky, and Richard Zemel. *Prototypical networks for few-shot learning.* Advances in Neural Information Processing Systems. 2017. [2] Fort, Stanislav. *Gaussian prototypical networks for few-shot learning on omniglot.* arXiv preprint arXiv:1708.02735 (2017). [3] Oh, Seong Joon, et al. *Modeling uncertainty with hedged instance embedding.* arXiv preprint arXiv:1810.00319 (2018).