Vectorize One-Nearest-Neighbor and Express It as a Neural Forward Pass
Quick Overview
Write NumPy-style pseudocode for vectorized one-nearest-neighbor inference, then express the same nearest-prototype decision as a neural-network forward pass.
Vectorize One-Nearest-Neighbor and Express It as a Neural Forward Pass
Company: OpenAI
Role: Machine Learning Engineer
Category: Machine Learning
Difficulty: medium
Interview Round: Technical Screen
## Prompt
You are given a matrix of training vectors `X_train` with one label per row and a matrix of query vectors `X_query`. Explain and write NumPy-style pseudocode for vectorized one-nearest-neighbor inference without a Python loop over training examples or queries. Then show how the same nearest-prototype decision can be represented as a neural-network forward pass.
### Constraints & Assumptions
- All vectors have the same feature dimension.
- Ties choose the smallest training-row index.
- The L2 portion may use squared Euclidean distance because square root does not change the nearest index.
- The implementation must discuss memory when the full query-by-training distance matrix is too large.
### Clarifying Questions to Ask
- Are labels arbitrary values or contiguous class IDs?
- Is the requirement nearest training example or nearest class centroid?
- Must the forward pass reproduce the nearest-example index or only the final label?
- What batch sizes and numeric precision are expected?
### Part 1 — Vectorized L2 One-Nearest Neighbor
Derive a broadcast or matrix-multiplication formulation for all pairwise squared distances, select one nearest training row per query, and map those indices to labels.
```hint Expand the square
Use `||q - x||^2 = ||q||^2 + ||x||^2 - 2 q dot x` to avoid materializing a three-dimensional difference tensor.
```
#### What This Part Should Cover
- Correct matrix shapes and an explicit tie rule.
- Numerical handling for small negative values caused by floating-point roundoff.
- Query or training blocking when the full distance matrix exceeds memory.
- The distinction between vectorized computation and reduced asymptotic complexity.
### Part 2 — Express L2 Nearest-Prototype Selection as a Forward Pass
Choose linear weights and biases whose largest logit identifies the nearest stored prototype. Explain why applying softmax does not change that selected index and how the prototype index maps to a label.
#### What This Part Should Cover
- A derivation from squared L2 distance, not a guessed parameter matrix.
- The query-only term that can be removed because it is identical across prototypes.
- Clear handling of multiple prototypes that share one class label.
- Why this is an exact representation of inference, not a claim that the parameters were learned.
### Part 3 — Replace L2 with L1 Distance
Explain whether one affine layer can exactly compute negative L1 distances. If not, construct a small piecewise-linear network that does so.
```hint Absolute value needs a nonlinearity
For one scalar difference `z`, `abs(z) = ReLU(z) + ReLU(-z)`.
```
#### What This Part Should Cover
- A correct explanation of why an affine map alone cannot represent absolute value globally.
- Hidden units for positive and negative coordinate differences relative to each prototype.
- A summation layer that creates one negative L1-distance score per prototype.
- The same deterministic nearest-index and label-mapping behavior as the L2 case.
### What a Strong Answer Covers
- Correct, shape-aware pseudocode and mathematical derivations for all three parts.
- Numeric stability, batching, and memory trade-offs rather than relying on “vectorized” as a performance guarantee.
- A distinction between nearest example, nearest prototype, and aggregated class probability.
### Follow-up Questions
1. How would cosine distance change the preprocessing and forward pass?
2. When can approximate nearest-neighbor indexing outperform the dense matrix method?
3. How would you make the model differentiable if nearest-neighbor labels were used during training?
Quick Answer: Write NumPy-style pseudocode for vectorized one-nearest-neighbor inference, then express the same nearest-prototype decision as a neural-network forward pass.
You are given a matrix of training vectors X_train with one label per row and a matrix of query vectors X_query. Explain and write NumPy-style pseudocode for vectorized one-nearest-neighbor inference without a Python loop over training examples or queries. Then show how the same nearest-prototype decision can be represented as a neural-network forward pass.
Constraints & Assumptions
All vectors have the same feature dimension.
Ties choose the smallest training-row index.
The L2 portion may use squared Euclidean distance because square root does not change the nearest index.
The implementation must discuss memory when the full query-by-training distance matrix is too large.
Clarifying Questions to Ask Guidance
Are labels arbitrary values or contiguous class IDs?
Is the requirement nearest training example or nearest class centroid?
Must the forward pass reproduce the nearest-example index or only the final label?
What batch sizes and numeric precision are expected?
Part 1 — Vectorized L2 One-Nearest Neighbor
Derive a broadcast or matrix-multiplication formulation for all pairwise squared distances, select one nearest training row per query, and map those indices to labels.
What This Part Should Cover Guidance
Correct matrix shapes and an explicit tie rule.
Numerical handling for small negative values caused by floating-point roundoff.
Query or training blocking when the full distance matrix exceeds memory.
The distinction between vectorized computation and reduced asymptotic complexity.
Part 2 — Express L2 Nearest-Prototype Selection as a Forward Pass
Choose linear weights and biases whose largest logit identifies the nearest stored prototype. Explain why applying softmax does not change that selected index and how the prototype index maps to a label.
What This Part Should Cover Guidance
A derivation from squared L2 distance, not a guessed parameter matrix.
The query-only term that can be removed because it is identical across prototypes.
Clear handling of multiple prototypes that share one class label.
Why this is an exact representation of inference, not a claim that the parameters were learned.
Part 3 — Replace L2 with L1 Distance
Explain whether one affine layer can exactly compute negative L1 distances. If not, construct a small piecewise-linear network that does so.
What This Part Should Cover Guidance
A correct explanation of why an affine map alone cannot represent absolute value globally.
Hidden units for positive and negative coordinate differences relative to each prototype.
A summation layer that creates one negative L1-distance score per prototype.
The same deterministic nearest-index and label-mapping behavior as the L2 case.
What a Strong Answer Covers Guidance
Correct, shape-aware pseudocode and mathematical derivations for all three parts.
Numeric stability, batching, and memory trade-offs rather than relying on “vectorized” as a performance guarantee.
A distinction between nearest example, nearest prototype, and aggregated class probability.
Follow-up Questions Guidance
How would cosine distance change the preprocessing and forward pass?
When can approximate nearest-neighbor indexing outperform the dense matrix method?
How would you make the model differentiable if nearest-neighbor labels were used during training?