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 proposes an ensemble-based video compression model to capture the predictive uncertainty of intermediate predictions. A loss is constructed to encourage diversity between ensemble members, and the paper investigates the benefit of incorporating adversarial training in video compression. The experimental result shows that the proposed model can save more than 20% compared to DVC Pro.