PACT: Parameterized Clipping Activation for Quantized Neural Networks

In plain words

Neural networks run on numbers stored with many bits. Storing and multiplying those numbers with fewer bits makes a model smaller and faster, but the model usually gets worse. Most earlier work shrank the weights (the learned parameters). The activations (the values passed from layer to layer while the network runs) were harder to shrink, because a common activation function, ReLU, lets them grow without a limit, and a few-bit scale then fits them poorly.

PACT's idea is to cap each layer's activations at a threshold that the network learns during training. The threshold starts large, and training pulls it down to the range where the activations matter, so the few available levels are spent where they count. The rest of training then adjusts the network to the rounding. The authors report that this works with 4-bit weights and activations at roughly full-precision accuracy, and that the smaller number formats let a hardware accelerator pack in more compute units.

TL;DR

• PACT swaps ReLU for a clipped activation with one learned clip level alpha per layer. §4 eq. 1
• Clipped activations are rounded to k bits in the forward pass; gradients reach alpha through a straight-through estimator. §4 eq. 2 eq. 3
• Author-reported: with 4-bit weights (DoReFa) and 4-bit activations (PACT), top-1 accuracy is within about 1 point of full precision on the five networks tested. tab. 1
• The hardware analysis is a model estimate (14 nm MAC area, a performance simulator), not a measured chip. §6 fig. 6

Problem

Quantizing activations is harder than quantizing weights. ReLU outputs are unbounded, so quantizing them needs a high dynamic range, and accuracy is significantly degraded §3 fig. 1. Clipping activations with an upper bound helps, but the authors say it is hard to find one globally optimal clip value across layers and models §3. The Gaussian-based scale of HWGQ is set offline and stays fixed during training §3. The authors state that such schemes do not fully use back-propagation to learn the clip level §3.

Method

PACT is an activation function with one trainable parameter alpha. The forward pass is a truncation followed by a uniform k-bit quantizer. The backward pass uses a straight-through estimator (STE) for the rounding. §4

Activation (forward pass). The ReLU in each activation position is replaced by the function in eq. 1. It is 0 below zero, the identity on [0, alpha), and the constant alpha above alpha. §4 eq. 1

Quantization (forward pass). The truncated output y is mapped onto 2^k levels spaced alpha / (2^k - 1) apart, using the formula in eq. 2. The authors say the result is linearly quantized to k bits for the dot products. §4 eq. 2

Gradient for alpha (backward pass). The STE sets the derivative of the rounding to 1. The derivative of the output with respect to alpha is then 1 where the input is at or above alpha, and 0 elsewhere eq. 3. The paper therefore learns alpha by gradient descent through the loss. §4 eq. 3

Regularization of alpha. The loss includes an L2 penalty on alpha so that alpha shrinks during training and limits the dynamic range. The penalty uses the same coefficient as the weight decay §4.2 app. A.2, and the paper calls it lambda_alpha §4.1. The exact form of the penalty term in the loss is not written out in the paper. §4 app. A.2

Effect of alpha during training. In the ResNet20/CIFAR10 run of fig. 2 (lambda_alpha = 0.01), alpha falls from its initial value and settles at values much smaller than that start, which limits the dynamic range of the activations. §4 fig. 2

Sharing scope. The best scope for alpha was one value per layer. Sharing per layer also lets alpha be applied once after all multiply-accumulate operations of the layer §4.2 app. A.1 fig. 7.

Initialization. Starting alpha at a large value and then regularizing it down worked better than starting it small §4.2 app. A.2. The ResNet20 run uses an initial value of 10 §4 fig. 2.

Layers kept at full precision. The paper follows common practice and does not quantize the first and last layers in its main comparisons §4.2. Table 2 and Figure 9 show how quantizing those layers changes accuracy tab. 2 fig. 9.

Pseudocode. The two functions below restate eq. 1 to eq. 3. The paper gives no algorithm box.

# PACT forward pass  [eq. 1] [eq. 2]
# Input:  x (activation tensor), alpha (scalar clip level, one per layer),
#         k (activation bit-width)
# Output: y_q (quantized activations, same shape as x)

def pact_forward(x, alpha, k):
    # Truncate to [0, alpha]  [eq. 1]
    truncated = 0.5 * (abs(x) - abs(x - alpha) + alpha)
    # Uniform k-bit grid with 2^k levels spaced alpha / (2^k - 1)  [eq. 2]
    step = alpha / (2 ** k - 1)
    quantized = round(truncated / step) * step
    return quantized

# Gradient of the output with respect to alpha, with STE  [eq. 3]
# Input:  x (pre-activation tensor), alpha, upstream (dL/dy_q, same shape as x)
# Output: grad_alpha (scalar)

def pact_grad_alpha(x, alpha, upstream):
    # 1 where x >= alpha, 0 where x < alpha  [eq. 3]
    inside_clip = (x >= alpha)
    # Sum over every element that shares this alpha (one alpha per layer)
    # [§4.2] [app. A.1]
    grad_alpha = sum(upstream * inside_clip)
    return grad_alpha

Hyperparameters

Values are as stated in the paper. Where the paper is silent, the cell says so. The full comparison tables are in the paper's appendix app. C.2.

