SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Frederik Warburg*
Technical University of Denmark  
frwa@dtu.dk
Marco Miani*
Technical University of Denmark mmia@dtu.dk
Silas Brack
Technical University of Denmark silasbrack@gmail.com
Søren Hauberg
Technical University of Denmark
sohau@dtu.dk
We propose a Bayesian encoder for metric learning. Rather than relying on neural amortization as done in prior works, we learn a distribution over the network weights with the Laplace Approximation. We first prove that the contrastive loss is a negative log-likelihood on the spherical space. We propose three methods that ensure a positive definite covariance matrix. Lastly, we present a novel decomposition of the Generalized Gauss-Newton approximation. Empirically, we show that our Laplacian Metric Learner (LAM) yields well-calibrated uncertainties, reliably detects out-of-distribution examples, and has state-of-the-art predictive performance.
Metric learning seeks data representations where similar observations are near and dissimilar ones are far. This elegantly allows for building retrieval systems with simple nearest-neighbor search. Such systems easily cope with a large number of classes, and new classes can organically be added without retraining. While these retrieval systems show impressive performance, they quickly, and with no raised alarms, deteriorate with out-of-distribution data [38]. In particular, in safety-critical applications, the lack of uncertainty estimation is a concern as retrieval errors may propagate unnoticed through the system, resulting in erroneous and possibly dangerous decisions.
We present the Laplacian Metric Learner (LAM) to estimate reliable uncertainties of image embeddings as demonstrated in Fig. 1. We learn a distribution over the network weights (weight posterior) from
Figure 1: Reliable stochastic embeddings. LAM (ours) estimates reliable uncertainties of latent image embeddings that reflect the amount of blur, noise, or occlusion in the input image.
which we obtain a stochastic representation by embedding an image through sampled neural networks. This Bayesian formulation has multiple benefits, namely (1) robustness to out-of-distribution examples, (2) calibrated in-distribution uncertainties, and (3) a slight improvement in predictive performance.
More specifically, our method extends the Laplace Approximation [27] for metric learning. We present a probabilistic interpretation of the contrastive loss [15] which justifies that it can be interpreted as an unnormalized negative log-likelihood. We propose three solutions to ensure a positive semidefinite Hessian for the contrastive loss and present two approaches to compute the Generalized Gauss-Newton [12] approximation for  $\ell_2$ -normalized networks. Finally, we boost our method with the online training procedure from [29] to achieve state-of-the-art performance.
We are not the first to consider uncertainty quantification in image retrieval. Seminal works [38, 33] have addressed the lack of uncertainties in retrieval with amortized inference [14], where a neural network predicts a stochastic embedding. The issues with this approach are that (1) it requires strong assumptions on the distribution of the embedding, (2) the networks are often brittle and difficult to optimize, and (3) out-of-distribution detection relies on the network's capacity to extrapolate uncertainties. As neural networks extrapolate poorly [50], the resulting predicted uncertainties are unreliable for out-of-distribution data [10] and are thus, in practice, of limited value.
In contrast, our method does not assume any distribution on the stochastic embeddings, is simple to optimize, and does not rely on a neural network to extrapolate uncertainties. Instead, our weight posterior is derived from the curvature of the loss landscape and the uncertainties of the latent embeddings deduced (rather than learned) with sampling. Empirically, we show that this leads to reliable out-of-distribution performance and calibrated uncertainties in both controlled toy experiments and challenging real-world applications such as bird, face, and place recognition.
Metric learning attempts to map data to an embedding space, where similar data are close together and dissimilar data are far apart. This is especially useful for retrieval tasks with many classes and few observations per class such as place recognition [46] and face recognition [37] or for tasks where classes are not well-defined, such as food tastes [48] or narratives in online discussions [4].
There exist many metric losses that optimize for a well-behaved embedding space. We refer to the excellent survey by [30] for an overview. We here focus on the contrastive loss [15]
$$
\mathcal {L} _ {\text {c o n}} (\theta) = \frac {1}{2} \| f _ {\theta} (x _ {a}) - f _ {\theta} (x _ {p}) \| ^ {2} + \frac {1}{2} \max  \left(0, m - \| f _ {\theta} (x _ {a}) - f _ {\theta} (x _ {n}) \| ^ {2}\right), \tag {1}
$$
which has shown state-of-the-art performance [30] and is one of the most commonly used metric losses. Here,  $f_{\theta}$  is a neural network parametrized by  $\theta$  which maps from the observation space to the embedding space. The loss consists of two terms, one that attracts observations from the same class (anchor  $x_{a}$  and positive  $x_{p}$ ), and one that repels observations from different classes (anchor  $x_{a}$  and negative  $x_{n}$ ). The margin  $m \geq 0$  ensures that negatives are repelled sufficiently far. We will later present a probabilistic extension of the contrastive loss that allows us to learn stochastic, rather than deterministic, features in the embedding space.
Uncertainty in deep learning is studied across many domains to mitigate fatal accidents and allow for human intervention when neural networks make erroneous predictions. Current methods can be divided into methods that apply amortized optimization to train a neural network to predict the parameters of the output distribution, and methods that do not. The amortized methods, best known from the variational autoencoder (VAE) [19, 36], seem attractive at first as they can directly estimate the output distribution (without requiring sampling), but they suffer from mode collapse and are sensitive to out-of-distribution data due to the poor extrapolation capabilities of neural networks [31, 10]. 'Bayes by Backprop' [2] learns a distribution over parameters variationally but is often deemed too brittle for practical applications. Alternatives to amortized methods include deep ensembles [23], stochastic weight averaging (SWAG) [28], Monte-Carlo dropout [13] and Laplace Approximation (LA) [24, 27] which all approximate the generally intractable weight posterior  $p(\theta | \mathcal{D})$  of a neural network. We propose to extend LA to metric learning.
Laplace approximations (LA) can be applied for every loss function  $\mathcal{L}$  that can be interpreted as an unnormalized log-posterior by performing a second-order Taylor expansion around a chosen weight
vector  $\theta^{*}$  such that
$$
\mathcal {L} (\theta) \approx \mathcal {L} ^ {*} + \left(\theta - \theta^ {*}\right) ^ {\top} \nabla \mathcal {L} ^ {*} + \frac {1}{2} \left(\theta - \theta^ {*}\right) ^ {\top} \nabla^ {2} \mathcal {L} ^ {*} \left(\theta - \theta^ {*}\right), \tag {2}
$$
where  $\mathcal{L}^*$  is the loss evaluated in  $\theta^{*}$ . Imposing the unnormalized log-posterior to be a second-order polynomial is equivalent to assuming the posterior to be Gaussian. If  $\theta^{*}$  is a MAP estimate, the first-order term vanishes, and the second-order term can be interpreted as a precision matrix, the inverse of the covariance. Assuming  $\theta^{*}$  is a MAP estimate, this second-order term is negative semi-definite for common (convex) supervised losses, such as the mean-squared error and cross-entropy. Recently, Daxberger et al. [8] demonstrated that post-hoc LA is scalable and produces well-behaved uncertainties for classification and regression. The Laplacian Autoencoder (LAE) [29] improves on the post-hoc LA with an online Monte Carlo EM training procedure to learn a well-behaved posterior. It demonstrates state-of-the-art uncertainty quantification for unsupervised representation learning.
Uncertainty in metric learning is not new [44], but the majority of recent methods apply amortized inference to predict distributions in the embedding space [47, 3, 5, 33, 38, 39, 41], making them sensitive to mode collapse and out-of-distribution data. Alternatives like deep ensembles [42] and Monte-Carlo dropout [43] suffer from increased training time, poor empirical performance, and limited Bayesian interpretation [8]. We explore LA in metric learning and attain state-of-the-art performance.
3 Laplacian Metric Learning
To perform Bayesian retrieval, we estimate the weight posterior of the embedding network  $f_{\theta}$  such that we can sample data embeddings to propagate uncertainty through the decision process. The embedding network is parametrized by  $\theta \in \Theta$  and trained with the contrastive loss. The network maps an image  $x \in \mathcal{X} \coloneqq \mathbb{R}^{HWC}$  to an embedding  $z \in \mathcal{Z}$ , which is restricted to be on a  $Z$ -dimensional unit sphere  $\mathcal{Z} \coloneqq S^Z$ . This spherical normalization is often done in retrieval to obtain faster retrieval and a slight performance boost [1, 35]. Fig. 2 illustrate our Bayesian mapping from image to latent space.
Figure 2: Model overview. We learn a distribution over parameters, such that we embed an image through sampled encoders  $f_{\theta}$  to points  $z_{i}$  (red dots) in a latent space  $\mathcal{Z}$ . We reduce these latent samples to a single measure of uncertainty by estimating the parameters of a von Mises-Fisher distribution.
Laplace Approximation (LA) recap. We apply LA to obtain the posterior over the weights  $\theta$ . LA comes in two flavors: (1) The post-hoc LA is found by first training a standard deterministic network through gradient steps with the contrastive loss to find the maximum a posteriori (MAP) parameters  $\theta^{*}$ . Since we are in a local optimum, the first-order term in the second-order Taylor expansion (Eq. 2) vanishes, and we can define the parameter distribution as
$$
p (\theta | \mathcal {D}) = \mathcal {N} \left(\theta \mid \theta^ {*}, \left(\nabla_ {\theta} ^ {2} \mathcal {L} _ {\text {c o n}} \left(\theta^ {*}; \mathcal {D}\right) + \sigma_ {\text {p r i o r}} ^ {- 2} \mathbb {I}\right) ^ {- 1}\right). \tag {3}
$$
The advantage of post-hoc LA is that the training procedure does not change, and already trained neural networks can be made Bayesian. In practice, however, we empirically observe post-hoc LA to be unstable. We refer to the large standard deviation in Table 1, Table 2, and Table 3 over random seeds. The instability stems from curvature differences between the local minimas through stochastic optimization. (2) Online LA [29] improves on this instability by marginalizing the LA during training with Monte Carlo EM. This helps the training recover a solution  $\theta^{*}$  where the Hessian reflects the loss landscape. Specifically, at each step  $t$  during training, we keep in memory a Gaussian distribution on the parameters  $q^{t}(\theta) = \mathcal{N}(\theta |\theta_{t},H_{\theta_{t}}^{-1})$ . The parameters are updated through an expected gradient step
$$
\theta_ {t + 1} = \theta_ {t} + \lambda \mathbb {E} _ {\theta \sim q ^ {t}} [ \nabla_ {\theta} \mathcal {L} _ {\text {c o n}} (\theta ; \mathcal {D}) ] \tag {4}
$$
and a discounted Laplace update
$$
H _ {\theta_ {t + 1}} = (1 - \alpha) H _ {\theta_ {t}} + \nabla_ {\theta} ^ {2} \mathcal {L} _ {\text {c o n}} (\theta ; \mathcal {D}), \tag {5}
$$
where  $\alpha$  describes an exponential moving average, similar to momentum-like training. The initialization follows the isotropic prior  $q^0 (\theta) = \mathcal{N}(\theta |0,\sigma_{\mathrm{prior}}^2\mathbb{I})$
The probabilistic contrastive likelihood. LA expects the loss function to be a log-posterior, that is log-likelihood plus log-prior. A simple choice of the log-prior is  $\| \theta \|_2^2$  (weight decay). Further, we find that the contrastive loss Eq. 1 has a probabilistic interpretation and is a negative log-likelihood on the spherical space. We define the attractive term  $\mathcal{P} \to (z|x,\theta) \sim \mathcal{N}^S(z|f_\theta(x),\kappa)$  and the repelling term  $\mathcal{P} \leftrightarrow (z|x,\theta) \sim \mathcal{N}^S(z| - f_\theta(x),\kappa)$  as Von Mises-Fisher distributions on the latent space spherical space, where the concentration parameter  $\kappa \geq 0$ . The product of these two likelihoods yields a contrastive likelihood, which is valid on the spherical space, and its negative log-likelihood is equivalent to the contrastive loss (proof and details in ??).
The intuition, helped by electrostatics, is that the log probability density of these Von Mises-Fisher distributions (i.e. spherical Gaussians) plays the role of a potential energy. The precision parameter  $\kappa$  controls the concentration of the potential energy and, consequently, the strength of the associated force. If  $\kappa = 0$  the distribution is uniform, the potential is constant and the associated force is zero. On the spherical space, moreover, a repulsive term from some point is equivalent to an attractive term from the antipodal point, which in turn is mathematically equivalent (see ??) to a negative precision  $\kappa$ . This is the fundamental reason behind the non positive-definiteness of the hessian.
Hessian of the contrastive loss. Both post-hoc and online LA require the Hessian of the contrastive loss  $\nabla_{\theta}^{2}\mathcal{L}_{\mathrm{con}}(\theta ;\mathcal{D})$ . The Hessian is commonly approximated with the Generalized Gauss-Newton (GGN) approximation [12, 8, 6, 11]. The GGN decomposes the loss into  $\mathcal{L} = g\circ f$ , where  $g$  is usually chosen as the loss function and  $f$  the model function, and only  $f$  is linearized [22].
However, in our case, this decomposition is non-trivial. Recall that the last layer of our network is an  $\ell_2$  normalization layer, which projects embeddings onto a hyper-sphere. This normalization layer can either be viewed as part of the model  $f$  (linearized normalization layer) or part of the loss  $g$  (non-linearized normalization layer). The former can be interpreted as using the Euclidean distance and the latter as using the Arccos distance for the contrastive loss (see ?? and ??). These two share the zero- and first-order terms for normalized embeddings but not the second-order derivatives due to the GGN linearization. The Euclidean interpretation leads to simpler derivatives and interpretations, and we will therefore use it for our derivations. We emphasize that the Arccos is theoretically a more accurate approximation, because the  $\ell_2$ -layer is not linearized, and we provide derivations in ??
The GGN matrix for contrastive loss with the Euclidean interpretation is given by
$$
\begin{array}{l} \nabla_ {\theta} ^ {2} \mathcal {L} _ {\mathrm {c o n}} (\theta ; \mathcal {I}) = \sum_ {i j \in \mathcal {I}} H _ {\theta} ^ {i j} = \sum_ {i j \in \mathcal {I} _ {p}} H _ {\theta} ^ {i j} + \sum_ {i j \in \mathcal {I} _ {n}} H _ {\theta} ^ {i j} \\ \approx \sum_ {i j \in \mathcal {I} _ {p}} J _ {\theta} ^ {i j} ^ {\top} \underbrace {\left( \begin{array}{l l} 1 & - 1 \\ - 1 & 1 \end{array} \right)} _ {:= H _ {p}} J _ {\theta} ^ {i j} + \sum_ {i j \in \mathcal {I} _ {n}} J _ {\theta} ^ {i j} ^ {\top} \underbrace {\left( \begin{array}{l l} - 1 & 1 \\ 1 & - 1 \end{array} \right)} _ {:= H _ {n}} J _ {\theta} ^ {i j}, \tag {6} \\ \end{array}
$$
where  $J_{\theta}^{ij} = \left(J_{\theta}f_{\theta}(x_i)^\top, J_{\theta}f_{\theta}(x_j)^\top\right)^\top$ , with  $J_{\theta}$  is the Jacobian wrt. the parameters and  $H_p$  and  $H_n$  are the Hessians of the contrastive loss wrt. the model output for positive and negative pairs. The first sum runs over positive pairs and the second sum runs over negative pairs within the margin. Negative pairs outside the margin do not contribute to the Hessian, and can therefore be ignored to reduce the computational load (???).
The eigenvalues of the Hessian wrt. to the output are  $(0,2)$  and  $(-2,0)$  for the positive  $H_{p}$  and negative  $H_{n}$  terms, so we are not guaranteed to have a positive semidefinite Hessian,  $H_{\theta}$ . To avoid covariances with negative eigenvalues, we propose three solutions to ensure a positive semidefinite Hessian. Proofs are in ??
Ensuring positive definiteness of the covariance matrix. We do not want to be restricted in the choice of the prior except to have non-zero precision, so we must ensure that  $\nabla_{\theta}^{2}\mathcal{L}_{\mathrm{con}}(\theta^{*};\mathcal{D})$  is positive semidefinite. Differently from the standard convex losses, this is not ensured by the GGN approximation [18]. Our main insight is that we can ensure a positive semidefinite Hessian  $H_{\theta}$  by only manipulating the Hessians  $H_{p}$  and  $H_{n}$  in Eq. 6.
1. Positive: The repelling term is ignored, such that only positive pairs contribute to the Hessian.
$$
H _ {p} = \left( \begin{array}{c c} 1 & - 1 \\ - 1 & 1 \end{array} \right), \quad H _ {n} = \left( \begin{array}{c c} 0 & 0 \\ 0 & 0 \end{array} \right) \tag {7}
$$
2. Fixed: The cross derivatives are ignored.
$$
H _ {p} = \left( \begin{array}{c c} 1 & 0 \\ 0 & 1 \end{array} \right), \quad H _ {n} = \left( \begin{array}{c c} - 1 & 0 \\ 0 & - 1 \end{array} \right) \tag {8}
$$
3. Full: Positive semidefiniteness is ensured with  $ReLU$ ,  $\max(0, \nabla_{\theta}^{2}\mathcal{L}_{con}(\theta; \mathcal{D}))$ , on the Hessian of the loss wrt. the parameters.
$$
H _ {p} = \left( \begin{array}{c c} 1 & - 1 \\ - 1 & 1 \end{array} \right), \quad H _ {n} = \left( \begin{array}{c c} - 1 & 1 \\ 1 & - 1 \end{array} \right) \tag {9}
$$
The positive approximation is inspired by [38], which only uses positive pairs to train an uncertainty module. The gradient arrows in Fig. 3a illustrate that negative pairs are neglected when computing the Hessian of the contrastive loss. The fixed approximation considers one data point at a time, assuming the other one is fixed. Thus, given a pair of data points, this can be interpreted as first moving one data point, and then the second (rather than both at the same time). We formalize this in ??. Fig. 3b illustrate this idea when all points except  $a$  are fixed. Lastly, we propose the full Hessian of the contrastive loss (Fig. 3c) and ensure positive semidefiniteness by computing the ReLU of the diagonal Hessian. This approximation can equivalently be interpreted as a projection into the space of psd matrices. In practice, the Hessian scales quadratically in memory wrt. the number of parameters. To mitigate this, we approximate this Hessian by its diagonal and only apply the LA on the last layer [25, 9]. We experimentally find that the fixed approximation yields the best performance (Section 4).
Hard negative mining. Most pairs, namely the negatives outside the margin, do not contribute to the Hessian, so it is wasteful to compute their Hessian. Therefore, we use hard negative mining [30] to only compute the Hessian of pairs that have non-zero Hessian, i.e. the negative sample lie within the margin (illustrated with the dotted line in Fig. 3).
Von Mises-Fisher distribution. To obtain a single measure of uncertainty from our sampled image embeddings (red dots in Fig. 4), we fit a von Mises-Fisher distribu
tion. The von Mises-Fisher distribution describes a normal distribution where all probability mass lies on a  $Z$ -dimensional hyper-sphere. It is parametrized with a directional mean  $\mu$  and a scalar concentration parameter  $\kappa$ , which can be interpreted as the inverse of an isotropic covariance  $\kappa = 1 / \sigma^2$ , i.e., small  $\kappa$  means high uncertainty and large  $\kappa$  means low uncertainty. There exist several methods to estimate  $\kappa$ . We opt for the simplest and most computationally efficient [40] (see ??).
(a) positives
(b) fixed
Figure 3: Hessian approximations. To ensure a positive semidefinite Hessian approximation we propose three approximations. In (a) only the positives  $p$  contribute to the Hessian as the negatives  $n$  are ignored. In (b) we consider one point at a time, e.g., only the anchor  $a$  contributes. In (c) we consider all interactions.
(c) full
We benchmark our method against strong probabilistic retrieval models. Probabilistic Face Embeddings (PFE) [38] and Hedge Image Embedding (HIB) [33] perform amortized inference and thus estimate the mean and variance of latent observation. We also compare against MC Dropout [13] and Deep Ensemble [23], two approximate Bayesian methods, which have successfully been applied in image retrieval [43, 42].
We compare the models' predictive performance with the recall (recall@k) and mean average precision  $(\mathrm{mAP}@\mathbf{k})$  among the  $k$  nearest neighbors [47, 30, 1]. We evaluate the models' abilities to interpolate and extrapolate uncertainties by measuring the Area Under the Sparsification Curve (AUSC), Expected Calibration Error (ECE) on in-distribution (ID) data, the Area Under Receiver Operator Curve (AUROC), and Area Under Precision-Recall Curve (AUPRC) on out-of-distribution (OoD) data. We provide more details on these metrics in ??
We extend StochMan [11] with the Hessian backpropagation for the contrastive loss, and the PyTorch [34] code is publicly available<sup>2</sup>. ?? details the experimental setup.
See https://github.com/FrederikWarburg/bayesian-metric-learning
(a) Out-of-distribution detection.
(b) Predictive performance.
Figure 4: Summary of experimental results. LAM consistently outperforms existing methods on OoD detection, as measured by AUROC, and matches or surpasses in predictive performance measured by mAP@  $k$ . Error bars show one std across five runs.
Experimental Summary. We first summarize our experimental results. Across five datasets, three network architectures, and three different sizes of the latent space (ranging from 3 to 2048), we find that LAM has well-calibrated uncertainties, reliably detects OoD examples, and achieves state-of-the-art predictive performance. Fig. 4a shows that the uncertainties from online LAM reliably identify OoD examples. Online LAM outperforms other Bayesian methods, such as post-hoc LAM and MC dropout, on this task, which in turn clearly improves upon amortized methods that rely on a neural network to extrapolate uncertainties. Fig. 4b shows that LAM consistently matches or outperforms existing image retrieval methods in terms of predictive performance. We find that the fixed Hessian approximation with the Arccos distance performs the best, especially on higher dimensional data.
Ablation: Positive definiteness covariance matrix. We experimentally study which method to ensure a positive semidefinite Hessian has the best performance measured in both predictive performance (mAP@5) and uncertainty quantification (AUROC, AUSC). We found that all methods perform similarly on simple datasets and low dimensional hyper-spheres, but the fixed approximation with Arcos distance performs better on more challenging datasets and higher dimensional hyper-spheres. We present results on one of these more challenging datasets, namely the LFW [17] face recognition dataset with the CUB200 [45] bird dataset as an OoD dataset. We use a ResNet50 [16] with a GeM pooling layer [35] and a 2048 dimensional embedding and diagonal, last-layer LA [8].
Table 1 shows the performance for post-hoc and online LA with fixed, positive, or full Hessian approximation using either Euclidean or Arccos distance. Across all metrics, the online LA with Arccos distance and the fixed Hessian approximation performs similarly or the best. We proceed to benchmark this method against sev