K-Means Skala Besar: Kenapa Gagal dan Solusi Modern

K-means adalah algoritma yang dipelajari semua orang di kelas ML pengantar, lalu sepanjang karier mereka dipakai dengan cara yang keliru. Kelihatannya sederhana: pilih K pusat klaster, tetapkan setiap titik data ke pusat terdekatnya, geser pusat ke rata-rata titik yang ditetapkan padanya, lalu ulangi. Cukup tiga baris pseudocode. Masalahnya, kesederhanaan ini menyembunyikan serangkaian jebakan yang meledak ketika diterapkan ke data nyata dalam skala besar.

Pada 10.000 titik, k-means selesai dalam hitungan milidetik dan tidak ada yang repot memikirkan detail implementasi. Pada 10 juta titik, pilihan metode inisialisasi bisa mengubah waktu eksekusi hingga 100 kali lipat. Pada 1 miliar titik, algoritma standar tidak muat di memori dan Anda butuh pendekatan yang berbeda secara fundamental. Sebuah paper baru tentang Flash-KMeans membahas tepat masalah ini, yaitu mencapai hasil k-means yang exact dengan memori yang jauh lebih kecil dan konvergensi yang lebih cepat. Mari kita bahas apa yang membuat k-means sulit di skala besar dan bagaimana varian modern menyelesaikannya.

Cara Kerja K-Means Standar

Algoritma Lloyd, yang dimaksud orang saat mengatakan 'k-means', punya dua langkah per iterasi. Langkah penetapan: untuk setiap titik data, hitung jaraknya ke semua K centroid lalu tetapkan ke yang paling dekat. Langkah pembaruan: hitung ulang setiap centroid sebagai rata-rata dari semua titik yang ditetapkan padanya.

import numpy as np
def kmeans_lloyd(X, K, max_iter=100):
n, d = X.shape
# Random initialization (bad, but we'll fix this later)
centroids = X[np.random.choice(n, K, replace=False)]
for _ in range(max_iter):
# Assignment: O(n * K * d) — this is the bottleneck
distances = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)
labels = np.argmin(distances, axis=1)
# Update: O(n * d)
new_centroids = np.array([
X[labels == k].mean(axis=0) for k in range(K)
])
if np.allclose(centroids, new_centroids):
break
centroids = new_centroids
return labels, centroids

Langkah penetapan membutuhkan biaya O(n × K × d) per iterasi, dengan n adalah jumlah titik, K adalah jumlah klaster, dan d adalah dimensi. Untuk n = 1 miliar, K = 1000, d = 128 (skenario klastering embedding yang realistis), itu berarti 128 triliun operasi floating-point per iterasi. Bahkan dengan compute satu teraflop, itu sudah 128 detik per iterasi, dan k-means biasanya butuh 10-50 iterasi untuk konvergen.

Masalah Inisialisasi

Sebelum k-means menjalankan satu iterasi pun, ia harus memilih posisi awal centroid. Pilihan ini jauh lebih penting daripada yang disadari kebanyakan orang. Inisialisasi yang buruk bisa membuat algoritma konvergen ke solusi yang jauh lebih buruk dari yang optimal.

Inisialisasi acak (memilih K titik data secara acak sebagai centroid awal) itu sederhana tetapi tidak bisa diandalkan. Jika dua centroid awal kebetulan jatuh di klaster yang sama, satu klaster tidak akan terwakili. Algoritmanya tetap konvergen, tetapi ke solusi yang suboptimal. Dengan K = 100, peluang terjadinya tabrakan minimal sekali sangat tinggi.

K-means++ (Arthur dan Vassilvitskii, 2007) adalah solusi standarnya. Ia memilih centroid pertama secara acak, lalu memilih setiap centroid berikutnya dengan probabilitas yang sebanding dengan jarak kuadratnya dari centroid terdekat yang sudah ada. Titik yang jauh dari centroid mana pun lebih berpeluang dipilih, sehingga centroid awal tersebar merata di seluruh data. Ini memberikan jaminan aproksimasi O(log K), artinya solusinya terbukti berada dalam faktor O(log K) dari yang optimal.

def kmeans_plus_plus_init(X, K):
n, d = X.shape
centroids = [X[np.random.randint(n)]]
for _ in range(1, K):
# Distance from each point to nearest existing centroid
dists = np.min([
np.sum((X - c) ** 2, axis=1) for c in centroids
], axis=0)
# Sample proportional to distance squared
probs = dists / dists.sum()
next_idx = np.random.choice(n, p=probs)
centroids.append(X[next_idx])
return np.array(centroids)

