xbe.at

Machine learning / Classification and clustering

KMeans Clustering with Approximate Nearest Neighbors in Python

10 min read, 15 sections

Contents
  1. Introduction to KMeans Clustering with Approximate Nearest Neighbors
  2. Understanding Approximate Nearest Neighbors (ANN)
  3. Implementing KMeans with ANN: The Algorithm
  4. Advantages of KMeans with ANN
  5. Choosing the Right ANN Algorithm
  6. Handling High-Dimensional Data
  7. Balancing Accuracy and Speed
  8. Real-Life Example: Image Segmentation
  9. Real-Life Example: Document Clustering
  10. Implementing KMeans with ANN using FAISS
  11. Evaluating Clustering Quality
  12. Handling Outliers and Noise
  13. Scalability and Parallel Processing
  14. Visualization Techniques for High-Dimensional Clusters
  15. Additional Resources

Introduction to KMeans Clustering with Approximate Nearest Neighbors

KMeans clustering is a popular unsupervised learning algorithm used for partitioning data into K distinct clusters. When dealing with large datasets, traditional KMeans can be computationally expensive. This is where Approximate Nearest Neighbors (ANN) comes into play, significantly speeding up the clustering process while maintaining accuracy.

import numpy as np
from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors

# Generate sample data
np.random.seed(42)
X = np.random.rand(1000, 2)

# Initialize KMeans
kmeans = KMeans(n_clusters=5, random_state=42)
kmeans.fit(X)

# Plot the results
import matplotlib.pyplot as plt
plt.scatter(X[:, 0], X[:, 1], c=kmeans.labels_, cmap='viridis')
plt.title('KMeans Clustering')
plt.show()

Understanding Approximate Nearest Neighbors (ANN)

Approximate Nearest Neighbors is a technique used to find the nearest neighbors of a given point in high-dimensional spaces. Unlike exact nearest neighbor search, ANN trades off some accuracy for improved speed, making it suitable for large-scale clustering tasks. ANN algorithms use various data structures and techniques to efficiently organize and search through the data points.

from sklearn.neighbors import NearestNeighbors
import numpy as np

# Generate sample data
np.random.seed(42)
X = np.random.rand(1000, 10)  # 1000 points in 10-dimensional space

# Create ANN model
ann = NearestNeighbors(n_neighbors=5, algorithm='ball_tree')
ann.fit(X)

# Find nearest neighbors for a query point
query_point = np.random.rand(1, 10)
distances, indices = ann.kneighbors(query_point)

print(f"Indices of 5 nearest neighbors: {indices[0]}")
print(f"Distances to 5 nearest neighbors: {distances[0]}")

Implementing KMeans with ANN: The Algorithm

To implement KMeans with ANN, we modify the standard KMeans algorithm by replacing the exact nearest centroid search with an approximate search. This approach significantly reduces the computational complexity, especially for large datasets with high dimensionality.

import numpy as np
from sklearn.neighbors import NearestNeighbors

def kmeans_ann(X, n_clusters, max_iter=100):
    # Initialize centroids randomly
    centroids = X[np.random.choice(X.shape[0], n_clusters, replace=False)]
    
    for _ in range(max_iter):
        # Build ANN index for current centroids
        ann = NearestNeighbors(n_neighbors=1, algorithm='ball_tree')
        ann.fit(centroids)
        
        # Assign points to nearest centroids
        distances, labels = ann.kneighbors(X)
        labels = labels.flatten()
        
        # Update centroids
        new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(n_clusters)])
        
        # Check for convergence
        if np.all(centroids == new_centroids):
            break
        
        centroids = new_centroids
    
    return labels, centroids

# Example usage
X = np.random.rand(1000, 2)
labels, centroids = kmeans_ann(X, n_clusters=5)

Advantages of KMeans with ANN

