Tree-Structured Data Clustering-Driven Neural Network for Intra Prediction in Video Coding

In H.265/HEVC, the quadtree-based block partition is adopted for intra prediction and transformation [3]. Each picture is divided into disjunct square coding tree units (CTUs) of the same size, each of which serves as the root of a block partition quadtree structure. Along with the coding tree structure, the CTUs can be subdivided into coding units (CUs). A CU can be further split along the coding tree structure into smaller prediction units (PUs) and transform units (TUs). In VVC, the concepts of CU, PU, and TU are replaced by CU. A more advanced block partition mechanism with embedded multi-type trees (MTT), including quadtree (QT), binary tree (BT), and ternary tree (TT), is adopted [27].

In HEVC, there are 35 directional prediction modes as shown in Fig. 1, including planar mode, DC mode, and 33 angular prediction modes. Planar mode is designed for areas with gradually changing contents. DC mode aims to predict smooth textures, and the predicted sample values are populated with a constant value representing the average of all surrounding reference samples. Angular prediction modes are designed to model structures with directional edges. To better adapt to the diverse contents and different block sizes, the number of angular prediction modes is increased to 65 in VVC by adding a new mode between every two neighboring angular prediction modes of HEVC as shown in Fig. 1. What’s more, aiming at the problem of asymmetric distribution of direction prediction in the non-square block, VVC adopts wide angular intra prediction (WAIP) [28] in which 14 angles using prediction from the shorter side are replaced by more extreme angles using prediction from the longer side, bringing the total number of intra directions supported in VVC to 93.

In addition to the increased directional prediction modes, multiple additional advanced prediction techniques are designed in VVC to make better use of surrounding reference samples, such as matrix-based intra prediction (MIP) [11], position-dependent prediction combination (PDPC) [29], multiple reference lines (MRL) [30], and cross-component linear model (CCLM) [31]. In PDPC, filtering with spatially varying weights is applied to blocks that use Planar mode, DC mode, and certain angular prediction modes. To benefit the prediction for sharp content, MRL is designed to use the non-adjacent reference line which is one, two or four lines away from the current block for prediction rather than only using the nearest line as in traditional directional prediction. CCLM is designed for the prediction of chroma components. The chroma components can be predicted from the collocated reconstructed luma samples by linear models whose parameters are derived from adjacent reconstructed samples.

Recently, neural network-based intra predictions in video coding have been developed rapidly and achieved impressive performance. In [7], a fully connected network-based intra prediction (IPFCN) is proposed to learn an end-to-end mapping from neighboring reconstructed pixels to the target block. The neural networks are fed by 8 reference lines above and on the left side of the target block. Two variants of IPFCN: IPFCN-S and IPFCN-D are designed. IPFCN-S has a single network that is trained utilizing an unclassified training data, while IPFCN-D has two networks that are trained by the data from the non-angular prediction modes (DC and Planar) and the other angular prediction modes separately. IPFCN-S and IPFCN-D could achieve 2.9% and 3.4% BD-rate improvement over HEVC, respectively.

In [10], different numbers of networks (up to 35) are designed for blocks of different sizes. All neural networks for a given block size share all layers but the last one. The network mode signaling is based on a sorted mode list, which is produced by another fully connected network. The bitrate saving in [10] is around 3.0% over HEVC. Compared with IPFCN, the network architecture is simplified. In a later version [11], each network has only one layer without non-linear activation, i.e. each network is reduced to a matrix plus a bias vector. Only one row and one column from the reconstructed picture are used as reference pixels, and only half of the pixels in the target block are generated by the selected network. The remaining half is generated by linear interpolation. This matrix-based intra prediction (MIP) has been adopted by VVC with a bitrate saving of 0.79% under all-intra configuration. In [13], multiple network modes (NM) are trained and each corresponds to a certain range of the directional prediction modes. A maximum 6.3% bitrate reduction is achieved by NM compared with HM-16.9 with 8$\times$8 constraint. In [12], a novel training loss function that reflects properties of the residual quantization and coding stages is designed by applying the $L1$-norm and a sigmoid-function to the prediction residual in DCT domain. It could achieve a bitrate saving of 3.79% compared with VTM 1.0. Different from the fully connected structure applied in above-mentioned works, a progressive spatial recurrent neural network (PS-RNN) which consists of three spatial recurrent units is designed in [14], and the prediction is progressively generated by passing information from preceding contents to blocks to be encoded. PS-RNN achieves on average 2.7% bitrate reduction over HEVC.

