Geometry-Aware Adaptation for Pretrained Models

TL;DR

This paper introduces Loki, a geometry-aware adaptation method for pretrained models that leverages label space metrics to improve zero-shot and few-shot predictions without additional training, supported by theoretical analysis and empirical results.

Contribution

It proposes a simple, metric-based prediction adjustment technique called Loki, with comprehensive theoretical insights and practical improvements for zero-shot and few-shot learning.

Findings

Up to 29.7% relative improvement over SimCLR on ImageNet.

Effective with external or self-derived label metrics.

Scales to hundreds of thousands of classes.

Abstract

Machine learning models -- including prominent zero-shot models -- are often trained on datasets whose labels are only a small proportion of a larger label space. Such spaces are commonly equipped with a metric that relates the labels via distances between them. We propose a simple approach to exploit this information to adapt the trained model to reliably predict new classes -- or, in the case of zero-shot prediction, to improve its performance -- without any additional training. Our technique is a drop-in replacement of the standard prediction rule, swapping argmax with the Fr\'echet mean. We provide a comprehensive theoretical analysis for this approach, studying (i) learning-theoretic results trading off label space diameter, sample complexity, and model dimension, (ii) characterizations of the full range of scenarios in which it is possible to predict any unobserved class, and (iii) an optimal active learning-like next class selection procedure to obtain optimal training classes for when it is not possible to predict the entire range of unobserved classes. Empirically, using easily-available external metrics, our proposed approach, Loki, gains up to 29.7% relative improvement over SimCLR on ImageNet and scales to hundreds of thousands of classes. When no such metric is available, Loki can use self-derived metrics from class embeddings and obtains a 10.5% improvement on pretrained zero-shot models such as CLIP.