PAMS: Quantized Super-Resolution via Parameterized Max Scale

Quantization function. As shown in [2], distributions on different layers of activations tend to be symmetric. This characteristic can help to improve the model accuracy of a quantized network with extremely low-bit weights and activations, as validated in [8]. Therefore, given a specific full-precision model with a parameter set $\mathcal{X}$ ($\mathcal{X}$ denotes either weights or activations of a specific layer), we quantize every element $x$ ($x\in\mathcal{X}$) by using the following point-wise quantization function $Q$ with a symmetric mode:

where $f(x)$ is the clamp function that limits the inputs range and $s(n)$ is the map function that scales the higher precision inputs to their lower bit reflections, which can be formulated by $f(x)=max(min(x,a),-a)$ and $s(n)=\frac{a}{2^{n-1}-1}$, respectively. $n$ denotes the number of quantization level. $a$ represents the maximum of the absolute value of $\mathcal{X}$ and $\lfloor\cdot\rceil$ rounds the value to the nearest integer.

Trainable upper bound. In the previous work [3], the dynamic range of activations can be partially alleviated by using a parameterized clipping activation function to replace the Rectified Linear Unit (ReLU), which limits the application scope. In this paper, we propose a novel activation quantization scheme, in which the clamp function $f^{*}(\cdot)$ has the trainable parameter $\alpha$ to dynamically adjust the upper bound of the quantization range. We can directly employ the stochastic gradient descent to update this parameter, which is able to minimize the performance degradation arising from the quantization. For a given activation, the corresponding value will be quantized to $n$ bits by:

where $\tilde{x}=f^{*}(x)=max(min(x,\alpha),-\alpha)$. The dynamic range is limited to $[-\alpha,\alpha]$. The advantages of our quantization function lie in that $Q(x,n)$ directly involves in the back-propagation process based on only one learnable parameter $\alpha$, based upon which we can train the quantized SR network in an end-to-end manner. Extensive experiments in Section 4 demonstrate that Eq.2 can introduce more effectiveness comparing to several state-of-the-art quantization methods [42, 17, 3].

Back-propagation with quantization. In back-propagation, $\frac{\partial x_{q}}{\partial\tilde{x}}$ can be approximated to $1$ based on the straight-through estimator (STE) [5]. Inspired by [3], the gradient of $\alpha$ is calculated as follows:

Initializing $\alpha$. To avoid gradients vanishing or exploding, non-convex optimization on DCNNs heavily depends on the initialization of parameters. Instead of manually designing the initial value of $\alpha$, initialization based on the statistics from a pre-trained network can achieve better performance. Therefore, we resort to task-related statistical methods based on the pre-trained network to calibrate the quantization error. In particular, given the $l$-th layer with $m$ input activations ${x_{1}^{(l,t)},...,x_{m}^{(l,t)}}$, $\alpha^{(l)}$ is calculated by the redefined exponential moving average (EMA) function at the start of training:

where $t$ is the iteration number and $\beta$ denotes the smoothing parameter of EMA, which is set to be 0.9997. Specially, we set $\beta$ to 0 when $t$ is 0.

Pixel-wise loss. Given a training dataset $D=\{I_{LR}^{i},I_{HR}^{i}\}_{i=1}^{n}$ with $n$ LR input images and their corresponding HR counterparts, SR models are commonly optimized by minimizing the conventional pixel-wise $L_{1}$ loss between the output $I_{SR}$ and the ground truth image $I_{HR}$:

where $||\cdot||_{1}$ denotes the $L_{1}$ norm. A better SR model needs to infer the high-frequency textures from a low-resolution input. However, it is hard to obtain by only using Eq. 5 based on low-bit quantization, which is due to the accumulated quantization error.

Structured knowledge transfer (SKT). Inspired by [38], we consider that the full-precision model has learned high-level representation, which provides knowledge to the low-precision one about where it concentrates. More specifically, instead of using the soft probability in the classification task, we align the structured features between the cumbersome network and the quantized one by minimizing their pixel-wise distance. Therefore, the loss function for our SKT is defined as:

