PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Bits: Unifying Quantization and Pruning

Mart van Baalen, Christos Louizos, Markus Nagel, Rana Ali Amjad, Ying Wang, Tijmen Blankevoort, Max Welling

arXiv:2005.07093v3cs.LGcs.CVstat.ML

TL;DR

Selecting bit widths for neural networks is difficult because tensor choices interact, exhaustive search is impractical, and efficient hardware typically supports power-of-two widths. Bayesian Bits uses gated quantized residuals to jointly learn bit widths and pruning, and the authors report more efficient networks than prior art.

  • Problem

    Bit-width choices interact across tensors, while the configuration space grows exponentially and learned widths may not match hardware-friendly power-of-two settings.

    Source: §1 · p. 1

  • Method

    Bayesian Bits recursively quantizes residuals and uses stochastic gates with sparsifying priors to select bit widths and optionally gate the whole quantized value for pruning.

    Source: §1 · p. 2, §2.2 · p. 4, §2.3 · p. 5

  • Results

    Bayesian Bits finds more efficient networks than prior art.

    Source: §5 · p. 9

  • Takeaways & Limitations

    The approach provides a hardware-friendly unified view of pruning and quantization.

    Source: §5 · p. 9

    The authors did not investigate how their method affects prediction fairness.

    Source: Broader Impact · p. 9

Abstract

from arXiv · show

We introduce Bayesian Bits, a practical method for joint mixed precision quantization and pruning through gradient based optimization. Bayesian Bits employs a novel decomposition of the quantization operation, which sequentially considers doubling the bit width. At each new bit width, the residual error between the full precision value and the previously rounded value is quantized. We then decide whether or not to add this quantized residual error for a higher effective bit width and lower quantization noise. By starting with a power-of-two bit width, this decomposition will always produce hardware-friendly configurations, and through an additional 0-bit option, serves as a unified view of pruning and quantization. Bayesian Bits then introduces learnable stochastic gates, which collectively control the bit width of the given tensor. As a result, we can obtain low bit solutions by performing approximate inference over the gates, with prior distributions that encourage most of them to be switched off. We experimentally validate our proposed method on several benchmark datasets and show that we can learn pruned, mixed precision networks that provide a better trade-off between accuracy and efficiency than their static bit width equivalents.

1 Introduction

Quantization and compression reduce inference cost, but selecting bit widths is difficult because tensor choices interact and exhaustive search is infeasible. Bayesian Bits decomposes quantization into gated residual additions, enabling joint learning of hardware-friendly bit widths, pruning, and model parameters.

1 Introduction

  • Quantization reduces weight and activation bit widths for integer arithmetic, while compression aims to reduce multiply-accumulate operations.Both techniques are applied before deployment to reduce neural-network inference cost.Source: §1 · p. 1
  • Bit-width choices can affect the quantization sensitivity of other tensors, so they cannot be selected independently across a network.Recent work found that selecting each tensor's bit width carefully can yield a better accuracy-complexity trade-off than assigning every layer the same width.Source: §1 · p. 1
  • The number of bit-width configurations grows exponentially with network layers, making exhaustive search impractical.Prior learning approaches do not account for commercially available hardware typically supporting efficient computation in power-of-two bit widths.Source: §1 · pp. 1–2
  • Bayesian Bits recursively quantizes residual errors and uses learnable gates to control effective bit widths, including a 0-bit option for pruning.The approach optimizes gates jointly with quantization scales and network parameters, using priors that favor low bit widths.Source: §1 · p. 2

2 Unifying quantization and pruning with Bayesian Bits

Bayesian Bits builds hardware-friendly quantization from sequential power-of-two residual additions, with stochastic gates controlling effective precision and a zero-bit option for pruning. A Bayesian objective encourages efficient gate configurations, with practical gradient-based optimization and test-time thresholding.

2.1 Mixed precision gating for quantization and pruning

  • Sequentially quantizing and adding residuals doubles precision at each step while retaining power-of-two bit widths.The residual at each new width is quantized using a grid compatible with the preceding grid.Source: §2.1 · p. 2
  • Turning off the lowest-bit gate assigns the tensor value zero, enabling pruning, including structured pruning by using a quantizer per convolutional filter.Higher-bit gates and grid step sizes can be shared so surviving elements use the same quantization grid.Source: §2.1 · p. 3

