A Deep Dive into K-Means Clustering with PySpark MLlib

ByJordan Brown

K-means clustering is one of the oldest and most widely used unsupervised machine learning algorithms. Dating back to the 1950s, it was first proposed by Stuart Lloyd at Bell Labs for pulse-code modulation and later refined and popularized by J. B. MacQueen in 1967. Today, k-means remains a workhorse algorithm used across numerous domains and applications.

In this in-depth guide, we‘ll dissect the k-means algorithm from a variety of angles – mathematical, computational, practical, and conceptual. We‘ll walk through how to implement scalable k-means clustering in PySpark MLlib and share expert tips to get the most out of this powerful technique.

Overview of K-Means

At a high level, the goal of k-means is to partition a set of data points into k clusters such that points within the same cluster are more similar to each other than points in different clusters. Formally, k-means aims to minimize the sum of squared distances between each data point and its closest cluster centroid:

$$\arg\minC \sum{i=1}^{k} \sum_{\mathbf{x} \in C_i} \lVert \mathbf{x} – \boldsymbol\mu_i \rVert^2$$

where $C = {C_1, \ldots, C_k}$ is the set of clusters, $\mathbf{x}$ is a data point, and $\boldsymbol\mu_i$ is the centroid (mean) of cluster $C_i$.

The algorithm alternates between two steps until convergence:

  1. Assignment step: Assign each data point to the cluster with the closest centroid.
  2. Update step: Recompute the cluster centroids as the mean of all points assigned to that cluster.

Mathematically, the assignment step can be written as:

$$C_i^{(t)} = \left{ \mathbf{x}_p : \lVert \mathbf{x}_p – \boldsymbol\mu_i^{(t)} \rVert^2 \leq \lVert \mathbf{x}_p – \boldsymbol\mu_j^{(t)} \rVert^2\ \forall j, 1 \leq j \leq k \right}$$

And the update step:

$$\boldsymbol\mu_i^{(t+1)} = \frac{1}{|Ci^{(t)}|} \sum{\mathbf{x}_p \in C_i^{(t)}} \mathbf{x}_p$$

The algorithm is guaranteed to converge to a local optimum, but not necessarily the global optimum. The final clusters depend on the initial placement of centroids, which is usually done randomly.

Implementing K-Means in PySpark MLlib

Let‘s dive into the code and see how to run k-means at scale using PySpark MLlib. We‘ll use the classic Iris dataset for demonstration purposes.

First, start a SparkSession and load the required libraries:

from pyspark.sql import SparkSession
from pyspark.ml.clustering import KMeans
from pyspark.ml.evaluation import ClusteringEvaluator
from pyspark.ml.feature import VectorAssembler, StandardScaler

spark = SparkSession.builder.appName(‘iris_kmeans‘).getOrCreate()

Load the Iris dataset into a Spark DataFrame:

data = spark.read.csv("iris.csv", header=True, inferSchema=True)
data.show(5)
+-----------+----------+-----------+----------+-----------+
|sepal_length|sepal_width|petal_length|petal_width|species    |
+-----------+----------+-----------+----------+-----------+
|        5.1|       3.5|        1.4|       0.2|Iris-setosa|
|        4.9|       3.0|        1.4|       0.2|Iris-setosa|
|        4.7|       3.2|        1.3|       0.2|Iris-setosa|
|        4.6|       3.1|        1.5|       0.2|Iris-setosa|
|        5.0|       3.6|        1.4|       0.2|Iris-setosa|
+-----------+----------+-----------+----------+-----------+

We‘ll ignore the species column since we‘re treating this as an unsupervised problem. Next, we need to create a feature vector by combining the individual feature columns:

assembler = VectorAssembler(inputCols=data.columns[:-1], outputCol=‘features‘)
data = assembler.transform(data)

It‘s important to standardize the features to have zero mean and unit variance, since k-means is sensitive to scale:

scaler = StandardScaler(inputCol=‘features‘, outputCol=‘scaledFeatures‘)
scalerModel = scaler.fit(data)
data = scalerModel.transform(data)

Now we‘re ready to build the k-means model with k=3 to match the true number of species in the dataset:

kmeans = KMeans(k=3, seed=1, featuresCol=‘scaledFeatures‘, predictionCol=‘cluster‘)
model = kmeans.fit(data)

Let‘s evaluate the clustering quality using the silhouette score:

predictions = model.transform(data)
evaluator = ClusteringEvaluator(predictionCol=‘cluster‘, featuresCol=‘scaledFeatures‘)
silhouette = evaluator.evaluate(predictions)
print(f‘Silhouette with squared euclidean distance = {silhouette:.3f}‘)
Silhouette with squared euclidean distance = 0.553

A silhouette score of 0.553 indicates reasonable clustering quality, but there may be some overlap or less cohesion within clusters. We can visualize the clusters to investigate:

import matplotlib.pyplot as plt
import pandas as pd

df = predictions.select(‘scaledFeatures‘, ‘cluster‘).toPandas()
pd.plotting.scatter_matrix(df, c=df[‘cluster‘], figsize=(10,10), hist_kwds={‘bins‘:20})
plt.show()

Scatter matrix of Iris k-means clusters

