Diffusion-Aided Joint Source Channel Coding For High Realism Wireless Image Transmission

The input image is denoted as $x\in\mathbb{R}^{C\times H\times W}$, with $C$ indicating the number of color channels, and $H$ and $W$ representing the height and width, respectively. Within the transmitter, a JSCC encoder, $E_{c}$, encodes the image $x$ into a sequence of complex-valued channel input symbols $y\in\mathbb{C}^{K}$. These $K$ symbols are subsequently transmitted through a noisy communication channel while each symbol utilizes the channel once. The rate of the transmission is quantified as the channel bandwidth ratio (CFR), which is referred as $\rho$ and is defined as:

To comply with power constraints, the final encoder layer normalizes the transmitted symbols $y$ to ensure an average power output of $\bar{P}$. Let $\tilde{y}$ represent the complex symbols prior to normalization; the normalization procedure is expressed mathematically as:

The wireless channel in our study is modeled as a conditional distribution, with the received symbols being drawn according to $\hat{y}\sim p(y|\varepsilon)$, where $\varepsilon$ encapsulates the channel parameters. In this work, we consider both the AWGN channel and Rayleigh Fading channel, and the channel transfer function can be represented as:

Here, the noise $n\in\mathbb{C}^{K}$ is assumed to have a circularly symmetric complex Gaussian distribution with $\sigma^{2}$ symbolizing the noise’s power level. $h\in\mathbb{C}$ denotes the channel gain. For AWGN channel, $h$ is set to 1. For the Rayleigh Fading channel, the channel gain $h$ is a random variable and it follows $h\sim\mathcal{CN}(0,H)$. The channel’s SNR, denoted as $\gamma$, is defined by the ratio $\gamma=\bar{P}\mathbb{E}[|h|^{2}]/\sigma^{2}$. In our experiments, we standardize the average power $\bar{P}$ to 1 for consistency. Following the prior works [21, 44], both the transmitter and receiver are assumed to have perfect knowledge of the channel gain $h$ and SNR $\gamma$. Thus, at the receiver side, we first perform equalization to the received symbols $\hat{y}$, which is:

where $h^{*}$ denotes the complex conjugate of $h$.

After that, a JSCC decoder ($D_{c}$) takes the equalized channel symbols $\tilde{y}$ as input and generate the initial reconstruction $\hat{x}\in\mathbb{R}^{C\times H\times W}$. In this paper, we employ a JSCC architecture akin to the one described in Yang et al.[23]. The encoder and decoder could take the CSI information $\gamma$ and $h$ and are trained jointly. The encoding and decoding processes are expressed as:

The aim is to minimize the MSE between the original image $x$ and the decoder’s output $\hat{x}$, with the following loss function deployed for optimization:

A key aspect of our approach is its model-agnosticism, facilitating seamless integration with more advanced deep JSCC frameworks. While we demonstrate the effectiveness of DiffJSCC using one JSCC structure, the design is inherently adaptable to other more recent architectures such as SwinJSCC[19] or NTSCC[24].

Rather than directly using the received symbols $\hat{y}$ as conditioning inputs for the conditional latent diffusion model, the initial JSCC reconstruction acts as an intermediate step and exhibits three primary benefits: First, it enables separate training of the JSCC encoder and the conditional denoiser, which allows more stable training compared to a joint training approach. Second, the initial reconstruction $\hat{x}$ harnesses the pre-trained LDM encoder to capture salient spatial features, thus avoiding the need to develop an encoder from scratch as originally proposed by the ControlNet framework[41]. Last, the initial reconstruction $\hat{x}$ facilitates the extraction of supplementary textual descriptors using pre-trained image captioning models such as BLIPv2[45] and the CLIP encoder[46], enriching the input conditions for the diffusion model.

Leveraging the initial reconstruction $\hat{x}$, our framework extracts multimodal features with the help of pre-trained neural networks, eliminating the need for additional training specific to our approach. The spatial features $f_{v}$ with a size of $4\times H/8\times W/8$ are extracted by the image encoder $E_{ldm}$ from the pre-trained Stable Diffusion model:

