Neural Video Coding using Multiscale Motion Compensation and Spatiotemporal Context Model

DNN-based image compressions usually utilize auto-encoder or VAE architectures, consisting of nonlinear transforms, attention or importance map, differentiable quantization, context model, and embedded loss functions, for end-to-end learning.

Recurrent Autoencoder. Toderic et al. [13] first proposed to utilize fully-connected recurrent auto-encoders for variable-rate thumbnail image compression. A serial improvements were then extended, including the full-resolution image support, learned entropy coding, unequal bits allocation, etc [14, 15] by the introductions of ConvLSTM or ConvGRU, for better coding efficiency. Variable bit rate is intrinsically enabled by such recurrent structure. It, however, suffers from the higher computational complexity at higher bit rates, because more recurrent processing are desired.

Convolutional Autoencoder. Alternatively, convolutional auto-encoders [16, 17, 18, 19], are extensively studied in past years where different bit rates are adapted by setting a variety of $\lambda$ in learning to optimize (1). Note that different network models may be required for individual bit rates, making it difficult for hardware implementation (e.g., model adaptation for various bit rates). Recently, conditional convolution [20] and scaling factor [21] were proposed to enable variable-rate compression using a single or a limited number of network parameters without noticeable coding efficiency loss. It makes the convolutional autoencoders more attractive for practical applications.

Attention/Importance Map. Li et al. [22] utilized a separate three-layer CNN to generate importance map for spatial-complexity-based adaptive bits allocation, leading to the noticeable subject quality improvement with well-preserved edges and textures. Instead, Mentzer et al. [18] further selected one channel from the bottleneck layer to unequally weigh features at different spatial locations for simplification. Such importance map embedding was lightweight and easy for training and end-to-end optimization. This approach was later improved with nonlocal attention mechanism to efficiently and implicitly capture both global and local important information for better compression [21].

Nonlinear Transform. To generate more compact feature representation, Balle et al. [16] suggested to replace the traditional nonlinear activation, e.g., ReLU, with the generalized divisive normalization (GDN) that was theoretically proven to be more consistent with image natural statistics for visual perception. A succeeding investigation by the same authors was given in [23], reporting that GDN outperformed other nonlinear rectifiers, such as ReLU, leakyReLU and tanh, in compression tasks. Several follow-up studies [24, 25] directly applied GDN in their networks for compression exploration.

Adaptive Contexts. Probabilistic model plays a vital role in data compression. Assuming the Gaussian distribution for feature elements, Balle et al. [17] utilized hyper priors to estimate the parameters of Gaussian Scale Model (GSM) for latent features. Later Hu et al. [26] used hierarchical hyper priors (coarse-to-fine) for improving the entropy models in multiscale representations. Minnen et al. [19] improved the context modeling using joint autoregressive spatial neighbors and hyper priors based on Gaussian Mixture Model (GMM). Autoregressive spatial priors were usually extracted by PixelCNNs or PixelRNNs [27], which have been widely adopted for natural image density modeling. Reed et al. [28] further introduced multiscale PixelCNNs, yielding competitive density estimation and great speedup (e.g., from $O(N)$ to $O(\log N)$). It was later extended from 2D architectures to 3D PixelCNNs [18]. Channel-wise weights sharing-based 3D implementations can greatly reduce network parameters with higher parallelization. Then Chen et al. [21] discussed parallel pipelines of 3D PixelCNNs for practical decoding. Previous methods accumulated all the accessible priors to estimate the probability based on a single Gaussian distribution for each element. Recent explorations have shown that weighted GMMs can further improvre the coding efficiency as reported by [29, 30].

Quantization. Quantization is a non-differentiable operation, basically converting continuous variables into discrete variables with a limited alphabet. This process has to be replaced by a differentiable operation when used in end-to-end learning framework for back propagation. A number of methods, such as uniform noise adding [16], stochastic rounding [13], soft-to-hard vector quantization [18] and universal quantization [20], were developed to approximate a continuous distribution for differentiation.

Loss Functions. Pixel-error, such as MSE, or MAE, was one of the most popular loss functions used. In the meantime, SSIM or MS-SSIM was also adopted because of its better consistency with visual perception. Simulations had revealed that SSIM-based loss can improve the perception quality, especially at low bit rates. Towards the perception-optimized encoding, perceptual losses that were measured by adversarial loss [31, 32, 33] and VGG loss [34] were embedded in learning to produce visually appealing results.

