DM-MIMO: Diffusion Models for Robust Semantic Communications over MIMO Channels

DMs achieve data generation through a progressive denoising procedure. With the forward diffusion process gradually corrupting data by adding noise, the reverse sampling process of DMs performs the opposite procedure, producing samples sharing the same distribution as the source data. Specifically, for a given data $\mathbf{X}_{0}$ with distribution $q(\mathbf{X}_{0})$, the forward diffusion process is derived as a Markov chain, generating a sequence of random variables $\mathbf{X}_{1},\mathbf{X}_{2},\cdots,\mathbf{X}_{T}$ modeled by

where $T$ is the number of diffusion steps, and $q(\mathbf{X}_{t}|\mathbf{X}_{t-1})$ denotes the conditional distribution of step $t$, formulated as

Here, $\alpha_{t}\in(0,1)$ is the noise schedule chosen ahead of model training. The reverse sampling process starts by sampling a pure Gaussian noise $\mathbf{X}_{T}$ consists of independent and identically distributed (i.i.d.) elements with distribution $\mathcal{N}(0,1)$, and then gradually generates the target data by

where ${\boldsymbol{\theta}}$ are the learnable parameters of DMs, and

Here $\mu_{\boldsymbol{\theta}}$ and $\Sigma_{\boldsymbol{\theta}}$ denotes the mean and variance of the conditional distribution $p_{\boldsymbol{\theta}}(\mathbf{X}_{t-1}|\mathbf{X}_{t})$ predicted by DMs. As such, DMs are capable of noise elimination. We can first simulate different noise addition process by selecting different diffusion steps, then adopt corresponding sampling steps to recover the source samples. Leveraging the denoising capability of DMs, we propose a plug-in denoising module DM-MIMO, detailed in Section III.

Consider a semantic communication system performing image transmission over MIMO channels. For simplicity, we assume both the transmitter and receiver are equipped with $M$ antennas. The input image signal is represented by $\mathbf{S}\in\mathbb{R}^{h\times w\times 3}$, where $h$ and $w$ denote the height and width of the input image respectively, while 3 is the quantity of color channels. The image is transmitted over $k$ channel uses, and the channel bandwidth ratio (CBR) is defined as $\text{CBR}=k/n$, where $n=3hw$.

As illustrated in Fig. 1, at the transmitter, the input image first adopts a JSCC encoding function $f_{{\boldsymbol{\phi}}}$ parameterized by ${\boldsymbol{\phi}}$ to output the encoded signal $\mathbf{Z}=[\mathbf{z}_{1},\cdots,\mathbf{z}_{M}]^{T}\in\mathbb{C}^{M\times k}$, expressed as $\mathbf{Z}=f_{{\boldsymbol{\phi}}}(\mathbf{S})$. The encoded signal $\mathbf{Z}$ is then mapped into the channel input signal $\mathbf{W}\in\mathbb{C}^{M\times k}$ via MIMO precoding. The power constraint $P_{s}$ is given by $\frac{1}{k}\left\|\mathbf{W}\right\|^{2}_{F}\leq P_{s}$.

We consider Rayleigh block fading MIMO channels. Let $\mathbf{H}\in\mathbb{C}^{M\times M}$ be the channel matrix which remains unchanged over $k$ channel uses. Each entry of $\mathbf{H}$ is i.i.d. random variables following the complex normal distribution $\mathcal{CN}(0,1)$. Then, the output signal of the MIMO channel can be written as

where $\mathbf{N}\in\mathbb{C}^{M\times k}$ is the additive noise term that consists of i.i.d. elements with distribution $\mathcal{CN}(0,\sigma^{2})$, in which $\sigma^{2}$ denotes the channel noise power.

As CSI is accessible to both the transmitter and receiver, we adopt SVD-based precoding and equalization to leverage spatial multiplexing and mitigate inter-channel interference. The channel matrix $\mathbf{H}$ can be decomposed as

where $\mathbf{U}=[\mathbf{u}_{1},\cdots,\mathbf{u}_{M}]\in\mathbb{C}^{M\times M}$ and $\mathbf{V}\in\mathbb{C}^{M\times M}$ are unitary matrices, and $\boldsymbol{\Sigma}={\rm diag}[\lambda_{1},\cdots,\lambda_{M}]\in\mathbb{C}^{M\times M}$ is a diagonal matrix with singular values $\lambda_{1}\geq\cdots\geq\lambda_{M}$. Applying $\mathbf{V}$ as the precoder at the transmitter, we have $\mathbf{W}=\mathbf{V}\mathbf{Z}$, thus (5) becomes