In parallel, given that the Stable Diffusion model is inherently a text-to-image generative model, it is natural to leverage textual information as a supplementary condition. Noticing the presence of an initial image reconstruction, we propose to utilize the pre-trained image captioning model BLIPv2 [45] denoted as $E_{blip}$ to generate textual descriptions for $\hat{x}$. Following the initial captioning step, these descriptions are further encoded by the pre-trained CLIP encoder $E_{clip}$[46] to obtain textual features $f_{t}\in\mathbb{R}^{77\times 768}$:

An alternative approach to leveraging textual data is by deploying image captioning models at the transmitter end and sending the generated captions along with the image. Nevertheless, since image captioning models often involve significant computational requirements, this method could substantially increase the transmitter’s processing load, making it impractical for mobile systems. It also incurs the overhead of sending text captions through the noisy channel. Within our method, despite the potential quality degradation in captions due to the distortion of initial image reconstructions, the inclusion of $f_{t}$ generated from $\hat{x}$ has been shown to be beneficial as it encapsulates content-related keywords.

The CSI information, described by the channel SNR $\gamma$ and the channel gain $h$, plays a crucial role alongside the multimodal features $f_{v}$ and $f_{t}$ as it provides insight into the extent of distortion that initial image reconstruction experiences. After integrating the CSI, the set of conditions $\bar{c}$ utilized by our conditional diffusion model is formulated as:

These diverse conditions work together to strengthen the conditional diffusion process, improving the model’s capability to reconstruct images with increased authenticity and accuracy, faithfully reconstructing the originally transmitted content.

The proposed DiffJSCC framework takes the set of conditions $\bar{c}$ as the input of the conditional latent diffusion process as depicted in Fig. 2. The diffusion process starts with a random input $z_{T}$ sampled from $\mathcal{N}(0,\mathbf{I})$ and iterates the denoising process with the multi-modal condition set $\bar{c}$ using the proposed denoiser (Fig. 3) until the denoised output $z_{0}$ is obtained. The final output image $\hat{x}_{diff}$ is generated using an LDM decoder on $z_{0}$: $\hat{x}_{diff}=D_{ldm}(z_{0})$.

For evaluations, we adopt standard distortion metrics such as PSNR and MS-SSIM[48] to measure image fidelity concerning pixel intensity and structural details, respectively. To evaluate the perceptual quality of reconstructed images, we employ more recently proposed metrics such as Learned Perceptual Image Patch Similarity (LPIPS)[49] and Fréchet Inception Distance (FID) scores[50]. These metrics provide crucial insights into visual similarity and how well the distribution of generated images aligns with that of real images. Given the Kodak dataset contains only 24 images, we transmit each image five times to measure our performance metrics reliably. For Kodak images, the FID score is calculated over half-overlapping patches of $256\times 256$ pixels. Our evaluation also includes the performance of the semantic segmentation task on the generated and original images measured by the mean Intersection over Union (mIoU) metric.

Figure 4 presents the performance comparison of our DiffJSCC framework against traditional separate source-channel coding schemes as well as the foundational DeepJSCC approach which is also used for DiffJSCC’s initial image reconstruction. The evaluation is across various transmission rates while fixing the SNR to 1dB or 10dB. The benchmark separate coding scheme utilizes the BPG codec for image compression paired with LDPC codes conforming to the IEEE 802.11n WiFi standard with a specific block length. Under the 1dB SNR scenario, the LDPC code adopts a $2/3$ coding rate with BPSK modulation, while the coding rate is $1/2$ with 16-QAM modulation for the 10dB SNR case. To adjust the transmission rate, we manipulate the BPG’s quantization parameter with the range from 51 to 0. The ‘BPG-Capacity’ is the upper bound of the separate source-channel coding scheme (i.e., LDPC is replaced by ideal channel coding) where channel coding is error-free when the rate is under the channel capacity calculated by $C=\text{log}_{2}(1+\gamma)$.

