Katkat
Back

Optimize Nearest-Center Assignment Without Materializing NxKxD Distances

Mistral AIMachine Learning Engineer · Onsite
MediumMachine Learning EngineerOnsiteMachine Learning

Last updated

This problem tests the ability to optimize memory usage in distance computations by leveraging algebraic expansion to avoid creating an intermediate (N, K, D) tensor.

Problem Statement

In a clustering task, you are given a set of $N$ data points $X$ and $K$ cluster centers $C$. Both $X$ and $C$ are represented as matrices of shape $(N, D)$ and $(K, D)$ respectively. Your task is to assign each point to the nearest center by minimizing the Squared Euclidean Distance: $||x_i - c_j||^2$.

Naively, one might compute the $(N, K, D)$ difference tensor, which is prohibitively expensive for large $N$ and $D$. Instead, implement a function that computes the assignment indices using the expanded form of the distance: $||x_i - c_j||^2 = ||x_i||^2 + ||c_j||^2 - 2(x_i \cdot c_j^T)$.

Function Signature

def assign_nearest_centers(X, C):
    # X: np.ndarray of shape (N, D)
    # C: np.ndarray of shape (K, D)
    # returns: np.ndarray of shape (N,) containing indices of the nearest center

Constraints

  • $N$ up to $10^6$, $D$ up to $512$, $K$ up to $10^3$.
  • Use NumPy broadcasting.
  • Avoid any intermediate structures larger than $(N, K)$.

Example

Input: X = [[1, 2], [5, 5]] C = [[0, 0], [6, 6]]

  • Distance to centers for X[0]: $||[1,2]-[0,0]||^2 = 5$; $||[1,2]-[6,6]||^2 = 41$. Assign index 0.
  • Distance to centers for X[1]: $||[5,5]-[0,0]||^2 = 50$; $||[5,5]-[6,6]||^2 = 2$. Assign index 1. Output: [0, 1]

Hints

Hint 1

Recall the algebraic expansion of the squared Euclidean distance: the distance is the sum of the squared norms of the vectors minus twice their dot product.

Hint 2

Compute the norms squared for X as an (N, 1) vector and the norms squared for C as a (1, K) vector to utilize broadcasting.

Hint 3

The matrix multiplication `np.dot(X, C.T)` will yield an (N, K) matrix, which is memory-efficient and sufficient to find the minimum for each row.

Follow-up questions

  • How would you modify this implementation if N and D are so large that the (N, K) distance matrix itself exceeds memory capacity?
  • What are the numerical stability implications of the expanded form $||x||^2 + ||c||^2 - 2(x \cdot c^T)$ compared to calculating $(x-c)^2$ directly?
  • If we were performing this in a distributed setting (e.g., across multiple GPUs), how would you partition X and C to minimize communication overhead?