Geometry-Aware Adaptation for Pretrained Models

Machine learning models often struggle to reliably predict new classes when trained on datasets where labels are only a small proportion of a larger label space, and this work proposes LOKI, an approach that exploits the geometric information (metric structure) between labels to adapt pretrained models without additional training.

The gist: LOKI is a simple adaptor for pretrained models that replaces the standard prediction rule with the Fréchet mean to exploit label space geometry for reliable prediction of unobserved classes.

How it works

LOKI functions by translating the soft outputs (probability vectors) produced by any supervised learning model into a more general model capable of exploiting geometric information. Specifically, LOKI uses the Fréchet mean estimator, defined as finding the label in the space that minimizes a weighted sum of squared distances to observed class embeddings:

yˆ ∈ mλ(Pyx) = arg min y∈Y X K i=1 Pλixd 2(y, λi).

This approach generalizes the standard prediction rule, which is typically arg max over the model's per-class probabilities. LOKI can be expressed as a fixed linear transformation of a model’s outputs. This transformation involves multiplying the model output vector by a fixed matrix D, where D is the matrix of negative squared distances between observed classes and the rest of the label space: yˆ ∈ arg maxj∈[N] (DPyx)j, where Dj,i = −d 2(yj, λi).

Theoretical Analysis and Results

The paper provides a comprehensive theoretical analysis covering three main areas:

1. Learning-theoretic results trading off label space diameter, sample complexity, and model dimension.

2. Characterizations of the full range of scenarios in which it is possible to predict any unobserved class.

3. An optimal active learning-like next class selection procedure to obtain optimal training classes when prediction of all unobserved classes is not possible.

Empirically, the approach demonstrates significant gains:

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 [external], LOKI can use self-derived metrics from class embeddings and obtains a 10.5% improvement on pretrained zero-shot models such as CLIP.

Key Theoretical Findings

The theoretical results address specific challenges:

1. Sample Complexity: A result is provided relating the sample complexity of estimating target y∗ from the test distribution Dtest with prediction yˆ∗ to parameters like dimensionality (d), number of observed classes (K), and sample size (n).

2. Optimal Label Subspaces: The work characterizes the minimum number (and structure) of subsets that enable reliable prediction for various relational structures, including tree graphs, grid graphs, and the complete graph. For phylogenetic trees, it shows that the leaves of any tree yield the minimum locus cover.

3. Locus Computation: Algorithms are provided to compute the set of all classes in the locus. The paper notes that for all considered graph types (tree and grid), any locus is pairwise decomposable, allowing for a polynomial-time algorithm to compute it efficiently.

Active Learning Strategy

When it is not possible to predict the entire range of unobserved classes, LOKI enables an active learning-like approach. The goal is to select the next observed class to maximize the total number of classes that can be predicted, which corresponds to maximizing the size of the locus, Π(Λ).

For tree graphs, Theorem 4.8 provides a specific optimization problem for selecting the next class:

arg max y ∈ Y∖Π(Λ) d(y, b) s.t. b ∈ ∂inT′ and Γ(y, b)∖ b ⊆ Y∖Π(Λ).

This procedure can be computed in polynomial time by searching over pairs of vertices, leading to an efficient active learning-based strategy to maximize the size of the locus for trees. This active approach was shown to yield a larger locus than the passive baseline and resulted in improved error on ImageNet.

Generalization and Calibration

LOKI demonstrates that it can improve zero-shot models even when no external metric is available, using internal class embedding distances. Furthermore, the paper shows that calibration can be improved through Softmax temperature scaling. In experiments with CLIP on CIFAR-100, the optimal Softmax temperature for LOKI is both close to the default temperature used by CLIP and to the optimally-calibrated temperature. The analysis confirms that tuning this parameter can lead to "improved accuracy with CLIP, even when no external metric space is available.

Here are the specific improvements to AI systems based on the proposed LOKI (Learning On Knowledge-Informed Inference) framework, and what these improved systems can achieve:

The LOKI framework provides a method for adapting pretrained models (like CLIP or SimCLR) to predict classes outside their training distribution by leveraging the geometric structure (metric space) of the label space. The core improvement is replacing the standard prediction rule,

ArgMax, with a Fréchet Mean estimator.

Here are the specific improvements and capabilities:

1. A new prediction mechanism that utilizes class relationships rather than just per-class confidence scores.

2. The ability to reliably predict unobserved or zero-shot classes without requiring additional fine-tuning or training data for those specific classes.

3. A method to quantify the geometric structure of label spaces (e.g., using phylogenetic trees, grid graphs, or complete graphs).

Specific Capabilities and Applications:

1. Predicting a Vast Number of Unobserved Classes:

The system can predict classes that were never seen during pretraining, provided the label space has a known metric structure (like WordNet or CIFAR-100 hierarchies).

2. Improving Zero-Shot Performance Without Fine-Tuning:

Pretrained models like CLIP can achieve up to 29.7% relative improvement in performance on ImageNet when adapted with LOKI, even when no external metric is explicitly provided, by deriving internal metrics from text/image embeddings.

3. Optimized Class Selection for Maximum Reach (Active Learning):

The system can employ an active learning-like strategy (Theorem 4.8) to intelligently select the next set of labels to observe from a large space, specifically targeting selections that maximize the locus (the set of classes reachable via LOKI). This prevents random sampling and ensures that observed data is used optimally to unlock the largest possible prediction frontier.

4. Robustness Across Different Label Structures:

The framework allows for explicit modeling of different label spaces (trees, grids, complete graphs). The system can be specifically tuned (e.g., by using a WordNet metric) to perform optimally on structured data like hierarchical classification or spatial data like maps.

5. Improved Calibration and Uncertainty Quantification:

By tuning the Softmax temperature parameter within the LOKI framework, the system can achieve optimal calibration (minimizing Expected Calibration Error, ECE), leading to more reliable confidence scores when making predictions, even in zero-shot settings without an external metric.

In summary, this improved AI system transforms a standard classification model into a geometrically aware predictor capable of expanding its knowledge base by seeing the underlying relationships between classes.

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échet 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.