A Low-Complexity Approach to Rate-Distortion Optimized Variable Bit-Rate Compression for Split DNN Computing

A significant amount of visual data being generated in visual IoT applications is intended for consumption by analytic algorithms such as classification, object detection/tracking, semantic segmentation, etc. In a typical framework, video data captured by a mobile device is compressed using conventional image and video codecs like JPEG, H.264, HEVC, etc. and transmitted over the network to an edge-server (or cloud). The bitstreams are first decompressed and then provided as input for various analytics tasks (Fig. (a)a), which are increasingly being performed by very complex deep neural networks (DNNs). The scale of these IoT scenarios creates unique challenges as large numbers of IoT devices with limited compute resources rely on shared edge/cloud compute, and access to that is via a shared, time-varying network resource. This places significant new demands on visual compression algorithms to scale beyond the compression efficiencies obtained with conventional video compression standards such as H.264/HEVC. These standard compression techniques optimize perceptual quality for human consumption, which may not necessarily be optimal for analytic tasks that rely on the semantics of content rather than its perceptual quality. Consequently, task performance can degrade severely in the presence of even relatively mild compression artifacts.

To address this, recent approaches rely on machine-learning models to learn features that are jointly optimized both for performance on an analytic task and for compression. Learnt representations from the front end or head of a DNN – consisting of an input layer and a number of subsequent layers – are compressed and transmitted to an edge server. The remaining layers (also referred to as the tail) then operate directly on these compressed representations (Fig. (b)b). Computationally, this is much more efficient than first reconstructing the image or video input and then performing the analytic task. Early approaches [7] explored the use of simple lossless and lossy compression techniques to compress the intermediate features. While lossless techniques result in only mild reduction in bandwidth, naively quantizing the intermediate features in lossy compression leads to a drop in the task performance. Inspired by the impressive results of ML-based image compression approaches [2, 3, 25], subsequent approaches [24, 9, 21] include the quantization operation in the end-to-end training of the model. These approaches yield significantly better rate-distortion performance (task accuracy vs compression level) than that obtained by traditional image compression methods such as JPEG and the more recent HEIC.

Early approaches [16] also attempted to determine an optimal splitting point between the device and the edge-server to minimize latency. These often result in the splitting point being placed towards the end of the model, which ends up placing most of the computational burden on a low-powered, compute limited mobile device. More recent contributions [13, 22, 19, 20] attempt to remedy this by introducing bottleneck layers at the split point. These are layers with small number of parameters and low computational overhead which reduce the dimensions of the features to be compressed and transmitted. This allows the model to be split at earlier layers. Crucially, though, in all these approaches, once the split point has been fixed, the model is then trained for that particular partition only. Changing the split point, therefore, cannot be done without retraining the entire model, or at least a significant portion of it if techniques such as head-network distillation [19] are used. This is a serious limitation that adversely impacts the applicability of the solution in conditions where dynamic partitioning of the workload between the mobile device and the cloud is required in response to changing compute and platform requirements at either end. Such capabilities are especially important when the edge-server needs to perform dynamic orchestration of workloads while servicing multiple mobile or client devices.

Another equally serious limitation is that the parameters of the trained pipeline are optimized for operation only at a certain compression level or bit-rate. Supporting different compression levels needed for variable bit-rate operation, therefore, also requires retraining the DNN and storing multiple copies of its weights, similar to the changing of the split point. Such multiple end-to-end trainings increase the overall training time and complexity. Moreover, during actual operation in inference mode, the entire set of parameters will have to be reloaded each time a different compression level or split point is desired. DNN parameters often run into several tens of millions, and hence storing and loading multiple copies of these parameters during runtime is expensive and can slow down overall operation. Some contributions avoid this problem by maintaining a learnable vector of scale parameters [24] or gain parameters [11] that are applied to the output of the encoder of a deep-autoencoder. While this is effective, its applicability has been demonstrated only at the output of the encoder network of a deep autoencoder, where the output dimension is already the lowest. It does not address the issue of high-dimensionality of intermediate features that arises in split computing and hence is not applicable in those scenarios.

