Hybrid Spatial-Temporal Entropy Modelling for
Neural Video Compression

The neural image codec has developed rapidly in recent years. The early work (Theis et al., 2017) uses a compressive autoencoder-based framework to achieve similar RD performance with JPEG 2000. Recently, many works focus on the design of entropy model. Ballé et al. (Ballé et al., 2017) proposed the factorized model and got better RD performance than JPEG 2000. The hyper prior (Ballé et al., 2018) introduces the hierarchical design and uses additional bits to estimate the distribution, which gets comparable results with H.265. Subsequently, the auto-regressive prior (Minnen et al., 2018) was proposed to explore the spatial correlation. It obtains higher compression ratio together with hyper prior. However, auto-regressive prior inferences in a serialized order for all spatial positions, and thus it is quite slow. The checkerboard context model (He et al., 2021) solves the complexity problem by introducing the two-step coding. In addition, to further boost the compression ratio, the mixture Gaussian model (Cheng et al., 2020) was proposed and its RD results are on par with H.266. Recently, the vision transformer attracts lots of attention, and the corresponding entropy model (Koyuncu et al., 2022) helps the neural image codec exceed H.266 intra coding.

The success of neural image codec also pushes the development of neural video codec. The pioneering work DVC (Lu et al., 2019) follows the traditional codec, and uses the residual coding-based framework where motion-compensated prediction is first generated and then the residual is coded using hyper prior (Ballé et al., 2018). With the help of auto-regressive prior (Minnen et al., 2018), its following work DVCPro achieves higher compression ratio.

Most recent works follow this motion estimation (Ranjan and Black, 2017; Hui et al., 2020) and residual coding-based framework. Then more advanced networks structures are proposed for generating the optimized residual or motion. For example, the residual is adaptively scaled by learned parameter in (Yang et al., 2021b). The optical flow estimation in scale space (Agustsson et al., 2020) was proposed to reduce the residual energy in fast motion area. In (Hu et al., 2020), the rate distortion optimization is applied to improve the coding of motion. The deformable compensation (Hu et al., 2021) is used to improve the prediction in feature space. Lin et al. (Lin et al., 2020) proposed using multiple reference frames to reduce the residual energy. In (Lin et al., 2020; Rippel et al., 2021), the motion prediction is introduced to improve the coding efficiency of motion.

Besides the residual coding, other coding frameworks are also investigated. For example, 3D autoencoder (Pessoa et al., 2020; Habibian et al., 2019; Sun et al., 2020) was proposed to encode the multiple frames simultaneously. It is a natural extension of neural image codec by expanding the input dimension. However, such kind of framework will bring significant encoding delay and is not suitable for real-time scenarios. Another emerging coding framework is the conditional coding whose entropy bound is lower than or equal to residual coding (Ladune et al., 2020). For example, Ladune et al. (Ladune et al., 2021b, a, 2020) used the conditional coding to code the foreground contents. In DCVC (Li et al., 2021), the condition is the extensible high-dimension feature rather than the 3-dimension predicted frame. The following work (Sheng et al., 2021) further boosts the compression ratio by introducing the feature propagation and multi-scale temporal contexts.

However, most existing neural video codecs focus on how to generate the optimized latent representation and design the network structures therein. As for the entropy model used for coding the latent representation, the ready-made solutions from neural image codec are directly used. In this paper, we focus on the entropy model design by efficiently leveraging both spatial and temporal correlations. Actually, some works also have begun to investigate it. For example, the conditional entropy coding was proposed in (Liu et al., 2020c). The temporal context prior extracted from temporal feature is used in (Li et al., 2021; Sheng et al., 2021). Yang et al. (Yang et al., 2021a) proposed the recurrent entropy model. However, these works (Liu et al., 2020c; Li et al., 2021; Sheng et al., 2021; Yang et al., 2021a) focus more on utilizing temporal correlation. Although the work in (Li et al., 2021) also investigates the spatial correlation, the auto-regressive prior is used and leads to a very slow coding speed. An entropy model which not only fully exploits the spatial-temporal correlation but also has low complexity is desired. To meet this requirement, we specially design the latent prior and time-efficient dual spatial prior to equip the entropy model, and push the compression ratio to a new height. In addition, all these methods (Liu et al., 2020c; Li et al., 2021; Sheng et al., 2021; Yang et al., 2021a) need to train separate model for each rate point. By contrast, we design a multi-granularity quantization, which is powered by our entropy model. This content-adaptive quantization mechanism helps our codec achieve smooth rate adjustment in single model. The rate adjustment in single model was investigated for neural image codec via the gain unit (Cui et al., 2021). And yet, to the best of our knowledge, our codec is the first neural video codec to obtain such capability via our multi-granularity quantization.

