Weighted k-NN Zero-Shot Transfer

Problem Statement

A strong training-free baseline (Tip-Adapter, and CLIP's own linear-probe cousins) classifies a query image by a similarity-weighted vote of its nearest neighbors in a labeled support set of embeddings. Implement it.

Background

Given a query embedding q and a support set of (embedding, label) pairs, score every support example by cosine similarity to q, take the top-k most similar, and let each vote for its label with a weight of exp(beta * (sim - 1)) (the sharpening used by Tip-Adapter; beta > 0). The predicted class is the label with the greatest total weight.

Tie-breaking, in order:

1. Highest total vote weight.
2. If tied, the label whose single best (highest-similarity) neighbor is more similar.
3. If still tied, the smaller label value.

Assume all embeddings are already L2-normalized, so cosine similarity is the dot product.

Your Task

Implement:

def knn_zero_shot(query, support_embs, support_labels, k, beta):

Return the predicted label (an int).

Input Format

• query: list of floats (unit vector).
• support_embs: M x D nested list (unit vectors).
• support_labels: list of M int labels.
• k (int), beta (float).

Output Format

• A single int label.

Sample

q = [1.0, 0.0]
se = [[1.0, 0.0], [0.9, 0.44], [0.0, 1.0]]
sl = [0, 0, 1]
print(knn_zero_shot(q, se, sl, 2, 5.0))

Output:

0