With the aid of SVD-based equalizer, the equalized signal $\mathbf{Y}^{\prime}=\boldsymbol{\Sigma}^{\dagger}\mathbf{U}^{H}\mathbf{Y}$ at the receiver can be expressed as

where $\mathbf{N}^{\prime}=\boldsymbol{\Sigma}^{\dagger}\mathbf{U}^{H}\mathbf{N}=[\mathbf{n}^{\prime}_{1},\cdots,\mathbf{n}^{\prime}_{M}]^{T}$, and $\boldsymbol{\Sigma}^{\dagger}={\rm diag}[\frac{1}{\lambda_{1}},\cdots,\frac{1}{\lambda_{M}}]$. Consequently, after employing SVD-based precoding and equalization, the MIMO channels are decomposed into $M$ parallel sub-channels. For sub-channel $i$, the effective noise power is $\sigma_{i}^{2}=\frac{\sigma^{2}}{\lambda_{i}^{2}}$. In order to remove noise from the equalized signal $\mathbf{Y}^{\prime}$, we introduce a DM-MIMO module $g_{{\boldsymbol{\theta}}}(\cdot)$ with parameters ${\boldsymbol{\theta}}$, represented as $\hat{\mathbf{Z}}=g_{{\boldsymbol{\theta}}}(\mathbf{Y}^{\prime})$, detailed in Section III. Finally the JSCC decoding function with parameter ${\boldsymbol{\varphi}}$ takes the denoised signal $\hat{\mathbf{Z}}\in\mathbb{C}^{M\times k}$ as input, reconstructing the input image $\hat{\mathbf{S}}=f_{{\boldsymbol{\varphi}}}(\hat{\mathbf{Z}})$.

Considering block fading, disparities in channel conditions occur over distinct sub-channels. Moreover, under varying channel conditions, the JSCC encoder and decoder jointly learn semantic feature allocation across sub-channels along with semantic feature extraction. The fluctuating channel conditions and the different significance of semantics across parallel sub-channels are crucial factors to be considered in the design of the denoising module.

As $\mathbf{U}$ is a unitary matrix, the noise of different sub-channels in (8) follows an independent Gaussian distribution. Moreover, the effective noise power $\sigma_{i}^{2}$ in different sub-channels varies significantly due to its dependence on the channel singular value $\lambda_{i}$. To illustrate the difference in effective noise power over sub-channels, we perform a Monte Carlo experiment with $1\times 10^{7}$ samples of $2\times 2$ MIMO Rayleigh fading channels. As shown in Fig. 2, the gap between the expectation of $\lambda_{i}^{2}$ in the two sub-channels reaches $10.37$ dB, presenting a challenge in joint denoising design.

With significant gap in effective noise power between different sub-channels, the naive method of applying one sampling step for all sub-channels fails to eliminate noise effectively. Therefore, an individual effective sampling step need to be employed for each sub-channel according to its effective noise power. Another naive method to remove noise is to employ a separate DM module for each sub-channel. This, however, fails to capture the joint distribution of the encoded signals over different sub-channels. Hence, based on the effective sampling steps, we propose a joint sampling algorithm, which is detailed in the following sub-sections.

We denote the output of each sub-channel as $\mathbf{y}^{\prime}_{i}$, which is the $i$-th row of $\mathbf{Y}^{\prime}$. To adapt to different channel matrices $\mathbf{H}$, for sub-channel $i$, we employ normalization with factor ${\sigma}_{i}$ as

We define $\mathbf{x}_{i,0}=\mathbf{z}_{i}$, $\mathbf{X}_{0}=[\mathbf{x}_{1,0},\cdots,\mathbf{x}_{M,0}]^{T}$ and design a forward diffusion process inspired by [7]

where $\bar{\alpha}_{t}=\prod_{l=1}^{t}\alpha_{l}$, $t\in\{1,2,\cdots,T\}$ and ${\boldsymbol{\epsilon}}=[{\boldsymbol{\epsilon}}_{1},\cdots,{\boldsymbol{\epsilon}}_{M}]^{T}$ denotes the additive noise term consisting of i.i.d. elements following $\mathcal{CN}(0,1)$. To handle signals over sub-channels with different effective noise power, DM-MIMO generates signals sharing the same distribution as $\bar{\mathbf{y}}^{\prime}_{i}$. Specifically, DM-MIMO utilizes the forward diffusion process with additive noise power matching the effective noise power, enabling the corresponding sampling process for denoising. Therefore, the effective sampling step $m_{i}$ is given by