To achieve higher compression ratio for the input frame $x_{t}$ ($t$ is the frame index), we adopt the conditional coding-based framework rather than the residual coding-based framework. Specifically, we follow the DCVC (Li et al., 2021) and its improved work (Sheng et al., 2021), and then redesign the core modules therein. The framework of our codec is presented in Fig. 1. As shown in this figure, the entire coding process can be roughly divided into three steps: temporal context generation, contextual encoding/decoding, and reconstruction.

Temporal context generation. To fully explore the temporal correlation, we generate the multi-scale contexts $\bar{C}^{0}_{t}$, $\bar{C}^{1}_{t}$,$\bar{C}^{2}_{t}$ at different resolutions via the temporal context mining module (more details of this module can be found in (Sheng et al., 2021)). To carry richer information, we also use the temporal feature $F_{t-1}$ rather than previous decoded frame $\hat{x}_{t-1}$ as the module input. As for motion estimation, we adopt the light-weight SPyNet (Ranjan and Black, 2017) for acceleration.

Contextual encoding/decoding. Conditioned by the multi-scale contexts, the current frame $x_{t}$ is transformed into latent representation $y_{t}$ by the contextual encoder. To achieve bitrate saving, the $y_{t}$ is quantized to $\ddot{y}_{t}$ before being sent to the arithmetic encoder which generates the bit-stream. During the decoding, $\ddot{y}_{t}$ is decoded from bit-steam by arithmetic decoder and inversely quantized to $\hat{y}_{t}$. Also conditioned on the multi-scale contexts, the contextual decoder decodes the high-resolution feature $\hat{F}_{t}$ from $\hat{y}_{t}$. In this encoding/decoding process, how to accurately estimate the distribution of $\ddot{y}_{t}$ by the entropy model is vital for bitrate reduction. To this end, we propose the hybrid spatial-temporal entropy model (Section 3.2). To support smooth rate adjustment in single model, the multi-granularity quantization powered by our entropy model is proposed (Section 3.3).

Reconstruction. After obtaining the high-resolution feature $\hat{F}_{t}$, our target is generating the high-quality reconstructed frame $\hat{x}_{t}$ via the frame generator. Different from DCVC (Li et al., 2021) and (Sheng et al., 2021) only using plain residual blocks (He et al., 2016), we proposing using the W-Net (Xia and Kulis, 2017) based structure which ties two U-Nets. Such kind of network design can effectively enlarge the receptive field of model with acceptable complexity. This results in stronger generation ability for model.

For arithmetic coding, it needs to know the probability mass function (PMF) of $\ddot{y}_{t}$ to code it. However, we do not know its true PMF $p(\ddot{y}_{t})$ and usually approximate it with an estimated PMF $q(\ddot{y}_{t})$. The cross-entropy $\mathbb{E}_{\ddot{y}_{t}\sim p}[-log_{2}q(\ddot{y}_{t})]$ captures the average number of bits needed by the arithmetic coding without considering the negligible overhead. In this paper, we follow the existing work (PyT, 2022) and assume that $q(\ddot{y}_{t})$ follows the Laplace distribution. Thus, our target is designing an entropy model which can accurately estimate the distribution parameter of $q(\ddot{y}_{t})$ to reduce the cross-entropy.