KMeans clustering with Approximate Nearest Neighbors offers several benefits over traditional KMeans. The primary advantage is improved scalability, allowing the algorithm to handle much larger datasets efficiently. This approach reduces the time complexity of nearest neighbor search from O(n) to approximately O(log n) in many cases, where n is the number of data points.

import time
import numpy as np
from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors

def benchmark_kmeans(X, n_clusters):
    start_time = time.time()
    kmeans = KMeans(n_clusters=n_clusters, random_state=42)
    kmeans.fit(X)
    end_time = time.time()
    return end_time - start_time

def benchmark_kmeans_ann(X, n_clusters):
    start_time = time.time()
    labels, centroids = kmeans_ann(X, n_clusters)
    end_time = time.time()
    return end_time - start_time

# Generate large dataset
X = np.random.rand(100000, 10)

# Benchmark traditional KMeans
traditional_time = benchmark_kmeans(X, n_clusters=10)

# Benchmark KMeans with ANN
ann_time = benchmark_kmeans_ann(X, n_clusters=10)

print(f"Traditional KMeans time: {traditional_time:.2f} seconds")
print(f"KMeans with ANN time: {ann_time:.2f} seconds")
print(f"Speedup: {traditional_time / ann_time:.2f}x")

Choosing the Right ANN Algorithm

Several ANN algorithms are available, each with its strengths and weaknesses. Common choices include K-d trees, Ball trees, and Locality-Sensitive Hashing (LSH). The choice of algorithm depends on factors such as data dimensionality, dataset size, and the desired trade-off between speed and accuracy.

from sklearn.neighbors import NearestNeighbors
import numpy as np
import time

def benchmark_ann(X, algorithm):
    start_time = time.time()
    ann = NearestNeighbors(n_neighbors=5, algorithm=algorithm)
    ann.fit(X)
    query_point = np.random.rand(1, X.shape[1])
    distances, indices = ann.kneighbors(query_point)
    end_time = time.time()
    return end_time - start_time

# Generate dataset
X = np.random.rand(10000, 10)

algorithms = ['ball_tree', 'kd_tree', 'brute']
for algo in algorithms:
    time_taken = benchmark_ann(X, algo)
    print(f"{algo} algorithm time: {time_taken:.4f} seconds")

Handling High-Dimensional Data

As the dimensionality of data increases, the performance of traditional ANN algorithms may degrade due to the "curse of dimensionality". To address this, specialized techniques like Product Quantization (PQ) or Hierarchical Navigable Small World (HNSW) graphs can be employed.

import numpy as np
from sklearn.neighbors import NearestNeighbors
import time

def benchmark_high_dim(n_samples, n_features):
    X = np.random.rand(n_samples, n_features)
    
    start_time = time.time()
    ann = NearestNeighbors(n_neighbors=5, algorithm='ball_tree')
    ann.fit(X)
    query_point = np.random.rand(1, n_features)
    distances, indices = ann.kneighbors(query_point)
    end_time = time.time()
    
    return end_time - start_time

dimensions = [10, 100, 1000]
for dim in dimensions:
    time_taken = benchmark_high_dim(10000, dim)
    print(f"Time for {dim} dimensions: {time_taken:.4f} seconds")

Balancing Accuracy and Speed

When implementing KMeans with ANN, it's crucial to find the right balance between clustering accuracy and computational speed. This trade-off can be adjusted by tuning parameters such as the number of neighbors considered or the approximation factor in the ANN algorithm.

import numpy as np
from sklearn.neighbors import NearestNeighbors
from sklearn.metrics import pairwise_distances

def accuracy_speed_tradeoff(X, n_neighbors_range):
    true_distances = pairwise_distances(X)[0]
    true_neighbors = np.argsort(true_distances)[1:n_neighbors_range[-1]+1]
    
    results = []
    for n_neighbors in n_neighbors_range:
        start_time = time.time()
        ann = NearestNeighbors(n_neighbors=n_neighbors, algorithm='ball_tree')
        ann.fit(X)
        distances, indices = ann.kneighbors(X[0].reshape(1, -1))
        end_time = time.time()
        
        accuracy = np.mean(np.isin(indices[0], true_neighbors[:n_neighbors]))
        time_taken = end_time - start_time
        results.append((n_neighbors, accuracy, time_taken))
    
    return results

