# IDKM: Memory Efficient Neural Network Quantization via Implicit, Differentiable k-Means

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 31 references that are in the map.
- First posted on arXiv 12 Dec 2023.

## Builds on (most cited first, 24 of 24)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices](https://mlsanity.com/paper/1707.01083/) (2017)
- [SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size](https://mlsanity.com/paper/1602.07360/) (2016)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [HAWQ: Hessian AWare Quantization of Neural Networks with Mixed-Precision](https://mlsanity.com/paper/1905.03696/) (2019)
- [Towards Accurate Binary Convolutional Neural Network](https://mlsanity.com/paper/1711.11294/) (2017)
- [Implicit Deep Learning](https://mlsanity.com/paper/1908.06315/) (2019)
- [Monotone operator equilibrium networks](https://mlsanity.com/paper/2006.08591/) (2020)
- [Soft Weight-Sharing for Neural Network Compression](https://mlsanity.com/paper/1702.04008/) (2017)
- [JFB: Jacobian-Free Backpropagation for Implicit Networks](https://mlsanity.com/paper/2103.12803/) (2021)
- [Post-Training Piecewise Linear Quantization for Deep Neural Networks](https://mlsanity.com/paper/2002.00104/) (2020)
- [And the Bit Goes Down: Revisiting the Quantization of Neural Networks](https://mlsanity.com/paper/1907.05686/) (2019)
- [An Overview of Neural Network Compression](https://mlsanity.com/paper/2006.03669/) (2020)
- [PROFIT: A Novel Training Method for sub-4-bit MobileNet Models](https://mlsanity.com/paper/2008.04693/) (2020)
- [Methods for Pruning Deep Neural Networks](https://mlsanity.com/paper/2011.00241/) (2020)
- [Recurrent Stacking of Layers for Compact Neural Machine Translation Models](https://mlsanity.com/paper/1807.05353/) (2018)
- [Deep k-Means: Re-Training and Parameter Sharing with Harder Cluster Assignments for Compressing Deep Convolutions](https://mlsanity.com/paper/1806.09228/) (2018)
- [DKM: Differentiable K-Means Clustering Layer for Neural Network Compression](https://mlsanity.com/paper/2108.12659/) (2021)
- [Robust Implicit Networks via Non-Euclidean Contractions](https://mlsanity.com/paper/2106.03194/) (2021)

## Abstract (arXiv)

Compressing large neural networks with minimal performance loss is crucial to enabling their deployment on edge devices. (Cho et al., 2022) proposed a weight quantization method that uses an attention-based clustering algorithm called differentiable k-means (DKM). Despite achieving state-of-the-art results, DKM's performance is constrained by its heavy memory dependency. We propose an implicit, differentiable k-means algorithm (IDKM), which eliminates the major memory restriction of DKM. Let t be the number of k-means iterations, m be the number of weight-vectors, and b be the number of bits per cluster address. IDKM reduces the overall memory complexity of a single k-means layer from O(t · m · 2b) to O( m · 2b). We also introduce a variant, IDKM with Jacobian-Free-Backpropagation (IDKM-JFB), for which the time complexity of the gradient calculation is independent of t as well. We provide a proof of concept of our methods by showing that, under the same settings, IDKM achieves comparable performance to DKM with less compute time and less memory. We also use IDKM and IDKM-JFB to quantize a large neural network, Resnet18, on hardware where DKM cannot train at all.