Learnt video compression is extended from the learnt image compression by further exploiting the temporal redundancy in a trainable way for efficient representations of temporal motion and displaced image difference (e.g., predictive residual). In most cases, network models for image compression were directly re-used for temporal displaced residuals, leaving a great deal of efforts devoted for better motion field compression.

Chen et al. [35] developed the DeepCoder where a simple convolutional autoencoder was applied for both intra and residual coding at fixed 32$\times$32 blocks, and block-based motion estimation in traditional video coding was re-used for temporal compensation. Lu et al.[8] introduced the optical flow for motion representation in their DVC work, which, together with the intra coding in [17], demonstrated similar performance compared with the HEVC. However, the coding efficiency suffer from a sharp loss at low bit rates. Liu et al. [10] extended their nonlocal attention optimized image compression (NLAIC) for coding the intra and residual frame, and applied spatiotempoal adaptive context model for more compact motion representation, showing consistent rate-distortion gains across different contents and bits rates.

Motion can be also implicitly inferred by temporal interpolations. For example, Wu et al. [36] applied recurrent neural network (RNN)-based frame interpolation. Together with the residual compensation, it offered comparable performance with H.264/AVC. Djelouah et al. [37] further imporved the interpolation-based video coding by utilizing the advanced optical flow estimation and feature domain residual coding. Yang et al. [38] also use a interpolation method with three hierarchical quality layers and a recurrent enhancement network for compression. However, temporal interpolation usually led to coding delay that may not be acceptable for low-latency applications.

Another interesting exploration made by Ripple et al. in [39] was to jointly encode the flow and residual signals using unified quantized features in an unsupervised way. A recurrent state was embedded to aggregate multi-frame information for efficient flow generation and residual coding.

Our NVC framework is designed for low-delay applications. As with all modern video encoders, the proposed NVC compresses the first frame in each group of pictures as an intra-frame using a VAE based compression engine (neuro-Intra). It codes remaining frames using motion compensated prediction. As shown in Fig. 1, it uses the VAE compressor (neuro-Motion) to generate the multiscale motion field between the current frame and the reference frame. Then, MS-MCN takes multiscale compressed flows, warps the multiscale features of the reference frame, and combines these warped features to generate the predicted frame. The prediction residual is then coded using another VAE-based compressor (neuro-Res).

Given a group of pictures (GOP) $\mathbb{X}$ = {${\bf X_{1}},{\bf X_{2}},...,{\bf X_{t}}$}, we first encode ${\bf X_{1}}$ using the neuro-Intra module, leading to the reconstructed frame $\hat{\bf X}_{1}$. The following frame ${\bf X}_{2}$ is encoded predictively, using neuro-Motion, MS-MCN, and neuro-Res together, shown in Fig. 1. Note that MS-MCN takes the multiscale optical flows $\left\{\vec{f}^{1}_{d},\vec{f}^{2}_{d},...,\vec{f}^{s}_{d}\right\}$ derived by the pyramid decoder in neuro-Motion, and then generates the predicted frame $\hat{\bf X}^{p}_{2}$ by multiscale motion compensation. Displaced inter-residual ${\bf r}_{2}={{\bf X}_{2}}-{\hat{\bf X}^{p}_{2}}$ is then compressed in neuro-Res, yielding the reconstruction $\hat{\bf r}_{2}$. The final reconstruction $\hat{\bf X}_{2}$ is given by ${\hat{\bf X}_{2}}={\hat{\bf X}^{p}_{2}}+{\hat{\bf r}_{2}}$. Encoding of the next P-frame follows the same procedure of ${\bf X}_{2}$ until all frames in the GOP are coded completely.

Table I summarizes relevant abbreviations used throughput this paper.

The general architecture of the VAE model is shown in Fig. 2, with main encoder-decoder pair for latent feature analysis and synthesis, and hyper encoder-decoder for hyper prior generation. The main encoder ${\bf E}_{m}$ uses four stacked CNN layers, where each convolutional layer employs stride convolutions to achieve downsampling (e.g., at a factor of 2 in this example) and cascaded convolutions (e.g., three ResNet-based residual blocks [40]¹¹1We apply cascaded ResNets for stacked CNNs because of its high-efficiency and reliable performance. Other efficient CNNs architectures can be applied as well.) for efficient feature extraction. We utilize two-layer hyper encoder ${\bf E}_{h}$ to further generate the subsequent hyper priors as side information, which are used for the entropy coding for the latent features.

