A Multi-Scale Spatial-Temporal Network for Wireless Video Transmission

The architecture of the proposed video wireless transmission scheme, noted as VDJSCC, is illustrated in Fig. 1. We consider the wireless transmission of videos over the additive white Gaussian noise (AWGN) channel. Denoting $T$ as the number of frames, $C$ as the number of color channels, $H$ and $W$ as the height and the width of the frame. The video sequences can be represented as $\mathbf{X}_{i}=\left\{\mathbf{x}_{1},\mathbf{x}_{2},...,\mathbf{x}_{T}\right\}$, where $\mathbf{X}_{i}\in\mathbb{R}^{T\times C\times H\times W}$. The proposed VDJSCC model includes a pair of trainable encoder $\mathbf{E_{\phi}}$ and decoder $\mathbf{D_{\theta}}$ and a non-trainable physical channel, where $\phi$ and $\theta$ are the parameters for the encoder and decoder, respectively.

Inspired by the patch embedding in ViT [8], video vision Transformer (ViViT) [10] first used the tubelet embedding to extract non-overlapping, spatial-temporal tubes from the input video sequence. The video sequence $\mathbf{X}_{i}$ is initially split into video tubes. Subsequently, these tubes are then flattened and converted to tokens $\mathbf{z}\in\mathbb{R}^{n_{t}\times{n_{h}\times n_{w}\times{K}}}$ through a trainable linear projection, where $K$ is the hidden dimension. For a video tube of dimension $t\times h\times w$, $n_{t}=\left\lfloor\frac{T}{t}\right\rfloor$, $n_{h}=\left\lfloor\frac{H}{h}\right\rfloor$, $n_{w}=\left\lfloor\frac{W}{w}\right\rfloor$, tokens are extracted from the temporal, height, and width dimensions respectively. The tokens $\mathbf{z}$ is processed by both spatial Transformer (ST) and temporal Transformer (TT) in sequence, resulting in the generation of feature $\mathbf{f}\in\mathbb{R}^{n_{t}\times(n_{h}\times n_{w})\times{K}}$. This approach allows for the effective capture of multi-scale features aimed at enhancing the representation of details. The down-scaled feature $\mathbf{f}_{s}\in\mathbb{R}^{n_{t}\times(\frac{n_{h}}{2}\times\frac{n_{w}}{2})\times{2K}}$ is obtained through a patch merging process, followed by further processing with ST and TT. Ultimately, the outputs from these separate branches are aggregated through averaging, resulting in the generation of feature $\mathbf{s}\in\mathbb{R}^{n_{t}\times(n_{h}\times n_{w})\times{K}}$.

In order to reduce the amount of tokens to be transmitted, we design a dynamic token selection module, which mask tokens with less semantic importance at a certain token keep ratio $\gamma$. We denote the pruned tokens that required to be transmitted as $\mathbf{\tilde{s}}$, and the mask matrix is represented as $\Omega\in\{0,1\}^{M}$, where $M$ is the amount of video tube. Additionally, the power normalization operation enables $\mathbf{\tilde{s}}$ to satisfy the average power constraint before transmitting into the channel. The Encoder can be formulated as:

The wireless channel can be formulated as $\mathbf{\hat{s}}=h\mathbf{\tilde{s}}+\mathbf{n}$, where $h$ denotes the channel gain coefficient. In this formula, $\mathbf{n}\in\mathcal{CN}\sim(0,\sigma^{2}\mathrm{I})$ denotes independent identically distributed (i.i.d) AWGN samples with power $\sigma^{2}$. Assuming that the dimension of the original data is $N=T\times C\times H\times W$, and $R=(\gamma\times K)/N$ is channel bandwidth ratio (CBR).

At the receiver, the VDJSCC Decoder $\mathbf{D_{\theta}}$ expands the received tokens $\mathbf{\hat{s}}$ with zero elements to maintain dimensional consistency. Finally, the TT and ST are utilized to reconstruct the input video sequence $\hat{\mathbf{X}}_{i}\in\mathbb{R}^{T\times C\times H\times W}$, which can be formulated as:

The VDJSCC model is trained in an end-to-end manner. We optimize the model by the mean square error (MSE) between $\mathbf{X}_{i}$ and $\hat{\mathbf{X}}_{i}$. The training loss function is formulated as:

In the video reconstruction task, accounting for temporal correlations is essential to minimize redundant information. The proposed VDJSCC scheme effectively captures multi-scale spatial-temporal representations, encompassing temporal information over long-term frames and spatial information within individual frames. While multi-scale feature extraction demands more computational resources, this paper focuses on two specific feature scales. Additionally, the dynamic token selection block generates an adaptive mask matrix $\Omega$ based on the video content to discard less significant tokens, thereby conserving transmission resources. By employing these approaches, VDJSCC can effectively reduce the information redundancy and save the transmission resources.