2.2 Bayesian Bits

  • Binary stochastic gates control residual additions, while variational inference uses priors favoring low-bit configurations and conditions higher-bit gates on lower-bit gates.The lower-bit inclusion probabilities scale the KL penalty for higher-bit gates, avoiding over-regularization when higher bits contribute only if lower gates are active.Source: §2.2 · pp. 3–4

2.3 A simple approximation for learning the bit width

  • The simplified objective penalizes the probability of including quantizer parameters and additional precision, using a large-prior approximation with prior strength scaled by dataset size.Under assumptions stated in the paper, the resulting objective is related to L0 regularization.Source: §2.3 · pp. 4–5

2.4 Practical considerations

  • Hard-concrete gate relaxations enable reparameterized gradient optimization; at test time, gates are pruned when their exact-zero probability exceeds a threshold.The method uses PACT clipping with learned ranges and a straight-through estimator for rounding.Source: §2.4 · p. 5

3 Related work

Related approaches learn mixed-precision assignments through quantization, sensitivity analysis, architecture search, reinforcement learning, and Bayesian methods. Bayesian Bits is distinguished by its recursively quantized residuals and fixed-point tensors by construction.

3 Related work

  • Differentiable Quantization learns ranges and scales jointly with weights, but its inferred widths must be rounded up to powers of two for hardware use, limiting control over resulting widths.Other methods use Hessian eigenvalues or trace for sensitivity [6], or posterior weight uncertainty to infer widths.Source: §3 · pp. 5–6
  • Mixed-precision search has also been framed as architecture search with separate weight tensors [38] or reinforcement learning with validation accuracy and target-hardware feedback.The architecture-search approach uses a stochastic version of DARTS.Source: §3 · p. 6
  • Prior work learns quantization scales [15] or activation clipping ranges [5] jointly with model parameters.Source: §3 · p. 6
  • Bayesian Bits relates to recursive residual vector quantization and neural-network compression, but produces valid fixed-point tensors by construction unlike.Concurrent work uses one decomposition step without handling scales or extending to multiple bit widths.Source: §3 · p. 6

4 Experiments

Across the reported experiments, Bayesian Bits finds better accuracy–complexity trade-offs than the compared baselines, including in ImageNet and toy-task evaluations. The experiments also examine post-training mixed precision and show how regularization and compute requirements shape the results.

4 Experiments

  • Across image-classification evaluations, weights and activations are quantized per tensor, and test accuracy is paired with BOPs as a hardware-agnostic complexity proxy.All quantizer bit widths are learned; the first and last layers are not held at higher precision as is common in some prior work.Source: §4 · p. 6
  • A global regularization parameter scales a gate prior whose inclusion penalty tracks each gate’s contribution to model BOPs.The experiments use this prior to encourage efficient gate configurations.Source: §4 · p. 6

4.1 Toy experiments on MNIST & CIFAR 10

  • On MNIST and CIFAR-10, Bayesian Bits improves accuracy–complexity trade-offs over the compared baselines; stronger CIFAR-10 regularization lowers both accuracy and complexity.Table 1 reports means and standard errors over three runs.Source: §4.1 · p. 7, Table 1 · p. 7
  • Toy-task models often retain higher precision in the first and last layers; under moderate VGG regularization, weights are mostly 2-bit while activations range from 2 to 16 bits.At aggressive regularization, nearly all tensors are quantized to 2-bit; moderate VGG regularization prunes only a few channels.Source: §4.1 · p. 7

4.2 Experiments on Imagenet

  • On ImageNet, Bayesian Bits yields better accuracy–complexity trade-offs than strong baselines, and joint pruning plus quantization beats either ablation.Figure 2a reports ResNet18 validation accuracy; its caption notes means over three Bayesian Bits runs and two runs for the pruning-only and quantization-only experiments.Source: §4.2 · pp. 7–8, Figure 2 · p. 8
  • A ResNet18 experiment takes approximately 70 hours on one Nvidia Tesla V100, about twice the runtime of 40 epochs of quantization-aware training.The reported run includes 30 Bayesian Bits training epochs and 10 fixed-gate fine-tuning epochs; results are not directly compared with, [6], or.Source: §4.2 · p. 8

