Katkat
Back

Represent k-means as an MLP

SealthMachine Learning Engineer · Onsite
EasyMachine Learning EngineerOnsiteMachine Learning

Last updated

This question tests the conceptual mapping between the distance-based assignment step of k-means clustering and a single-layer neural network architecture.

Problem Statement

Given a set of $k$ cluster centroids $C = {c_1, c_2, ..., c_k}$ in a $d$-dimensional space, the assignment step of the k-means algorithm assigns an input vector $x$ to the centroid $c_i$ that minimizes the squared Euclidean distance: $||x - c_i||^2$.

Your task is to design a single-layer neural network (an MLP with no hidden layers) that computes the soft assignment of an input $x$ to these centroids. Specifically, define the weights $W$ and bias $b$ of a linear layer followed by a Softmax activation such that the output represents a probability distribution over the $k$ clusters, where the probability mass concentrates on the nearest centroid as a temperature parameter $T$ approaches 0.

Requirements:

  1. Show how to expand the squared Euclidean distance $||x - c_i||^2$ to isolate the terms involving $x$ and $c_i$.
  2. Determine the weight matrix $W$ and bias vector $b$ such that the input to the Softmax function is proportional to $-||x - c_i||^2 / (2T)$.
  3. Explain the relationship between the weight matrix $W$ and the centroids $C$.

Example:

If $x = [1, 2]$ and $c_1 = [0, 0]$, $c_2 = [2, 2]$, the network should produce higher values for the index closer to $x$ when passed through the Softmax function with a small temperature $T$.

Hints

Hint 1

Expand the quadratic form ||x - c_i||^2 = ||x||^2 - 2(x^T * c_i) + ||c_i||^2.

Hint 2

Observe that ||x||^2 is constant across all centroids and does not affect the ranking of the Softmax output.

Hint 3

Set the weight matrix $W$ as a matrix where each row is a centroid $c_i$, and define the bias $b$ based on the norm of each centroid.

Follow-up questions

  • How does increasing the temperature parameter T affect the output distribution of the Softmax layer?
  • If we use a ReLU activation instead of Softmax, what happens to the interpretability of the cluster assignments?
  • Can this approach be extended to perform online learning of the centroids using backpropagation?