Tapi ada jebakannya: inisialisasi k-means++ sendiri berbiaya O(n × K × d), karena membutuhkan K kali pemindaian seluruh dataset. Untuk K dan n yang besar, inisialisasi bisa lebih lama daripada beberapa iterasi algoritma Lloyd. Varian yang lebih scalable seperti k-means|| (Bahmani dkk., 2012) mengurangi jumlah pemindaian dengan mengambil sampel kandidat secara berlebih secara paralel, lalu menggabungkannya menjadi K centroid.

Tembok Memori

K-means standar membutuhkan seluruh dataset berada di memori sekaligus. Pada langkah penetapan, Anda harus mengakses setiap titik data dan membandingkannya dengan setiap centroid. Dengan 1 miliar vektor float32 berdimensi 128, datanya saja memerlukan 512 GB, lebih dari yang dimiliki sebagian besar mesin tunggal.

Solusi naifnya adalah mini-batch k-means: setiap iterasi mengambil sampel acak sebagian titik, lalu memperbarui centroid berdasarkan sampel tersebut. Cara ini berhasil, tetapi konvergensinya lebih lambat dan menghasilkan solusi yang sedikit berbeda (biasanya sedikit lebih buruk) dibanding k-means exact. Untuk banyak aplikasi, penurunan kualitasnya masih bisa diterima. Untuk yang lain, misalnya codebook kuantisasi di vector search, selisih kualitas itu penting.

Flash-KMeans mengambil pendekatan yang berbeda. Alih-alih memproses seluruh dataset di memori, ia menggunakan pemindaian streaming dengan caching yang cerdas. Insight utamanya: sebagian besar titik tidak berubah penetapan klasternya antar iterasi. Jika sebuah titik jelas berada di klaster 7 (jauh lebih dekat ke centroid 7 dibanding centroid lainnya), memeriksa semua K jarak adalah kerja yang sia-sia. Dengan menyimpan batas jarak, Flash-KMeans bisa melewati perhitungan jarak penuh untuk sebagian besar titik di sebagian besar iterasi.

Optimasi Pertidaksamaan Segitiga

Algoritma Elkan (2003) menggunakan pertidaksamaan segitiga untuk melewati perhitungan jarak yang tidak perlu. Pertidaksamaan segitiga menyatakan: jarak dari titik P ke centroid A paling besar sama dengan jarak P ke centroid B ditambah jarak dari B ke A. Jika Anda tahu P saat ini ditetapkan ke centroid B, dan Anda tahu jarak antara centroid A dan B, Anda kadang bisa membuktikan bahwa A terlalu jauh dari P tanpa menghitung jarak sebenarnya.

Dalam praktiknya, ini menghilangkan 80-95% perhitungan jarak setelah beberapa iterasi pertama, ketika sebagian besar titik sudah dekat dengan centroid yang benar. Sisa perhitungan jarak hanya untuk titik di dekat batas klaster, yaitu satu-satunya titik yang mungkin benar-benar berubah penetapan.

Trade-off-nya: algoritma Elkan membutuhkan memori tambahan O(n × K) untuk menyimpan batas jarak (batas bawah dari setiap titik ke setiap centroid, ditambah batas atas ke centroid yang ditetapkan). Untuk K yang besar, biaya memori ini bisa sangat signifikan. Algoritma Hamerly menguranginya menjadi O(n) dengan hanya menyimpan satu batas bawah per titik, dengan konsekuensi lebih sedikit perhitungan yang bisa dipangkas.

Akselerasi GPU

K-means sangat mudah diparalelkan pada langkah penetapan, karena perhitungan jarak setiap titik bersifat independen. Ini membuatnya cocok sekali untuk akselerasi GPU. Library cuML milik NVIDIA dan FAISS milik Facebook sama-sama menyertakan implementasi k-means di GPU yang mencapai percepatan 10-50 kali lipat dibanding implementasi CPU untuk dataset besar.

Masalahnya, memori GPU terbatas. A100 punya memori 80 GB, cukup untuk sekitar 150 juta vektor berdimensi 128. Dataset yang lebih besar memerlukan partisi multi-GPU atau pendekatan streaming, di mana data dimuat per potongan, diproses di GPU, lalu hasilnya diagregasi di CPU.