To capture the spatial locality, we apply convolutional layers with limited receptive field (e.g., 3$\times$3) that are stacked altogether to simulate the layer-wise feature extraction. These same ideas are used in many relevant studies [17, 19]. We utilize the simplest ReLU as the nonlinear activation function. Other nonlinear activation functions can be used as well, such as the GDN (Generalized Divisive Normalization) in [16].

Attention mechanism is leveraged to intelligently allocate bit resource (e.g., via unequal feature quantization) for image/video compression [41, 18]. It basically enables more accurate reconstruction of salient areas. We adopt the nonlocal attention module (e.g., NLAM) at the bottleneck layers of both the main encoder and hyper encoder, prior to quantization to include both global and local information for more accurate importance selection. This module is motivated by the behaviour of HVS, where we often promptly scan the entire viewing scene to have the complete understanding of the field of vision, and then fixate to the salient regions.

To enable more accurate conditional probability density modeling for entropy coding of the latent features, we introduce the Prior Aggregation (PA) engine which fuses the inputs from the hyper priors, spatial neighbors and temporal context (if applicable)²²2Intra and residual coding only use joint spatial and hyper priors without temporal inference.. The more accurate context modeling requires less resource (e.g., bits) for information representation as suggested in information theory [42]. For the sake of simplicity, we assume the latent features (e.g., motion, image pixel, residual) following the Gaussian distribution as in [19, 26], and use the PA engine to derive the mean and standard deviation of the distribution for each feature.

Our neuro-Intra is a simplified version of the NLAIC that was originally proposed in [21].

NLAIC: One major distinction of NLAIC from the VAE model using autoregressive spatial context in [19] is the introduction of a nonlocal attention module (NLAM) inspired by [43]. NLAM is used to capture the global and local importance for saliency aggregation, by which we can assign different importance to various spatial-channel elements implicitly. Such nonlocal attention mechanism is inspired by the visual information processing of spatial perception (e.g., coarse global structure plus fine-grain local details). We have found that with the addition of NLAM, we can achieve similar performance as [19] without using GDN nonlinearity.

In addition, we have applied 3D 5$\times$5$\times$5 masked CNN³³3This 5$\times$5$\times$5 convolutional kernel shares the same parameters for all channels, offering great model complexity reduction compared with 2D CNN-based solution in [19]. to extract spatial priors, which are fused with hyper priors in PA for entropy context modeling (e.g., bottom part of Fig. 6). Here, we have assumed the single Gaussian distribution for the context modeling of entropy coding. More details of NLAIC can be found in [21]. Note that temporal priors are not adopted for intra-pixel and inter-residual in this paper.

Improvements to NLAIC. Original NLAIC applies multiple NLAMs in both main and hyper coders, leading to excessive memory consumption at a large spatial scale. In NVC, NLAMs are only used at the bottleneck layers for both main and hyper encoder-decoder pairs, achieving adaptive bits allocation implicitly.

To overcome the non-differentiability of the quantization operation, quantization is simulated by adding uniform noise in [17]. However, such noise augmentation is not exactly consistent with the rounding in inference, yielding the performance loss as reported by [20]. Thus, we apply the universal quantization (UQ) [20] in neuro-Intra, i.e.,:

where $\hat{x}$ is a quantized symbol and $u$ represents the random uniform variable ranging from $-\frac{1}{2}$ to $\frac{1}{2}$. Statistically, we define the gradient as 1 for back propagation since UQ can be approximated as a linear function. Such UQ (2) is used for neuro-Motion and neuro-Res as well. When applying to common Kodak dataset, neuro-Intra achieved similar performance as NLAIC [21], outperforming Minnen (2018) [19], BPG (4:4:4) and JPEG2000, as shown in Fig. 3.

Inter-frame coding plays a vital role in video coding. The key is how to efficiently represent motion in a compact format for effective compensation. In comparison to the pixel-domain block-based motion estimation and compensation in conventional video coding, we rely on the optical flow to accurately capture the temporal information for motion compensation.

Inter-frame residual coding is another significant module contributing to the overall system efficiency. It is used to compress the temporal prediction error pixels. It affects the efficiency for next frame prediction since errors usually propagate temporally.

Here we use the VAE architecture in Fig. 2 for encoding the residual ${\bf r}_{t}$. The rate-constrained loss is used:

where $\mathbb{D}_{2}$ is the $\ell_{2}$ loss between a residual compensated frame ${\bf X}^{p}_{t}+{\hat{\bf r}_{t}}$ and ${\bf X}_{t}$. neuro-Res will be first pretrained using the predicted frames by the pretrained neuro-Motion and MS-MCN, and a loss function in (8) where the rate $R$ only consider the bits for residual. Then we refine it jointly with neuro-Motion and MS-MCN, using a loss where $R$ considers the bits for both motion and residual with two frames.

