Sliced-Wasserstein Distribution Alignment Loss Improves the Ultra-Low-Bit Quantization of Large Language Models

TL;DR

This work tackles the resource and performance trade-offs of ultra-low-bit quantization for large language models by introducing a sliced-Wasserstein (SW) distribution-alignment loss. The SW loss complements standard MSE by aligning full-precision and quantized activation distributions under random projections, enabling more faithful representations with minimal inference overhead. Integrated into frontier PTQ frameworks OmniQuant and TesseraQ, it yields consistent gains in perplexity and downstream accuracy across LLaMA-2 and OPT models, and is demonstrated with publicly available code. Overall, distributional calibration emerges as a practical, plug-and-play strategy to push ultra-low-bit quantization closer to full-precision fidelity, facilitating energy- and memory-efficient deployment of LLMs.

Abstract

The benefits of most large language models come with steep and often hidden economic and environmental costs due to their resource usage inefficiency during deployment. Model quantization improves energy and memory efficiency through representing model parameters by lower-precision values. However, compression below 4-bits often distorts activation distributions and degrades performance. We address this challenge by introducing a sliced Wasserstein loss function for distribution-aware calibration in ultra-low-bit post-training quantization. The proposed loss aligns the output distributions of full-precision and quantized models under random linear projections, complementing standard mean-squared error loss without adding any computational overhead during inference. Our proposed loss function can be incorporated with any post-training quantization framework that has a retraining component. We demonstrate the performance gains of our proposed model by incorporating it with two frontier methods known as OmniQuant and TesseraQ. Compared to these two baselines, the proposed loss consistently improves both perplexity and downstream task accuracy across multiple ultra-low-bit settings. Our proposed loss function recovers 4.12-20.37% of the OmniQuant's lost accuracy on the language model LLaMA-2-7B, 0.93-7.65% on OPT-6.7B, and 2.26-6.20% on LLaMA-2-13B. TesseraQ's accuracy degradation is recovered by 3.63-7.63% in relative terms when augmented by our proposed loss function. Taken together, these results demonstrate that distributional alignment provides a simple yet effective performance boost that can push the limits of frontier quantization methods. Our method is available on GitHub to facilitate future progress in ultra-low-bit quantization.

Figures (5)

• Figure 1: Perplexity comparison between the OmniQuant Baseline and the Proposed SW Loss Quantization on LLaMA-2-7B.
• Figure 2: Perplexity comparison between the OmniQuant Baseline and the Proposed SW Loss Quantization on LLaMA-2-13B. Left: LLaMA-2-13B W2A16g128 (mixed calibration). Right: LLaMA-2-13B W4A4 (wikitext2 calibration).
• Figure 3: Perplexity comparison between the OmniQuant Baseline and the Proposed SW Loss Quantization on OPT-6.7B. Left: OPT-6.7B W2A16g128 (wikitext2 calib.). Right: OPT-6.7B W2A16g128 (mixed calib.).
• Figure 4: Perplexity comparison for TesseraQ with and without the SW loss on LLaMA-2-7B and LLaMA-2-13B based on WikiText-2 (aggregated panels).
• Figure 5: Impact of the SW hyperparameters on quantization performance of LLaMA-2-7B under the W2A16g128 configuration with mixed calibration. Figure (a) examines how varying the SW loss weight ($w_{\text{SW}}$) affects the perplexity when $n_{\text{proj}}=128$, while Figure (b) analyzes the influence of the number of projection samples ($n_{\text{proj}}$) when $w_{\text{SW}}=0.1$. The dashed lines indicate the OmniQuant baseline, showing that the proposed SW-loss quantization achieves lower perplexity with minimal sensitivity to these hyperparameters.