Soft Quantization: Model Compression Via Weight Coupling

TL;DR

Soft quantization introduces a layer-wise short-range attractive coupling among weights during fine-tuning, implemented via a compression term $\mathcal{L}_C(\boldsymbol{\theta})=\sum_l h_l \sum_{i\neq j} U_{w_l}(\theta_i^{(l)}-\theta_j^{(l)})$, where $U_w$ is a short-range triangular well and the two global hyperparameters $h$ and $w$ control strength and range. The coupling drives weights into discrete clusters, yielding layer-dependent effective bit-widths that are computed from the resulting cluster counts, with an efficient histogram-based approximation to form the effective potential $V_{w_l}$. On ResNet-20 trained on CIFAR-10, softly quantized models outperform histogram-equalized post-training quantization (HEQ) and achieve mixed-precision representations with low variability across runs, particularly for settings that beat 4-bit HEQ. The work links compression to loss-landscape geometry, provides a scalable compression pipeline, and suggests directions for applying soft quantization to larger models and contexts such as transfer and continual learning.

Abstract

We show that introducing short-range attractive couplings between the weights of a neural network during training provides a novel avenue for model quantization. These couplings rapidly induce the discretization of a model's weight distribution, and they do so in a mixed-precision manner despite only relying on two additional hyperparameters. We demonstrate that, within an appropriate range of hyperparameters, our "soft quantization'' scheme outperforms histogram-equalized post-training quantization on ResNet-20/CIFAR-10. Soft quantization provides both a new pipeline for the flexible compression of machine learning models and a new tool for investigating the trade-off between compression and generalization in high-dimensional loss landscapes.

Figures (6)

• Figure 1: Weight distribution of a ResNet-20 layer during the soft quantization procedure of Table \ref{['tab:soft_quant_algorithm']}. (A) Distribution after pretraining on CIFAR-10; the inset shows a schematic of the potential between weight pairs. (B) Distribution after one epoch of soft quantization; the black curve shows the effective compression potential, Eq \ref{['eq:effective_potential']}, acting as a local L1 regularizer. (C) Distribution at the end of the soft quantization; small clusters (in green) are reassigned to nearby larger ones, further reducing the bit-width of the layer.
• Figure 2: Performance of soft quantization compared with HEQ. (A,B) Test accuracy degradation (A) and resulting bit-width (B) of softly quantized models as functions of the potential strength $h$ and range $w$. Colormap midpoints correspond to a 4-bit HEQ-compressed model. The black envelope outlines $(h,w)$ pairs that outperform 4-bit HEQ, achieving higher accuracy at lower bit-width. Each pixel represents an average over 15 independent simulations; standard deviations are small (see Fig. \ref{['fig:coeff_var']} in the End Matter). (C,D) Scatter plots of bit-width versus accuracy degradation for a range of models, compared with 4-bit and 3-bit HEQ benchmarks. Smaller $h$ values approach the Pareto-optimal frontier (D), while smaller $w$ values yield higher bit-width and lower degradation (C). Poor performance at $h = 0.001$ coincides with failure of the clustering refinement step, which reassigns small ($\leq 10$ weights) clusters when the coupling is too weak. Consistent with this, Figure \ref{['fig:clamped_sq']} in the End Matter shows that soft quantization effects are minimal for small $h$. (E) Dynamics of final-layer weights during soft quantization. (F) Loss curves during soft quantization: compression loss (blue) rapidly decreases and saturates, while task loss (red) initially increases before gradually stabilizing.
• Figure 3: Relative task loss change $\frac{\Delta\mathcal{L}_{\mathrm{task}}}{\Delta\boldsymbol{\theta}}$, for compressed and randomly perturbed models. Here $\Delta \theta$ denotes the metric distance from the pretrained model and $\Delta \mathcal{L}_{\mathrm{task}}$ the corresponding change in task loss. Randomly perturbed models are obtained by adding layer-wise Gaussian noise, matching the parameter-space distance of the quantized models from the pretrained one. (A) Soft quantization produces a significantly smaller $\frac{\Delta\mathcal{L}_{\mathrm{task}}}{\Delta\boldsymbol{\theta}}$ than random perturbations at comparable distances. (B) In contrast 4-bit HEQ gives values of $\frac{\Delta\mathcal{L}_{\mathrm{task}}}{\Delta\boldsymbol{\theta}}$ within the tail of the distribution. The softly quantized model shows a bit-width of $\approx3.79$ while being $50\%$ farther from the pretrained model than the 4-bit HEQ model.
• Figure 4: Power law between layer size and associated extensive ($h_l=1$) per-layer compression potential for the ansatz $w_l\sim\sigma_l$ for three distinct pretrained models. The exponent of this power law is approximately $1.66$. Average mean-squared error in this power law is 0.065.
• Figure 5: Coefficients of variation for (A) test accuracy degradation and (B) average bit-width.