Results

Author-reported. Setup: CIFAR10, SVHN, ImageNet (AlexNet, ResNet18, ResNet50), trained from scratch in TensorFlow §5 app. B. Baselines are DoReFa-Net for the same bit-widths §5.2 tab. 1 and the prior schemes of app. C.1 for ImageNet. "PACT" below means DoReFa weights with PACT activations, unless noted tab. 1 tab. 1 caption.

Top-1 accuracy on CIFAR10, SVHN, AlexNet, ResNet18, ResNet50 tab. 1

Validation error (%) for CIFAR10-ResNet20 with first and last layers at different bit-widths tab. 2. Column groups are the bit-width of the other layers (BIT-M). Within each group the columns are the bit-width of the first and last layers (BIT-FL = 2, 4, 8, 32). FL/M/FL quantizes first, middle and last layers. FL/M/NQ leaves the last layer unquantized. NQ/M/FL leaves the first layer unquantized. NQ means no quantization.

ImageNet, PACT rows only (full comparisons with DoReFa, BalancedQ, WRPN, FGQ, WEP, LPBN and HWGQ are in the cited tables app. C.2). The first row of each block is the bit configuration as BitW / BitA. The full-precision baseline is given per table.

AlexNet top-1, full precision 55.1 tab. 3:

AlexNet top-5, full precision 77.0 tab. 4:

ResNet18 top-1, full precision 70.4 tab. 5:

ResNet18 top-5, full precision 89.6 tab. 6:

ResNet50 top-1, full precision 76.9 tab. 7:

ResNet50 top-5, full precision 93.1 tab. 8:

Hardware (modelled). Over 16 to 2 bits, the MAC density improves by about 14x when weight and activation bit-widths are reduced together §6 fig. 6. With external bandwidth limited to 64 GB/s, going from 16 to 4 bits gives a 4x increase in peak FLOPs and about 4.5x in performance. The authors attribute the extra gain to data fitting in on-chip memory §6. The first and last layers stayed at 8 bits in these runs; the authors estimate about 1.24x more if they were also quantized to 4 or 2 bits §6.

Paper vs code

No official code. The arXiv source is a wrapper around the PDF, and the PDF has no code link [raw/src/1805.06085/main2.tex]. Searches with the GitHub API: "PACT quantization" returned six repositories, none owned by an author or by IBM; "pact ibm org:IBM" returned none; "parameterized clipping activation" returned none. The closest match, KwangHoonAn/PACT (Apache-2.0), is a third-party reimplementation and is not used.

Pitfalls

• Alpha's L2 penalty is unspecified. The paper says alpha has an L2 penalty with coefficient lambda_alpha §4 app. A.2, but does not write the term in the loss (for example whether it is lambda * alpha^2 or with a 1/2 factor). No official code is available to settle it.
• Gradient with respect to the input is not stated. eq. 3 covers only alpha. The gradient with respect to x, and the handling of the STE for the rounding in eq. 2, are not written out. An implementation has to choose them. The usual choice is the straight-through identity for rounding and the step of eq. 1 for the input; this is an inference, not stated.
• Initial alpha is only given for one run. The value 10 is from the ResNet20/CIFAR10 study §4 fig. 2. For other networks, the initial value is not stated. app. A.2 says a very small start puts activations where the gradient is nonzero, so alpha adapts unstably in early epochs. A very large start makes the gradient small, so alpha may stay large and quantization error suffers.
• Quantization of the first and last layers in Table 1 is not stated. The paper's default is to keep them full precision §4.2; tab. 2 shows that quantizing them to low bits costs accuracy. Whether the Table 1 PACT runs did so is not stated.
• Weight precision in the activation-only runs is not stated. Figure 4 and Figures 7 to 9 vary activation bit-width. Whether weights were full precision in those runs is not stated.
• "Full precision" versus "near full precision." The text says 4-bit runs reach "full-precision accuracy consistently across the networks tested", then says "near (<= 1%)" §5.2. In tab. 1, ResNet18 with 4-bit PACT is 1.0 point below full precision (0.692 against 0.702). Use the table numbers.
• Accuracy degradation curves are read from plots. Figures 5 and 4(f) give degradation values only as bar heights, so they are not reproduced here.
• Hardware numbers are a model. The area and performance results come from a 14 nm MAC area estimate and the DeepMatrix framework §6 fig. 6. They are not measured silicon. The memory bandwidth of the modelled system is set per configuration §6.
• Scope of alpha drives the hardware benefit. Per-layer scope lets alpha be applied once after all MAC operations §4.2. Per-channel and per-neuron scopes give different training and validation error app. A.1 fig. 7, and the paper prefers per-layer scope partly for hardware complexity app. A.1.

Lineage

• Used with: DoReFa-Net (arXiv:1606.06160) for weights in tab. 1 §5.2.
• Builds on: clipping activation functions with an upper bound (Hubara et al., Zhou et al., as cited in §3), with the clip level learned rather than set.
• Compared against: DoReFa, BalancedQ, WRPN, FGQ, WEP, LPBN, HWGQ, with their reported numbers and settings app. C.1 app. C.2.
• Superseded by: not stated.