Neural network could also be utilized for prediction result refinement. In [8], a convolutional neural network-based intra prediction is proposed. The block size is constrained as 8$\times$8. After HEVC intra prediction, the target block, along with its three nearest reconstruction blocks, forms a 16$\times$16 block as the input of CNN. The output of the neural network is a residual block of the same size as the input. By subtracting the generated residual block from the original HEVC prediction result, the prediction quality can be enhanced. MSCNN proposed in [9] further improves this method by increasing the number of supported block sizes. In addition, the multi-scale feature extraction is proposed to take advantage of feature maps in different scales to improve the performance. MSCNN could achieve an average of 3.4% bitrate saving with all intra configuration over HEVC.

In this paper, we propose a tree-structured data clustering-driven neural network (TreeNet) for intra prediction. TreeNet builds the networks and clusters the training data in a tree-structured manner, rather than the parallel networks as in [7] and [10] - [13].

The network structure in TreeNet is shown in Fig. 2 (a). The network is composed of four parts: the input layer, the output layer, the fully connected layers, and the non-linear activation layers. The network input is the reconstructed reference samples surrounding the target block. To utilize more spatial contextual information for more accurate prediction, the multi-reference-line scheme [32] is employed, thus the network input is a $2k\times\left(w+h\right)+k^{2}$ vector, where $w$ and $h$ are the width and height of the target block and $k$ is the number of reference lines. In TreeNet, $w=h$. If some required reference pixels are not available, the vacancies are filled in by a padding mechanism which is similar to the method used in HEVC [33]. The network output is an $w\times h$ vector, corresponding to the predicted pixels of the target block in raster order. Since overmuch non-square shapes are supported in VVC, we transform non-square shapes to square shapes as follows. Firstly, the shorter reference line is padded to the same length as the longer one. Then the prediction is obtained after cutting off the redundant part of the network output. The prediction of non-square blocks is shown in Figs. 2 (b) and (c).

All linear layers of the network in TreeNet are fully connected, which is convenient for the network split. Except for the last fully connected layer, all fully connected layers are with the same output dimension. In addition, except for the last fully connected layer, each fully connected layer is followed by a non-linear activation layer. If the depth of the neural network is represented by $d$, the output of the $i_{th}$ layer can be calculated as follows:

where $W_{i}$ and $B_{i}$ represent the weights and bias, respectively. $Y_{0}$ and $Y_{d}$ are the input and output of the network. $\sigma$ represents the non-linear activation function. In this paper, Parametric Rectified Linear Unit (PReLU) [34] is selected as the activation function.

In TreeNet, the number of reference lines $k$ is set as 8, and the network depth $d$ is set as 4, as recommended in [7]. Different networks are designed for square blocks of different sizes. The layer dimensions of networks are set to be 512, 1024, 1024, 2048, 2048 for 4$\times$4, 8$\times$8, 16$\times$16, 32$\times$32, and 64$\times$64 blocks,respectively.

In traditional Linde-Buzo-Gray (LBG) data clustering algorithm [35], a perturbation is produced and applied as an addition or subtraction term for splitting one clustering centroid into two. Inspired by [35], we design a network split method for TreeNet. In each network split, every parent network on a leaf node is split into two networks by adding or subtracting Gaussian random noise to its each layer. The variance of the introduced Gaussian random noise is decided by the weights of each layer so that the mapping from the input to the output in the subsequent activation function won’t be changed. The two derived child networks have the same structure as their parent network. Fig. 3 (a) is an example of splitting one parent network into two child networks, where $n_{i}$ represents the Gaussian random noise introduced to the $i_{th}$ layer.

The networks in TreeNet are built in a tree-structured manner. An example of constructing TreeNet with $Depth$=4 is shown in Fig. 3 (b). TreeNet is a perfect binary tree network and the number of networks at level $l$ is $2^{l-1}$. The total number of networks in TreeNet with $Depth$=$L$ is $2^{L}-1$. After each network split, a data clustering-driven training is applied to train the every two derived child networks using the clustered training data of their parent.

