Differentiable Soft Min-Max Loss to Restrict Weight Range for Model
Quantization

Arnav Kundu * 1 Chungkuk Yoo * 1 Minsik Cho 1 Saurabh Adya 1

Abstract

the full-precision model’s accuracy slightly as the modelhas to operate under new constraints, but it would have aquantization friendly weight distribution removing outlier
weights. Therefore, lower bit quantization accuracy can be
improved as shown in Figure 1 (right). In case of higher bit
quantization such as 4bit or 8bit, a model might already haveenough bits to properly represent wide ranges of weights.
Therefore the benefit of DSMM loss could be limited.

The range of weights in a model disrupts effective lower bit quantization. Penalizing the range
of weights improve quantization accuracy, but
calculation of range (max-min) is not differentiable. In this work, we propose Differentiable
Soft Min-Max Loss (DSMM) to restrict weight
ranges so that we can get a quantization-friendly
model which has narrow weight ranges. We apply DSMM with a learnable parameter which canadjust hardness of DSMM without requiring a special hyper-parameter. DSMM improves lower bit
quantization accuracy with state-of-the-art posttraining quantization (PTQ), quantization-aware
training (QAT), and weight clustering across various domains and model sizes.

Soft max operation has been proposed in previous works [1],but we bring it to the quantization domain for the first timeby modifying the function. The key innovation of our work
is DSMM loss for reducing a range of weights by operating on outliers more preferably to quantization. To thisend, we make the degree of smoothness (temperature α inEquation (1)) in the function as a learnable parameter sothat it automatically fits to an optimal value without further
hyper-parameter search for the temperature factor. We thenadd a new term e−α term into the function to encourage thetemperature to increase. Without the additional term, thetemperature would move towards negative infinite direction
to minimize the DSMM always and it disrupts the purposeof DSMM, making quantization-friendly pre-trained model.

1. Introduction
Quantization bit-resolution is inversely proportional to therange of weights and affects accuracy of the quantized models. Since outliers tend to increase range, outliers are detrimental for quantization friendly models.

We show that DSMM loss works well with state of theart post-training quantization (PTQ), quantization aware
training (QAT) and weight clustering algorithms. We also
show that our method is applicable to multiple domains likecomputer vision and natural language processing.

As an example, lets assume we want to quantize the weight
distributions shown in Figure 1 (left) into 4 bins. For theoriginal distribution in red most of the weights will be quantized to the central 2 bins and the model accuracy would
drop significantly. This problem gets worse for low bit
quantization such as 1 or 2 bit quantization. To this end,
we introduce Differentiable Soft Min-Max Loss (DSMM),
a simple yet powerful method that helps to reduce weight
range during training without severely affecting full precision accuracy and provides a quantization friendly checkpoint. Using DSMM loss we intend to trim the edges of theblack distribution and convert it to the red distribution asshown in Figure 1 (left). Such a restriction might regress*
Equal contribution 1 AIML, Apple.
Saurabh Adya <sadya@apple.com>.

2. Related Works
In this paper we have applied our DSMM loss with various training time quantization (quantization-aware training,
QAT) algorithms like LSQ [2] and DoReFa [3] used in
PACT [4]. PACT clips activation values with a trainable
parameter for activation quantization and uses DoReFa forweight quantization. LSQ quantizes weights and activations
with learnable step size (scale or bin size).
Also, we have compared our DSMM loss with a state-ofthe-art post-training quantization (PTQ) methods. DFQ[5] equalizes per-channel weight ranges by applying perchannel scaling factors. It resolves the wide weight range
problem across channels, but still the weight range would

Correspondence to:

Published at the 2 nd Differentiable Almost Everything Workshop
at the 41 st International Conference on Machine Learning, Vienna, Austria. July 2024. Copyright 2024 by the author(s).

1Differentiable Soft Min-Max Loss to Restrict Weight Range for Model Quantization

Bin 0

Bin 2

