---
title: 'DiffQ: Differentiable Model Compression'
url: https://www.emergentmind.com/topics/diffq
type: topic
---

# DiffQ: Differentiable Model Compression

DiffQ is a differentiable method for neural-network model compression by weight quantization that replaces hard quantization during training with additive pseudo quantization noise, thereby avoiding gradient approximations such as the Straight-Through Estimator (STE). Introduced in “Differentiable Model Compression via Pseudo Quantization Noise” [2104.09987], it is designed to optimize both the underlying real-valued weights and the bitwidth allocation used for quantization, including mixed-precision settings in which different weights or groups of weights use different numbers of bits. The method targets end-to-end training under a single tradeoff parameter balancing task loss against compressed model size, and is evaluated on image classification, language modeling, and audio source separation.

## 1. Concept and problem formulation

DiffQ addresses model compression by weight quantization, especially when low-precision weights are required with minimal accuracy loss and when mixed precision is desirable. In the paper’s formulation, standard uniform quantization maps a real weight \(w \in [0,1]\) to one of \(2^B\) levels:
\[
Q(w,B)=\frac{\mathrm{round}\!\left(w(2^B-1)\right)}{2^B-1}.
\]
For a weight vector \(w\in\mathbb{R}^d\), quantization is preceded by layerwise normalization,
\[
\hat w=\frac{w-\min(w)}{\max(w)-\min(w)},
\]
with \(\min(w)\) and \(\max(w)\) stored separately. The ideal quantization-aware objective is
\[
\min_{w\in\mathbb{R}^d} L\bigl(f_{Q(w,B)}\bigr),
\]
but the rounding operator has zero gradient almost everywhere, which makes direct gradient-based optimization difficult [2104.09987].

The paper positions DiffQ against the common STE-based approach to quantization-aware training. In that setup, \(Q(w,B)\) is used in the forward pass while the quantizer’s Jacobian is replaced by the identity in backpropagation. The paper characterizes this as a usable but biased surrogate-gradient method and gives a one-dimensional least-squares example in which STE causes the learned quantized value to oscillate between adjacent quantization levels rather than converge to the optimal quantized point. This motivates replacing hard quantization during training with a differentiable approximation whose scale depends continuously on the bitwidth.

A central feature of DiffQ is that bitwidth is itself optimized. Rather than fixing \(B\) as an integer hyperparameter, the method relaxes bit allocations to continuous variables during training and penalizes model size directly. This makes mixed-precision allocation part of the same optimization problem as task fitting.

## 2. Pseudo quantization noise and differentiability

The technical core of DiffQ is pseudo quantization noise (PQN). The paper starts from the signal-processing approximation that true quantization noise can be modeled as an independent random variable uniformly distributed in half a quantization bin. If
\[
\Delta=\frac{1}{2^B-1},
\]
then the true quantization noise is
\[
N(w,B)=Q(w,B)-w,
\]
and the training-time pseudo-quantizer is defined as
\[
\widetilde Q(x,B)=x+\frac{\Delta}{2}\cdot U,\qquad U\sim \mathcal U[-1,1].
\]
Because the training graph contains additive noise rather than hard rounding, \(\widetilde Q(x,B)\) is differentiable with respect to both the real-valued weight \(x\) and the bitwidth parameter \(B\). The paper notes in particular that
\[
\frac{\partial \widetilde Q}{\partial x}=1,
\]
and that the dependence on \(B\) enters through \(\Delta=(2^B-1)^{-1}\), so gradients can push the bitwidth upward or downward according to the loss [2104.09987].

The paper’s toy analysis compares this with STE on the objective
\[
\min_{w\in[0,1]} L(w):=\mathbb{E}\left[\frac12\left(XQ(w,B)-Xw_*\right)^2\right].
\]
Under STE, the expected gradient is
\[
G_n=\sigma^2\bigl(Q(w_n,B)-w_*\bigr), \qquad \sigma^2=\mathbb{E}[X^2],
\]
which depends on the quantized value and can induce oscillation between neighboring quantization levels. Replacing \(Q\) by \(\widetilde Q\) gives
\[
G_n=\sigma^2(w_n-w_*),
\]
because \(\mathbb{E}[U]=0\). The paper presents this as an unbiased approximation in the toy setting and as the main analytic argument for PQN.

Although the theoretical exposition uses uniform pseudo noise, the implementation used in the main experiments employs Gaussian noise with unit variance. The paper states that this empirically improved robustness to the final mismatch introduced by hard rounding of both weights and bitwidths at inference time. On ResNet-18/CIFAR-10 at 2.7MB, the paper reports \(86.9\%\) with uniform noise and \(93.6\%\) with Gaussian noise, and all main experiments use Gaussian noise.

## 3. Compression objective and bitwidth parameterization