DeepJSCC and DiffJSCC methods support four distinct rates: $\rho=1/384,1/192,1/96,$ and $1/48$. Notably, the basic DeepJSCC method surpasses traditional separate coding methods across multiple metrics, validating previous studies. Though the diffusion-based refinements from our DiffJSCC may induce degradation in traditional distortion metrics such as PSNR and MS-SSIM similar to [30, 39], DiffJSCC exhibits remarkable improvements in visual quality assessed by LPIPS and FID metrics. These metrics capture the capability of DiffJSCC to produce better photorealistic reconstructions even at compression levels unachievable by the other schemes (including theoretical ‘BPG-Capacity’).

Notice, at a 1dB SNR, DiffJSCC attains a low FID score of $117.9$ at the transmission rate of merely $\rho=1/384=0.0026$, outperforming the ‘BPG+Capacity’ baseline’s FID score of $165.4$ at a substantially higher rate of $\rho=0.0377$. This translates to DiffJSCC achieving superior image quality by a margin of $28.7\%$ while utilizing $14\times$ less symbols. In scenarios where the SNR is increased to 10dB, DiffJSCC continues to hold a significant advantage over the other methods although the performance gap in terms of LPIPS and FID is reduced. When transmitting at a rate of $\rho=0.0026$, the FID score of DiffJSCC demonstrates an impressive $64.3\%$ enhancement over DeepJSCC and a $70.7\%$ over ‘BPG+Capacity’. This confirms the effectiveness of the DiffJSCC framework in reconstructing highly realistic images under stringent rate and channel constraints.

In Figure 5 and 6, we assess DiffJSCC in comparison to other baseline approaches across various SNR conditions for AWGN channel and Rayleigh Fading channel respectively. The DeepJSCC-$adv$ model from Yang et al.[32] is included for comparison, which uses the hyperparameter $\alpha$ to adjust the weight of adversarial loss. Various $\alpha$ parameters were tested to explore the trade-off between distortion and visual quality. Similarly, we experimented with different $\lambda$ values in DiffJSCC with the same goal. These evaluations were conducted at transmission rates of $\rho=1/384$ and $\rho=1/96$, with SNR values set at $\gamma=1,4,7,10,13$ dB. All benchmarks including ours are based on the same deep JSCC network architecture. The results in Figure 5 demonstrate that although the DeepJSCC-$adv$ baseline successfully lowers LPIPS, it struggles to effectively minimize FID scores regardless of the $\alpha$ value used. The FID score remained above $190$ at $\rho=1/384$ and exceeded $110$ at $\rho=1/96$. Conversely, DiffJSCC significantly improves both LPIPS and FID scores over the entire range of SNR values, while maintaining similar PSNR and MS-SSIM distortion metrics. At $\lambda=0$ and $\rho=1/384$, DiffJSCC shows a notable average improvement of $41.4\%$ in LPIPS and $56.7\%$ in FID scores compared to the standard DeepJSCC. When compared with DeepJSCC-$adv$ using $\alpha=5\times 10^{-3}$, it achieves $21.1\%$ better LPIPS and $52.2\%$ enhanced FID metrics. At $\rho=1/96$, DiffJSCC consistently demonstrates a significant gain in performance, improving average LPIPS and FID scores by $43\%$ and $53.2\%$ over DeepJSCC, respectively, and $18.9\%$ and $48.4\%$ compared to DeepJSCC-$adv$. Increasing the scale factor $\lambda$ from 0 to 100 gradually aligns the generated image with the initial JSCC reconstruction, enhancing PSNR and MS-SSIM but reducing perceptual quality. However, this trade-off still leads to better visual perception metrics (LPIPS and FID) than the baselines. This significant improvement in the perceived visual quality highlights the effectiveness of our fine-tuning control module along with the Stable Diffusion model. A similar trend can be observed under the Rayleigh Fading channel as shown in Figure 6.