Bin 1

Bin 3

Accuracy

Accuracy benefit from DSMM loss
at lower bit quantization/compression

Frequency

Without DSMM
min ofweights

max ofweights

Accuracy drop by adding
DSMM loss for FP 32 model

Reducing weight range
by DSMMWith DSMM
Without DSMM
With DSMM

0

FP32

Weight value

Higher bit

Lower bit

Figure 1. (Left) Example weight distribution for quantization using 4 bins. (Right) Benefit of DSMM loss for low bit quantization.

approach hard-min-max loss towards the end of training.
This loss smoothly penalizes not only outliers but also nearoutlier weights together. We allow α to be learnable because
in our experiments we found that fixing it introduces a new
hyper-parameter to tune while worsening the accuracy of themodel. A trainable α also allows us to control the smoothness of the loss per layer therefore, introducing more degrees
of freedom to the loss.

remain wide for lower bit quantization like 4bit as DFQ
does not target outliers within a channel. AdaRound [6]proposed adaptive rounding for quantization bin assignment
instead of nearest rounding. SQuant [7] decomposes a layer
by the Hessian-based optimization objective into sub-items,
then composes them in a quantized domain. PD-Quant [8]quantizes weights by comparing model prediction result
before and after quantization of each layer. Additionally,
we show DSMM loss is effective not only for integer quantization but also for weight clustering. DKM [9] introduces
differentiable k-means clustering for weights to represent
them in n-bit centroids which have arbitrary float values.

Σ(W ⊙ eα×(W −Wmax ) )Σeα×(W −Wmax )Σ(W ⊙ e−α×(W −Wmin ) )smin =
Σe−α×(W −Wmin )LDSM M = (smax − smin ) + e−αsmax =

In our extensible experiments, we show our DSMM loss improves accuracies with cutting-edge QAT, PTQ and weightclustering for lower bit quantization like 1bit, 2bit and 4bit.

(1)

DSMM loss was employed during training time of the basemodel itself and not during quantization. This was done
because the purpose of DSMM loss is to provide effective
initial weights for quantization. This ensures extensibilityof DSMM loss to any quantization technique.

3. Differentiable Soft Min-Max Loss
We introduce Differentiable Soft Min-Max Loss as an auxiliary loss to reduce the range of weights for every layer to getbetter pre-trained models for further quantization or compression. Just like L1 and L2 regularization our approach isinvariant to the quantization or compression technique used.
But as opposed to L1 or L2 regularization, Differentiable
Soft Min-Max Loss only affects the range of the distribution
and not the absolute magnitude of it. In our experiments,
we demonstrate that L2 regularization (1x(baseline) and10x(heavy L2)) only affects the magnitude of the weights
but does not remove outliers from the distribution. As evident from Figure 1, heavier L2 regularization just reduces
the scale of weights but does affect the shape of the distribution whereas DSMM loss reshapes the distribution.
To capture the range we intend to calculate the difference
between the maximum and minimum value of the weights
in a easily differentiable form as illustrated below. The loss
for a given weight W is described in Equation (1).

4. Experiment4.1. Post-Training Quantization with DSMM loss
We compare models trained using DSMM loss and otherweight regularization, L2 and heavy L2, using PTQ methods such as DFQ [5], AdaRound [6], SQuant [7], and
PD-Quant [8]. As shown in Table 1, models trained with
DSMM loss are more quantization friendly than other regularization. DSMM loss shows the best accuracy as it reduces
outliers as well as weight range. On the other hand, heavy
L2 regularization makes weight ranges smaller, but it does
not remove outliers, therefore prove to be ineffective here.
Compared to FP32 accuracy of baseline models, models
trained with DSMM loss, have slight accuracy regression in
full-precision inference as we expected in Figure 1 (right).
However, after quantizing, the models trained with DSMM
loss shows better accuracies.

Here temperature α is a learnable parameter per layer. e−αterm in the auxiliary loss Ldsmm , encourages temperature
α to increase during training time optimization process to
2