4.2.1 Post-training mixed precision

  • Post-training Bayesian Bits occupies a middle ground between push-button post-training methods and full-model fine-tuning, with relatively minor data and compute requirements.Figure 3 compares its Pareto front with an iterative baseline and a fixed 8/8 baseline.Source: §4.2.1 · p. 9, Figure 3 · p. 9

5 Conclusion

Bayesian Bits learns efficient bit widths through gradient descent, using gated quantized residuals and sparsifying priors to unify pruning and quantization.

5 Conclusion

  • Bayesian Bits uses gated additions of quantized residuals and sparsifying priors to learn efficient bit widths through gradient descent, unifying pruning and quantization.The paper reports that the approach finds more efficient networks than prior art.Source: §5 · p. 9

Broader Impact

The broader-impact discussion describes possible energy benefits and unstudied fairness effects, while the appendix relates the objective to L0 regularization and details gate optimization and quantization variants.

Broader Impact

  • The method may reduce inference-time energy consumption, but its effects on prediction fairness were not investigated.The paper notes that pruning can affect classes differently even when aggregate statistics are unchanged.Source: Broader Impact · p. 9

A.1 Relationship to L0 norm regularization

  • When quantization gates are fixed on, the objective corresponds to regularizing the expected L0 norm of the parameter-inclusion vector.This equivalence assumes parameters are not quantized.Source: §A.1 · p. 12

A.2 Optimizing the Bayesian bits objective

  • Hard-concrete relaxation gives gates a mixture with exact zero and one values, and test-time threshold t = 0.34 favors zero when its mixture component is most probable.The relaxation also permits gate values between zero and one during optimization.Source: §A.2 · p. 12

A.3 Alternative gating approaches

  • Deterministic gates can lose gradients when saturated and create a training–inference mismatch; REINFORCE was abandoned because its gradient estimates had high variance.The reported deterministic-gate issue was especially evident on ImageNet, while CIFAR-10 results were not significantly hurt.Source: §A.3 · p. 13

A.4 Bayesian Bits algorithm

  • The appendix provides pseudo-code for the Bayesian Bits quantizer’s forward pass.Source: §A.4 · p. 13

A.5 Decomposed quantization for non-doubling bit widths

  • For non-doubling bit-width transitions, the resulting bin count differs from the desired count, so the quantization range endpoints must be rescaled.The discrepancy can be an excess or deficit of bins depending on whether the new width is above or below twice the previous width.Source: §A.5 · pp. 13–14

B.1 Experimental setup

The appendix records model-training setups and defines BOP-aware regularization and counting, including adjustments for sparsity and ResNet18 residual connections.

B.1 Experimental setup

  • LeNet-5 and VGG experiments used MNIST and CIFAR-10 training schedules, while ResNet18 used separate SGD and Adam optimization followed by gate fixing and fine-tuning.Source: §B.1 · p. 14

B.2 BOP and MAC count

  • A layer’s BOP count is based on its weight bit width and input-activation bit width.Source: §B.2 · p. 14

B.2.1 BOP-aware regularization

  • Each gate’s regularization strength is weighted by its bit width and layer MAC count; in practice, a range of strengths is used to select a task-performance/BOP trade-off.The authors leave targeting a specific BOP count for future work.Source: §B.2.1 · p. 15

B.2.2 BOP and MAC count under sparsity

  • For dense convolutions, MACs depend on channel counts, output-map dimensions, and filter dimensions; input and output pruning ratios then adjust the count and permit BOP estimation without recomputing pruned-layer MACs.Source: §B.2.2 · p. 15

B.2.3 ResNet18 MAC count computation

  • ResNet18 BOP counts are upper bounds because the counting method handles input pruning conservatively around residual connections.It only accounts for input pruning at the second convolution in each residual block; actual BOP counts may be lower.Source: §B.2.3 · p. 15

B.2.4 ResNet18 regularization

  • For ResNet18 activation quantizers feeding both a downsample layer and a block’s first convolution, the regularization weight accounts for both layers’ MAC counts.Source: §B.2.4 · pp. 15–16

C.1 Differences between baselines and our experimental setup

The appendix explains setup differences relevant to comparisons with LSQ and PACT, describes extensive LSQ baseline tuning, and identifies limits on comparing with other published results.