As TreeNet goes deeper, the encoder complexity and memory burden will increase exponentially, and the overhead to indicate the networks will also increase. If the network split can’t bring enough prediction improvement, the network split will be terminated.

Inspired by the K-means-like training strategy in [17], a data clustering-driven training is applied to train every two derived child networks using the clustered training data of their parent. Fig. 4 is an example of the data clustering and the network training.

1) Clustering and Training: In each iteration of clustering and training, the training data of a parent network is first clustered according to the recovery qualities of the two derived child networks. In this paper, the squared error between original pixels and the prediction is used to measure the recovery quality. Then the clustered training data are fed into the corresponding derived child networks for training. Given a collection of $M$ training sample pairs, the loss function is formulated as:

where $\mathcal{F}\left(x_{m}|\Theta\right)$ represents the network output and $y_{m}$ represents the ground truth. $\Theta$ is the parameter set. $\gamma$ represents the weight of the $L2$ regularization term, and it is set to be $10^{-4}$ in the experiments. $W$ is the weight of fully connected layers. The optimization algorithm to update all parameters is the stochastic gradient descent ($SGD$) with a momentum term [36]. The momentum of $SGD$ is set as 0.9.

After several iterations of clustering and training, the training loss gradually converges, which indicates that the clustering centers no longer move. Then the data clustering-driven training will stop.

2) Training Details: Firstly, the parameter set $\Theta=\left\{W,B,A\right\}$ is initialized. The weight $W$ is randomly generated from a Gaussian distribution with a mean of 0 and a standard deviation of 1. The bias $B$ is initialized as 0. The scale factor $A$ in PReLU layers is initialized as 0.25. Then a root network of TreeNet is pre-trained for 80 epochs. The base learning rate is set to 0.1 as recommended in [7] and [9]. The network is trained with a fixed learning rate for 20 epochs and then uses the exponentially decaying learning for another 20 epochs until its convergence. The child networks are trained with exponentially decaying learning rates from 0.01 to 0.00001 for 40 epochs. The learning rate settings are listed in Table I.

After TreeNet is well-trained, it can be used for intra prediction. However, the full search of TreeNet is very time consuming. Since the training of derived child networks only use the clustered training data from their parent network, the predictions between a parent network and its two child networks are correlated. Therefore, a fast termination strategy for TreeNet is proposed by exploring the correlation, which is shown in Algorithm 1. The rate-distortion (R-D) cost of a network can be expressed as:

where $J$, $D$, and $R$ represent the R-D cost, the sum of squared error (SSE), and costed bits, respectively. $\lambda_{QP}$ is the Lagrange multiplier decided by a quantization parameter (QP). The R-D cost difference $\Delta J$ between a child network and its parent network can be calculated as:

where the subscripts: $C$ and $P$ represent the child network and the parent network. If $\Delta J_{1}$ and $\Delta J_{2}$ are not smaller than 0, indicating that the parent network is better than both the two derived child networks, the search of TreeNet will be early terminated. Otherwise, the derived child network with less R-D cost will be searched, while the other child network is ignored. The fast termination strategy could help save the encoding running time greatly because at most $2\times L-1$ networks are searched in TreeNet with $Depth$=$L$ compared with $2^{L}-1$ networks for the full search.

To test the performance of TreeNet, we integrate TreeNet into VVC and HEVC reference softwares. In VVC, the best prediction mode is decided at the CU level. Fig. 5 shows an example of integrating TreeNet into the VVC framework. RDO means the rate-distortion optimization [37]. Similar to MRL, TreeNet doesn’t work with Intra Sub-Partitions (ISP) coding mode [38]. A new flag $S$ has to be coded and transmitted in the bitstream to indicate whether TreeNet is used or not. In HEVC, the best prediction mode is decided at the PU level. The integration of TreeNet in HEVC is exactly the same as in VVC.

The codeword of each TreeNet mode is shown in Table. II. At each branch, one bit is needed to represent if the parent network is better than child networks, which is indicated by $p$. If a child network works better, another bit is needed to represent which child network is better, which is indicated by $c$. Therefore, 1, 3, 5, and 6 bits are needed for coding TreeNet modes at level $1$, $2$, $3$, and $4$, respectively. The Context-Based Adaptive Binary Arithmetic Coding (CABAC) [39] is used to encode the codewords and a context model is built for a corresponding bit.