where $F_{T}^{\prime}$, $F_{S}^{\prime}$ are a pair of structure features after the spatial mapping of activations from the full-precision network and the correspond low-precision one, respectively. The spatial mapping defined by $F^{\prime}=\sum_{i=1}^{C}\left|F_{i}\right|^{2}\in\mathbb{R}^{H\times W}$, where $F\in\mathbb{R}^{C\times H\times W}$ denotes the activations after the last layer in the high-level feature extractor. We set $p=2$ for $p$-norm in our experiments. In sum, SKT enhances the learning process of spatial correlation in the low-precision model which effectively improves the performance of the quantized network and provides an additional constraint to avoid producing over-smoothed images.

The overall loss function. Given an SR model, as consistent with the distillation term mentioned above, the whole objective function is given as:

Datasets and metrics. DIV2K [36] contains 800 training images, 100 validating images and 100 testing images. We train all models with DIV2K training images. For testing, we use four standard benchmark datasets: Set5 [1], Set14 [21], BSD100 [30] and Urban100 [15]. For the evaluation metrics, we use PSNR and SSIM [37] over the Y channel between the output quality image and the original HR image.

SR models and alternative approaches. Both residual block and dense block are widely used in SR models, like VDSR [18], EDSR [23] and RDN [40]. To validate the superiority of our approach, we choose EDSR and RDN as backbones and use 8-bit and 4-bit quantization on them. As most parameters exist in the high-level feature extraction module, we do not quantize weights and activations in low-level feature extraction and reconstruction modules, which ensures a trade-off between performance and model size. The qualitative comparisons are generated by the publicly available source code in EDSR [23].

Training setting. The model is implemented by using PyTorch [31]. Following the setting of [23], we pre-process all images in the DIV2K training dataset by subtracting the mean RGB and adopt a normal data augmentation during training, which includes random horizontal flips and vertical rotations. The mini-batch size is set to 16. We deploy the ADAM optimizer with $\beta_{1}=0.9$, $\beta_{2}=0.999$ and $\epsilon=10^{-8}$ to the model, which is trained for 30 epochs. The learning rate is initialized by $10^{-4}$ and is halved at every 10 epochs.

Effect of BN in SR models. To investigate the effect of quantizing normalized features, we use PACT to quantize EDSR with BN and without BN. As shown in Table. 4, the performance gap between the quantized EDSR without BN is larger than the quantized EDSR with BN. For example, The gap of 8-bit PACT-EDSR without BN is 0.790dB PSNR on Urban100, which is larger than PACT-EDSR with BN (0.354dB PSNR). It shows that the performance degradation of unnormalized features is more pronounced in lower-precision SR models, Moreover, PAMS-EDSR can save more important information for unnormalized weights and activations which largely decrease the performance gaps.

Effect of the learnable $\alpha$. We compare our learnable max scale (PAMS) with the fixed maximum (TF Lite) for quantizing activations. Quantitative and qualitative results are represented in Table. 2 and Fig. 4, respectively. Compared to TF Lite-EDSR, PAMS-EDSR achieves a better score as it produces sharper images and more realistic textures. It indicates that our method can learn a more suitable quantization range which contains more information about the full-precision model and reduces the quantization error.

Effect of the initialization of $\alpha$. We evaluate our EMA initialization with random initialization on EDSR with a scale factor of $\times 4$. For the random mode, we initialize $\alpha$ in the activation quantization layer with a random number ranges from 0 to 128, which ensures that $\alpha$ can be initialized to a larger value in different layers independently. As illustrated in Table. 5, EMA initialization achieves better performance on all benchmark datasets. To explain, EMA achieves better statistical distribution by $\alpha$ that can further help improve SR performance.

Investigating SKT loss. To investigate the effectiveness of SKT, we further compare the quantized model with and without SKT. As shown in Table. 6, PAMS-EDSR which is optimized with the SKT outperforms the corresponding counterpart. Especially, our method obtains much better performance on lower bits. For instance, compared to the PAMS-EDSR without $L_{SKT}$ on Urban100, 4-bit PAMS-EDSR with $L_{SKT}$ gains $0.071$dB PSNR while 8-bit PAMS-EDSR with the same optimization gains only $0.032$dB PSNR. It also indicates that the feature maps from the full-precision model can help the low-precision model to better capture the spatial correlation from images.