By employing tailored effective sampling steps, DM-MIMO is able to handle different signals over different sub-channels in one sampling process, utilizing the inter-sub-channel signal correlations. Moreover, considering the channel noise power varying in a wide range, additive channel noise $\mathbf{N}^{\prime}$ introduces dramatic power fluctuations to the equalized signal $\mathbf{Y}^{\prime}$. DM-MIMO reduces the power fluctuations by learning the distribution of the encoded signal $\mathbf{Z}$ and eliminating noise through the reverse sampling process, thus enhancing the robustness of the receiver.

In the training process, as $\mathbf{x}_{i,m_{i}}$ and $\mathbf{y}^{\prime}_{i}$ share the same distribution, DM-MIMO can be trained on the forward diffusion process of $\mathbf{Z}$ instead of $\mathbf{Y}^{\prime}$. Aiming to recover ${\mathbf{Z}}$ through learning its distribution, the loss function $L$ is defined by the variational bound on the negative log likelihood function of $\mathbf{X}_{0}$, as given in (12). Analyzing the components of loss function, $L_{\mathbf{Y}^{\prime}}$ and $L_{t}$ can be ignored during training as they are not related to ${{\boldsymbol{\theta}}}$. Therefore, we focus on $L_{0}$ and $L_{t-1}$, revealing the training goal of approximating the distribution of $q(\mathbf{X}_{t-1}|\mathbf{X}_{t},\mathbf{X}_{0},\boldsymbol{\Sigma})$ with $p_{{\boldsymbol{\theta}}}(\mathbf{X}_{t-1}|\mathbf{X}_{t},\boldsymbol{\Sigma})$. After re-parameterization and re-weighting, the loss function $L_{t-1}$ can be simplified as

Finally, to optimize (15) for all $t\in\{1,\cdots,T\}$, the loss function of DM-MIMO can be expressed as

The training process of DM-MIMO is detailed in Algorithm 1. With different effective sampling steps chosen according to different effective noise power $\sigma_{i}^{2}$, DM-MIMO demonstrates the ability of denoising under diverse channel conditions while utilizing fixed parameters.

Inspired by [11], we design a joint sampling algorithm for DM-MIMO. Specifically, in the $t$-th sampling step of DM-MIMO, in order to guide the denoising process of equalized signals with high effective noise power, we add noise to the equalized signals with low effective noise power and send them to the reverse sampling step. In a word, we employ either noise addition or the reverse sampling process for each sub-channel based on the value of its effective sampling step $m_{i}$.

We begin the sampling process from sampling step $t=\max\{m_{1},\cdots,m_{M}\}$. As illustrated in Fig. 3, if $m_{i}\leq(t-1)$, to keep the correct properties of the distribution of $\mathbf{X}_{t-1}$, we derive $\mathbf{x}_{i,t-1}$ by adding noise to $\bar{\mathbf{y}}^{\prime}_{i}$, given by

On the other hand, if $m_{i}>(t-1)$, assume the knowledge of $\mathbf{x}_{i,t}$ and $\mathbf{x}_{i,0}$ to be available, we can derive the sampling process of $\mathbf{x}_{i,t-1}$ as

where $\mathbf{x}_{i,0}$ can be acquired by re-writing (10) as

In the reverse sampling process, with only $\mathbf{x}_{i,t}$ and ${\boldsymbol{\epsilon}}_{{\boldsymbol{\theta}},i}(\mathbf{X}_{t},\boldsymbol{\Sigma},t)$ known to the receiver, ${\boldsymbol{\epsilon}}_{i}$ is replaced with ${\boldsymbol{\epsilon}}_{{\boldsymbol{\theta}},i}(\mathbf{X}_{t},\boldsymbol{\Sigma},t)$. Therefore, the sampling process can be expressed as

As for the last sampling step $t=1$, DM-MIMO only predicts $\mathbf{x}_{i,0}$ with $\mathbf{x}_{i,1}$, given by

As outlined in Algorithm 2, for sub-channel $i$, by comparing $m_{i}$ with the sampling step $t$, the joint sampling algorithm either adds noise or performs the reverse sampling process, thus leveraging all the semantic information while addressing different effective noise power over different sub-channels.

Since DM-MIMO is trained to learn the distribution of the encoded signal $\mathbf{Z}$, a three-stage training algorithm is proposed. In the first stage, the JSCC encoder and decoder are jointly trained to minimize the reconstruction distortion. With MSE adopted as performance metric, the loss function of the first training stage can be written as

With the well-trained and fixed parameters of the semantic encoder, DM-MIMO is trained in the second stage using Algorithm 1. Specifically, DM-MIMO adopts the whole encoded signal $\mathbf{Z}$ as input, leveraging all the semantic information received. Benefiting from the different effective sampling steps and the joint sampling algorithm, DM-MIMO is capable of noise elimination over different sub-channels.

