# DKM: Differentiable K-Means Clustering Layer for Neural Network Compression

arXiv: https://arxiv.org/abs/2108.12659 · HTML: https://mlsanity.com/paper/2108.12659/

## Key facts

- Cited by 31 papers in the map (data through 14 Oct 2026).
- Builds on 41 of its 46 references that are in the map.
- First posted on arXiv 28 Aug 2021.

## Builds on (most cited first, 30 of 41)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks](https://mlsanity.com/paper/1905.11946/) (2019)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [MobileNetV2: Inverted Residuals and Linear Bottlenecks](https://mlsanity.com/paper/1801.04381/) (2018)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://mlsanity.com/paper/1704.04861/) (2017)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [EfficientNetV2: Smaller Models and Faster Training](https://mlsanity.com/paper/2104.00298/) (2021)
- [MnasNet: Platform-Aware Neural Architecture Search for Mobile](https://mlsanity.com/paper/1807.11626/) (2018)
- [XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks](https://mlsanity.com/paper/1603.05279/) (2016)
- [Up or Down? Adaptive Rounding for Post-Training Quantization](https://mlsanity.com/paper/2004.10568/) (2020)
- [MobileBERT: a Compact Task-Agnostic BERT for Resource-Limited Devices](https://mlsanity.com/paper/2004.02984/) (2020)
- [HAQ: Hardware-Aware Automated Quantization with Mixed Precision](https://mlsanity.com/paper/1811.08886/) (2018)
- [Model compression via distillation and quantization](https://mlsanity.com/paper/1802.05668/) (2018)
- [LQ-Nets: Learned Quantization for Highly Accurate and Compact Deep Neural Networks](https://mlsanity.com/paper/1807.10029/) (2018)
- [HAWQ-V2: Hessian Aware trace-Weighted Quantization of Neural Networks](https://mlsanity.com/paper/1911.03852/) (2019)
- [Understanding Straight-Through Estimator in Training Activation Quantized Neural Nets](https://mlsanity.com/paper/1903.05662/) (2019)
- [NISP: Pruning Networks using Neuron Importance Score Propagation](https://mlsanity.com/paper/1711.05908/) (2017)
- [Towards Accurate Binary Convolutional Neural Network](https://mlsanity.com/paper/1711.11294/) (2017)
- [Additive Powers-of-Two Quantization: An Efficient Non-uniform Discretization for Neural Networks](https://mlsanity.com/paper/1909.13144/) (2019)
- [Training with Quantization Noise for Extreme Model Compression](https://mlsanity.com/paper/2004.07320/) (2020)
- [Well-Read Students Learn Better: On the Importance of Pre-training Compact Models](https://mlsanity.com/paper/1908.08962/) (2019)
- [GOBO: Quantizing Attention-Based NLP Models for Low Latency and Energy Efficient Inference](https://mlsanity.com/paper/2005.03842/) (2020)
- [Soft Weight-Sharing for Neural Network Compression](https://mlsanity.com/paper/1702.04008/) (2017)
- [ESPNetv2: A Light-weight, Power Efficient, and General Purpose Convolutional Neural Network](https://mlsanity.com/paper/1811.11431/) (2018)

## Built on by (most cited first, 30 of 31)

- [GPTVQ: The Blessing of Dimensionality for LLM Quantization](https://mlsanity.com/paper/2402.15319/) (2024)
- [Stable and low-precision training for large-scale vision-language models](https://mlsanity.com/paper/2304.13013/) (2023)
- [VPTQ: Extreme Low-bit Vector Post-Training Quantization for Large Language Models](https://mlsanity.com/paper/2409.17066/) (2024)
- [Revisiting the Parameter Efficiency of Adapters from the Perspective of Precision Redundancy](https://mlsanity.com/paper/2307.16867/) (2023)
- [Q-GaLore: Quantized GaLore with INT4 Projection and Layer-Adaptive Low-Rank Gradients](https://mlsanity.com/paper/2407.08296/) (2024)
- [Model Compression in Practice: Lessons Learned from Practitioners Creating On-device Machine Learning Experiences](https://mlsanity.com/paper/2310.04621/) (2023)
- [PixOOD: Pixel-Level Out-of-Distribution Detection](https://mlsanity.com/paper/2405.19882/) (2024)
- [Recycling Model Updates in Federated Learning: Are Gradient Subspaces Low-Rank?](https://mlsanity.com/paper/2202.00280/) (2022)
- [A Review of Pseudo-Labeling for Computer Vision](https://mlsanity.com/paper/2408.07221/) (2024)
- [PDP: Parameter-free Differentiable Pruning is All You Need](https://mlsanity.com/paper/2305.11203/) (2023)
- [Reducing Computational Complexity of Neural Networks in Optical Channel Equalization: From Concepts to Implementation](https://mlsanity.com/paper/2208.12866/) (2022)
- [LO-BCQ: Block Clustered Quantization for 4-bit (W4A4) LLM Inference](https://mlsanity.com/paper/2502.05376/) (2025)
- [Talaria: Interactively Optimizing Machine Learning Models for Efficient Inference](https://mlsanity.com/paper/2404.03085/) (2024)
- [Hyperspherical Quantization: Toward Smaller and More Accurate Models](https://mlsanity.com/paper/2212.12653/) (2022)
- [eDKM: An Efficient and Accurate Train-time Weight Clustering for Large Language Models](https://mlsanity.com/paper/2309.00964/) (2023)
- [R2 Loss: Range Restriction Loss for Model Compression and Quantization](https://mlsanity.com/paper/2303.08253/) (2023)
- [ViM-VQ: Efficient Post-Training Vector Quantization for Visual Mamba](https://mlsanity.com/paper/2503.09509/) (2025)
- [Learning Low-Rank Representations for Model Compression](https://mlsanity.com/paper/2211.11397/) (2022)
- [SPT: Fine-Tuning Transformer-based Language Models Efficiently with Sparsification](https://mlsanity.com/paper/2312.10365/) (2023)
- [Polytopic Autoencoders with Smooth Clustering for Reduced-order Modelling of Flows](https://mlsanity.com/paper/2401.10620/) (2024)
- [MetaCluster: Enabling Deep Compression of Kolmogorov-Arnold Network](https://mlsanity.com/paper/2510.19105/) (2025)
- [Separate, Dynamic and Differentiable (SMART) Pruner for Block/Output Channel Pruning on Computer Vision Tasks](https://mlsanity.com/paper/2403.19969/) (2024)
- [Deep learning model compression using network sensitivity and gradients](https://mlsanity.com/paper/2210.05111/) (2022)
- [Hyperspherical Loss-Aware Ternary Quantization](https://mlsanity.com/paper/2212.12649/) (2022)
- [CrAFT: Compression-Aware Fine-Tuning for Efficient Visual Task Adaptation](https://mlsanity.com/paper/2305.04526/) (2023)
- [IDKM: Memory Efficient Neural Network Quantization via Implicit, Differentiable k-Means](https://mlsanity.com/paper/2312.07759/) (2023)
- [EncodingNet: A Novel Encoding-based MAC Design for Efficient Neural Network Acceleration](https://mlsanity.com/paper/2402.18595/) (2024)
- [MAS-Attention: Memory-Aware Stream Processing for Attention Acceleration on Resource-Constrained Edge Devices](https://mlsanity.com/paper/2411.17720/) (2024)
- [SSVQ: Unleashing the Potential of Vector Quantization with Sign-Splitting](https://mlsanity.com/paper/2503.08668/) (2025)
- [Efficient VQ-QAT and Mixed Vector/Linear quantized Neural Networks](https://mlsanity.com/paper/2604.23172/) (2026)

## Abstract (arXiv)

Deep neural network (DNN) model compression for efficient on-device inference is becoming increasingly important to reduce memory requirements and keep user data on-device. To this end, we propose a novel differentiable k-means clustering layer (DKM) and its application to train-time weight clustering-based DNN model compression. DKM casts k-means clustering as an attention problem and enables joint optimization of the DNN parameters and clustering centroids. Unlike prior works that rely on additional regularizers and parameters, DKM-based compression keeps the original loss function and model architecture fixed. We evaluated DKM-based compression on various DNN models for computer vision and natural language processing (NLP) tasks. Our results demonstrate that DKM delivers superior compression and accuracy trade-off on ImageNet1k and GLUE benchmarks. For example, DKM-based compression can offer 74.5% top-1 ImageNet1k accuracy on ResNet50 DNN model with 3.3MB model size (29.4x model compression factor). For MobileNet-v1, which is a challenging DNN to compress, DKM delivers 63.9% top-1 ImageNet1k accuracy with 0.72 MB model size (22.4x model compression factor). This result is 6.8% higher top-1accuracy and 33% relatively smaller model size than the current state-of-the-art DNN compression algorithms. Additionally, DKM enables compression of DistilBERT model by 11.8x with minimal (1.1%) accuracy loss on GLUE NLP benchmarks.