xbe.at

Machine learning / Classification and clustering

Optimizing KMeans Clustering with Approximate Nearest Neighbors

8 min read, 14 sections

Contents
  1. KMeans Clustering Optimization
  2. Traditional KMeans Algorithm
  3. Approximate Nearest Neighbors
  4. Faiss: Fast Library for Approximate Nearest Neighbors
  5. Accelerated KMeans with Faiss
  6. Performance Comparison
  7. Inverted Index in Faiss
  8. Tuning Faiss KMeans
  9. Real-Life Example: Image Color Quantization
  10. Real-Life Example: Customer Segmentation
  11. Limitations and Considerations
  12. Scaling to Larger Datasets
  13. Future Directions and Research
  14. Additional Resources

KMeans Clustering Optimization

KMeans clustering is a popular unsupervised learning algorithm used for partitioning data into K clusters. While effective, it can be computationally expensive for large datasets. This presentation explores an optimization technique using approximate nearest neighbors to significantly speed up the KMeans algorithm.

from sklearn.cluster import KMeans

# Generate sample data
X = np.random.rand(10000, 2)

# Standard KMeans
kmeans = KMeans(n_clusters=5)
kmeans.fit(X)

Traditional KMeans Algorithm

The traditional KMeans algorithm follows an iterative process: initialize centroids, assign points to nearest centroids, update centroids, and repeat until convergence. The bottleneck in this process is the nearest centroid search, which becomes time-consuming for large datasets.

    centroids = X[np.random.choice(X.shape[0], k, replace=False)]
    for _ in range(max_iters):
        # Assign points to nearest centroid
        distances = np.sqrt(((X - centroids[:, np.newaxis])**2).sum(axis=2))
        labels = np.argmin(distances, axis=0)
        
        # Update centroids
        new_centroids = np.array([X[labels == i].mean(axis=0) for i in range(k)])
        
        if np.all(centroids == new_centroids):
            break
        centroids = new_centroids
    return centroids, labels

# Run simple KMeans
centroids, labels = simple_kmeans(X, k=5)

Approximate Nearest Neighbors

Approximate Nearest Neighbors (ANN) algorithms provide a faster alternative to exact nearest neighbor search. These algorithms trade a small amount of accuracy for significant speed improvements, making them ideal for large-scale clustering tasks.


def ann_search(query, database, k):
    d = database.shape[1]
    index = faiss.IndexFlatL2(d)
    index.add(database)
    distances, indices = index.search(query, k)
    return distances, indices

# Example usage
query = np.random.rand(1, 2).astype('float32')
database = np.random.rand(10000, 2).astype('float32')
distances, indices = ann_search(query, database, k=5)
print(f"Indices of 5 nearest neighbors: {indices[0]}")

Faiss: Fast Library for Approximate Nearest Neighbors

Faiss is a library developed by Facebook AI Research for efficient similarity search and clustering of dense vectors. It provides optimized implementations of various indexing methods, including the Inverted File Index (IVF) which is particularly useful for KMeans acceleration.


# Create a Faiss index
d = 2  # dimension of the data
nlist = 100  # number of clusters for coarse quantization
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)

# Train and add vectors to the index
index.train(X.astype('float32'))
index.add(X.astype('float32'))

# Perform a search
k = 5  # number of nearest neighbors to retrieve
D, I = index.search(X[:5].astype('float32'), k)
print(f"Distances and indices for first 5 points: {D}, {I}")

Accelerated KMeans with Faiss

By leveraging Faiss, we can significantly speed up the KMeans algorithm. The key idea is to use Faiss for efficient nearest centroid search, replacing the computationally expensive brute-force approach in traditional KMeans.


def faiss_kmeans(X, k, niter=20):
    d = X.shape[1]
    kmeans = faiss.Kmeans(d, k, niter=niter, verbose=True, gpu=False)
    kmeans.train(X.astype(np.float32))
    centroids = kmeans.centroids
    _, labels = kmeans.index.search(X.astype(np.float32), 1)
    return centroids, labels.ravel()

# Run Faiss KMeans
centroids, labels = faiss_kmeans(X, k=5)
print(f"Centroids shape: {centroids.shape}, Labels shape: {labels.shape}")

Performance Comparison

To demonstrate the speedup achieved using Faiss, let's compare the execution time of traditional KMeans and Faiss KMeans on datasets of varying sizes. The results show a significant reduction in computation time, especially for larger datasets.

from sklearn.cluster import KMeans

def benchmark(X, k):
    # Traditional KMeans
    start = time.time()
    KMeans(n_clusters=k).fit(X)
    trad_time = time.time() - start
    
    # Faiss KMeans
    start = time.time()
    faiss_kmeans(X, k)
    faiss_time = time.time() - start
    
    return trad_time, faiss_time

