Weighted k-NN Zero-Shot Transfer

1. Background Knowledge

k-Nearest Neighbors (k-NN) is a non-parametric classification method where a query point is assigned the majority label among its $k$ closest neighbors in a feature space. In vision-language models (VLMs), embeddings from models like CLIP live in a high-dimensional space where semantic similarity corresponds to geometric proximity. Because these embeddings are typically L2-normalized, cosine similarity reduces to a simple dot product, making neighbor retrieval computationally efficient.

Tip-Adapter is a training-free inference-time adaptation technique that refines standard k-NN voting. Instead of treating all $k$ neighbors equally, it assigns each neighbor a weight based on its similarity to the query. The weight function $\exp(\beta \cdot (s - 1))$ (where $s$ is cosine similarity and $\beta > 0$) acts as a sharpening mechanism: neighbors with similarity close to 1 receive weight near 1, while less similar neighbors receive exponentially smaller weights. This effectively down-weights noisy or distant neighbors, improving classification accuracy without any gradient updates.

The problem also introduces a multi-level tie-breaking strategy, which is common in competitive programming and robust ML pipelines. When two labels receive identical total vote weights, the system must resolve ambiguity using secondary criteria (best individual neighbor similarity) and a final deterministic fallback (smaller label value).

2. Algorithm Approach

The solution follows a retrieve-and-aggregate pattern:

1. Compute similarities: Calculate the dot product between the query vector and every support embedding.
2. Select top-k: Identify the $k$ support examples with the highest similarity scores.
3. Weighted voting: For each of the top-$k$ neighbors, compute its vote weight using the Tip-Adapter formula and accumulate weights per label.
4. Resolve ties: Apply the three-level tie-breaking rule to select the final label.

This is not a graph or tree problem; it is a straightforward linear scan with sorting and aggregation.

3. Step-by-Step Strategy

• Compute similarity scores: Iterate over all $M$ support embeddings and compute dot(query, emb) for each. Store as a list of (similarity, index) tuples.
• Sort and slice: Sort the list in descending order of similarity and take the first $k$ entries.
• Accumulate votes: Initialize a dictionary mapping each label to its total weight. For each of the $k$ neighbors:
• Compute weight = exp(beta * (sim - 1)).
• Add weight to the label's total.
• Also track the maximum similarity seen for each label among its top-$k$ neighbors (for tie-breaking).
• Determine winner:
• Find the label with the highest total weight.
• If multiple labels share the max weight, compare their best individual neighbor similarities.
• If still tied, return the smallest label value.

import math

def knn_zero_shot(query, support_embs, support_labels, k, beta):
 # Step 1: Compute similarities
 sims = []
 for i, emb in enumerate(support_embs):
 dot = sum(q * e for q, e in zip(query, emb))
 sims.append((dot, i))
 
 # Step 2: Top-k
 sims.sort(key=lambda x: -x)
 top_k = sims[:k]
 
 # Step 3: Weighted voting
 total_weight = {}
 best_sim_per_label = {}
 for sim, idx in top_k:
 label = support_labels[idx]
 w = math.exp(beta * (sim - 1))
 total_weight[label] = total_weight.get(label, 0.0) + w
 if label not in best_sim_per_label or sim > best_sim_per_label[label]:
 best_sim_per_label[label] = sim
 
 # Step 4: Tie-breaking
 max_w = max(total_weight.values())
 candidates = [l for l, w in total_weight.items() if w == max_w]
 if len(candidates) == 1:
 return candidates
 
 max_best_sim = max(best_sim_per_label[l] for l in candidates)
 candidates = [l for l in candidates if best_sim_per_label[l] == max_best_sim]
 if len(candidates) == 1:
 return candidates
 
 return min(candidates)

4. Common Pitfalls

• Floating-point equality: Comparing floats with == for tie-breaking can fail due to precision errors. In practice, use a small epsilon (e.g., abs(a - b) < 1e-9) if the problem allows, though many competitive programming problems expect exact equality since inputs are controlled.
• Forgetting to track best similarity per label: The second tie-breaker requires knowing the highest similarity among a label's top-$k$ neighbors, not just the last one processed.
• Incorrect weight formula: The exponent is beta * (sim - 1), not beta * sim. Since $s \leq 1$, the exponent is non-positive, ensuring weights are in $(0, 1]$.
• Sorting stability: When similarities are equal, the order of indices in the top-$k$ slice may affect which neighbors are included. The problem does not specify a tie-breaker for neighbor selection, so ensure your sort is deterministic (e.g., sort by (-sim, index)).

5. Time & Space Complexity

• Time: $O(M \cdot D + M \log M)$, where $M$ is the number of support examples and $D$ is the embedding dimension. The dot product computation is $O(M \cdot D)$, and sorting $M$ elements is $O(M \log M)$.
• Space: $O(M)$ for storing similarity scores and $O(C)$ for the vote dictionaries, where $C$ is the number of distinct labels.