C.1 Differences between baselines and our experimental setup

  • The LSQ and PACT comparison covers activation placement, batch-normalization handling, input/output quantization, full-precision accuracy, and gradient scaling.These setup differences can affect accuracy, storage, or deployment latency.Source: Table 3 · p. 16, §C.1 · p. 16
  • Direct comparison with several prior results is constrained by missing BOP counts or mixed-precision configurations and by differences in hardware efficiency and optimization targets.The cited studies include, [6], and.Source: §C.1 · p. 16
  • Because published LSQ results used a different setup, the authors ran extensive baseline experiments and grid-searched learning rates, finding Adam for both parameters and scales performed consistently better than their tested SGD alternative.Source: §C.1 · p. 16

D.1 Updated ImageNet results

The updated ImageNet experiments leave activations feeding residual connections unquantized, and the paper reports that Bayesian Bits more clearly outperforms baselines in this setting. The appendix also describes learned architectures, fine-tuning evaluations, and gate evolution across regularization settings.

D.1 Updated ImageNet results

  • With residual-connection inputs left unquantized, Bayesian Bits more clearly outperforms the baselines in this arguably more realistic scenario.The change addresses over-regularization of activation quantizers whose bit widths do not affect BOP count.Source: §D.1 · p. 17, Figure 5 · p. 17

D.2 MNIST and CIFAR 10

  • Learned LeNet-5 and VGG architectures show bit allocation and sparsity; structural sparsity on weights implicitly applies to activation tensors.The figure also includes a more aggressively regularized VGG with fewer BOPs and more quantization and sparsity.Source: §D.2 · p. 17, Figure 6 · p. 18

D.3 Effect of fine-tuning

  • Fine-tuning effects on final model accuracy are presented for Bayesian Bits ImageNet ResNet18 experiments.Figure 7 compares validation accuracy before and after the final 10 fine-tuning epochs across four µ settings and three runs per setting.Source: §D.3 · p. 17, Figure 7 · p. 18

D.4 ImageNet

  • Learned ResNet18 architectures tend to assign higher bit precision to the first and last layers and to weights on some shortcut connections.The appendix points to learned architectures for four regularization strengths.Source: §D.4 · p. 18

D.4.1 ImageNet ResNet18 gate evolution

  • For three experiments with µ = 0.05, Figure 10 plots the evolution of the gates; the cited passage does not identify the model architecture or dataset.Source: §D.4.1 · p. 18

D.4.2 ImageNet post-training

The post-training quantization experiments assess accuracy as quantizers are successively set to lower bit widths, and report that accuracy does not decline monotonically as more quantizers are changed. The appendix also provides comparisons across regularization settings and visualizations of training and learned architectures.

D.4.2 ImageNet post-training

  • Accuracy does not monotonically decrease as more quantizers are turned on in the post-training quantization experiment.The results are presented as a Pareto front of model efficiency versus accuracy.Source: §D.4.2 · p. 18
  • The post-training baseline measures each quantizer’s sensitivity in an INT16 network, then lowers bit widths in increasing-sensitivity order and measures accuracy cumulatively.The procedure measures accuracy after each quantizer is set to the lower bit width.Source: §D.4.2 · p. 18
  • The ImageNet ResNet18 validation-accuracy plots cover µ values 0.03, 0.05, 0.07, and 0.2; the PACT [5] comparison is marked hypothetical under an 8-bit input/output assumption.The hypothetical comparison assumes that changing PACT’s full-precision input and output layers to 8 bits would not change its results.Source: Figure 8 · p. 19, Figure 9 · p. 19, Table 4 · p. 20
  • Table 5 compares learning only gates with learning both gates and scales across regularization strengths, reporting means and standard errors over three runs per setting..Source: Table 5 · p. 21
  • Training visualizations track validation accuracy and cross-entropy loss, gate probabilities, and the Pareto-front and gate-loss trajectories across the stated µ settings.The gate-probability plots show mean probability with one-standard-deviation shading for µ values 0.03, 0.05, 0.07, and 0.2.Source: Figure 11 · p. 22, Figure 12 · p. 22, Figure 13 · p. 22
  • Learned ResNet18 architecture diagrams are provided for first-run experiments at µ = 0.03, 0.05, 0.07, and 0.2..Source: Figure 15 · p. 23, Figure 16 · p. 24, Figure 17 · p. 25

Related papers

Loading 2005.07093v3…