e-means: supera k-means (in un decimo delle epoche)

Abstract degli autori

L’algoritmo k-means di Lloyd è da tempo uno standard per il clustering, con applicazioni nell’analisi dei dati, nel recupero, nella quantizzazione e nella curatela dei dati. In questo articolo presentiamo una procedura di ottimizzazione basata sulla media mobile esponenziale per k-means con aggiornamenti a mini-batch, che sfrutta gli strumenti moderni del deep learning e nuove intuizioni. Stabiliamo inoltre leggi di scaling che descrivono come il tasso di decadimento e la dimensione del batch influenzino le prestazioni finali. Ne ricaviamo una regola basata su una formula per fissare questi parametri per diverse dimensioni dei problemi e diversi dataset di embedding, senza dover effettuare alcuna scansione degli iperparametri. L’algoritmo di clustering risultante, battezzato e-means, è molto più veloce di Lloyd k-means e converge a loss inferiori, superando spesso persino k-means++. Su un’ampia gamma di dataset di embedding, raggiunge in 1,4-2,6 epoche una loss inferiore a quella di Lloyd k-means dopo 25 epoche, con accelerazioni da 10 a 45 volte nei tempi effettivi di esecuzione su GPU rispetto alle implementazioni di k-means all’avanguardia. Per esempio, e-means suddivide DINO10B, un dataset di 10 miliardi di embedding DINO, in 16 milioni di cluster in sole 2,5 epoche, raggiungendo una loss inferiore a quella di Lloyd dopo 25 epoche.