A Non-isotropic Probabilistic Take on Proxy-based Deep Metric Learning

TL;DR

This paper introduces a probabilistic framework for proxy-based deep metric learning using non-isotropic von Mises-Fisher distributions, enhancing uncertainty modeling, gradient stability, and generalization in embedding spaces.

Contribution

It proposes a novel non-isotropic probabilistic approach that models images and class proxies with vMF distributions, improving class-internal structure learning and uncertainty representation.

Findings

Improved generalization performance on standard benchmarks.

Enhanced uncertainty-awareness and gradient behavior during training.

Favorable comparison with existing proxy-based DML methods.

Abstract

Proxy-based Deep Metric Learning (DML) learns deep representations by embedding images close to their class representatives (proxies), commonly with respect to the angle between them. However, this disregards the embedding norm, which can carry additional beneficial context such as class- or image-intrinsic uncertainty. In addition, proxy-based DML struggles to learn class-internal structures. To address both issues at once, we introduce non-isotropic probabilistic proxy-based DML. We model images as directional von Mises-Fisher (vMF) distributions on the hypersphere that can reflect image-intrinsic uncertainties. Further, we derive non-isotropic von Mises-Fisher (nivMF) distributions for class proxies to better represent complex class-specific variances. To measure the proxy-to-image distance between these models, we develop and investigate multiple distribution-to-point and distribution-to-distribution metrics. Each framework choice is motivated by a set of ablational studies, which showcase beneficial properties of our probabilistic approach to proxy-based DML, such as uncertainty-awareness, better-behaved gradients during training, and overall improved generalization performance. The latter is especially reflected in the competitive performance on the standard DML benchmarks, where our approach compares favorably, suggesting that existing proxy-based DML can significantly benefit from a more probabilistic treatment. Code is available at github.com/ExplainableML/Probabilistic_Deep_Metric_Learning.