Progressive Training. Training all NVC components directly is difficult and unreliable, because these modules are interdependent with each other. For example, inter prediction depends on the former reconstructed frame and meanwhile a better predicted frame will reduce the residual energy. In our model development, we first train the neuro-Intra for multiple bit rates, using multiple $\lambda$ values; Then we pretrain and fine-tune the neuro-Motion and MS-MCN through Step 1-3 described in Sec. III-D2, to minimize a training loss that consider both the prediction error and the rate for motion features. We also pretrain neuro-Res using the predicted frames generated by the trained neuro-Motion and MS-MCN so far, with a loss that considers both bit rate of residual and the final reconstruction of the current frame. Finally, we further refine all the modules for inter-coding, including neuro-Motion, MS-MCN and neoro-Res, using a training loss including the reconstruction error and the total rate for the motion features and residual features. For each target bit rate, we set the $\lambda$ for training the inter-coding model to be proportional to the $\lambda$ used for training the neuro-Intra module.

Training with Multi-frame Loss. One way to train the inter-coding model (including neuro-Motion, MS-MCN and neuro-Res modules) is to use only two frames as a training sample and use the neuro-Intra module to code the first frame, and use the inter-coding model to code the second frame. This training approach can only learn the intra-to-inter variations, yielding instability and poor reconstruction for the succeeding inter frames distant from the first intra frame. To overcome the quality degradation in future inter-frames, we adopt a multi-frame training strategy.

We first pretrain the inter-coding model using pairs of two successive frames as training samples, where the first frame is encoded and decoded using the neuro-Intra. This step follows the progressive training procedure. We then refine the inter-coding model by using groups of four successive frames as training samples, where the first frame is encoded and decoded by a pretrained neuro-Intra, and each following frame is coded using the inter-coding model and the decoded frame is then used recursively as the reference frame for coding the next frame. We use a loss function that considers the sum of the distortions (in MSE or negative MS-SSIM) in all three P-frames, and the sum of the rate for these frames, so that the model update considers the impact of the propagation of P-frame reconstruction errors. Note that once the inter-coding model is trained, it is applied to all inter frames in a GOP during testing. Although it is possible to improve the performance by training a different inter-coding model for each possible distance between a P-frame and the intra-frame, we choose to train a single model that is used repeatedly for all P-frames, to reduce the implementation complexity.

Datasets. The neuro-Intra is trained on COCO [47] and CLIC [48] with training samples randomly cropped into 256$\times$256$\times$3. The neuro-Motion, MS-MCN and neuro-Res are joint trained using Vimeo 90k [49] with sample size of 192$\times$192$\times$3.

Our NVC is evaluated using the standard HEVC test sequences and ultra video group (UVG) dataset. HEVC test sequences include different classes, covering the contents with a variety of motion, frame rate, resolution and texture; UVG dataset has seven 1080p videos which are often selected for testing video applications.

Loss Function. We have used either MSE or negative MS-SSIM loss for training the intra-coding and inter-coding modules. For pretraining the neuro-Motion and MS-MCN modules, we use the $\ell_{1}$ loss on the prediction error instead which is similar to the mean absolute error (MAE) used in traditional motion estimation. We do not use MS-SSIM loss on the predicted frame because MS-SSIM mainly cares about the structure similarity and may ignore the background noise to some extent. Through experiments, we have found that using MS-SSIM loss for neuro-Motion can lead to temporal error accumulation, which not only increases the bits consumption but also potentially leads to unreliable training.

Hyperparameters and Platform. The initial learning rate (LR) is set to 10e-4 and is halved for every 10 epochs, and final models are obtained using a LR of 10e-5 for both pretraining and overall training. We apply the distributed training on 4 GPUs (Titan Xp) for 3 days for each bit rate model.

Evaluation Criteria. We apply the same low-delay coding setting as DVC in [8] for our method and traditional H.264/AVC, and HEVC for fair comparison. We encode 100 frames and use GOP of 10 on HEVC test sequences, and 600 frames with GOP of 12 on UVG dataset. Both PSNR and MS-SSIM results are offered to understand the efficiency of NVC.

