Prediction and Reference Quality Adaptation for Learned Video Compression

Temporal contexts are generated by feature-domain temporal prediction. Therefore, the quality of temporal prediction affects the temporal context quality. Since the complexities of video content and motion patterns are unevenly distributed in space, spatial-wise prediction quality difference exists in temporal contexts. For example, as shown in Fig. 3 (a), the prediction error of the horsetail in the red rectangle is large since it has complex textures and motion, while that of the horse body is smaller since their textures are simpler and their motion pattern is mainly simple translational motion. Using temporal context with large spatial prediction errors may increase coding costs. Therefore, it is necessary to give explicit discrimination to the spatial-wise temporal context quality difference to help the context encoder and decoder decide whether to use the temporal context in a certain region.

In addition to the spatial-wise prediction quality difference, channel-wise prediction quality difference also exists in temporal contexts. Commonly, temporal contexts have multiple feature channels, and each channel is compensated by different motion vectors [46]. Therefore, temporal information is unevenly distributed across context channels. Different context channels may have different prediction qualities and contain different temporal information. For example, as shown in Fig. 3 (b), the $9^{th}$ and $47^{th}$ channels of the temporal context have more temporal information and their temporal prediction qualities are higher, while the $17^{th}$ channel of temporal context has less temporal information and its temporal prediction quality is lower. Therefore, it is also necessary to give explicit discrimination to the channel-wise temporal context quality difference to help the contextual encoder and decoder decide which channel of temporal information to use.

To give explicit discrimination to the spatial and channel-wise prediction quality difference of temporal contexts, we design a simple yet effective prediction quality adaptation module as illustrated in Fig. 4. Given a certain scale temporal context $C_{t}^{i}$, we first concatenate (concat) it with an intermediate feature $I_{t,e}^{i}$ of the contextual encoder or an intermediate feature $I_{t,d}^{i}$ of the contextual decoder with the same scale. For the biggest temporal context $C_{t}^{0}$ in the encoder, the input frame $x_{t}$ serves as $I_{t,e}^{0}$. Then, we feed them into a convolutional layer to obtain a feature with the same spatial and channel dimension as $C_{t}^{i}$. The feature is passed into a Sigmoid function to obtain a group of confidence maps for each channel of the temporal context $C_{t}^{i}$. Considering that the temporal contexts for reducing temporal redundancy and video reconstruction may be different, given a temporal context $C_{t}^{i}$, we calculate different confidence maps for the contextual encoder and decoder.

where $w_{t,e}^{i}$ is the confidence map group for contextual encoder and $w_{d,t}^{i}$ is the confidence map group for contextual decoder. $i$ is the scale index of temporal contexts. $\mathcal{F}_{e}$ and $\mathcal{F}_{d}$ denote 2D convolution operations with $3\times 3$ filter size. $\sigma$ refers to the Sigmoid function. Each element of $w_{t,e}^{i}$ and $w_{d,t}^{i}$ is in the range of 0 to 1. The greater the value of one element, the more important the temporal context of the corresponding position is.

After obtaining the confidence map for each channel of $C_{t}^{i}$, we perform element-wise product operation between the confidence map and the corresponding channel of temporal context.

where $\odot$ is the element-wise product operation, $m$ is the channel index of the temporal context $C_{t}^{i}$, and $M$ is the number of channels of $C_{t}^{i}$.

In addition to prediction quality, reference quality also needs to be adapted to reduce reconstruction error propagation. As shown in Fig. 5, the continuous degradation of reconstruction quality is relatively limited in traditional video codecs, such as VTM. Their transforms like DCT are lossless and their reconstruction distortion only comes from the quantization. Therefore, it is easier for them to achieve a target reconstruction quality by adjusting a quantization step. However, in learned video codecs, the reconstruction distortion not only comes from quantization but also comes from non-linear transform networks. The non-linear transform networks generate an implicit quantization for the input frame. The degree of the implicit quantization depends on the $\lambda$ in the rate-distortion loss function and the quality of reference that is used to predict temporal contexts. In other words, the transform networks will adapt to the reference quality to achieve the target reconstruction quality controlled by $\lambda$.

However, current learned video compression schemes are difficult to achieve this goal. There are two main reasons for this problem. One is their transform networks cannot adapt to different reference qualities well. Although some learned video compression schemes proposed to add a periodically varying weight before $\lambda$ in the loss function [39, 46] to adapt to different reference qualities, it is still difficult to adapt various reference qualities with a limited number of weights (1.2, 0.5, 0.9) to achieve the target quality. As shown in Fig. 5, for our baseline, which has used this kind of reference quality adaptive loss function, the qualities of its reconstructed frames still gradually decrease as the reference qualities decrease. Another reason is that the transform networks only “see” a limited range of reference qualities during training, resulting in an inability to adapt to unseen reference qualities during testing.