As is shown in Fig. 2, the spatial-temporal Transformer encoder consists of two separate Transformer encoders, the ST encoder and the TT encoder. Firstly, spatial position embedding is added to retain spatial positional information. After $L$-layer ST, we rearrange the tokens from $n_{t}\times({n_{h}\times n_{w})\times{K}}$ to $({n_{h}\times n_{w})\times n_{t}\times{K}}$ to make the model pay more attention to temporal connection. For the TT, the temporal position embedding is also added to obtain temporal positional information. Finally, the tokens are processed through $L$-layer TT to obtain the final output.

In the spatial-temporal Transformer encoder, each layer $l$ contains the multi-head self-attention (MSA), layer normalization (LN), and the multilayer perceptron (MLP) blocks. Taking feature $\mathbf{f}$ as an example, this process can be written as follows:

In the spatial-temporal Transformer decoder, we incorporate both spatial and temporal position embedding to capture positional information. Initially, the tokens $\mathbf{\hat{s}}$ and the corresponding downscaled features are processed by TT, after which they are rearranged and fed into ST. Subsequently, the two multi-scale branches are merged and processed by another TT and ST, resulting in the generation of reconstructed tokens $\mathbf{\hat{z}}$.

In order to enhance the feature representation capability, we down sample the video frames to achieve multi-scale features. Inspired by the downsampling method in Swin Transformer [11], we use patch merging to downsample the video frames. The patch merging layer concatenates the features of $2\times 2$ neighboring patches, and applies a linear layer on the $4K$-dimensional concatenated features. This patch merging process reduces the number of tokens by a multiple of $2\times 2=4$ ($2\times$ downsampling of resolution), and the output dimension is set to $2K$. Additionally, the patch reverse merging process reverts the downsampled tokens to their original dimensions, thereby preserving the dimensional consistency between the input and output tokens. At the receiver, we continue to use patch merging to downsample the tokens $\mathbf{\hat{s}}$ and employ patch reverse merging to restore the original dimensions.

In long-term video sequences, the static background in each frame remains largely consistent, leading to substantial redundancy during transmission. Consequently, this background can be regarded as content with lower semantic significance. In contrast, it is essential to prioritize the capture of motion information within the video, as this represents content with higher semantic importance. A token selection network can be developed to dynamically identify and prioritize the transmission of semantically significant content, optimizing transmission resources.

In VDJSCC, we design a dynamic token selection module to reduce the redundancy, as illustrated in Fig. 3. This module dynamically generates a mask matrix $\Omega$ based on the input video content and adjusts the encoding length according to the token keep ratio $\gamma$, enabling content-adaptive variable-length coding. To enhance the capability of the mask decisions, we take into account the influence of both local features and global features. The local features contains the information of a certain token while the global feature contains the context of the whole video sequence. We first utilize an MLP to divide the features $\mathbf{s}$ into the local features $\mathbf{s}^{local}\in\mathbb{R}^{n_{t}\times(n_{h}\times n_{w})\times{K^{\prime}}}$ and global features $\mathbf{s}^{global}\in\mathbb{R}^{n_{t}\times 1\times{K^{\prime}}}$. The process can be formulated as:

where $Agg$ is the function which aggregates all tokens and can be simply implemented as the average pooling. $K^{\prime}$ is the dimension associated with token splitting. In this paper, we set $K^{\prime}$ as $K/2$.

Subsequently, we combine local features $\mathbf{s}^{local}$ and global features $\mathbf{s}^{global}$ to obtain the concatenated features $\mathbf{S}\in\mathbb{R}^{n_{t}\times(n_{h}\times n_{w})\times{K}}$,

Ultimately, we calculate the mask matrix $\Omega$ using ${P_{keep}}$ and $\gamma$ to select more semantically significant tokens from $\mathbf{S}$ for transmission. The process is written as:

where $\odot$ is elementwise production, and $\mathbb{I}$ is indicator function. Specifically, $\Omega$ is clearly determined based on the input video content, with “0” representing the discarded features of lower semantic importance, and “1” indicating the retained features of higher semantic importance.

Moreover, the token keep ratio $\gamma$ can be adjusted to control the number of retained features, thereby enabling dynamic control over the encoding rate. Lower $\gamma$ results in lower coding rate since more tokens are discarded. At the receiver, zero padding is employed to restore the tokens to their original dimensions, resulting in the generation of feature $\mathbf{\hat{s}}$.

The VDJSCC model is trained and evaluated on the UCF101 dataset [12], which consists of 13320 video clips across 101 action classes. All clips have fixed frame rate and resolution of 25fps and $320\times 240$ respectively. In this paper, we use the first train/test list to split the dataset into training dataset and test dataset. For each video frame, we first resize the frame to $300\times 225$, and then randomly crop into $224\times 224$. Additionally, we randomly select $16$ consecutive frames in a video to serve as the input for VDJSCC model.

