Uncertainty-Aware Deep Video Compression with Ensembles

1 INTRODUCTION . Video contents are reported to account for 82 % percent of all consumer Internet traffic by 2021 and are growing rapidly with an increasing demand for high-resolution videos ( e.g. , 4K movies ) and live streaming services ( Cisco , 2020 ) . It is therefore critical for us to improve the video compression performance to transmit video with higher quality given limited Internet bandwidth . In recent years , there has been a surge of deep learning-based video compression models ( Rippel et al. , 2019 ; Lu et al. , 2020b ; Agustsson et al. , 2020 ) and some of them have achieved comparable or even better performance than previous traditional video codecs , such as x.264 ( Tomar , 2006 ) . Despite previous deep learning-based video codecs have achieved improved performance on many challenging datasets , most of the state-of-the-art models estimate deterministic predictions for intermediate representations , such as optical flows and residuals . These models fail to represent the aleatoric uncertainty inherent in the model inputs or the epistemic uncertainty in the model parameters and would blindly assume the predictions to be accurate , which is not always the case ( Der Kiureghian & Ditlevsen , 2009 ; Kendall & Gal , 2017 ) . In terms of video compression , such models produce deterministic motion vectors ( or optical flows ) and residuals for each pixel location , ignoring the fact that optical flows may not be estimated accurately in occluded regions and around object boundaries , and the quantization operation for lossless entropy coding also introduces additional noises to the inputs of the decoders . Underlying errors in such overconfident intermediate predictions are propagated to later stages of the P-frame model and even to subsequent frames for models built on temporal correlation , leading to suboptimal performance of the compression system . Predictive uncertainty is crucial for us to understand how certain the model is about the predictions , especially for out-of-distribution data . However , most neural networks do not offer such information and tend to produce overconfident predictions ( Gal , 2016 ; Lakshminarayanan et al. , 2017 ) . Bayesian neural networks ( MacKay , 1992 ; Hinton & Neal , 1995 ) are widely used to quantify predictive uncertainty , but lack practicality due to greatly increased computation complexity and do not scale well to high dimensional data . Gal & Ghahramani ( 2016 ) proposed Monte Carlo dropout that performs test-time dropout . It is simple to implement but is not suitable for deep learning-based compression since it requires multiple decoding-time inferences and yields non-deterministic outputs . In terms of deep learning-based video compression , two non-Bayesian approaches are considered to represent the predictive uncertainty : ( 1 ) modeling the uncertainty explicitly by regressing the empirical uncertainty of the model outputs ( Nix & Weigend , 1994 ) ; and ( 2 ) using ensembles for predictive uncertainty estimation ( Lakshminarayanan et al. , 2017 ) . Agustsson et al . ( 2020 ) took the first approach and proposed scale-space flow for motion compensation . Specifically , their model predicts a scale field besides the standard 2-dimension flow field , representing the uncertainty associated with each MV . Gaussian blurring is then applied to the reference frame and the scale parameter is used to control the scale of the Gaussian kernel . Experimental results show that the self-supervised scale parameter increases around the boundaries of the objects and in regions where the motion is large , which confirms that the scale field learns to estimate the uncertainty of the optical flow prediction . In this work , we consider the second approach and represent the underlying uncertainty with an ensemble of outputs . Ensembles of models have been shown to improve predictive performance ( Dietterich , 2000 ) , and ensembles of neural networks have also been used to boost model performance on many challenging benchmark datasets ( Szegedy et al. , 2015 ) . Instead of producing a deterministic prediction , ensemble methods perform model combination and reflect the uncertainty of out-of-distribution data . Here we propose an ensemble-based multi-head decoding module that can generate an ensemble of intermediate outputs , such as motion vectors and residuals , and implicitly represent the predictive uncertainty with the variance of the Gaussian mixture prediction . We no longer assume that any individual intermediate prediction is accurate and explicitly consider multiple candidates of them . Such uncertainty is then propagated to later stages and all modules in our framework are optimized in an end-to-end fashion . To further improve the performance of our ensemble-based video compression model , we propose an ensemble-aware loss to encourage diversity between different branches and incorporate an adversarial training strategy , fast gradient sign method ( FGSM ) ( Goodfellow et al. , 2014 ) , to smooth the intermediate latent representations . Our experiments show that our model achieves an average bitrate saving of 15 % when the distortion is measured in PSNR . The contributions of this work can be summarized as follows : • We are the first to leverage deep ensembles in a video compression model and demonstrate that our approach can be used to capture the underlying uncertainty of intermediate representations and prevent the model from being overconfident with intermediate predictions . • We demonstrate that our approach can effectively improve the performance of deeplearning based video compression and can be widely applied to optical flow-based video codecs with negligible complexity increase . • Our experiments show that our model outperforms previous state-of-the-art models such as DVC Pro ( Lu et al. , 2020b ) and Lu et al . ( 2020a ) , and our ablation study proves the effectiveness of each module . 2 RELATED WORK . Learned video compression . Previous learning-based video compression methods can be categorized into two groups : ( i ) one-stage models , such as methods based on 3D autoencoders ( Pessoa et al. , 2020 ; Habibian et al. , 2019 ) ; and ( ii ) two-stage models , which is adopted by most previous state-of-the-art methods , consist of predicted frame generation and residual coding . Lu et al . ( 2019 ) proposed an end-to-end trainable video codec , DVC , that utilizes an optical-flow network ( Ranjan & Black , 2017 ) for motion compensation and then compresses the residuals . An extension of DVC , known as DVC Pro ( Lu et al. , 2020b ) , improves the compression performance by introducing refinement modules and auto-regressive entropy models . Feng et al . ( 2020 ) considered residual coding in the feature domain and choose between pixel-level residuals and feature-level residuals with RDO at encoding time . Agustsson et al . ( 2020 ) proposed scale-space flow to blur the intermediate reconstructions when motion vectors are not estimated well . Chen et al . ( 2021 ) proposed to represent videos with neural networks and achieved promising results for video compression . Model uncertainty . Predictive uncertainty can be grouped into aleatoric uncertainty and epistemic uncertainty ( Der Kiureghian & Ditlevsen , 2009 ) . Aleatoric uncertainty captures the noises inherent in the observations and can not be explained away with more data , while epistemic uncertainty accounts for uncertainty in the model structure or parameters and can be reduced with more training data . Bayesian neural networks ( MacKay , 1992 ; Hinton & Neal , 1995 ) is a widely used approach for modeling predictive uncertainty that extends the traditional neural networks by learning a posterior distribution of model parameters from the observed data . Various non-Bayesian approaches have also been proposed , such as utilizing the probabilities of softmax distributions ( Hendrycks & Gimpel , 2016 ) and Specialists+1 Ensemble for representing the predictive uncertainty for adversarial samples ( Abbasi & Gagné , 2017 ) . Gal & Ghahramani ( 2016 ) proposed Monte Carlo dropout by performing multiple inferences with Dropout at test time . Lakshminarayanan et al . ( 2017 ) proposed to use an ensemble of neural networks for quantifying predictive uncertainty . Deep Ensembles . The neural networks community has been investigating ensembles of deep networks since the early 1990s ( Hansen & Salamon , 1990 ; Wolpert , 1992 ; Perrone & Cooper , 1993 ) . Krogh & Vedelsby ( 1995 ) proved the bias-variance trade-off for ensemble models which suggested the importance of the diversity among ensemble members . Lee et al . ( 2015 ) investigated several training strategies to train an ensemble and proposed ensemble-aware oracle loss to encourage diversity . GoogLeNet ( Szegedy et al. , 2015 ) , one of the best-performing models on ILSVRC 2014 , is an ensembles of CNNs . Lakshminarayanan et al . ( 2017 ) proposed to estimate predictive uncertainty by training multiple stand-alone neural networks . Garipov et al . ( 2018 ) ; Fort et al . ( 2020 ) showed that deep ensembles can learn different modes of function with ensemble members that only differ in initialization weights . 3 UNCERTAINTY-AWARE DEEP VIDEO COMPRESSION . This section presents our main contributions . We introduce the theoretical background and the motivation of our proposed approach in Section 3.1 . Then we introduce the ensemble-based multihead structure to decode multiple candidates of motion vectors and residuals in Section 3.2 . In order to encourage the diversity among the ensemble members and to improve the overall performance , we propose an ensemble-aware loss for multi-head decoders in Section 3.3 . Finally , we introduce an adversarial training strategy that we find beneficial for the learning-based video compression task in Section 3.4 . 3.1 UNCERTAINTIES IN DEEP VIDEO COMPRESSION . We consider the predictive coding-based framework , DVC , proposed by Lu et al . ( 2020b ) ( Fig . 1a ) . Let the current frame be xi and the reconstructed previous frame from the buffer be x̂i−1 . We estimate a motion vector ( MV ) map fi with a motion estimation network . The optical flow is then sent to a motion auto-encoder for transform coding , yielding quantized bits âi and the reconstructed optical flow f̂i . Bilinear warping is used for motion compensation ( MC ) and a MC prediction x̃i with residual ri = xi−x̃i is obtained . The residual ri is then compressed with a residual encoder and decoder , outputting quantized residual bit stream b̂i and the decoded residual r̂i . The reconstructed current frame is the sum of the MC prediction and the decoded residual written as x̂i = BilinearWarp ( x̂i−1 , f̂i ) + r̂i . ( 1 ) Aleatoric uncertainty . Although at encoding time we have full information necessary to decode F̂i , for lossy compression at certain bit rates , we quantize the bit stream that is passed to the decoder , and inevitably introduce aleatoric uncertainty at decoding time . Since the aleatoric uncertainty can not be reduced with more training data , a well-trained codec can not reduce the quantization noise or fully recover the estimated MV fi . Consider the MV auto-encoder in the DVC framework above . The lossy compression of the motion vectors can be summarized as ai = MVEncoder ( fi ) âi = q ( ai ) = ai + η f̂i = MVDecoder ( âi ) ( 2 ) If the MV decoder is implemented with a linear model parameterized by w , the impact of the quantization noise η on the decoded MV is given by w > âi = w > ( ai + η ) = w > ai + w > η ( 3 ) Since η is introduced by the quantization operation , we have ‖η‖∞ ≤ 1/2 = ε and it follows that the upper bound of the effects from the quantization operation is given by ‖w > η‖1 ≤ ε ∥∥∥∥w > sign ( ∂w > ai∂ai ) ∥∥∥∥ 1 = 1 2 ‖w > sign ( w ) ‖1 ( 4 ) While in practice the MV decoder is usually implemented with a stack of convolution layers and nonlinear activation layers , such as leaky ReLUs , the transformation of the MV decoder may be too linear to reject the quantization noise ( Goodfellow et al. , 2014 ) . We visualize the aleatoric uncertainty by adding a small perturbation ‖η0‖ ≈ 0.2 · ‖η‖∞ to the bit stream ( see Appendix A.2 ) and the impact on the decoded MVs is depicted in Fig . 2d . It shows that there is more aleatoric uncertainty in regions where the motion is complicated and the impact due to the aleatoric uncertainty is nontrivial . Epistemic uncertainty . Due to limited observed data during training , epistemic uncertainty accounts for the uncertainty in the estimated motion vectors . Motion vectors near the object boundaries and occluded regions tend not to be estimated well , and warping erroneous motion vectors would propagate errors to the residual coding . We may roughly visualize such uncertainty by optimizing an motion estimation network with regards to the mean squared errors ( MSE ) between the current frame xi and the warped frame L = MSE ( xi , BilinearWarp ( x̂i−1 , fi ) ) ( 5 ) We depict the results in Fig . 2e . Since the motion estimation is optimized to minimize the MSE , regions where the MSE is large are likely to have a larger epistemic uncertainty and the corresponding motion vectors can not be estimated well given the limited training data . Ideally , we could save MV bits by not encoding MVs that are not estimated well and save residual bits by not warping MVs that would not help to reduce residuals . This is often difficult to implement as an end-to-end trained deep neural network . In the next section , we will show that with the help of multi-head decoders , the model could learn to exploit available information in the bit stream and handle those predictions with larger uncertainty .

Summary: This paper aims to capture the predictive uncertainty and proposes an ensemble-based video compression approach. In the proposed approach, an ensemble-aware loss is also introduced to seek diversity among ensemble members. Experimental results show its performance.