sizes = [1000, 10000, 100000]
for size in sizes:
    X = np.random.rand(size, 100).astype('float32')
    trad_time, faiss_time = benchmark(X, k=5)
    speedup = trad_time / faiss_time
    print(f"Dataset size: {size}")
    print(f"Traditional KMeans: {trad_time:.2f}s")
    print(f"Faiss KMeans: {faiss_time:.2f}s")
    print(f"Speedup: {speedup:.2f}x\n")

Inverted Index in Faiss

The Inverted File Index (IVF) in Faiss is a key component for achieving high performance. It organizes data points into clusters and creates an inverted list structure, allowing for efficient nearest neighbor search without exhaustive comparisons.


def create_ivf_index(X, nlist):
    d = X.shape[1]
    quantizer = faiss.IndexFlatL2(d)
    index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
    index.train(X)
    index.add(X)
    return index

# Create and use IVF index
X = np.random.rand(10000, 100).astype('float32')
index = create_ivf_index(X, nlist=100)
D, I = index.search(X[:5], k=5)
print(f"Distances and indices for first 5 points: {D}, {I}")

Tuning Faiss KMeans

To achieve optimal performance with Faiss KMeans, it's important to tune parameters such as the number of lists (nlist) in the IVF index and the number of probes during search. These parameters allow for a trade-off between speed and accuracy.

    d = X.shape[1]
    best_time = float('inf')
    best_params = None
    
    for nlist in nlist_values:
        for nprobe in nprobe_values:
            kmeans = faiss.Kmeans(d, k, niter=20, verbose=False)
            kmeans.train(X)
            index = faiss.IndexIVFFlat(kmeans.index, d, nlist)
            index.train(X)
            index.add(X)
            index.nprobe = nprobe
            
            start = time.time()
            D, I = index.search(X, 1)
            end = time.time()
            
            if end - start < best_time:
                best_time = end - start
                best_params = (nlist, nprobe)
    
    return best_params, best_time

X = np.random.rand(100000, 100).astype('float32')
nlist_values = [10, 100, 1000]
nprobe_values = [1, 10, 100]
best_params, best_time = tune_faiss_kmeans(X, k=5, nlist_values=nlist_values, nprobe_values=nprobe_values)
print(f"Best parameters: nlist={best_params[0]}, nprobe={best_params[1]}")
print(f"Best search time: {best_time:.4f}s")

Real-Life Example: Image Color Quantization

One practical application of KMeans clustering is image color quantization. By applying KMeans to the pixel colors of an image, we can reduce the number of colors while maintaining the overall appearance. Using Faiss can significantly speed up this process for large images.

import numpy as np
import faiss

def quantize_colors(image_path, k):
    # Load image and reshape to 2D array of pixels
    img = Image.open(image_path)
    pixels = np.array(img).reshape(-1, 3).astype('float32')
    
    # Perform KMeans clustering using Faiss
    kmeans = faiss.Kmeans(3, k, niter=20, verbose=False)
    kmeans.train(pixels)
    
    # Assign each pixel to nearest centroid
    _, labels = kmeans.index.search(pixels, 1)
    
    # Replace pixel colors with centroids
    quantized = kmeans.centroids[labels.ravel()].reshape(img.size[1], img.size[0], 3)
    
    return Image.fromarray(quantized.astype('uint8'))

# Example usage
original_image_path = "path_to_your_image.jpg"
quantized_image = quantize_colors(original_image_path, k=16)
quantized_image.save("quantized_image.jpg")

Real-Life Example: Customer Segmentation

Another common application of KMeans clustering is customer segmentation in marketing. By clustering customers based on features like purchase history, demographics, and behavior, businesses can tailor their marketing strategies. Faiss can accelerate this process for large customer databases.

import faiss
import pandas as pd

def segment_customers(data, k):
    # Normalize the data
    mean = np.mean(data, axis=0)
    std = np.std(data, axis=0)
    normalized_data = (data - mean) / std
    
    # Perform KMeans clustering using Faiss
    kmeans = faiss.Kmeans(data.shape[1], k, niter=20, verbose=False)
    kmeans.train(normalized_data.astype('float32'))
    
    # Assign customers to segments
    _, labels = kmeans.index.search(normalized_data.astype('float32'), 1)
    
    return labels.ravel()

# Example usage
customer_data = pd.DataFrame({
    'age': np.random.randint(18, 80, 10000),
    'income': np.random.randint(20000, 200000, 10000),
    'spending_score': np.random.randint(1, 100, 10000)
})

segments = segment_customers(customer_data.values, k=5)
customer_data['segment'] = segments

print(customer_data.groupby('segment').mean())

Limitations and Considerations