Generate Models for Different Target Rates. To generate a model for a particular bit rate, we first train the neuro-Intra using a certain $\lambda$ value, denoted by $\lambda_{\rm intra}$. Then we train the inter-frame models (neuro-motion, MS-MCN and neuro-Res) using a proportionally reduced $\lambda$ value, $\lambda_{\rm inter}=\lambda_{intra}/4.$ This simple way of setting the $\lambda_{\rm intra}$ and $\lambda_{\rm inter}$ values yielded good results. We chose not to further optimize $\lambda_{\rm inter}$ for given $\lambda_{\rm intra}$.

Rate distortion Performance. We compare the coding perfomance of different methods in Fig. 7 and 8 using respective PSNR and MS-SSIM measures, across HEVC and UVG test sequences. Note that when we report PSNR or MS-SSIM values, the models are trained using PSNR and MS-SSIM, respectively, as the distortion metric. In Table III, by setting the same anchor using H.264/AVC, our NVC presents 35% BD-Rate gains when the distortion is measured by PSNR, while HEVC offers 30% gains. Gains are even larger, if the distortion is measured by the MS-SSIM. In this case, NVC can achieve 50% improvement, while HEVC is around 25%.

Wu et al. [36] proposed a interpolation based video coding framework and could not get better performance than H.265/HEVC. Our NVC performs significantly better than DVC [8], which has 22.26% and 23.08% BD-Rate reductions under PSNR and MS-SSIM metrics, respectively. From Fig. 7 and Fig. 8, DVC [8] has mainly improved the coding efficiency against HEVC at high bit rates. However, DVC is not competitive at low bit rate (e.g., having performance even worse than H.264/AVC at some rate regimes.). We have also observed that DVC’s performance varies for different test sequences. An improved version of DVC, known as DVC_pro, has been recently reported, which has shown the state of the art performance [9] by using [19] for intra and residual coding and $\lambda$ tuning. NVC is slightly better that DVC Pro in both PSNR and MS-SSIM, at 34.5% and 45.88%, respectively.

Visual Comparison. We provide the visual quality comparison with H.264/AVC and H.265/HEVC in Fig. 9. Generally, NVC leads to a higher quality reconstruction at slightly lower bit rate. For example, for the “RaceHorse” which has nontranslational motion and complex background, NVC uses 7% less bits for more than 1.5 dB PSNR improvement compared with H.264/AVC. For other cases, our method also shows robust improvement. Traditional codec usually suffers from blocky artifacts and motion-induced noise close to object edges. One can clearly observe block partition boundaries with severe pixel discontinuity in reconstructed frames by H.264/AVC. Our results have much less visible noise and artifacts.

This section examines modular components in NVC to further understand its capacity for application.

Spatiotemporal Context Modeling. To evaluate the gain from using temporal priors generated by ConvLSTM for the entropy coding of the motion features, we have also trained a model when the PA engine in neuro-Motion does not use the temporal priors. Table IV shows that 2% to 7% rate increase are incurred if the temporal priors are not used. This reveals that temporal priors help to make probability prediction more accurate, leading to less bit consumption for compressing the motion features.

Generally, bits consumed by motion information vary across different content and total target rates. More saving is reported for low bit rates, and for motion intensive content. In these cases, motion bits usually occupy a higher percentage of the total rate. For stationary content, such as HEVC Class E, motion bits contribute less percentage, thus the gain from using temporal priors is less significant.

Comparison with Cascaded Denoising Networks. Linear warping often brings artifacts (e.g., ghosting edges), especially when flow estimation is not accurate and occlusion happens. To remove such artifacts, a denoising network is often added after motion compensated warping as shown in Fig. 4a and 4b. Our MS-MCN, instead, rely on multiscale motion fields to generate the predicted frame. As a comparison, we also trained a U-net like denoising network on the warped frames based on the single-scale decoded motion field. As shown in Fig. 10, MS-MCN achieves 0.2 to 0.3 dB PSNR gains across a large bit rate ranges over using a cascaded denoising network for various content for overall performance.

Temporal Stability with Multi-frame Training. Compression Errors (e.g., lossy compression noise) often propagate from one frame to another due to predictive coding structure. It is critical to limit the temporal quality degradation during training. To evaluate the gain from multi-frame training, we compare the performance of the NVC model trained by minimizing reconstruction errors for only one P-frame (single frame training) and the model further refined by minimizing the reconsturction loss for multiple consecutive P-frames. We have varied the number of total frames to use while training, and found that using 4 frames reaches a good compromise. As shown in Fig. 11, multi-frame training could well capture temporal quality variations from intra to inter, and from inter to inter, leading to a more generalized model with consistent performance for all the frames in a GOP with slower quality degradation, and improved stability.