Here, we provide some background information on intermediate feature compression and training DNNs for split computing. As explained earlier, we split a DNN model that has been designed for a particular task like classification or segmentation into a front-end or head-network and a back-end or tail-network as shown in Fig. (b)b. The intermediate features are to be transmitted to the tail that resides on an edge-server or cloud. The dimensions of the intermediate features, especially towards the earlier layers of the DNN, tend to be very high making them unsuitable for efficient compression. Therefore, the use of bottleneck units has been proposed. The unit comprises two modules –- a Bottleneck Encoder and a Bottleneck Decoder –- as shown in Fig. 2. The bottleneck unit can hence be viewed as a deep autoencoder that has been designed for the feature space of a DNN rather than for its input. The Bottleneck Encoder transforms the high-dimensional intermediate feature tensor into an appropriate lower dimensional space. This lower dimensional tensor is compressed and transmitted to the edge-server where the Bottleneck Decoder restores it to its original dimension.

Early works explored the use of lossless of simple lossless compression techniques for intermediate DNN features [7], but these don’t yield a significant reduction in bandwidth nor offer flexibility in adapting the compression bitrate. For this reason, lossy compression is applied, which involves first quantizing the intermediate features by some quantization step-size, $Q$, followed by lossless compression of the resulting discrete-valued signal by algorithms such as Huffman Coding or Arithmetic Coding. Simply quantizing the intermediate features during inference, however, leads to a drop in the task performance. To overcome this, the quantization operation is included in the training of the network. Since quantization is a non-differentiable operation, it cannot be directly incorporated during a gradient-based training regimen. This can be circumvented by either replacing quantization by additive noise [2], or using a stochastic form of binarization [25], or using a smooth approximation of the gradient [24]. In this work, we adopt the third approach [24], using an identity function as an approximation for quantization during backpropagation. Note that this approximation is applied only to the gradients during backpropagation, and during the forward pass, true quantization is performed.

The model itself is trained with a loss function of the form:

where, $L_t$ is a task-loss to maximize task performance, $L_r$ is a rate-loss term to minimize bit-rate of the encoded data, and $\alpha$ is a term that controls the relative weighting of the two terms. Higher values of $\alpha$ assign higher weight to $L_t$ relative to $L_r$, i.e. emphasize task-accuracy more than bit-rate, resulting in higher accuracy but at higher bit-rates. Conversely, lower values of $\alpha$ allow the network to achieve lower bit-rates, but at the cost of task performance. Note that both $L_r$ and $L_t$ are functions of the quantization step-size, $Q$, since true quantization is performed in the forward pass during training and this directly affects the values of the loss terms.

For $L_t$, the usual loss functions such as cross-entropy loss for classification, and sum of per-pixel cross-entropy losses for segmentation are used. Various approaches have been explored in literature to choose an appropriate $L_r$. Since the feature values are discrete-valued following quantization, the rate is the expected code length which is lower-bounded by the entropy of the probability distribution of the quantized alphabet [10]. In [24], a continuous, differentiable function is used to approximate this probability distribution. In [2, 3], methods of variational inference are used to approximate the true distribution by parametric variational densities. Other contributions in literature adopt indirect approaches. In [8], the ${\ell }_2$ norm of the compressed feature is taken as a proxy for the number of bits required to encode the data. In [18], the ${\ell }_1$ norm of the DCT coefficients of the original image was used as an estimate for rate. [1] improved upon this by incorporate spatial prediction into the calculation of the loss to get a better estimate of the rate. In our work, we adopt the indirect approach and use the ${\ell }_1$-norm of the output of the bottleneck encoder as a proxy for the rate. As will be demonstrated shortly, even this simple approach yields impressive results.

In this section, we describe a systematic approach to the design and training of low-complexity bottleneck layers for flexible workload partitioning and enabling variable bit-rate compression in the split computing paradigm.

We have tested our approach on two analytic tasks: image classification and semantic segmentation. For the classification task, we test on the ImageNet dataset [12] and use a Resnet50 [14] model pretrained on the same. The images are scaled to $224\times 224$ resolution prior to classification. For segmentation, we test on the validation split of the MS-COCO 2017 dataset [17], and use a DeepLab v3 model [6] pretrained on its train split. For this task, the images are scaled to a resolution of $513\times 513$. For both tasks, we test our approach at four different split points shown in see Table II within their respective models. The layer names for each model are the same as those used in their respective implementations in Torchvision [26].