DiffQ formulates compression either as constrained optimization of task loss under a model-size budget or as constrained minimization of model size under a target loss. The relaxed training objective is written as
\[
\min_{w,b} L\bigl(f_{\widetilde Q(w,b)}\bigr)+\lambda M(b),
\]
where \(\lambda\) is the main tradeoff hyperparameter and \(M(b)\) is a train-time estimate of model size [2104.09987].

The paper uses a group-based representation of bit allocation. A tensor of \(d\) weights is partitioned into groups of size \(g\), each group \(s\) being assigned a bitwidth \(b_s\). This supports per-weight allocation when \(g=1\), small-group mixed precision for finite \(g\), and effectively layerwise allocation when \(g=\infty\). Most experiments use group-wise allocation with default group size \(g=8\), which the paper describes as the best tradeoff between flexibility and metadata overhead.

The train-time model-size penalty is
\[
M(b)=\frac{g\sum_{s=1}^{d/g} b_s}{2^{23}}.
\]
At evaluation time, the actual representation also includes per-layer min/max scaling values and metadata for encoding the selected bitwidth of each group. The paper gives the true compressed layer size as
\[
\tilde M(b)=\frac{1}{2^{23}} \left( 2\cdot 32 + 8 + \frac{d}{g}\max(C) + g\sum_{s=1}^{d/g} b_s \right),
\]
where
\[
C_s=\log_2\!\bigl(1+\tilde b_s-b_{\min}\bigr),
\]
and \(\tilde b_s=\mathrm{round}(b_s)\) is the final integer bitwidth.

Continuous bitwidths are parameterized through logits \(l\) via
\[
b = b_{\min} + \sigma(l)(b_{\max}-b_{\min}),
\]
with default settings
\[
b_{\min}=2,\qquad b_{\max}=15,\qquad b_{\text{init}}=8.
\]
This makes each \(b_s\) differentiable and keeps it within a bounded interval. The same \(b_s\) determines both the noise scale
\[
\Delta_s=\frac{1}{2^{b_s}-1}
\]
and the size penalty \(M(b)\), so the learned bitwidth allocation directly reflects the compression–accuracy tradeoff induced by \(\lambda\).

The paper also emphasizes that DiffQ optimizes naive bit count rather than entropy-coded size. It notes, for example, that on Wikitext-103 a QAT 8-bit model compresses from 236MB naive to 150MB after ZLib, whereas a DiffQ model changes from 113MB naive to 122MB after ZLib. This suggests that DiffQ encourages fuller use of the available discrete code space rather than low-entropy weight distributions.

## 4. Training procedure and implementation

The training loop described in the paper is conceptually simple. A standard full-precision model is augmented with trainable bit logits for each weight group. During each forward pass, a group-specific step size
\[
\Delta_s=\frac{1}{2^{b_s}-1}
\]
is computed from the current continuous bitwidth, pseudo quantization noise is added to the weights, the task loss is evaluated with the noisy weights, and the model-size penalty is added:
\[
L_{\text{total}} = L\bigl(f_{\widetilde Q(w,b)}\bigr)+\lambda M(b).
\]
Backpropagation then updates both the real-valued weights and the bitwidth parameters, and at inference time the continuous bitwidths are rounded,
\[
\tilde b_s=\mathrm{round}(b_s),
\]
followed by standard hard quantization with the resulting integer precisions [2104.09987].

The paper states that DiffQ can optimize the number of bits used per individual weight or groups of weights. It also reports architectural patterns in learned allocation: for EfficientNet-B3, earlier layers tend to keep higher bitwidth while deeper layers use lower bitwidth except for the final classifier, and for DeiT the pattern is described as similar but less pronounced and more symmetric across depth.

Implementation details matter for reproducibility. The main experiments use a separate Adam optimizer for the bit logits, typically with learning rate \(10^{-3}\), except \(10^{-2}\) for language modeling, and \(\beta_1=0.9\), \(\beta_2=0.999\). Very small tensors are not quantized: below 0.01MB in general and below 0.1MB for DeiT. The paper also notes that tied embeddings and output weights in the language model require shared bit parameters and shared sampled pseudo noise across reuses; otherwise validation performance degrades significantly.

Code is stated to be available at `github.com/facebookresearch/diffq`. The method is described as simple to implement because it amounts to injecting noise with bit-dependent scale and adding a model-size penalty, rather than introducing a separate surrogate-gradient mechanism.

## 5. Empirical performance

DiffQ is evaluated on image classification, language modeling, and audio source separation, and the paper’s overall claim is that it is competitive with or better than STE-based and related baselines in the compression–accuracy tradeoff [2104.09987].

A concise summary of several headline results is given below.