DSMM

Baseline

Heavy L2

Differentiable Soft Min-Max Loss to Restrict Weight Range for Model Quantization

Figure 2. Weight distribution of the first five 3x3 convolution layers of MobileNet-V2 using L2 norm (baseline), heavy L2 norm (10xheavy L2 norm than baseline) and the proposed Differentiable Soft Min-Max Loss (the red dots correspond outliers).
Table 1. Top-1 accuracies (%) of MobileNet-V1 and V2 on ImageNet-1K using PTQ methods with 4bit weight and 8bit activation
quantization. Heavy L2: applied 10x heavy L2 regularization than baseline. Naı̈ve: quantizing without any advanced PTQ techniques.
DFQ [5], AR [6], SQ [7], PD-Q [8]. DSMMα : fixed α (10.0). DSMM: trainable α.

FP32

NA ÏVE

M OBILE N ET-V1
DFQ
AR

SQ

PD-Q

FP32

NA ÏVE

M OBILE N ET-V2
DFQ
AR

SQ

PD-Q

BASELINEH EAVY L2

74.12
72.67

2.67
13.41

54.06
57.68

70.42
69.23

63.85
66.51

71.87
69.86

73.08
71.00

2.57
4.17

56.56
0.09

71.29
68.06

59.30
57.30

71.54
68.92

DSMMα=10
DSMM( OURS )

74.21
73.95

8.52
44.24

42.48
59.21

69.71
71.35

67.23
67.86

71.92
72.39

72.80
72.81

13.77
36.69

56.19
51.46

70.79
71.77

62.16
67.16

71.27
71.98

M ETHOD

Table 2. BLEU score for machine translation task on Transformer
Base model [10] with 8bit weight quantization.

M ETHOD

FP32

NA ÏVE

BASELINE

28.2

2.9

DSMM

27.9

27.8

Table 4. Top-1 accuracies (%) of ResNet50 and ResNet101 on
ImageNet-1K with 4bit weight and 8bit activation PTQ.

Table 3. Top-1 accuracies (%) of 2bit weight and 8bit activation PTQ using PD-Quant [8]. MNV1: MobileNet-V1, MNV2:
MobileNet-V2, RN50: ResNet-50, RN101: ResNet-101. FP32
accuracy is in Table 1 and Table 4. DSMM1 : applied DSMM fromscratch. DSMM2 : applied DSMM during fine-tuning

M ETHOD

MNV1

MNV2

RN50

RN101

BASELINE

47.62

50.66

62.92

66.50

54.20
53.10

57.53
56.31

69.31
66.85

71.24
69.66

1

DSMM
DSMM2

FP32

R ES N ET 50
AR
SQ

PD-Q

BASELINE

78.04

74.73

74.68

76.60

DSMM

78.22

75.61

75.75

77.21

FP32

R ES N ET 101
AR
SQ

PD-Q

BASELINE

79.45

75.18

75.23

78.16

DSMM

79.63

76.71

77.20

78.49

M ETHODM ETHOD

we finetuned pre-trained models with a low learning rate
while applying DSMM as an additional loss to the existingtask loss. Our experimental results are shown in Table 3
where we applied DSMM on pre-trained models, followedby PD-Quant (PD-Q) [8] for PTQ. It can be observed that on
quantization, finetuning a pre-trained model does not show
the same performance gains as training from scratch but
it is still significantly better than the baseline. This showsthat our method is also applicable to pre-trained models if
finetuned with the additional loss.

Even without advanced PTQ approaches, models trainedwith DSMM loss can be reasonably quantized without any
further fine-tuning (See Naı̈ve in Table 1). This proves thatmodels with DSMM loss have good weight distribution sothey can be quantized with fairly high quantization accuracies. This approach scales to other applications like machinetranslation as illustrated in Table 2. It is also applicable tolarger models like ResNet{50,101} as shown in Table 4.