In the third stage, the JSCC decoder is retrained to adapt to the JSCC encoder and the DM-MIMO. Though only the parameters of the JSCC decoder are trained, the entire system operates under real MIMO channels. The loss function is

To investigate the robustness of semantic communication systems, we consider a universal transmission scenario with random transmit power. Therefore, to enhance robustness of the communication system, the channel noise power is randomly chosen from a regime during both the first and the third training stage.

We consider DIV2K dataset in the experiments. This dataset contains $1000$ diverse 2K images from a wide range of real-world scenes, $800$ of which are used for training, $100$ for validating and the rest $100$ for testing. We randomly crop the images into $256\times 256$ patches in the training process. We consider block-fading MIMO channels with antenna number $M=2$ and channel SNRs ranging from $0$ dB to $20$ dB, where the channel SNR is defined as

We adopt the latest Swin-Transformer based JSCC [12] for the JSCC in our DM-MIMO-JSCC. For simplicity, channel adaptation is not considered in the JSCC encoder and decoder. The proposed DM-MIMO is established on U-Net architecture [13]. We choose hyper-parameter $T=1000$ and employ a noise schedule $\alpha_{t}$ that linearly decreases from $\alpha_{1}=0.9999$ to $\alpha_{T}=0.98$. DM-MIMO is trained with an Adam optimizer for 800 epochs, which adopts a cosine warm-up learning rate schedule with initial learning rate of $1\times 10^{-4}$. Besides, the JSCC is trained for $800$ epochs during the first training stage, with a learning rate of $1\times 10^{-4}$. The retraining epochs in the third training stage is set to $20$. We implement DM-MIMO-JSCC on one NVIDIA A40 GPU using Pytorch. For fair comparison, the same Swin-Transformer-based JSCC but without DM-MIMO is considered as a benchmark. It is trained with the same setup as the first training stage of DM-MIMO-JSCC.

We first evaluate the effectiveness of DM-MIMO by comparing the MSE between the denoised signal and the encoded signal in DM-MIMO-JSCC and the MSE between the equalized signal and the encoded signal in universal JSCC without DM-MIMO. ${\rm MSE}_{avg}$ denotes the average MSE across all sub-channels, while ${\rm MSE}_{i}$ denotes the average MSE of the $i$-th sub-channel. As shown in Fig. 4, with DM-MIMO adopted, both ${\rm MSE}_{1}$ and ${\rm MSE}_{2}$ decrease over the SNR regime of $[0,20]$ dB. The lower the channel SNR, the higher the MSE gain achieved by DM-MIMO on both sub-channels. The MSE gain of ${\rm MSE}_{1}$ and ${\rm MSE}_{2}$ achieve over $0.233$ dB and $2.482$ dB respectively in the SNR regime of $[0,20]$ dB. As such, our proposed DM-MIMO, with effective sampling step adaptation and joint sampling algorithm, demonstrates effectiveness in noise elimination and signal quality enhancement.

Fig. 5 shows the PSNR performance versus the channel SNR. Our DM-MIMO-JSCC outperforms universal JSCC without DM-MIMO in SNR regime from $0$ dB to $20$ dB. The higher the channel SNR, the higher the PSNR gain achieved by DM-MIMO-JSCC. Specifically, DM-MIMO-JSCC achieves a PSNR gain of $0.488$ dB at SNR = $20$ dB. For comparison, the JSCC scheme where training SNR matches testing SNR is also ploted in Fig. 5. It can be seen that the DM-MIMO-JSCC achieves comparable performance in SNR regime of $[0,15]$ dB. Moreover, we compare the reconstructed samples of different methods under channel SNR of $20$ dB. As shown in Fig. 6, the samples reconstructed by DM-MIMO-JSCC show better visual quality compared with those reconstructed by universal JSCC without DM-MIMO, as the first one shows a sharper edge and the second one shows a clearer fur detail. In a word, by eliminating noise of the equalized signal, DM-MIMO enhances the robustness of the semantic communication system and achieves better performance in image recovery over various channel conditions.

We analyze the computational complexity of the proposed DM-MIMO in terms of multiply accumulate operations (MACs). Table I shows the number of MACs for one step of sampling in DM-MIMO at different CBRs. It can be seen that the amount of MACs of DM-MIMO increases proportionally with the square of CBR. This is expected as the number of channels in U-Net adopted by DM-MIMO directly depends on the length of the received signal. Meanwhile, the number of MACs in JSCC only increases slightly when CBR increases. As such, our DM-MIMO is more suitable at lower CBR region in terms of computational complexity.