# Using FAISS for GPU-accelerated k-means
import faiss
import numpy as np
# 10 million 128-dimensional vectors
n, d, K = 10_000_000, 128, 1000
X = np.random.randn(n, d).astype('float32')
# CPU k-means (for comparison)
kmeans_cpu = faiss.Kmeans(d, K, niter=20, verbose=True)
kmeans_cpu.train(X)  # ~120 seconds
# GPU k-means (single GPU)
kmeans_gpu = faiss.Kmeans(d, K, niter=20, verbose=True, gpu=True)
kmeans_gpu.train(X)  # ~8 seconds — 15x faster

Kapan K-Means Adalah Pilihan yang Salah

Sebelum mengoptimalkan implementasi k-means Anda, pertimbangkan dulu apakah k-means memang algoritma yang tepat. Algoritma ini membuat asumsi kuat yang tidak berlaku untuk banyak dataset nyata.

• Klaster berbentuk bola. K-means mengasumsikan klaster kurang lebih berbentuk bola dan berukuran sama. Jika klaster Anda memanjang, berbentuk tidak beraturan, atau ukurannya sangat berbeda, k-means akan memecah klaster besar dan menggabungkan klaster kecil. Gaussian Mixture Models (GMM) menangani klaster berbentuk elips. DBSCAN menangani bentuk arbitrer.
• K sudah diketahui. Anda harus menentukan jumlah klaster di awal. Jika tidak tahu K, Anda perlu menjalankan k-means beberapa kali dengan nilai berbeda dan memakai metrik (silhouette score, metode elbow) untuk memilih yang terbaik. Ini melipatgandakan total compute sebanyak jumlah nilai K yang dicoba.
• Jarak Euclidean. K-means secara default memakai jarak Euclidean. Untuk embedding teks, cosine similarity biasanya lebih tepat. Anda bisa mengakalinya dengan melakukan normalisasi L2 pada vektor (sehingga jarak Euclidean setara dengan jarak cosine), tetapi ini mudah terlupakan.
• Sensitif terhadap outlier. Satu outlier yang jauh dari klaster mana pun akan menarik centroid yang ditetapkan ke arahnya. Varian yang lebih robust seperti k-medoids (yang memakai median alih-alih mean) lebih tahan terhadap outlier, tetapi lebih mahal.

Saran Praktis

Setelah memakai k-means pada dataset dari ribuan hingga miliaran titik, berikut yang menurut saya paling penting:

1. Selalu gunakan inisialisasi k-means++. Inisialisasi acak tidak pernah sepadan dengan risikonya. Selisih kualitas klaster akhir sering mencapai 10-30%, dan k-means++ hanya menambah overhead yang kecil untuk K kecil hingga menengah.
2. Jalankan beberapa kali. K-means menemukan optimum lokal, bukan global. Jalankan 5-10 kali dengan random seed berbeda dan ambil hasil terbaik (total jarak dalam-klaster terkecil). Ini asuransi murah untuk menghindari run yang buruk.
3. Normalisasi fitur Anda. Jika satu fitur punya rentang [0, 1000000] dan fitur lain punya rentang [0, 1], fitur dengan rentang besar akan mendominasi perhitungan jarak. Lakukan standarisasi (kurangi mean, bagi dengan standar deviasi) atau min-max normalization sebelum klastering.
4. Gunakan FAISS untuk klastering skala besar. Jika Anda punya lebih dari satu juta titik, k-means di scikit-learn akan lambat. Implementasi FAISS sangat dioptimalkan dan mendukung akselerasi GPU langsung.
5. Pertimbangkan metode aproksimasi untuk eksplorasi. Mini-batch k-means 10 kali lebih cepat daripada k-means exact dan biasanya menghasilkan hasil yang cukup baik untuk eksplorasi. Gunakan k-means exact untuk klastering produksi di mana kualitas penting.

K-means adalah salah satu algoritma yang mudah dipakai, sulit dipakai dengan baik, dan layak dipahami secara mendalam. Jurang antara implementasi naif dan yang dioptimalkan, baik dari sisi kecepatan maupun kualitas hasil, sangat besar. Di skala kecil, semua ini tidak penting. Di skala di mana ini benar-benar berpengaruh, memahami inisialisasi, manajemen memori, pemangkasan jarak, dan akselerasi GPU adalah pembeda antara klastering yang butuh berjam-jam dan yang selesai dalam hitungan menit.