4.3. Quantization-Aware Training with DSMM loss

In Table 3, we can clearly see the benefit of DSMM lossfor
lower bit PTQ. The accuracy of ResNet101 trained with
DSMM loss only regressed by 8.39% (79.63% → 71.24%),
while the baseline model trained without DSMM loss showshigher absolute regression, 12.95% (79.45% → 66.50%).

We apply state-of-the-art quantization techniques like PACT
[4] while training the models from scratch using DSMM
loss. For LSQ [2] we initialize the model to pre-trained
ResNet-18 (RN) [11] with DSMM loss.
As shown in Table 5, DSMM loss helps the quantization
techniques in improving their accuracy, especially for extremely low bit quantization such as at 2 bit while it shows
similar accuracies with 4 bit. For example, DSMM loss improves 2 bit quantization accuracy with LSQ to over than

4.2. PTQ for models fine-tuned with DSMM loss
Training models from scratch might not always be feasible especially given the cost and time taken. Therefore,3

Differentiable Soft Min-Max Loss to Restrict Weight Range for Model Quantization70

M ETHOD

2W2A QATFP32 PACT LSQBASELINE

69.76

51.97

58.33

DSMM

69.84

55.64

62.47

M ETHOD

4W4A QATFP32 PACT LSQBASELINE

69.76

66.90

69.90

DSMM

69.84

68.36

69.45

Top-1 Accuracy (%)

Table 5. Top-1 accuracies (%) of ResNet18 with various QAT methods. Weights and activations are quantized with the same bit(2W2A: 2bit, 4W4A: 4bit).

65
60
ResNet18 Baseline
ResNet18 DSMM
MobileNet-V1 Baseline
MobileNet-V1 DSMM

55
50

Improvement
by DSMM

45
1B/1D

2B/2D

4B/4D

2B/1D

4B/1D

Figure 3. Top-1 accuracy for various compression ratios and models using weight clustering. nB/mD: n-bit/m-dim weight clustering.

training task of MobileBERT from scratch so that DSMM
loss would have more chances to get effective weight distribution for model compression.

62% from 58%, but there is no noticeable difference in 4
bit LSQ accuracies with and without DSMM loss. The reason why DSMM loss would not help much for higher bit
like 4 bit quantization is that QAT can effectively represent
outliers using many bits as we expected in Figure 1 (right).

Table 6. Question-answering NLI (QNLI) accuracies of MobileBERT using single dimension DKM

M ETHOD

4.4. Weight Clustering with DSMM loss
We evaluate the effectiveness of DSMM loss with thestate-of-the-art weight clustering technique, DKM [9], for
ResNet-18 and MobileNet-V1. The bit-dim ratio, db is an
important factor in the DKM algorithm which effectively
defines the kind of compression a DKM palettized model
would see. We ran these experiments for both scalar andvector palettization. For scalar palettization(dim = 1) weran 1 bit, 2 bit and 4 bit compression. Figure 3 shows that
DSMM loss significantly improves accuracy from DKM 1
bit and 2 bit models. As we discussed, there is no significant difference for higher bit like 4 bit because many bitcompression can also cover outliers even without DSMM.

P RE - TRAIN

1- BIT

2- BIT

DKM BASELINE

90.41

61.34

80.12

DKM + DSMM

90.83

61.49

80.87

5. Discussion
As we discussed earlier, DSMM loss would be effective
for lower bit quantization, because higher bit quantization can represent outlier weights and wide weight range.
Also most of state-of-the-art quantization techniques alreadyachieved reasonable accuracy with higher bit comparing its
full-precision models, so there would not be a room for
improvement further in higher bit quantization.
We are planning to conduct more studies with DSMM, suchas expending to large language model quantization and providing theoretical grounding from current empirical studies.