While Faiss provides significant speedups for KMeans clustering, it's important to consider some limitations. The approximate nature of the algorithm may lead to slightly different results compared to exact KMeans. Additionally, the memory requirements can be substantial for very large datasets.

import faiss
from sklearn.cluster import KMeans

def compare_kmeans(X, k):
    # Exact KMeans
    kmeans = KMeans(n_clusters=k, n_init=1, max_iter=20)
    kmeans.fit(X)
    exact_labels = kmeans.labels_
    
    # Faiss KMeans
    faiss_kmeans = faiss.Kmeans(X.shape[1], k, niter=20, verbose=False)
    faiss_kmeans.train(X.astype('float32'))
    _, faiss_labels = faiss_kmeans.index.search(X.astype('float32'), 1)
    
    # Compare results
    agreement = np.mean(exact_labels == faiss_labels.ravel())
    print(f"Agreement between exact and Faiss KMeans: {agreement:.2%}")

# Example usage
X = np.random.rand(10000, 100).astype('float32')
compare_kmeans(X, k=5)

Scaling to Larger Datasets

For extremely large datasets that don't fit in memory, Faiss provides GPU support and out-of-core indexing options. These features allow for efficient clustering of datasets with billions of points across multiple GPUs or machines.


def large_scale_kmeans(X, k, gpu_id=0):
    d = X.shape[1]
    
    # Create GPU resource
    res = faiss.StandardGpuResources()
    
    # Configure the index
    config = faiss.GpuIndexFlatConfig()
    config.device = gpu_id
    
    # Create the index on GPU
    index = faiss.GpuIndexFlatL2(res, d, config)
    
    # Create and train KMeans object
    kmeans = faiss.Kmeans(d, k, niter=20, gpu=True)
    kmeans.train(X)
    
    return kmeans.centroids, kmeans.obj

# Example usage (Note: This requires a CUDA-enabled GPU)
X = np.random.rand(1000000, 100).astype('float32')
centroids, obj = large_scale_kmeans(X, k=100)
print(f"Final objective: {obj[-1]:.2f}")

Future Directions and Research

The field of approximate nearest neighbor search and clustering continues to evolve. Recent research focuses on improving the accuracy-speed trade-off, developing new indexing structures, and exploring applications in areas such as recommendation systems and anomaly detection.

import faiss

def hybrid_kmeans(X, k, exact_fraction=0.1):
    d = X.shape[1]
    n = X.shape[0]
    n_exact = int(n * exact_fraction)
    
    # Perform exact KMeans on a subset
    kmeans_exact = faiss.Kmeans(d, k, niter=20, verbose=False)
    kmeans_exact.train(X[:n_exact])
    
    # Use the result to initialize Faiss KMeans for the full dataset
    kmeans_full = faiss.Kmeans(d, k, niter=20, verbose=False)
    kmeans_full.centroids = kmeans_exact.centroids
    kmeans_full.train(X)
    
    return kmeans_full.centroids, kmeans_full.obj

# Example usage
X = np.random.rand(100000, 100).astype('float32')
centroids, obj = hybrid_kmeans(X, k=10)
print(f"Final objective: {obj[-1]:.2f}")

Additional Resources

For those interested in delving deeper into approximate nearest neighbor search and its applications in clustering, the following resources are recommended:

  1. "Billion-scale similarity search with GPUs" by Johnson, Douze, and Jégou (2017) ArXiv: https://arxiv.org/abs/1702.08734
  2. "Faiss: A Library for Efficient Similarity Search" by Facebook AI Research GitHub: https://github.com/facebookresearch/faiss
  3. "Approximate Nearest Neighbor Search in High Dimensions" by Andoni and Indyk (2006) ACM: https://dl.acm.org/doi/10.1145/1132

More in Classification and clustering

  1. Practical Limitations of KMeans Clustering13 min

    KMeans is a popular clustering algorithm, but it has several limitations that are often overlooked. Understanding these limitations is crucial for effective...

  2. Naive Bayes Classification Use Cases9 min

    Naive Bayes classifiers are probabilistic algorithms based on Bayes' theorem, assuming feature independence. They excel in text classification, spam...

  3. Predicting Premier League Winner with Random Forest7 min

    Predicting the winner of the Premier League using machine learning has become an exciting application of data science in sports. This presentation focuses on...

  4. Mastering Support Vector Machines with Python13 min

    Support Vector Machines (SVMs) are powerful supervised learning models used for classification and regression tasks. They aim to find the optimal hyperplane...

  5. Preventing Overfitting in Decision Trees with CCP10 min

    Decision trees are powerful machine learning algorithms, but they have a tendency to overfit the training data. This means they can capture noise and create...

  6. Mastering Support Vector Machines for Classification and Regression12 min

    Support Vector Machines rely on fundamental mathematical principles to find the optimal hyperplane separating data classes. The primary objective is...