To test the performance of TreeNet, it is implemented into VVC reference software VTM-4.0 and HEVC reference software HM-16.9. Libtorch is used to perform the feed-forward of networks. All experiments comply with the common test conditions specified in [40] and [41] with QP set as 22, 27, 32, 37, and the all-intra main configuration is deployed. 24 video sequences with different resolutions, grouped as $ClassA$, $B$, $C$, $D$, and $E$, are utilized for experiments, and the first frame of each sequence is tested. Bj$\phi$ntegaard distortion rate (BD-rate) [42] is used to evaluate the performance.

In our experiments, TreeNet is trained using the deep learning framework Pytorch [44] 1.8.0 on NVIDIA GeForce RTX 3060 GPU. The models for small blocks (4$\times$4 and 8$\times$8) are trained with a batch size of 512, while the models for large blocks (16$\times$16 to 64$\times$64) are trained with a batch size of 256.

The training data is generated from pictures in the New York city library image set [43]. Firstly, these images are converted to YUV 4:2:0 color format. Then all training images are encoded with all intra configuration. The QP is set to 22, 27, 32, and 37. For all QPs, the encoder generates 5 different bitstreams for different block sizes from 4$\times$4 to 64$\times$64 separately, which will be used to train the models in TreeNet for different block sizes. The reconstructed reference samples of a block are extracted from the bitstreams on the decoder side as the network input. The original pixels of the block are taken as the ground truth.

It should be noted that, before the recursive training, all training data should be pre-processed by zero-centering to get rid of the low-frequency component and make the training easier. In [7] and [9], blocks with complex textures are excluded from the training data since such data would degrade the training efficiency and practical performance, while in TreeNet training, all blocks are kept in the training data.

In the experiments, we first test the performance of applying TreeNet with different depths to the VTM-4.0. Then the performance of the fast termination strategy is provided.

The experimental results are shown in Table III. As observed, TreeNet with $Depth$=4, 3, 2, and 1 can bring an average BD-rate saving of 2.9%, 2.9%, 2.5%, and 1.9% for the luma component, respectively. For all test sequences, TreeNet with $Depth$=4 brings similar gains as TreeNet with $Depth$=3. Therefore, TreeNet with $Depth$=3 is adopted and used in the following experiments. For simplicity, we call TreeNet with $Depth$=3 as TreeNet in the following. The performance of TreeNet with the fast termination strategy (fast TreeNet) is degraded to an average of 2.8% BD-rate saving (up to 8.1%), however the encoding time is much reduced, which is shown in Table IV. Compared with MIP [11], which has simpler network architecture and more network modes than TreeNet, the BD-rate improvement brought by TreeNet is significant, reaching 2.0% improvement on average.

The running times of these tests are summarized in Table IV. All tests for TreeNet are done with an AMD Ryzen7 5800H CPU and an NVIDIA GeForce RTX 3060 GPU. MIP is conducted on an Intel Xeon cluster (E5-2697A v4). As TreeNet goes deeper, the encoding time increases exponentially. The encoding time of TreeNet with $Depth$=3 is around 41 times of the anchor. With the fast termination strategy, the encoding time is reduced to around 25 times. The decoding time of fast TreeNet is roughly 50 times of the anchor. Compared with MIP, both the encoding and decoding time increase greatly due to the complex network structure in TreeNet. How to simplify the network structure in TreeNet is our future work.

In addition to the BD-rate saving, we also provide some visual examples of prediction results of TreeNet and compare them with VVC intra prediction. The examples are taken from the first frame of Kimono encoded with QP=27. The comparison results are shown in Figs. 6, 7, and 8. As observed, TreeNet could produce more accurate predictions than the VVC directional prediction modes. As shown in Fig. 6, the prediction of TreeNet distinguishes the black and white areas more accurately and clearly than the directional prediction of VVC. In Fig. 7, TreeNet generates the predictions that are closer to the original textures compared with the VVC intra prediction. In Fig. 8, there is an obvious directional pattern in the original block. TreeNet could still generate better predictions than the VVC directional prediction mode.