X = np.random.rand(1000, 10)
n_neighbors_range = [5, 10, 20, 50, 100]
results = accuracy_speed_tradeoff(X, n_neighbors_range)

for n_neighbors, accuracy, time_taken in results:
    print(f"n_neighbors: {n_neighbors}, Accuracy: {accuracy:.2f}, Time: {time_taken:.4f} seconds")

Real-Life Example: Image Segmentation

KMeans clustering with ANN can be applied to image segmentation tasks, where the goal is to partition an image into multiple segments or objects. This technique is particularly useful for processing large, high-resolution images efficiently.

import numpy as np
from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors
import matplotlib.pyplot as plt
from PIL import Image

# Load and preprocess the image
image = np.array(Image.open('sample_image.jpg'))
pixels = image.reshape(-1, 3)

# Perform KMeans clustering with ANN
n_clusters = 5
ann = NearestNeighbors(n_neighbors=1, algorithm='ball_tree')
kmeans = KMeans(n_clusters=n_clusters, random_state=42)
kmeans.fit(pixels)

# Assign labels using ANN
ann.fit(kmeans.cluster_centers_)
labels = ann.kneighbors(pixels)[1].flatten()

# Reshape the result back to the original image shape
segmented_image = labels.reshape(image.shape[:2])

# Display the results
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
ax1.imshow(image)
ax1.set_title('Original Image')
ax2.imshow(segmented_image, cmap='viridis')
ax2.set_title('Segmented Image')
plt.show()

Real-Life Example: Document Clustering

KMeans with ANN can be used for efficient document clustering, helping to organize large collections of text documents into thematic groups. This application is particularly useful in information retrieval and text mining tasks.

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.neighbors import NearestNeighbors
import numpy as np

# Sample documents
documents = [
    "Machine learning is a subset of artificial intelligence",
    "Natural language processing deals with text and speech",
    "Computer vision focuses on image and video analysis",
    "Deep learning uses neural networks with multiple layers",
    "Reinforcement learning involves agents and environments"
]

# Create TF-IDF vectors
vectorizer = TfidfVectorizer()
X = vectorizer.fit_transform(documents)

# Perform KMeans clustering with ANN
n_clusters = 3
ann = NearestNeighbors(n_neighbors=1, algorithm='ball_tree')
kmeans = KMeans(n_clusters=n_clusters, random_state=42)
kmeans.fit(X)

# Assign labels using ANN
ann.fit(kmeans.cluster_centers_)
labels = ann.kneighbors(X)[1].flatten()

# Print results
for doc, label in zip(documents, labels):
    print(f"Cluster {label}: {doc[:50]}...")

Implementing KMeans with ANN using FAISS

FAISS (Facebook AI Similarity Search) is a library that provides efficient implementations of ANN algorithms, particularly suited for high-dimensional data. Let's see how to use FAISS to implement KMeans with ANN.

import numpy as np
import faiss

def kmeans_faiss(X, n_clusters, max_iter=100):
    d = X.shape[1]
    kmeans = faiss.Kmeans(d, n_clusters, niter=max_iter, verbose=True)
    kmeans.train(X.astype(np.float32))
    
    # Get cluster assignments and centroids
    _, labels = kmeans.index.search(X.astype(np.float32), 1)
    centroids = kmeans.centroids
    
    return labels.flatten(), centroids

# Generate sample data
np.random.seed(42)
X = np.random.rand(10000, 128).astype(np.float32)

# Perform KMeans clustering with FAISS
n_clusters = 10
labels, centroids = kmeans_faiss(X, n_clusters)

