# eDKM: An Efficient and Accurate Train-time Weight Clustering for Large Language Models

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 22 references that are in the map.
- First posted on arXiv 2 Sep 2023.

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

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding](https://mlsanity.com/paper/1510.00149/) (2015)
- [GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers](https://mlsanity.com/paper/2210.17323/) (2022)
- [XNOR-Net: ImageNet Classification Using Binary Convolutional Neural Networks](https://mlsanity.com/paper/1603.05279/) (2016)
- [HAQ: Hardware-Aware Automated Quantization with Mixed Precision](https://mlsanity.com/paper/1811.08886/) (2018)
- [LLM-QAT: Data-Free Quantization Aware Training for Large Language Models](https://mlsanity.com/paper/2305.17888/) (2023)
- [Model compression via distillation and quantization](https://mlsanity.com/paper/1802.05668/) (2018)
- [HAWQ-V2: Hessian Aware trace-Weighted Quantization of Neural Networks](https://mlsanity.com/paper/1911.03852/) (2019)
- [NISP: Pruning Networks using Neuron Importance Score Propagation](https://mlsanity.com/paper/1711.05908/) (2017)
- [Training with Quantization Noise for Extreme Model Compression](https://mlsanity.com/paper/2004.07320/) (2020)
- [Soft Weight-Sharing for Neural Network Compression](https://mlsanity.com/paper/1702.04008/) (2017)
- [And the Bit Goes Down: Revisiting the Quantization of Neural Networks](https://mlsanity.com/paper/1907.05686/) (2019)
- [Lookahead: A Far-Sighted Alternative of Magnitude-based Pruning](https://mlsanity.com/paper/2002.04809/) (2020)
- [Value-aware Quantization for Training and Inference of Neural Networks](https://mlsanity.com/paper/1804.07802/) (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)
- [Neural Epitome Search for Architecture-Agnostic Network Compression](https://mlsanity.com/paper/1907.05642/) (2019)

## Built on by (most cited first, 4 of 4)

- [CompGS: Smaller and Faster Gaussian Splatting with Vector Quantization](https://mlsanity.com/paper/2311.18159/) (2023)
- [KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation](https://mlsanity.com/paper/2405.05329/) (2024)
- [On-Device LLMs for SMEs: Challenges and Opportunities](https://mlsanity.com/paper/2410.16070/) (2024)
- [HybridEP: Scaling Expert Parallelism to Cross-Datacenter Scenario via Hybrid Expert/Data Transmission](https://mlsanity.com/paper/2510.19470/) (2025)

## Abstract (arXiv)

Since Large Language Models or LLMs have demonstrated high-quality performance on many complex language tasks, there is a great interest in bringing these LLMs to mobile devices for faster responses and better privacy protection. However, the size of LLMs (i.e., billions of parameters) requires highly effective compression to fit into storage-limited devices. Among many compression techniques, weight-clustering, a form of non-linear quantization, is one of the leading candidates for LLM compression, and supported by modern smartphones. Yet, its training overhead is prohibitively significant for LLM fine-tuning. Especially, Differentiable KMeans Clustering, or DKM, has shown the state-of-the-art trade-off between compression ratio and accuracy regression, but its large memory complexity makes it nearly impossible to apply to train-time LLM compression. In this paper, we propose a memory-efficient DKM implementation, eDKM powered by novel techniques to reduce the memory footprint of DKM by orders of magnitudes. For a given tensor to be saved on CPU for the backward pass of DKM, we compressed the tensor by applying uniquification and sharding after checking if there is no duplicated tensor previously copied to CPU. Our experimental results demonstrate that prjname can fine-tune and compress a pretrained LLaMA 7B model from 12.6 GB to 2.5 GB (3bit/weight) with the Alpaca dataset by reducing the train-time memory footprint of a decoder layer by 130×, while delivering good accuracy on broader LLM benchmarks (i.e., 77.7% for PIQA, 66.1% for Winograde, and so on).