To improve the estimation, for each element $\ddot{y}_{t,i,j,k}$ ($i$, $j$, and $k$ are the height, width, and channel indexes) therein, we need to fully mine the correlation between it and its known information, e.g., the previous decoded latent representation, temporal context feature, and so on. Theoretically, $\ddot{y}_{t,i,j,k}$ may correlate to all these information from all previous decoded positions. For traditional codec, it is unable to explicitly exploit such correlation due to the huge space. Thus, traditional codec usually uses simple handcrafted rules to use context from a few neighbour positions. By contrast, the deep learning enables the capability of automatically mining the correlation in huge space.

Thus, we propose feeding manifold inputs to the entropy model, and let the model extract the complementary information from the rich high-dimensional inputs. The illustration of our entropy model is shown in Fig. 2. As shown in this figure, the inputs not only contain the commonly-used hyper prior $\hat{z}_{t}$ and the temporal context $\bar{C}^{2}_{t}$, but also include the $\hat{y}_{t-1}$ as the latent prior. Although both $\bar{C}^{2}_{t}$ and $\hat{y}_{t-1}$ are from temporal direction, they have different characteristics. For example, $\bar{C}^{2}_{t}$ at 4x down-sampled resolution usually contains lots of motion information (more details can be found in (Sheng et al., 2021)). By contrast, $\hat{y}_{t-1}$ is in the latent representation domain at 16x down-sampled resolution, and has more similar characteristics with $y_{t}$. Thus, $\bar{C}^{2}_{t}$ and $\hat{y}_{t-1}$ can provide the complementary auxiliary information to improve estimation. In addition, it is noted that we adopt the cascaded training strategy (Chan et al., 2021; Sheng et al., 2021), where the gradients will back propagate to multiple frames. Under such training strategy, the propagation chain of latent representation is formed. It means that the connection between the latent representation of the current frame and that of long-range reference frame is also built. Such connection is very helpful for extracting the correlation across the latent representation of multiple frames, and then results in more accurate prediction on distribution.

Besides using the latent prior $\hat{y}_{t-1}$ to enrich the input, our entropy model also adopts the dual spatial prior to exploit the spatial correlation. To pursue a time-efficient mechanism, we follow the checkerboard model (He et al., 2021) rather than the commonly-used auto-regressive model (Minnen et al., 2018) which seriously slows down the coding speed. However, all channels in the original checkerboard model use the same coding order, namely the even positions are always first coded and then used as context for the coding of odd positions. Such coding order cannot handle various videos because sometimes coding the even positions first has worse RD performance than coding the odd positions first. Therefore, we design the dual spatial prior, where half channels in both even and odd positions are first coded at the same time.

As shown in Fig. 2, the latent representation is split into two chunks along the channel dimension. During the first-step coding, the above branch in Fig. 2 will code the even positions $\ddot{y}_{t,(i+j)\%2=0,k<\frac{C}{2}}$ in the first chunk. Simultaneously, the below branch codes the odd positions $\ddot{y}_{t,(i+j)\%2=1,k\geq\frac{C}{2}}$ in the second chunk. The uncoded positions (i.e., the white regions in Fig. 2) in these two chunks are set as zero. After the first-step coding, the coded $\ddot{y}_{t,(i+j)\%2=0,k<\frac{C}{2}}$ and $\ddot{y}_{t,(i+j)\%2=1,k\geq\frac{C}{2}}$ are fused together and then further generate the contexts for the second-step coding. During the second-step coding, the left positions in each chunk are coded. The above branch codes the odd positions $\ddot{y}_{t,(i+j)\%2=1,k<\frac{C}{2}}$ in the first chunk and the below branch codes the even positions $\ddot{y}_{t,(i+j)\%2=0,k\geq\frac{C}{2}}$ in the second chunk. As shown in Fig. 2, this coding manner enables that the second-step coding can benefit from the contexts from all positions. When compared with original checkerboard model (He et al., 2021), the scope of spatial context is doubled and results in more accurate prediction on distribution. It is noted that, for both steps, the first chunk and second chunk will be added and then sent to the arithmetic encoder. Thus, our dual spatial prior will not bring any additional coding delay when compared with (He et al., 2021).

