Vectorize One-Nearest-Neighbor and Express It as a Neural Forward Pass

Vectorize One-Nearest-Neighbor and Express It as a Neural Forward Pass

Nearest-neighbor classification is one of the foundational non-parametric learning algorithms. In deep learning architectures, metric learning and prototype-based classifiers (e.g., Learning Vector Quantization, ProtoNets) bridge nearest-neighbor principles with neural network operations.

This problem has two parts:

Part 1: Vectorized 1-NN Inference

Implement a function predict_1nn(X_query, X_ref, y_ref) in NumPy/PyTorch-style pseudocode that computes predictions for a batch of query vectors against a reference set using squared Euclidean distance.

• X_query: array of shape (M, D) representing M query vectors of dimension D.
• X_ref: array of shape (N, D) representing N reference vectors of dimension D.
• y_ref: 1D array of shape (N,) containing the label for each reference vector.

Constraint: Do not use explicit Python loops over samples (M or N). The implementation must be fully vectorized.

Part 2: Nearest Prototype as a Linear Layer

Assume we have K class prototypes C of shape (K, D), where each row c_k is the representative vector for class k in {0, 1, ..., K-1}. An unnormalized query x of shape (D,) is assigned to class k* = argmin_k ||x - c_k||^2.

Show that this decision rule can be expressed as a standard linear feed-forward layer followed by an argmax:

logits = x @ W.T + b
pred = argmax(logits)

1. Derive the mathematical expressions for weight matrix W (shape (K, D)) and bias vector b (shape (K,)) in terms of the prototype matrix C.
2. Explain why the data-dependent norm ||x||^2 can be omitted from the logits.

Example

# Reference dataset (N=3, D=2)
X_ref = np.array([
    [0.0, 0.0],
    [1.0, 1.0],
    [2.0, 2.0]
])
y_ref = np.array([0, 1, 1])

# Query points (M=2, D=2)
X_query = np.array([
    [0.1, 0.2],
    [1.8, 1.9]
])

# Nearest ref for X_query[0] is X_ref[0] -> label 0
# Nearest ref for X_query[1] is X_ref[2] -> label 1
# Output: array([0, 1])