Figure 7 provides visual examples of the superior perceptual quality achieved by DiffJSCC at a challenging SNR of 1dB under the AWGN channel. At a transmission rate of $\rho=1/384$, the reconstructions produced by the conventional DeepJSCC framework exhibit significant distortion and have difficulty recovering important image details. Conversely, DiffJSCC demonstrates a remarkable capacity to restore crucial elements of the image, like the structures of the lighthouse and house, offering significantly enhanced perceptual quality despite slight variances from the original. With the transmission rate raised to $1/96$, DiffJSCC’s capabilities become more evident, resulting in reconstructions that accurately reflect the original images. Furthermore, we demonstrate the validity of captions generated by the BLIPv2 model based on the initial JSCC reconstructions. These captions prove to be robust against variations in reconstruction quality, consistently capturing the important content keywords of each image (highlighted in red). This confirms the ability of DiffJSCC to convey essential semantic information even at reduced transmission rates.

In Figure 8, we assess the DiffJSCC framework using the ADE20K dataset at transmission rates $\rho=1/384$ and $\rho=1/96$ under the AWGN channel, examining FID scores and mIoU for semantic segmentation performance. For mIoU evaluations, semantic masks are inferred from the reconstructed images utilizing a pre-trained ViT-adaptor [51]. The data indicates that while the DeepJSCC-$adv$ approach can effectively reduce FID scores, it comes at the cost of decreased mIoU, thus limiting its practicality for subsequent image processing tasks. In contrast, DiffJSCC notably improves both FID and mIoU. When $\rho=1/384$, it demonstrates an average improvement (reduction) in FID by $69.2\%$ and a $52.5\%$ increase in mIoU over DeepJSCC, alongside the notable improvement of $75.6\%$ in FID and $161\%$ in mIoU compared to DeepJSCC-$adv$. For the higher transmission rate $\rho=1/96$, DiffJSCC consistently shows superior performance, with a mean enhancement of $42.9\%$ in FID and $8.4\%$ in mIoU over DeepJSCC, along with gains of $33.3\%$ in FID and $19.1\%$ in mIoU in comparison to DeepJSCC-$adv$. With increasing SNR values, the mIoU obtained by DiffJSCC begins to align closely with that of the original undistorted images. Figure 9 presents a visual example from the ADE20K dataset. The example shows that, with DeepJSCC, elements such as grass and buildings are somewhat recognizable, yet roads and lamps appear unclear. The reconstructions obtained from DeepJSCC-$adv$ significantly lack authenticity, and the ViT-adaptor fails to generate accurate segmentation masks. On the other hand, images produced by DiffJSCC demonstrate higher visual quality to produce more accurate segmentation labels in most regions, thus highlighting the framework’s applicability for practical JSCC scenarios.

Figure 10 presents a comparison of our DiffJSCC framework with DeepJSCC and other generative JSCC methods, including InverseJSCC [35] and the approach presented by Yilmaz et al. [37] on CelebAHQ dataset. For DiffJSCC, we showcase results with varying $\lambda$ values of 0, 10, and 50. The outcomes of InverseJSCC and Yilmaz et al. are taken from their original papers. We only evaluate the PSNR and LPIPS metrics since MS-SSIM and FID evaluations were not conducted in their works. Overall, DiffJSCC shows superior performance, surpassing Inverse-JSCC regarding both PSNR and LPIPS metrics. In comparison with the Yilmaz et al. method, DiffJSCC has a slightly lower PSNR but achieves superior LPIPS scores. Unlike the Kodak image results, where $\lambda=0$ provided the optimal LPIPS, this configuration introduces artifacts because of the significant domain gap between the Stable Diffusion training dataset and the CelebAHQ dataset. However, by utilizing direct guidance with the initial JSCC reconstructions (with $\lambda=50$), these artifacts are effectively mitigated, resulting in enhanced LPIPS performance. Quantitatively, at $\lambda=50$, DiffJSCC accomplishes an average enhancement of $15.9\%$ and $28.6\%$ over InverseJSCC for $\rho=1/768$ and $\rho=1/192$, respectively, and showcases an improvement of $10.3\%$ and $9.4\%$ for $\rho=1/768$ and $\rho=1/192$ against the Yilmaz et al. method. Figure 11 visualizes examples from the CelebA dataset, highlighting DiffJSCC’s ability to reconstruct high-quality images in extremely challenging conditions such as a very low bit rate ($\rho=1/768$) and a low SNR value ($\text{SNR}=-5\text{dB}$). As the SNR increases, the quality of the reconstructions improves, making them more closely resemble the original images.