To make the transform networks adapt to reference qualities well, we propose a reference quality adaptation module, as presented in Fig. 6. We first feed the reference frame $\hat{x}_{t-1}$ into a convolutional layer $\mathcal{F}_{RQA}$ to obtain spatially variant filters $W_{t}^{\theta}\in R^{H\times W\times k\times k}$ for each position of an intermediate feature $I_{t}$ of the contextual encoder or decoder, where $H$ and $W$ are height and width of the intermediate feature, and $k$ refers to the size of filters:

Then, each filter $W_{t}^{\theta}(i,j)$ is applied to a $k\times k$ window of the intermediate feature $I_{t}$ centered at position $(i,j)$ to conduct spatially variant filtering.

To make the RQA module “see” a wide range of reference qualities during training, we add a repeat-long training step that combines repeating compressing and long-sequence cascaded training in the commonly-used multiple-step training strategy, which will be introduced in Section IV-C.

We propose a step-by-step training strategy to train our learned video compression scheme. The training details for encoding RGB videos when the reconstruction distortion is measured by RGB PSNR are listed in Table. I. According to the training loss function, the training strategy can be classified into 7 classes: $L_{t}^{meD}$, $L_{t}^{meRD}$, $L_{t}^{recD}$, $L_{t}^{recRD}$, $L_{t}^{all}$, $L_{T}^{all}$, and $L_{T}^{all-r-l}$.

• •

  $L_{t}^{meD}$ is comprised of $D_{t}^{me}$. We denote $D_{t}^{me}$ to the distortion between $x_{t}$ and its predicted frame $\tilde{x}_{t}$. MSE is used to measure the distortion. The predicted frame $\tilde{x}_{t}$ is generated by warping the reference frame $\hat{x}_{t-1}$ using the decoded motion vectors $\hat{v}_{t}$. Using $L_{t}^{meD}$ as the loss function is to make the motion decoder generate accurate motion vectors:

  $$L_{t}^{meD}=w_{t}\cdot\lambda\cdot D_{t}^{me}.$$

  (5)

• •

  $L_{t}^{meRD}$ is comprised of $D_{t}^{me}$ and $R_{t}^{me}$. We denote $R_{t}^{me}$ to the joint bit rate used for encoding the quantized motion latent representation $\hat{m}_{t}$ and its associated hyperprior. Using $L_{t}^{meRD}$ as the loss function is to achieve a trade-off between the accuracy and the consumed bit rate of motion vectors:

  $$L_{t}^{meRD}=w_{t}\cdot\lambda\cdot D_{t}^{me}+R_{t}^{me}.$$

  (6)

• •

  $L_{t}^{recD}$ is comprised of $D_{t}^{rec}$. We denote $D_{t}^{rec}$ to the distortion between $x_{t}$ and its reconstructed frame $\hat{x}_{t}$. MSE is used to measure the distortion. Using $L_{t}^{recD}$ as the loss function is to make the contextual decoder and frame generator generate a high-quality reconstructed frame:

  $$L_{t}^{recD}=w_{t}\cdot\lambda\cdot D_{t}^{rec}.$$

  (7)

• •

  $L_{t}^{recRD}$ is comprised of $D_{t}^{rec}$ and $R_{t}^{rec}$. We denote $R_{t}^{rec}$ to the joint bit rate used for encoding the quantized contextual latent representation $\hat{y}_{t}$ and its associated hyperprior of $x_{t}$. Using $L_{t}^{recRD}$ as the loss function is to achieve a trade-off between the quality of reconstructed frame $\hat{x}_{t}$ and the consumed bit rate of contextual latent representation $\hat{y}_{t}$:

  $$L_{t}^{recRD}=w_{t}\cdot\lambda\cdot D_{t}^{rec}+R_{t}^{rec}.$$

  (8)

• •

  $L_{t}^{all}$ is comprised of $D_{t}^{rec}$ and $R_{t}^{all}$. We denote $R_{t}^{all}$ to all the bit rates used for encoding the quantized contextual latent representation $\hat{y}_{t}$, the quantized motion latent representation $\hat{m}_{t}$, and their hyperprior. Using $L_{t}^{all}$ as the loss function is to achieve a trade-off between the quality of the reconstructed frame and all the consumed bit rates of the coded frame:

  	$\displaystyle L_{t}^{all}$	$\displaystyle=w_{t}\cdot\lambda\cdot D_{t}^{rec}+R_{t}^{all}$		(9)
  		$\displaystyle=w_{t}\cdot\lambda\cdot D_{t}^{rec}+R_{t}^{me}+R_{t}^{rec}.$