In addition, from the perspective of channel dimension, our dual spatial prior also mines the correlation across channels. The coded $\ddot{y}_{t,(i+j)\%2=0,k<\frac{C}{2}}$ during the first-step coding process can be also used as the condition for the coding of $\ddot{y}_{t,(i+j)\%2=0,k\geq\frac{C}{2}}$ during the second-step coding process. It is similar for the coding of $\ddot{y}_{t,(i+j)\%2=1,k\geq\frac{C}{2}}$ and $\ddot{y}_{t,(i+j)\%2=1,k<\frac{C}{2}}$, where the coding direction of channel is inverse. In a summary, our proposed dual spatial prior further squeezes the redundancy in $\ddot{y}_{t}$ by more efficiently exploiting the correlation across the spatial and channel positions.

It is a big pain that most existing neural codecs cannot handle rate adjustment in single model. To achieve different rates, the model needs to be retrained by adjusting the weight in the RD loss. It will bring large training cost and model storage burden. Such shortcoming calls for a mechanism supporting neural codec to achieve wide rate range in single model. To this end, we propose an adaptive quantization mechanism to enable this capability.

As shown in Fig. 3, our multi-granularity quantization involves three different kinds of quantization step (QP): the global QS $qs^{global}$, the channel-wise QS $qs^{ch}$, and the spatial-channel-wise QS $qs^{sc}$. The $qs^{global}$ is only a single value and is set from the user input for controlling the target rate. As all positions take the same QS, the $qs^{global}$ brings a coarse quantization effect. Thus, motivated by the channel attention mechanism (Hu et al., 2018), we also design a modulator $qs^{ch}$ to scale the QS at different channels because different channels carry information with different importance. However, the different spatial positions also have different characteristics due to the various video contents. Thus, we follow (Huang et al., 2022) and learn the spatial-channel-wise $qs^{sc}$ to achieve the precise adjustment on each position.

The $qs^{sc}$ is generated by our entropy model, as shown in Fig. 2. It is noted that, for each frame, it is dynamically changed to adapt the video contents. Such design not only helps us achieve smooth rate adjustment but also improves the final RD performance by content-adaptive bit allocation. The important information which is vital for the reconstruction or is referenced by the coding of the subsequent frames will be allocated with smaller QS. We show a visualization example of $qs^{sc}$ in Fig. 4. In this example, the model learns that the moving players are more important and produces smaller QS for these regions. By contrast, the backgrounds have larger QS, which brings considerable bitrate saving (as shown in region indicated by the yellow line). The ablation study in section 4.3 shows $qs^{sc}$ can achieve substantial improvements on multiple datasets.

During the decoding, the corresponding inverse quantization is applied. It is noted that, the $qs^{global}$ needs to be transmitted to the decoder along with the bit-stream. However, this overhead is negligible as only single number is transmitted for each frame or video (flexible setting for user). The modulator $qs^{ch}$ belongs to a part of our neural codec and is learned during the training.

Datasets. We use Vimeo-90k (Xue et al., 2019) for training. The videos are randomly cropped into 256x256 patches. For testing, we use the same test videos as (Sheng et al., 2021). All these test videos are widely used in the evaluation of traditional and neural video codecs, including HEVC Class B, C, D, E, and RGB. In addition, the 1080p videos from UVG (Mercat et al., 2020) and MCL-JCV (Wang et al., 2016) datasets are also tested.

