[Paper Review] Stochastic Layer-Wise Precision in Deep Neural Networks

Griffin Lacey, Graham W. Taylor+1 more|arXiv (Cornell University)|Jul 3, 2018

Adversarial Robustness in Machine Learning1 references3 citations

TL;DR

This paper proposes a stochastic, layer-wise precision allocation method in deep neural networks where each layer dynamically explores different bit-widths during training using Gumbel-Softmax sampling. The approach learns optimal precision configurations under a fixed budget, improving generalization and outperforming uniform precision baselines on MNIST and ILSVRC12, with gains up to 1.41% error reduction on MNIST and 2.22% on ILSVRC12.

ABSTRACT

Low precision weights, activations, and gradients have been proposed as a way to improve the computational efficiency and memory footprint of deep neural networks. Recently, low precision networks have even shown to be more robust to adversarial attacks. However, typical implementations of low precision DNNs use uniform precision across all layers of the network. In this work, we explore whether a heterogeneous allocation of precision across a network leads to improved performance, and introduce a learning scheme where a DNN stochastically explores multiple precision configurations through learning. This permits a network to learn an optimal precision configuration. We show on convolutional neural networks trained on MNIST and ILSVRC12 that even though these nets learn a uniform or near-uniform allocation strategy respectively, stochastic precision leads to a favourable regularization effect improving generalization.

Motivation & Objective

• To investigate whether heterogeneous precision allocation across layers improves deep neural network performance compared to uniform precision.
• To develop a differentiable, stochastic mechanism for learning optimal precision configurations per layer under a fixed total bit budget.
• To evaluate whether stochastic precision exploration acts as a regularizer, improving generalization without sacrificing hardware acceleration predictability.
• To demonstrate that precision learning via Gumbel-Softmax sampling leads to better accuracy than fixed, uniform precision settings.

Proposed method

• The method uses a Gumbel-Softmax relaxation to sample different precision configurations (bit-widths) per layer during training, enabling differentiable optimization over discrete precision choices.
• Each layer's precision is modeled as a categorical distribution over possible bit-widths, with learnable logits parameterized by the network.
• The total precision budget is constrained, ensuring predictable hardware acceleration and enabling efficient deployment.
• Gradient estimation via reparameterization allows backpropagation through the stochastic precision choices, even though precision is discrete.
• The model is trained end-to-end with standard backpropagation, using stochastic bit-width assignments per forward pass.
• Hard assignments are made at inference using the most likely precision configuration from the learned Gumbel-Softmax distribution.

Experimental results

Research questions

• RQ1 Can learning heterogeneous precision configurations across layers improve deep neural network performance compared to uniform precision allocation?
• RQ2 Does stochastic exploration of precision configurations act as a regularizer, improving generalization?
• RQ3 How does the learned precision allocation compare to uniform allocation under the same total bit budget on image classification tasks?
• RQ4 Does the model learn meaningful precision distributions, such as assigning fewer bits to early layers and more to deeper layers?
• RQ5 Can this method maintain hardware acceleration determinism while improving model accuracy?

Key findings

• On MNIST, the learned precision model achieved a test error of 2.32% with a 10-bit budget, outperforming the uniform baseline of 3.73%.
• With a 40-bit budget on MNIST, the learned model reached 1.14% error, compared to 1.18% for the uniform baseline.
• On ILSVRC12, the learned precision model reduced top-1 error from 49.68% (uniform) to 48.54% under a 14-bit budget.
• The ILSVRC12 model learned non-uniform allocations, assigning 9 bits to later layers and 3–4 bits to early layers, indicating architectural awareness.
• The learned precision configuration consistently converged faster and achieved lower test error across all budgets, demonstrating regularization benefits.
• The method achieved better accuracy with the same bit budget, showing that precision allocation is a tunable, learnable hyperparameter that improves model efficiency and performance.