We also expand the application of DSMM loss to vector palettization(2-bit/2-dim and 4-bit/4-dim) DKM [9] asdemonstrated in Figure 3. For these experiments, we kept
the effective bit-dim ratio, db equivalent to 1 so as to see variation across the models almost compressed to 32x. Sincea vector palettized model will require range constraining
for all dimensions, we applied multi-dimensional DSMM
loss for all layers. For vector palettized ResNet-18 there
is an average absolute improvement of > 1% using models trained with DSMM loss, and for vector palettized
MobileNet-V1, the gain ranges from 2% to 5%.

6. Conclusion
In this paper, we introduced Differentiable Soft Min-Max
Loss as an effective technique to reduce the weight range
for quantization for multiple tasks across domains. This
serves as a good initialization for PTQ, QAT and weightclustering methods, and hence can be coupled with any ofthem. This helps to augment the accuracy gained from suchtechniques and is invariant to the quantization algorithm.
We demonstrated how DSMM loss converts a wide weight
range distribution to a more densely-packed distribution
for model quantization. While full-precision accuracy with
DSMM loss can be slightly regressed as it penalizes outlier
weights, it significantly improves quantization accuracy,
especially for lower bit.

Finally we also validated that DSMM loss for weight clustering scales to other domains as well by applying it incompressing MobileBERT [12]. For Question Answering (QNLI) [13] using MobileBERT, DSMM loss slightly
improved the performance of the model as demonstrated
in Table 6. Note that we applied DSMM loss to a QNLI
fine-tuning task based on a pre-trained MobileBERT [14].
It might be necessary to apply DSMM loss to the entire4Differentiable Soft Min-Max Loss to Restrict Weight Range for Model Quantization

References
[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]

[10]

[11]

[12]

[13]

[14]

[15]

[16]

K. Asadi and M. L. Littman, “An alternative softmax operator for reinforcement learning,” in Proceedings of the34th International Conference on Machine Learning, D.
Precup and Y. W. Teh, Eds., ser. Proceedings of Machine
Learning Research, vol. 70, PMLR, Jun. 2017, pp. 243–
252. [Online]. Available: https : / / proceedings .
mlr.press/v70/asadi17a.html.
S. K. Esser, J. L. McKinstry, D. Bablani, R. Appuswamy,
and D. S. Modha, “Learned step size quantization,” in International Conference on Learning Representations, 2019.
S. Zhou, Y. Wu, Z. Ni, X. Zhou, H. Wen, and Y. Zou,
“Dorefa-net: Training low bitwidth convolutional neuralnetworks with low bitwidth gradients,” arXiv preprint
arXiv:1606.06160, 2016.
J. Choi, Z. Wang, S. Venkataramani, P. I.-J. Chuang, V.
Srinivasan, and K. Gopalakrishnan, “Pact: Parameterizedclipping activation for quantized neural networks,” arXiv
preprint arXiv:1805.06085, 2018.
M. Nagel, M. v. Baalen, T. Blankevoort, and M. Welling,“Data-free quantization through weight equalization andbias correction,” in Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019, pp. 1325–
1334.
M. Nagel, R. A. Amjad, M. Van Baalen, C. Louizos, andT. Blankevoort, “Up or down? adaptive rounding for posttraining quantization,” in International Conference on Machine Learning, PMLR, 2020, pp. 7197–7206.
C. Guo, Y. Qiu, J. Leng, et al., “SQuant: On-the-fly datafree quantization via diagonal hessian approximation,” inInternational Conference on Learning Representations,2022. [Online]. Available: https : / / openreview .
net/forum?id=JXhROKNZzOc.
J. Liu, L. Niu, Z. Yuan, D. Yang, X. Wang, and W. Liu,
“Pd-quant: Post-training quantization based on prediction
difference metric,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023,
pp. 24 427–24 437.
M. Cho, K. Alizadeh-Vahid, S. Adya, and M. Rastegari,“Dkm: Differentiable k-means clustering layer for neuralnetwork compression,” in International Conference on
Learning Representations, 2021.
A. Vaswani, N. Shazeer, N. Parmar, et al., “Attention is
all you need,” Advances in neural information processing
systems, vol. 30, 2017.
K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in Proceedings of the IEEE
conference on computer vision and pattern recognition,
2016, pp. 770–778.
Z. Sun, H. Yu, X. Song, R. Liu, Y. Yang, and D. Zhou,“Mobilebert: A compact task-agnostic bert for resourcelimited devices,” arXiv preprint arXiv:2004.02984, 2020.
P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang, “Squad:
100,000+ questions for machine comprehension of text,”
arXiv preprint arXiv:1606.05250, 2016.
T. Wolf, L. Debut, V. Sanh, et al., “Transformers: State-ofthe-art natural language processing,” in Proceedings of the2020 conference on empirical methods in natural language
processing: system demonstrations, 2020, pp. 38–45.
A. G. Howard, M. Zhu, B. Chen, et al., “Mobilenets: Efficient convolutional neural networks for mobile visionapplications,” arXiv preprint arXiv:1704.04861, 2017.