In all experiments, the image patch size ($h$ and $w$) is set to 16, and the frame patch size ($t$) is set to 2, the channel dimension $K$ is set to 768. The spatial-temporal Transformer encoder, with a depth of $L=5$ and 12 heads, is used to extract high-dimensional features via multi-head self-attention. In the dynamic token selection module, the token keep ratio $\gamma$ is set to 0.8 during the training process. Furthermore, we train the model at different channel SNRs and evaluate each model at the same SNR, with SNR is sampled uniformly from $[1,4,7,10,13]$ dB. For each model, we use the Adam optimizer [13] with a learning rate of $10^{-4}$. The batch size is set to 4, and it takes about 1.5 week to train the model on the single RTX 3090 GPU.

In this paper, we qualify the end-to-end video transmission performance of the proposed VDJSCC model using pixel-wise metric PSNR and the perceptual metric MS-SSIM [14]. MS-SSIM evaluates the model at multiple scales, providing more comprehensive similarity information.

We compare the performance of proposed VDJSCC model with the classic video coded transmission schems and a deep-learning based JSCC model. Specifically, we utilize the standard video codecs H.264 [15] for source coding and practial LDPC [16] for channel coding. FFmpeg [17] is used to achieve H.264 encoding.

In addition, we also compare the model with a state-of-the-art DeepJSCC model named WITT [5]. WITT uses Swin Transformers as the backbone to extract the long-range information, which shows the great performance in the image reconstruction tasks. In our comparison with WITT, we first downsample the video at intervals of 10 frames and save as images, and then input the all images into the WITT model.

Fig. 4 shows the PSNR and MS-SSIM performance versus SNR over the AWGN channel. For the H.264+LDPC, we compare the impact of different modulation methods and code rates. The proposed VDJSCC model significantly outperforms classical video transmission schemes, avoiding the cliff-effect. When the SNR is 1dB, VDJSCC achieves a PSNR of 31.17dB, compared to WITT’s 29.07dB, indicating a notable improvement in reconstruction quality. However, WITT slightly surpasses VDJSCC in PSNR when the SNR exceeds 10dB due to its focus on individual frame processing, which neglects temporal correlations and increases computational resources. In terms of MS-SSIM, VDJSCC consistently outperforms other schemes, particularly at low SNR levels, demonstrating superior perceptual quality with higher visual similarity.

Fig. 5 vividly demonstrates the visual comparison of VDJSCC and the classical video transmission scheme under AWGN channel at SNR=13dB. The proposed VDJSCC preserves clearer details and achieves higher reconstruction quality.

Next, we explore the impact of CBR on the performance of the proposed VDJSCC scheme. In this part, we set different token keep ratio to alter the CBR, adjusting the CBR to 0.031 when the token keep ratio $\gamma$=1.0. Fig. 6 provides the visual results of token selection of the same input video under AWGN channel at SNR=13dB. The model dynamically generates the mask matrix based on the input videos and token keep ratio, effectively masking less semantically important background information while retaining key portrait details. These results demonstrate the model’s ability to achieve high-quality video reconstruction.

Furthermore, in Fig. 7, we present the performance of PSNR and MS-SSIM of different CBR. The models are trained under the AWGN channel at SNR=13dB, SNR=7dB, and SNR=1dB. It is shown that, the PSNR and MS-SSIM increase along with the increase of CBR. Specifically, with a CBR of 0.019, the PSNR exceeds 30dB, indicating good video reconstruction quality even with nearly half of the tokens masked. Notably, token selection allows for dynamic adjustment of coding length based on the token keep ratio, enabling variable-length encoding and efficient bandwidth utilization.

Last but not least, we present the results of ablation study to evaluate the influence of different modules. Table \@slowromancapi@ illustrates the performance of base models, all trained under an AWGN channel at SNR=13dB, with CBR=0.031. For models with the token selection module, the token keep ratio $\gamma$ is set to 0.8. ”VDJSCC w/o multi-scale” refers to the model without patch merging, where video frames are transmitted directly into the 10-layer spatial-temporal Transformer encoder. ”VDJSCC w/o token selection” means the model without dynamic token selection, where all tokens are transmitted. Comparative results clearly show that the proposed multi-scale method can improve the reconstruction quality by 3.52dB. However, the token selection module slightly reduces PSNR by 0.62dB, an acceptable trade-off given its bandwidth-saving benefits.

With respect to complexity, the average inference time of the proposed VDJSCC is approximately 77.42ms using a RTX 3090 GPU. By comparison, the classical H.264 scheme runs at speeds ranging from 1fps to over 100fps, depending on coding settings, with typical encoding rates between 10-20fps. Although the multi-scale method requires more computational resources, it results in an approximately $10\%$ increase in PSNR. The ablation study indicates that the proposed VDJSCC method achieves high-quality video reconstruction while optimizing bandwidth and computational resources.