Implement 1NN Embeddings and Forward Pass

Problem Statement

In metric learning and retrieval-augmented systems, a common pattern is to project raw inputs into a learned latent space via a feed-forward layer and assign labels based on the nearest neighbor found in an embedding memory bank.

You are tasked with implementing a vectorized forward-pass and 1-Nearest Neighbor (1NN) classification module in NumPy.

Given a batch of query feature vectors X, a linear projection layer parameterized by weight matrix W and bias vector b, a reference embedding matrix R, and the corresponding reference labels y_ref:

1. Project the inputs to obtain query embeddings: E = X @ W + b.
2. Compute the pairwise Euclidean distance between each query embedding E[i] and every reference embedding R[j].
3. For each query, find the 1-nearest neighbor in R. In the event of a tie in distance, break the tie by choosing the candidate with the smallest reference index.
4. Return the predicted labels and the corresponding Euclidean distances.

Your solution must be fully vectorized across the query dimension N and reference dimension M without explicit Python loops (for or while).

Function Signature

import numpy as np

def knn_embedding_forward(
    X: np.ndarray,
    W: np.ndarray,
    b: np.ndarray,
    R: np.ndarray,
    y_ref: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """
    Computes 1NN classification in a projected embedding space.

    Parameters:
    - X: np.ndarray of shape (N, D_in), query input features.
    - W: np.ndarray of shape (D_in, D_out), projection weights.
    - b: np.ndarray of shape (D_out,), projection bias.
    - R: np.ndarray of shape (M, D_out), reference embedding bank.
    - y_ref: np.ndarray of shape (M,), labels for the reference embeddings.

    Returns:
    - y_pred: np.ndarray of shape (N,), predicted labels.
    - min_dist: np.ndarray of shape (N,), Euclidean distances to the nearest neighbor.
    """
    pass

Constraints

• 1 <= N <= 5,000
• 1 <= M <= 20,000
• 1 <= D_in <= 512
• 1 <= D_out <= 256
• All floating-point arrays are float32 or float64.
• Labels in y_ref can be integers or strings.

Examples

Example 1

X = np.array([[1.0, 2.0]], dtype=np.float32)        # N=1, D_in=2
W = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32)  # D_in=2, D_out=2
b = np.array([0.0, 0.0], dtype=np.float32)          # D_out=2

# E = [[1.0, 2.0]]
R = np.array([
    [1.0, 1.0],  # dist = sqrt((1-1)^2 + (2-1)^2) = 1.0
    [4.0, 6.0],  # dist = sqrt((1-4)^2 + (2-6)^2) = 5.0
], dtype=np.float32)
y_ref = np.array([10, 20])

y_pred, min_dist = knn_embedding_forward(X, W, b, R, y_ref)
# y_pred -> array([10])
# min_dist -> array([1.0], dtype=np.float32)

Example 2 (Tie Breaking)

X = np.array([[0.0]], dtype=np.float32)
W = np.array([[1.0]], dtype=np.float32)
b = np.array([0.0], dtype=np.float32)

# E = [[0.0]]
R = np.array([[-2.0], [2.0]], dtype=np.float32) # Both have dist = 2.0
y_ref = np.array([100, 200])

y_pred, min_dist = knn_embedding_forward(X, W, b, R, y_ref)
# Index 0 is selected due to smallest-index tie-breaking
# y_pred -> array([100])
# min_dist -> array([2.0], dtype=np.float32)

Edge Cases to Consider

• Numerical precision: Expanding ||u - v||^2 = ||u||^2 + ||v||^2 - 2(u . v) can produce tiny negative values (e.g., -1e-16) due to floating-point rounding. Clamping distances to 0.0 before computing the square root prevents NaNs.
• Memory consumption: Direct 3D broadcasting (N, 1, D_out) - (1, M, D_out) allocates O(N * M * D_out) memory, which may cause out-of-memory errors for large N and M. Using the matrix expansion identity requires only O(N * M) intermediate space.
• Tie-breaking behavior: np.argmin inherently preserves the lowest index on duplicate minimum values along the specified axis.