We benchmark our method against standard image compression algorithms - JPEG and HEIC (high-efficiency image compression) - and also against some of the recent ML-based image compression methods which include the Entropic Student [20] and variational image compression [3]. To compare and benchmark the effectiveness of our approach, we derive the Pareto optimal accuracy-vs-compression curve, which is a plot of a task-specific accuracy metric against the compression level. For classification, the metric is simply the classification accuracy; for segmentation, it is the Jaccard index as measured by the mean intersection-over-union (mIOU). The compression level is represented by bits-per-pixel (bpp), which, as the name indicates is the number of bits of information that needs to be transmitted divided by the number of pixels in the input image. The results for [20] were originally reported in terms of file sizes, but we convert those into equivalent bpp values using $\text{bpp}=\text{file\_size}/(W\times H)$. The results are plotted in Fig. 4 for classification and Fig. 5 for segmentation. For our method, results at all the different split points tested are plotted in the same figure.

Our method shows significant reduction in bit-rates compared to other methods across all accuracy levels for classification, and across all mIOU levels for segmentation. This is true for all the split points that were tested. Sample bpp values at specific accuracy and mIOU levels are shown in Table IV. At an accuracy of $70\%$, the bit-rate of our method is 3 to 27 times lower than the best performing amongst the benchmarks that we tested. For segmentation, the results are even more impressive with our method achieving bit-rates that are 9 to 392 times lower. The peak classification accuracy with our approach is slightly lower (within 0.5%-1.5%) than that from other methods, but for segmentation, the peak mIOU is higher than all other methods.

Note that each point on the Pareto optimal performance curves of our method represents a different point in the 4-dimensional parameter space $(C_r,S,\alpha ,Q)$. This implies that the size (number of learnable weights) and computational complexity of the bottleneck layer, which depend on $C_r$ and $S$, differs for each point. In Table III, therefore, we report the size and complexity of the largest bottleneck unit resulting from our search procedure. As shown, even this largest bottleneck layer is a negligibly small fraction of the size and complexity of the overall model. In the supplementary material, we document in detail the 4-d parameter values for each point on the Pareto optimal curve for all splits. One point to note in a variable bit-rate setting is that the information about which hyperparameter set has been used for bottleneck encoding must be sent over to the edge-server, so that the corresponding correct bottleneck decoder can be used. The impact to bit-rate, however, is vanishingly small. In our results, for example, a maximum of 7 different hyperparameter sets have been used to generate a Pareto optimal performance curve, which implies a maximum of 3-bits ($\lceil {\log }_{2}7\rceil =3)$ of additional overhead, which translates to a bpp increase of less than $6\times {10}^{-5}$.

Finally, we evaluate the performance of the method if all the parameters of the pipeline - both the original model and the bottleneck layers - are trained rather than only the bottleneck parameters. Recall that not only does this increase the number of parameters to be trained by over a 100-fold (from Table III), it also requires more epochs to train (as explained in section III-B). The results are shown in Figures 6 and 7 for classification and segmentation, respectively. The performance curves for full-model training are shown in solid lines, while the curves for bottleneck-only training are shown in dashed lines. The full-model curves show mild improvement over the bottleneck-only curves, but this comes at the cost of dramatically higher training time. Equally importantly, as explained in Section I, our method enables practical variable bit-rate compression, as only the parameters of the bottleneck layers have to be reloaded to switch to a different compression level, and these are a small fraction (less than 1%, see Table III) of the number of parameters in the original model.

In this paper, we have presented an approach that enables flexible workload partitioning and enables practical variable bit-rate compression in distributed media analytics pipeline. Our approach is remarkably lightweight, both during training and inference, highly effective and achieves excellent rate-distortion performance at a small fraction of the compute and storage overhead compared to existing methods.

Finally, we note that in this initial study of ours, we have intentionally made simple, low-complexity choices for various parts of our method. These include the architecture of bottleneck layers (single depthwise separable convolutional layer); the rate-loss term (${\ell }_1$-norm of the intermediate feature); and for the compression scheme (simple uniform quantization followed by Huffman encoding). For each of these, there are more sophisticated alternatives that could help improve the performance further, such as increasing the number of layers in the bottleneck units or using different topologies for the bottleneck layers; using arithmetic coding during lossless compression; and using entropy estimation for the rate-loss term during training. Hence, although our approach already yields impressive bit-rate reduction over state-of-the-art methods, we believe that there is still room for further improvement via exploring the approaches just outlined for future work.