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.

翻译：机器学习模型——包括著名的零样本模型——通常基于数据集进行训练，而数据集的标签仅涵盖更大标签空间中的一小部分。此类标签空间通常配备有度量，通过标签之间的距离建立关联。我们提出一种简单方法，利用这一信息适配训练好的模型，使其能够可靠地预测新类别——或在零样本预测场景下提升性能——且无需额外训练。该技术可直接替换标准预测规则，将"argmax"替换为"Fréchet均值"。我们对该方法进行了全面的理论分析，研究了：（i）权衡标签空间直径、样本复杂度和模型维度的学习理论结果；（ii）预测任意未观测类别可行性的完整场景特征描述；（iii）当无法预测全部未观测类别时，用于获取最优训练类别的类主动学习式最优下一类别选择方案。实验表明，利用易获取的外部度量，我们提出的方法Loki在ImageNet上相较于SimCLR获得最高29.7%的相对性能提升，并可扩展至数十万类别。当不存在此类外部度量时，Loki能通过类别嵌入自推导度量，在CLIP等预训练零样本模型上获得10.5%的性能提升。