Test conditions. We test 96 frames for each video. The intra period is set to 32 rather than 10 or 12. The reason is that intra period 32 gets closer to the practical usage in the real applications. For example, when compared with intra period 12, intra period 32 has an average of 23.8% (Sheng et al., 2021) bitrate saving for HM. We follow the low delay encoding settings as most existing works (Lu et al., 2019, 2020b; Li et al., 2021; Sheng et al., 2021). The compression ratio is measured by BD-Rate (Bjontegaard, 2001), where negative numbers indicate bitrate saving and positive numbers indicate bitrate increase.

Besides x265 (x26, 2022) (veryslow preset is used), our benchmarks include HM-16.20 (HM, 2022) and VTM-13.2 (VTM, 2022), which represent the best encoder of H.265 and H.266, respectively. For HM and VTM, we follow (Sheng et al., 2021) and use the configuration with the highest compression ratio. We also compare previous SOTA neural video codecs including DVCPro (Lu et al., 2020b), MLVC (Lin et al., 2020), RLVC (Yang et al., 2021a), DCVC (Li et al., 2021), as well as Sheng 2021 (Sheng et al., 2021).

Implementation and training details. The entropy model and quantization for the latent representation of motion vector follow those of ${y}_{t}$. The only difference is the input of entropy model. In the coding of motion vector, the inputs are the corresponding hyper prior and the latent prior, i.e., the quantized latent representation of motion vector from the previous frame. There is no temporal context prior for the coding of motion vector because the generation of temporal context depends on the decoded motion vector. In addition, as we target at single model handling multiple rates, we also train a neural image codec supporting such capability for intra coding.

During the training, the loss function includes the distortion and rate: $Loss=\lambda\cdot D+R$. $D$ refers to the distortion between the input frame and the reconstructed frame. The distortion can be $L_{2}$ loss or MS-SSIM (Wang et al., 2003) for different visual targets. $R$ represents the bits used for encoding $\ddot{y}_{t}$ and the quantized latent representation of motion vector, both associated with the bits used for encoding their corresponding hyper prior.

We adopt the multi-stage training, same as (Sheng et al., 2021). In addition, to train single model supporting rate adjustment, we use different $\lambda$ values in different optimization steps. For simplifying the training process, 4 $\lambda$ values (85, 170, 380, 840) are used. During the training, 4 $qs^{global}$ values will be learned via the RD loss with each corresponding $\lambda$ value. We trained image and video models separately and they have different $qs^{global}$ values. It is noted that, although we only use 4 $\lambda$ values during the training, the model still can achieve smooth rate adjustment by manually adjusting the $qs^{global}$ during the testing, and the corresponding study is presented in Section 4.4.

Table 1 and Table 2 show the BD-Rate (%) comparisons in terms of PSNR and MS-SSIM, respectively. The best traditional codec VTM is used as anchor. From Table 1, we can find that our neural codec achieves an average of 4.7% bitrate saving over VTM on all datasets. By contrast, the second best method Sheng 2021 (Sheng et al., 2021) is far behind than VTM and has 35.7% bitrate increase. To the best of our knowledge, this is the first end-to-end neural video codec that outperforms VTM using the highest compression ratio configuration, which is an important milestone in the development of neural video codec. In particular, our neural codec performs better for 1080p videos (HEVC B, HEVC RGB, UVG, MCL-JCV). Fig. 5 shows RD curves on these datasets. We can find our codec consumes the least bits under the same quality. These results verify the effectiveness of our entropy model on exploiting the correlation among the volumed video data. In addition, the high-resolution video will be more popular in the future and the advantage of our neural codec will be more obvious. When oriented to MS-SSIM, our neural video codec has larger improvement. As shown in Table 2, we achieve an average of 46.4% bitrate saving over VTM on all datasets.

To verify the effectiveness of each component, we conduct a comprehensive ablation study. We analyze the effect of entropy model input, our dual spatial prior, the multi-granularity quantization, as well as the design on frame generator. For simplification, we only use HEVC testsets in ablation study. The comparisons are measured by BD-Rate (%).