print(f"Number of clusters: {n_clusters}")
print(f"Shape of centroids: {centroids.shape}")
print(f"Cluster assignments: {labels[:10]}...")

Evaluating Clustering Quality

When using KMeans with ANN, it's important to evaluate the quality of the resulting clusters. Common metrics include the silhouette score and the Davies-Bouldin index. These metrics help assess how well-separated the clusters are and whether the approximation introduced by ANN significantly affects the clustering quality.

from sklearn.metrics import silhouette_score, davies_bouldin_score
import numpy as np
from sklearn.cluster import KMeans
from sklearn.neighbors import NearestNeighbors

def evaluate_clustering(X, labels):
    silhouette = silhouette_score(X, labels)
    davies_bouldin = davies_bouldin_score(X, labels)
    return silhouette, davies_bouldin

# Generate sample data
np.random.seed(42)
X = np.random.rand(1000, 10)

# Perform traditional KMeans
kmeans = KMeans(n_clusters=5, random_state=42)
kmeans_labels = kmeans.fit_predict(X)

# Perform KMeans with ANN
ann = NearestNeighbors(n_neighbors=1, algorithm='ball_tree')
ann.fit(kmeans.cluster_centers_)
ann_labels = ann.kneighbors(X)[1].flatten()

# Evaluate both methods
kmeans_scores = evaluate_clustering(X, kmeans_labels)
ann_scores = evaluate_clustering(X, ann_labels)

print("Traditional KMeans:")
print(f"Silhouette Score: {kmeans_scores[0]:.4f}")
print(f"Davies-Bouldin Index: {kmeans_scores[1]:.4f}")

print("\nKMeans with ANN:")
print(f"Silhouette Score: {ann_scores[0]:.4f}")
print(f"Davies-Bouldin Index: {ann_scores[1]:.4f}")

Handling Outliers and Noise

KMeans clustering, including when used with ANN, can be sensitive to outliers and noise in the data. To address this issue, we can implement a robust version of KMeans that is less affected by anomalous data points.

import numpy as np
from sklearn.neighbors import NearestNeighbors

def robust_kmeans_ann(X, n_clusters, max_iter=100, outlier_threshold=2.0):
    centroids = X[np.random.choice(X.shape[0], n_clusters, replace=False)]
    
    for _ in range(max_iter):
        ann = NearestNeighbors(n_neighbors=1, algorithm='ball_tree')
        ann.fit(centroids)
        
        distances, labels = ann.kneighbors(X)
        labels = labels.flatten()
        
        # Identify outliers
        outlier_mask = distances.flatten() > outlier_threshold * np.mean(distances)
        
        # Update centroids excluding outliers
        new_centroids = np.array([
            X[~outlier_mask & (labels == k)].mean(axis=0) 
            for k in range(n_clusters)
        ])
        
        if np.allclose(centroids, new_centroids):
            break
        
        centroids = new_centroids
    
    return labels, centroids

# Example usage
X = np.random.rand(1000, 2)
X[0] = [10, 10]  # Add an outlier
labels, centroids = robust_kmeans_ann(X, n_clusters=5)

print("Cluster centroids:")
print(centroids)

Scalability and Parallel Processing

As datasets grow larger, it becomes crucial to leverage parallel processing techniques to maintain efficient clustering. We can implement a parallel version of KMeans with ANN using Python's multiprocessing module to distribute the workload across multiple CPU cores.

import numpy as np
from sklearn.neighbors import NearestNeighbors
from multiprocessing import Pool, cpu_count

def assign_points_to_clusters(args):
    chunk, centroids = args
    ann = NearestNeighbors(n_neighbors=1, algorithm='ball_tree')
    ann.fit(centroids)
    distances, labels = ann.kneighbors(chunk)
    return labels.flatten()

