Modeling Uncertainty with Hedged Instance Embeddings

1 INTRODUCTION . An instance embedding is a mapping f from an input x , such as an image , to a vector representation , z ∈ RD , such that “ similar ” inputs are mapped to nearby points in space . Embeddings are a versatile representation that support various downstream tasks , including image retrieval ( Babenko et al. , 2014 ) and face recognition ( Schroff et al. , 2015 ) . Instance embeddings are often treated deterministically , i.e. , z = f ( x ) is a point in RD . We refer to this approach as a point embedding . One drawback of this representation is the difficulty of modeling aleatoric uncertainty ( Kendall & Gal , 2017 ) , i.e . uncertainty induced by the input . In the case of images this can be caused by occlusion , blurriness , low-contrast and other factors . To illustrate this , consider the example in Figure 1a . On the left , we show an image composed of two adjacent MNIST digits , the first of which is highly occluded . The right digit is clearly a 7 , but the left digit could be a 1 , or a 4 . One way to express this uncertainty about which choice to make is to map the input to a region of space , representing the inherent uncertainty of “ where it belongs ” . We propose a new method , called hedged instance embedding ( HIB ) , which achieves this goal . Each embedding is represented as a random variable , Z ∼ p ( z|x ) ∈ RD . The embedding effectively spreads probability mass across locations in space , depending on the level of uncertainty . For example in Figure 1b , the corrupted image is mapped to a two-component mixture of Gaussians covering both the “ 17 ” and “ 47 ” clusters . We propose a training scheme for the HIB with a learnable-margin contrastive loss and the variational information bottleneck ( VIB ) principle ( Alemi et al. , 2016 ; Achille & Soatto , 2018 ) . ∗Work performed during an internship with Google Research . To evaluate our method , we propose a novel dataset , N-digit MNIST , which we will open source . Using this dataset , we show that HIB exhibits several desirable properties compared to point embeddings : ( 1 ) downstream task performance ( e.g . recognition and verification ) improves for uncertain inputs ; ( 2 ) the embedding space exhibits enhanced structural regularity ; and ( 3 ) a per-exemplar uncertainty measure that predicts when the output of the system is reliable . 2 METHODS . In this section , we describe our method in detail . 2.1 POINT EMBEDDINGS . Standard point embedding methods try to compute embeddings such that z1 = f ( x1 ) and z2 = f ( x2 ) are “ close ” in the embedding space if x1 and x2 are “ similar ” in the ambient space . To obtain such a mapping , we must decide on the definition of “ closeness ” as well as a training objective , as we explain below . Contrastive loss Contrastive loss ( Hadsell et al. , 2006 ) is designed to encourage a small Euclidean distance between a similar pair , and large distance of margin M > 0 for a dissimilar pair . The loss is Lcon = { ||z1 − z2||22 if match max ( M − ||z1 − z2||2 , 0 ) 2 if non-match ( 1 ) where zi = f ( xi ) . The hyperparameter M is usually set heuristically or based on validation-set performance . Soft contrastive loss A probabilistic alternative to contrastive loss , which we will use in our experiments is defined here . It represents the probability that a pair of points is matching : p ( m|z1 , z2 ) : = σ ( −a||z1 − z2||2 + b ) ( 2 ) with scalar parameters a > 0 and b ∈ R , and the sigmoid function σ ( t ) = 11+e−t . This formulation calibrates Euclidean distances into a probabilistic expression for similarity . Instead of setting a hard threshold likeM , a and b together comprise a soft threshold on the Euclidean distance . We will later let a and b be trained from data . Having defined the match probability p ( m|z1 , z2 ) , we formulate the contrastive loss as a binary classification loss based on the softmax cross-entropy ( negative log-likelihood loss ) . More precisely , for an embedding pair ( z1 , z2 ) the loss is defined as Lsoftcon = − log p ( m = m̂|z1 , z2 ) = { − log p ( m|z1 , z2 ) if m̂ = 1 , − log ( 1− p ( m|z1 , z2 ) ) if m̂ = 0 , ( 3 ) where m̂ is the indicator function with value 1 for ground-truth match and 0 otherwise . Although some prior work has explored this soft contrastive loss ( e.g . Bertinetto et al . ( 2016 ) ; Orekondy et al . ( 2018 ) ) , it does not seem to be widely used . However , in our experiments , it performs strictly better than the hard margin version , as explained in Appendix B . 2.2 STOCHASTIC EMBEDDINGS . In HIB , we treat embeddings as stochastic mappings x 7→ Z , and write the distribution as Z ∼ p ( z|x ) . In the sections below , we show how to learn and use this mapping . Match probability for probabilistic embeddings The probability of two inputs matching , given in Equation 2 , can easily be extended to stochastic embeddings , as follows : p ( m|x1 , x2 ) = ∫ p ( m|z1 , z2 ) p ( z1|x1 ) p ( z2|x2 ) dz1dz2 . ( 4 ) We approximate this integral via Monte-Carlo sampling from z ( k1 ) 1 ∼ p ( z1|x1 ) and z ( k2 ) 2 ∼ p ( z2|x2 ) : p ( m|x1 , x2 ) ≈ 1 K2 K∑ k1=1 K∑ k2=1 p ( m|z ( k1 ) 1 , z ( k2 ) 2 ) . ( 5 ) In practice , we get good results using K = 8 samples per input image . Now we discuss the computation of p ( z|x ) . Single Gaussian embedding The simplest setting is to let p ( z|x ) be a D-dimensional Gaussian with mean µ ( x ) and diagonal covariance Σ ( x ) , where µ and Σ are computed via a deep neural network with a shared “ body ” and 2D total outputs . Given a Gaussian representation , we can draw K samples z ( 1 ) , · · · , z ( K ) iid∼ p ( z|x ) , which we can use to approximate the match probability . Furthermore , we can use the reparametrization trick ( Kingma & Welling , 2013 ) to rewrite the samples as z ( k ) = diag ( √ Σ ( x ) ) · ( k ) + µ ( x ) , where ( 1 ) , · · · , ( K ) iid∼ N ( 0 , I ) . This enables easy backprop- agation during training . Mixture of Gaussians ( MoG ) embedding We can obtain a more flexible representation of uncertainty by using a mixture of C Gaussians to represent our embeddings , i.e . p ( z|x ) = ∑Cc=1N ( z ; µ ( x , c ) , Σ ( x , c ) ) . To enhance computational efficiency , the 2C mappings { ( µ ( x , c ) , Σ ( x , c ) ) } Cc=1 share a common CNN stump and are branched with one linear layer per branch . When approximating Equation 5 , we use stratified sampling , i.e . we sample the same number of samples from each Gaussian component . Computational considerations The overall pipeline for computing the match probability is shown in Figure 2 . If we use a single Gaussian embedding , the cost ( time complexity ) of computing the stochastic representation is essentially the same as for point embedding methods , due to the use of a shared network for computing µ ( x ) and Σ ( x ) . Also , the space requirement is only 2× more . ( This is an important consideration for many embedding-based methods . ) 2.3 VIB TRAINING OBJECTIVE . For training our stochastic embedding , we combine two ingredients : soft contrastive loss in Equation 3 and the VIB principle Alemi et al . ( 2016 ) ; Achille & Soatto ( 2018 ) . We start with a summary of the original VIB formulation , and then describe its extension to our setting . Variational Information Bottleneck ( VIB ) A discriminative model p ( y|x ) is trained under the information bottleneck principle ( Tishby et al. , 1999 ) by maximizing the following objective : I ( z , y ) − βI ( z , x ) ( 6 ) where I is the mutual information , and β > 0 is a hyperparameter which controls the tradeoff between the sufficiency of z for predicting y , and the minimality ( size ) of the representation . Intuitively , this objective lets the latent encoding z capture the salient parts of x ( salient for predicting y ) , while disallowing it to “ memorise ” other parts of the input which are irrelevant . Computing the mutual information is generally computationally intractable , but it is possible to use a tractable variational approximation as shown in Alemi et al . ( 2016 ) ; Achille & Soatto ( 2018 ) . In particular , under the Markov assumption that p ( z|x , y ) = p ( z|x ) we arrive at a lower bound on Equation 6 for every training data point ( x , y ) as follows : −LVIB : = Ez∼p ( z|x ) [ log q ( y|z ) ] − β · KL ( p ( z|x ) ||r ( z ) ) ( 7 ) where p ( z|x ) is the latent distribution for x , q ( y|z ) is the decoder ( classifier ) , and r ( z ) is an approximate marginal term that is typically set to the unit Gaussian N ( 0 , I ) . In Alemi et al . ( 2016 ) , this approach was shown ( experimentally ) to be more robust to adversarial image perturbations than deterministic classifiers . It has also been shown to provide a useful way to detect out-of-domain inputs ( Alemi et al. , 2018 ) . Hence we use it as the foundation for our approach . VIB for learning stochastic embeddings We now apply the above method to learn our stochastic embedding . In particular , we train a discriminative model based on matching or mismatching pairs of inputs ( x1 , x2 ) , by minimizing the following loss : LVIBEmb : =− Ez1∼p ( z1|x1 ) , z2∼p ( z2|x2 ) [ log p ( m = m̂|z1 , z2 ) ] + β · [ KL ( p ( z1|x1 ) ||r ( z1 ) ) + KL ( p ( z2|x2 ) ||r ( z2 ) ) ] ( 8 ) where the first term is given by the negative log likelihood loss with respect to the ground truth match m̂ ( this is identical to Equation 3 , the soft contrastive loss ) , and the second term is the KL regularization term , r ( z ) = N ( z ; 0 , I ) . The full derivation is in appendix F. We optimize this loss with respect to the embedding function ( µ ( x ) , Σ ( x ) ) , as well as with respect to the a and b terms in the match probability in Equation 2 . Note that most pairs are not matching , so the m = 1 class is rare . To handle this , we encourage a balance of m = 0 and m = 1 pair samples within each SGD minibatch by using two streams of input sample images . One samples images from the training set at random and the other selects images from specific class labels , and then these are randomly shuffled to produce the final batch . As a result , each minibatch has plenty of positive pairs even when there are a large number of classes .

Summary: While most works consider embedding as the problem of mapping an input into a point in an embedding space, paper 1341 considers the problem of mapping an input into a distribution in an embedding space. Computing the matching score of two inputs (e.g. two images) involves the following steps: (i) assuming a Gaussian distribution in the embedding space, computing the mean and standard deviation for each input, (ii) drawing a set of samples from each distribution, (3) computing the normalized distances between the samples and (iv) averaging to obtain a global score.