The inputs of entropy model. As shown in Fig. 2, our entropy model contains three different kinds of inputs: hyper prior, temporal context prior, and latent prior. Table 3 compares the effectiveness of these inputs. When removing our latent prior, there is 10.1% bitrate increase. If only enabling one prior input, the first and second most important inputs are hyper prior (15.6%) and our latent prior (19.3%). From these comparisons, we can find that the hyper prior is still the most important. But enriching the entropy model input via our latent prior also brings significant bitrate saving.

Different spatial priors. Table 4 compares the effect of different spatial prior modules on exploring the inner correlation in $\ddot{y}_{t}$. The tested spatial context models include our proposed dual spatial prior, checkerboard prior (He et al., 2021), and the parallel-unfriendly auto-regressive prior (Minnen et al., 2018). From this table, we can find that our dual spatial prior could bring 14.1% bitrate saving. And there is 6.2% improvement over the checkerboard prior. This verifies the benefit of enlarging the scope of spatial context and exploiting the cross-channel correlation. In addition, we also find that the auto-regressive prior could further save the bitrate by a large margin (12.0%). However, considering the very slow encoding and decoding speed, we still do not adopt it in our neural codec.

Multi-granularity quantization. As introduced in Section 3.3, the multi-granularity quantization powered by our entropy model not only helps us achieve rate adjustment in single model but also improves the final RD performance by content-adaptive bit allocation. Table 5 shows the BD-rate comparison. From this table, we can find that there is 11.8% loss if disabling the whole multi-granularity quantization. When only removing the spatial-channel-wise quantization step $qs_{t}^{sc}$, there is a significant 8.9% bit increase. It shows that the $qs_{t}^{sc}$ play an important role in multi-granularity quantization and it is necessary to design a dynamic bit allocation which can adapt to various video contents.

Frame generator design. To improve the generation ability of neural codec, our frame generator uses W-Net based structure to enlarge the receptive field. To verify its effectiveness, we compare the frame generator using different networks, as shown in Table 6. We compare W-Net, U-Net, and residual block (Li et al., 2021; Sheng et al., 2021) based structures. From this table, we can find that W-Net has more than 10% bitrate saving than the plain network using residual blocks. The W-Net which ties two U-Nets is also 7.9% better than the single U-Net. Thus, it is possible that larger bitrate saving can be achieved by using more U-Nets. However, considering the complexity, we currently use W-Net.

Table 5 shows the RD improvement of multi-granularity quantization. This sub-section investigates whether our multi-granularity quantization can achieve smooth rate adjustment. For our codec, the global quantization step $qs^{global}$ can be flexibly adjusted during the testing. It serves the similar role of quantization parameter in traditional video codecs. Fig. 6 shows the results of our codec using the learned 4 $qs^{global}$ values which are guided by RD loss using the 4 $\lambda$ values during the training. In addition, we also manually generate 30 $qs^{global}$ values by the interpolation between the maximum and minimum of the learned $qs^{global}$ values. From the results at the 30 rate points, we can find that our single model could achieve smooth rate adjustment without any outlier. By contrast, DCVC (Li et al., 2021) and Sheng 2021 (Sheng et al., 2021) need different models for each rate point.

We compare the model complexity in model size, MACs (multiply–accumulate operations), encoding time, and decoding time with previous SOTA neural video codecs, as shown in Table 7. We use 1080p frame as input to measure these numbers. For the encoding/decoding time, we measure the time on NVIDIA V100 GPU, including the time of writing to and reading from bitstream as in (Sheng et al., 2021). As aforementioned, our model supports rate adjustment in single model. Thus, we significantly reduce the model training and storage burden. Because DCVC (Li et al., 2021) uses the parallel-unfriendly auto-regression prior model, its encoding/decoding time is quite slow. By contrast, (Sheng et al., 2021) and our codec are much faster. Compare with (Sheng et al., 2021), our encoding/decoding time is increased a little. However, the compression ratio is pushed into a new height (from Sheng 2021 surpassing HM to ours surpassing VTM). We believe this is a price worth paying.