| Benchmark | DiffQ result | Comparison stated in the paper |
|---|---:|---|
| DeiT-B / ImageNet | 81.5%, 33.02MB | 81.8%, 371.4MB uncompressed; 81.6%, 82.9MB for QAT 8-bit |
| Wikitext-103 / 16-layer Transformer | PPL 18.6, 113MB | PPL 18.1, 942MB uncompressed; PPL 29.9, 118MB for QAT 4-bit |
| Demucs / MusDB | 6.28 SDR, 120MB | 6.31 SDR, 1014MB uncompressed; 5.99 SDR, 130MB for QAT 4-bit |

The abstract’s most prominent example is ImageNet with DeiT-B, a 12-layer transformer-based model. The paper reports that DiffQ compresses it from 371.4MB to 33.02MB with top-1 accuracy changing from \(81.8\%\) to \(81.5\%\), which the abstract summarizes as more than an \(8\times\) compression at lower than 4 bits precision per weight on average, with about \(0.3\%\) accuracy loss. The same table reports QAT 8-bit at \(81.6\%\) and 82.9MB, and QAT 4-bit at \(79.2\%\) and 41.7MB.

On Wikitext-103 with a 16-layer Transformer, the paper reports: uncompressed perplexity \(18.1\) at 942MB; QAT 8-bit plus 8-bit activations at \(19.7\) and 236MB; QAT 4-bit plus 8-bit activations at \(29.9\) and 118MB; LSQ 4-bit plus 8-bit activations at \(18.9\) and 118MB; DiffQ at \(\lambda=5, g=16\) with \(18.1\) and 130MB; and DiffQ at \(\lambda=10, g=16\) with \(18.6\) and 113MB. The paper interprets this as showing that DiffQ reaches virtually baseline perplexity at about 4.4 bits per weight on average and that QAT “breaks down” at 4 bits in this task.

For audio source separation with Demucs on MusDB, the reported results are 6.31 SDR and 1014MB uncompressed, 6.27 SDR and 162MB for QAT 5-bit, 5.99 SDR and 130MB for QAT 4-bit, and 6.28 SDR and 120MB for DiffQ at \(\lambda=3\times 10^{-4}\). The paper interprets this as a smaller model than QAT 4-bit with only 0.03dB loss from baseline, compared with 0.3dB loss for QAT 4-bit.

The paper also includes fixed-bit comparisons intended to isolate PQN versus STE rather than mixed-precision allocation. On ResNet-18/CIFAR-10, uncompressed accuracy is \(95.3\), QAT 2-bit gives \(87.2\), and DiffQ 2-bit fixed gives \(92.3\). On ResNet-18/CIFAR-100, QAT 2-bit gives \(58.7\) while DiffQ 2-bit fixed gives \(66.6\). On WideResNet/CIFAR-100, QAT 2-bit gives \(46.5\) while DiffQ 2-bit fixed gives \(72.8\). The paper presents these as direct evidence for the benefit of pseudo quantization noise itself.

## 6. Position among related methods and limitations

DiffQ is presented as a differentiable alternative to STE-based quantization-aware training, and the paper repeatedly emphasizes that it avoids surrogate gradients. It is also compared with LSQ, DQ, NICE/UNIQ, Quant-Noise, and Product Quantization in various experiments, with the general conclusion that it is competitive or superior in the compression–accuracy tradeoff on the reported benchmarks [2104.09987].

The method also has explicit limitations. First, pseudo quantization noise is an approximation to real quantization noise, and the paper notes that this is less reliable in the extreme low-bit regime, especially around 1 bit. It states that \(b_{\min}<2\) is often unstable in practice, although some task-specific experiments alter \(b_{\min}\). Second, the train-time size penalty models raw bit count rather than entropy-coded size, so it does not directly optimize final compressed entropy. Third, training incurs overhead from the extra bit parameters and the stochastic noise injection. The paper reports batch-time increases such as 115ms to 125ms for language modeling, 120ms to 150ms for ResNet-18 on CIFAR-10, 0.9s to 1.1s for Demucs, and nearly doubling for EfficientNet-B3.

A further limitation is that \(\lambda\) must still be swept to reach a desired size–accuracy operating point, even though it is the only sensitive tradeoff hyperparameter highlighted by the method. The paper also notes the need for careful treatment of shared or tied weights, since bit parameters and sampled noise must be shared consistently across reuses.

Finally, the name “DiffQ” is not unique across the research literature. The data block also includes “DiffQ: Unified Parameter Initialization for Variational Quantum Algorithms via Diffusion Models” [2509.17324], which uses the same name for a distinct method in variational quantum algorithms. In the context of model compression and pseudo quantization noise, however, DiffQ refers specifically to the differentiable compression method introduced in 2021 [2104.09987].

Source: https://www.emergentmind.com/topics/diffq