Learning Unified Distance Metric Across Diverse Data Distributions with Parameter-Efficient Transfer Learning

TL;DR

This work addresses the problem of learning a unified distance metric across multiple heterogeneous data distributions, a scenario common in real world applications. It introduces PUMA, a parameter efficient approach that freezes a pre trained Vision Transformer and adds stochastic adapters plus a conditional prompt pool to capture both shared and dataset specific information without bias toward dominant distributions. The authors certify their method on a new eight dataset UML benchmark, showing superior universal accuracy and harmonic mean performance while using up to 69x fewer trainable parameters than dataset specific or naive unified baselines. The results indicate that parameter efficient transfer learning can achieve scalable, cross distribution retrieval with strong performance, enabling a single model deployment across diverse data regimes.

Abstract

A common practice in metric learning is to train and test an embedding model for each dataset. This dataset-specific approach fails to simulate real-world scenarios that involve multiple heterogeneous distributions of data. In this regard, we explore a new metric learning paradigm, called Unified Metric Learning (UML), which learns a unified distance metric capable of capturing relations across multiple data distributions. UML presents new challenges, such as imbalanced data distribution and bias towards dominant distributions. These issues cause standard metric learning methods to fail in learning a unified metric. To address these challenges, we propose Parameter-efficient Unified Metric leArning (PUMA), which consists of a pre-trained frozen model and two additional modules, stochastic adapter and prompt pool. These modules enable to capture dataset-specific knowledge while avoiding bias towards dominant distributions. Additionally, we compile a new unified metric learning benchmark with a total of 8 different datasets. PUMA outperforms the state-of-the-art dataset-specific models while using about 69 times fewer trainable parameters.

Figures (3)

• Figure 1: Comparison between conventional and unified metric learning methods. (a) Conventional metric learning employs separate models for individual datasets, incurring significant computational and memory costs as data diversity grows. (b) A naïve solution is to fine-tune the model on a merged dataset, but this often leads to a severe bias toward major data distributions. (c) In contrast, our method excels on all datasets with just one model. This is highly resource-efficient as it enables one-time learning and evaluation on diverse data distributions using a single model.
• Figure 2: An overview of PUMA. PUMA consists of two learnable modules: stochastic adapters (Sec. \ref{['sec:adapter']}) and a prompt pool (Sec. \ref{['sec:prompt']}). Using the output of the transformer's embedding layer as a query, and it creates a conditional prompt by integrating relevant prompts through an attention mechanism. The conditional prompt is combined with image embeddings and class token, and then fed into the transformer. The modified input is embedded through transformer blocks, each coupled with a stochastic adapter, a learnable bottleneck module that turns on stochastically during training.
• Figure 3: The average similarity between input queries and prompts for each dataset. The $x$-axis represents prompt indices.