From the scatter matrix, we see that the clusters do a fairly good job of separating the three Iris species, with some overlap between Iris-versicolor and Iris-virginica.

We can also examine the cluster centroids:

centers = pd.DataFrame(model.clusterCenters(), columns=data.columns[:-1])
print(centers)
sepal_length  sepal_width  petal_length  petal_width
0    -0.896180    1.032057     -1.341272    -1.313444
1     0.376393   -0.331490      0.529985     0.515151
2     1.122200   -1.476885      1.096662     1.053749

The centroids align with the known characteristics of the Iris species, with Iris-setosa having the smallest petal dimensions and Iris-virginica the largest.

Choosing the Number of Clusters

One of the key hyperparameters in k-means is k, the number of clusters to form. While there are mathematical ways to estimate the optimal k, it often requires domain knowledge and experimentation to determine the right number of clusters for the task at hand.

Some common heuristics include:

  • Elbow method: Plot the within-cluster sum of squared distances (WSSD) vs. k and look for an "elbow" point where increasing k further leads to diminishing reductions in WSSD.
  • Silhouette analysis: Calculate the average silhouette score for different values of k and choose the k that maximizes the score.
  • Gap statistic: Compare the WSSD to the expectation under a null reference distribution for each k.

In PySpark MLlib, we can use the KMeansModel.computeCost method to calculate the WSSD for the elbow method:

costs = []
for k in range(2,10):
    kmeans = KMeans(k=k, seed=1, featuresCol=‘scaledFeatures‘)
    model = kmeans.fit(data)
    costs.append(model.computeCost(data))

plt.plot(range(2,10), costs)
plt.xlabel(‘Number of clusters‘)
plt.ylabel(‘WSSD‘)
plt.title(‘K-means clustering elbow method‘)
plt.show()

K-means elbow curve for Iris data

Here we see an "elbow" around k=3, matching our prior knowledge of the Iris dataset.

Handling Outliers and High-Dimensional Data

K-means is sensitive to outliers, since it aims to minimize squared distances. A single outlier can dramatically shift a centroid and lead to skewed clusters. Some strategies to handle outliers include:

  • Remove outliers in a preprocessing step using standard deviation cutoffs, isolation forests, or expert rules
  • Use a variant like k-medoids that is more robust to outliers
  • Assign outliers to a separate "junk" cluster

High-dimensional data can also pose challenges for k-means, due to the curse of dimensionality. As the number of dimensions increases, Euclidean distances between points tend to become less meaningful. Some tips:

  • Carefully select relevant features and remove noise dimensions
  • Use dimensionality reduction techniques like PCA or t-SNE before clustering
  • Try a cosine distance metric instead of Euclidean distance
  • Use a variant like spherical k-means that operates in the unit hypersphere

Comparing K-Means with Alternatives

While k-means is a go-to clustering algorithm, it‘s not always the best choice. Let‘s compare it with two popular alternatives on the Iris dataset.

Gaussian Mixture Models (GMM) assume the data is generated from a mixture of Gaussian distributions and can handle clusters with different sizes and correlation structures. In MLlib:

from pyspark.ml.clustering import GaussianMixture

gmm = GaussianMixture(k=3, featuresCol=‘scaledFeatures‘, seed=1)
gmm_model = gmm.fit(data)
gmm_preds = gmm_model.transform(data)

gmm_silhouette = evaluator.evaluate(gmm_preds)
print(f‘Silhouette for GMM: {gmm_silhouette:.3f}‘)
Silhouette for GMM: 0.559

GMM performs slightly better than k-means here, likely because the Iris clusters have different variances that GMM can model.

DBSCAN is a density-based algorithm that finds clusters of high density separated by areas of low density. It does not require specifying the number of clusters upfront. In MLlib:

from pyspark.ml.clustering import DBSCAN

dbscan = DBSCAN(eps=0.5, minPoints=5, featuresCol=‘scaledFeatures‘)
dbscan_model = dbscan.fit(data)
dbscan_preds = dbscan_model.transform(data)

DBSCAN ends up partitioning the Iris data into two clusters – one containing primarily Iris-setosa and the other containing a mix of Iris-versicolor and Iris-virginica. It fails to separate the latter two species, likely because they have significant overlap in feature space.

Conclusion

K-means is a versatile and scalable clustering algorithm that is a core part of the data scientist‘s toolkit. In this guide, we covered:

  • The mathematical formulation and key steps of the k-means algorithm
  • How to implement k-means clustering in PySpark MLlib, including preprocessing and evaluation
  • Strategies for choosing the number of clusters, handling outliers, and clustering high-dimensional data
  • Comparisons of k-means with Gaussian mixture models and DBSCAN on the Iris dataset

We also highlighted some of the strengths and weaknesses of k-means:

Strengths Weaknesses
Simple and easy to implement Sensitive to initialization and outliers
Computationally efficient and scalable Assumes spherical clusters of equal size
Guaranteed to converge May get stuck in local optima
Tends to produce balanced clusters Need to specify k beforehand

Overall, k-means is a powerful tool when you need to partition a large dataset into distinct groups. By understanding its assumptions and limitations, you can wield it effectively and know when to reach for other methods.

Leave a comment and let me know your experience with k-means clustering in MLlib. Happy clustering!

References

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts