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\'echet均值。我们对该方法进行了全面理论分析，包括：(i) 权衡标签空间直径、样本复杂度和模型维度的学习理论结果；(ii) 可预测任意未观测类别的完整场景刻画；(iii) 当无法预测全部未观测类别时，获取最优训练类别的主动学习式下一类别选择流程。实验表明，利用易获取的外部度量，我们提出的方法Loki在ImageNet上相较SimCLR获得最高29.7%的相对提升，并可扩展至数十万类别。当无此类度量时，Loki可利用类别嵌入自推度量，在CLIP等预训练零样本模型上取得10.5%的性能提升。