def parallel_kmeans_ann(X, n_clusters, max_iter=100, n_jobs=None):
    if n_jobs is None:
        n_jobs = cpu_count()
    
    centroids = X[np.random.choice(X.shape[0], n_clusters, replace=False)]
    
    for _ range(max_iter):
        # Split data into chunks for parallel processing
        chunks = np.array_split(X, n_jobs)
        
        # Assign points to clusters in parallel
        with Pool(n_jobs) as pool:
            results = pool.map(assign_points_to_clusters, 
                               [(chunk, centroids) for chunk in chunks])
        
        labels = np.concatenate(results)
        
        # Update centroids
        new_centroids = np.array([X[labels == k].mean(axis=0) 
                                  for k in range(n_clusters)])
        
        if np.allclose(centroids, new_centroids):
            break
        
        centroids = new_centroids
    
    return labels, centroids

# Example usage
X = np.random.rand(100000, 10)
labels, centroids = parallel_kmeans_ann(X, n_clusters=10)

print("Final cluster centroids:")
print(centroids)

Visualization Techniques for High-Dimensional Clusters

Visualizing high-dimensional clusters can be challenging. We can use dimensionality reduction techniques like t-SNE or UMAP to project the data and clusters into a 2D or 3D space for visualization.

import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

def visualize_clusters(X, labels, n_components=2):
    # Perform t-SNE
    tsne = TSNE(n_components=n_components, random_state=42)
    X_tsne = tsne.fit_transform(X)
    
    # Plot the results
    plt.figure(figsize=(10, 8))
    scatter = plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=labels, cmap='viridis')
    plt.colorbar(scatter)
    plt.title('t-SNE visualization of clusters')
    plt.xlabel('t-SNE component 1')
    plt.ylabel('t-SNE component 2')
    plt.show()

# Generate sample data and perform clustering
X = np.random.rand(1000, 50)  # 1000 points in 50-dimensional space
labels, _ = kmeans_ann(X, n_clusters=5)  # Assuming kmeans_ann is defined

# Visualize the clusters
visualize_clusters(X, labels)

Additional Resources

For those interested in diving deeper into KMeans clustering with Approximate Nearest Neighbors, the following resources provide valuable insights and advanced techniques:

  1. "Scalable Nearest Neighbor Algorithms for High Dimensional Data" by Muja and Lowe (2014) ArXiv: https://arxiv.org/abs/1403.1231
  2. "Billion-scale similarity search with GPUs" by Johnson et al. (2017) ArXiv: https://arxiv.org/abs/1702.08734
  3. "Approximate Nearest Neighbor Search in High Dimensions" by Andoni and Indyk (2018) ArXiv: https://arxiv.org/abs/1806.09823

These papers provide in-depth discussions on various ANN algorithms, their applications in clustering, and strategies for handling high-dimensional data efficiently.

More in Classification and clustering

  1. Mastering K-Means Clustering with Python9 min

    K-Means is an unsupervised machine learning algorithm used for clustering data points into groups based on similarity. It's widely used in data analysis...

  2. K-Means Clustering Algorithm Explained14 min

    The fundamental k-means clustering algorithm implements iterative refinement to partition n observations into k clusters. Each cluster is represented by the...

  3. Mastering K-Nearest Neighbors Hyperparameters in Python6 min

    K-Nearest Neighbors is a simple yet powerful machine learning algorithm used for classification and regression tasks. It works by finding the K closest data...

  4. K-means & K-modes for numerical data9 min

    The K-means algorithm partitions n observations into k clusters by minimizing within-cluster variances. Each cluster is represented by its centroid, which is...

  5. Mastering Random Forest A Comprehensive Guide11 min

    Random Forest operates on the principle of ensemble learning by constructing multiple decision trees during training. Each tree is built using a bootstrap...

  6. Introduction to Tree-Based Machine Learning Algorithms9 min

    Tree-based machine learning algorithms are powerful tools for decision-making and prediction. These algorithms use tree-like structures to make decisions...