[17]

5

M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L.-C.
Chen, “Mobilenetv2: Inverted residuals and linear bottlenecks,” in Proceedings of the IEEE conference on computer
vision and pattern recognition, 2018, pp. 4510–4520.
J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. FeiFei, “Imagenet: A large-scale hierarchical image database,”in 2009 IEEE conference on computer vision and patternrecognition, Ieee, 2009, pp. 248–255.

Differentiable Soft Min-Max Loss to Restrict Weight Range for Model Quantization

A. Experiment settings
A.1. Pre-training from scratch with and without DSMM loss
We train ResNet-18 [11], MobileNet-V1 [15] and MobileNet-V2 [16] on ImageNet 1K [17] with proposed Differentiable
Soft Min-Max Loss on a x86 Linux machine with eight GPUs to get pre-trained models before model compression andquantization-aware training. We set initial learning rates to 1.0, 0.4 and 0.4 for ResNet-18, MobileNet-V1 and MobileNet-V2
respectively. We use SGD with 0.9 of momentum with Nesterov. We apply 1e-4 of weight decay (L2 norm weight
regularization) for ResNet-18 and 4e-5 for MobileNet-V1 and V2. For heavy L2-regularization, in Figure 2, we use 4e-4 ofweight decay (10x heavier than baseline) for MobileNet-V2 to see whether heavy L2-regularization helps quantization or
not as a naive solution for range restriction. Strength of DSMM loss is set to 0.01. The learnable parameter α is initially set
to 0.1. For comparison, we use pre-trained models of Resnet-18 from Torchvision. As we are using modified version of
ResNet-50, and ResNet-101, MobileNet-V1 and V2 for better FP32 performance, we trained those models from scratch
without DSMM loss using the same settings above.
A.2. Quantization and weight clustering
DSMM loss is not a model compression nor quantization method. It penalizes outlier weights during training of thebase model from scratch. To evaluate the effectiveness of DSMM loss with model compression and quantization, weapply state-of-the-art compression/quantization techniques, DKM [9], LSQ [2], DFQ [5], AdaRound [6], SQuant [7], and
PD-Quant [8] to the pre-trained model with and without DSMM loss. Except SQuant 1 , PD-Quant 2 , DFQ and AdaRound 3 ,since other works do not provide official implementation, we implement those techniques ourselves.
We follow the same hyper-parameters used in the works, but we apply compression and quantization for all layers including
the first and last layers. It is important to compress/quantize all layers including first and last layers considering computation
burden at the first layer with a large convolutional filter size such as 7x7 convolutions in the first layer of ResNet and thelarge number of weights in the last linear layer, e.g., 1.2M of weights in the last layer of MobileNet-V1 which has 4.2M ofweights in total.

1

https://github.com/clevercool/SQuant
https://github.com/hustvl/PD-Quant3
https://quic.github.io/aimet-pages/
26