Partitioning-Guided K-Means: Extreme Empty Cluster Resolution for Extreme Model Compression

TL;DR

This paper introduces Partitioning-Guided K-Means, a novel method to effectively resolve empty clusters in iterative product quantization, significantly improving model accuracy and efficiency in extreme deep learning model compression.

Contribution

It presents a new partitioning-guided k-means approach with strategies for initial cluster assignment, empty cluster resolution, and cluster consolidation, enhancing quantization accuracy.

Findings

Reduces empty clusters by 100x on average

Uses 8x fewer iterations for empty cluster resolution

Improves model accuracy by up to 12% on GLUE tasks

Abstract

Compactness in deep learning can be critical to a model's viability in low-resource applications, and a common approach to extreme model compression is quantization. We consider Iterative Product Quantization (iPQ) with Quant-Noise to be state-of-the-art in this area, but this quantization framework suffers from preventable inference quality degradation due to prevalent empty clusters. In this paper, we propose several novel enhancements aiming to improve the accuracy of iPQ with Quant-Noise by focusing on resolving empty clusters. Our contribution, which we call Partitioning-Guided k-means (PG k-means), is a heavily augmented k-means implementation composed of three main components. First, we propose a partitioning-based pre-assignment strategy that ensures no initial empty clusters and encourages an even weight-to-cluster distribution. Second, we propose an empirically superior empty cluster resolution heuristic executed via cautious partitioning of large clusters. Finally, we construct an optional optimization step that consolidates intuitively dense clusters of weights to ensure shared representation. The proposed approach consistently reduces the number of empty clusters in iPQ with Quant-Noise by 100x on average, uses 8x fewer iterations during empty cluster resolution, and improves overall model accuracy by up to 12%, when applied to RoBERTa on a variety of tasks in the GLUE benchmark.