We compare fast TreeNet with Intra Prediction Fully-Connected Networks (IPFCN) [7], Progressive Spatial Recurrent Neural Network (PS-RNN) [14], and multiple neural network model-based enhanced intra prediction (NM) [13]. Among these approaches, NM in [13] only allows 8$\times$8 intra coding, and such a manner is also applied in PS-RNN. In [13], IPFCN-D integrated into HM with 8$\times$8 constraint is also reimplemented. To make a more detailed comparison, fast TreeNet is implemented with and without 8$\times$8 constraint, respectively. The comparison results are presented in Table V.

As shown in Table V, for the luma component, fast TreeNet brings an average BD-rate saving of 4.9% (up to 8.2%). Compared with IPFCN-D, fast TreeNet achieves 1.5% BD-rate reduction on average. Fast TreeNet can also bring an average of 1.7% BD-rate reduction over PS-RNN on Class B, C, D, and E. When only 8$\times$8 intra coding is allowed, the improvement of fast TreeNet is more obvious, reaching an average of 7.5% BD-rate reduction. Compared with IPFCN-D and NM, fast TreeNet achieves 3.5% and 1.0% BD-rate reduction on average, respectively. Fast TreeNet can also bring 1.7% BD-rate reduction on average over PS-RNN on Class B, C, D, and E.

Furthermore, we replace all the directional prediction modes in HEVC with fast TreeNet and test the performance on the first frame of some high-resolution sequences. The results are presented in Table VI. As shown in Table VI, fast TreeNet with $Depth$=4 and 3 brings an average of 2.3% and 2.1% BD-rate reduction respectively, which indicates fast TreeNet replacing all the directional prediction modes still performs well for high-resolution sequences. The result also indicates that more networks should be adopted when TreeNet is used to replace the directional prediction modes than combining with them.

We calculate the usage rate of fast TreeNet at the pixel level when combining with the VVC directional prediction modes, which can be calculated by:

where $\eta$ denotes the usage rate. $W$ and $H$ denote the width and height of each frame. $n_{i}$ represents the number of blocks with a size of $w_{i}\times h_{i}$ coded by fast TreeNet in each frame. Since MTT [27] is enabled in VTM, there are 14 blocks sizes supported. The QPs are set as $\left\{22,27,32,37\right\}$. As shown in Table VII, the usage rates of fast TreeNet in all sequences are remarkable, especially in Class A. The average usage rate of fast TreeNet in all test sequences is 51.3% (up to 71.7%), which demonstrates its efficiency.

Furthermore, we provide an example of the usage rates of fast TreeNet at different levels for different block sizes, which is shown in Fig. 9. We can observe that fast TreeNet has a significant usage rate in predicting blocks of all sizes. For blocks of size 4$\times$4, 4$\times$8, 8$\times$4, and 8$\times$8, the usage rates of fast TreeNet are more than 80%, which demonstrates that fast TreeNet has more advantage in predicting small blocks than the directional prediction modes.

The visual result is shown in Fig. 10. The background parts, such as the tree leaves, are mostly coded by fast TreeNet with large blocks. The foreground parts and background parts with sharp edges, such as the woman’s head, the belt, and the trunks, are mostly coded by fast TreeNet with small blocks.

The size of each Libtorch model in TreeNet is listed in Table VIII. The model size for block of 64$\times$64, 32$\times$32, 16$\times$16, 8$\times$8, and 4$\times$4 are 80.5 MB, 48.5 MB, 11.2 MB, 9.5 MB, and 2.4 MB respectively. There are 7 models for each block size in TreeNet, so overall $7\times(80.5+48.5+11.2+9.5+2.4)=1064.7$ MB is needed to store TreeNet with $Depth$=3 in memory. In addition, another $2\times(64+64+8)\times(64+64+8)\div(1024)=36.125$KB should be allocated in memory to store the 8-line reference samples.

To assess the data clustering-driven training, we first give some investigations on the prediction losses during the network split and iterative training. The average training losses of each network mode at level $2$, $3$, and $4$ in TreeNet for 8$\times$8 blocks are provided in Tables IX. The average training losses of TreeNet modes gradually decrease and converge during the iterative data clustering-driven training. Then we conduct a statistic of the rate of samples remaining in the same cluster after each round of data clustering to see whether the data clustering could finally converge as shown in Table X. In addition, the ratio of the training data to the entire training dataset after the last iteration is also provided. We can observe that less and less training data is clustered to another cluster during the iterative training, which indicates that the data clustering and network training are gradually stabilized.