• •

  $L_{T}^{all}$ is the average $L_{t}^{all}$ loss of $T$ consecutive frames. Calculating the average loss of multiple frames is to achieve a cascaded fine-tuning for reducing the error propagation [37, 39, 42, 46]:

  	$\displaystyle L_{T}^{all}$	$\displaystyle=\frac{1}{T}\sum_{t}L_{t}^{all}$		(10)
  		$\displaystyle=\frac{1}{T}\sum_{t}\left\{w_{t}\cdot\lambda\cdot D_{t}^{rec}+R_{t}^{all}\right\}$
  		$\displaystyle=\frac{1}{T}\sum_{t}\left\{w_{t}\cdot\lambda\cdot D_{t}^{rec}+R_{t}^{me}+R_{t}^{rec}\right\},$

  

• •

  $L_{T}^{all-r-l}$ is the RQA-associated repeat-long training loss proposed in our paper to make RQA module “see” more reference qualities during training, as described in Section IV-B. To generate reference frames with various qualities, we first repeat compressing the first P-frame for certain times (randomly selected from 0 to $N$). Then, we regard its reconstructed frame as the reference of the second P-frame and calculate the cascaded loss $L_{T}^{all}$ of the remaining $T$ frames. A similar repeating compressing strategy is proposed by the coding experts of [53] and has been accepted into the developing AVS end-to-end video coding standard reference model AVS-EEM-v2.0¹¹1AVS video coding group is exploring a lightweight end-to-end learned video coding standard. AVS-EEM is the corresponding reference model. Its training and testing codes can be accessed at https://gitlab.com/xhsheng/avs-eem once authorized. . However, we find that directly applying the repeating compressing strategy only brings a little compression performance improvement in our model. This is because this strategy can only generate reference with various qualities for the first P-frame, and cannot affect the reference qualities of subsequent P-frames. Therefore, we combine this strategy with long-sequence training. We increase the number of frames $T$ ($T=17$) used to calculate the cascaded loss $L_{T}^{all}$ to generate various reference qualities for subsequent P-frames. With the repeating compressing and long-sequence training strategies, the RQA module can adapt to a wider range of reference qualities.

We use the Lagrangian multiplier $\lambda$ to control the rate-distortion (R-D) trade-off. We also add a periodically varying weight $w_{t}$ for each P-frame before the Lagrangian multiplier $\lambda$ to implement the hierarchical quality [39, 46]. The detailed setting of $\lambda$ and $w_{t}$ can be found in the Section V-A3.

When using $L_{t}^{meD}$ as the loss function, we only train the motion estimation module, motion encoder, and motion decoder (Inter). When using $L_{t}^{meRD}$ as the loss function, we add the motion entropy model into the training loop (Inter). When using $L_{t}^{recD}$ as the loss function, we only train the temporal context mining module, contextual encoder, contextual decoder, and frame generator (Rec). When using $L_{t}^{recRD}$ as the loss function, we add the contextual entropy model into the training loop (Rec). When using $L_{t}^{all}$ or $L_{T}^{all}$ as the loss function, we train all modules.

When the quality of reconstructed RGB videos is measured by MS-SSIM, we use the RGB-MSE model before being trained with $L_{T}^{all}-repeat-long$ as the pre-trained model. Then, we replace the distortion metric from MSE to 1–MS-SSIM to fine-tune the pre-trained model. The detailed fine-tuning strategy is listed in Table. II.

When encoding YUV420 videos, we use the RGB-MSE model before being trained with $L_{T}^{all}$ in RGB colorspace as the pre-trained model. Then, we fine-tune the pre-trained model with the loss functions calculated in the YUV colorspace. The detailed fine-tuning strategy is listed in Table. III.

Following the setting of previous learned video coding schemes [37, 39, 46, 42], we include the time for model inference, entropy modeling, entropy coding, and data transfer between CPU and GPU when calculating the encoding and decoding time. Table. VIII lists the detailed encoding and decoding time for a 1920$\times$1080 video frame of different learned video codecs. We run all these codecs on a NVIDIA 3090 GPU. The comparison results show that our proposed technologies lead to only 0.17s encoding time and 0.11s decoding time increase compared with our baseline [37]. In addition, the number of trainable parameters of our proposed scheme is 19.3M, which is only 0.5M more than our baseline.