Rate-Adaptive Generative Semantic Communication Using Conditional Diffusion Models
Pujing Yang Guangyi Zhang Yunlong Cai

	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	
	Rate-Adaptive Generative Semantic Communication Using Conditional Diffusion Models
    Pujing Yang, Student Member, IEEE, Guangyi Zhang, Student Member, IEEE, and Yunlong Cai, Senior 
Member, IEEE
		P. Yang, G. Zhang, and Y. Cai are with the College of Information Science and Electronic Engineering, Zhejiang University, Hangzhou 310027, China (e-mail: yangpujing@zju.edu.cn; zhangguangyi@zju.edu.cn;  ylcai@zju.edu.cn). 
	 
    September 9, 2024
======================================================================================================================================================================================================================================================================================================================================================



	§ ABSTRACT 
		Recent advances in deep learning-based joint source-channel coding (DJSCC) have shown promise for end-to-end semantic image transmission.
		However, most existing schemes primarily focus on optimizing pixel-wise metrics,  which often fail to align with human perception, leading to lower perceptual quality.
		In this letter, we propose a novel generative DJSCC approach using conditional diffusion models to enhance the perceptual quality of transmitted images. 
		Specifically, by utilizing entropy models, we effectively manage transmission bandwidth based on the estimated entropy of transmitted symbols. These symbols are then used at the receiver as conditional information to guide a conditional diffusion decoder in image reconstruction.
		Our model is built upon the emerging advanced mamba-like linear attention (MLLA) skeleton, which excels in image processing tasks while also offering fast inference speed. 
		Besides, we introduce a multi-stage training strategy to ensure the stability and improve the overall performance of the model.
		Simulation results demonstrate that our proposed method significantly outperforms existing approaches in terms of perceptual quality. 
		
		
		
		
		
	
	
	
	
		Semantic communications, conditional diffusion models, joint source-channel coding, image transmission.
	
	
	

§ INTRODUCTION

	
	
	
	
	
	
	The rapid development of sixth-generation (6G) communication systems has driven the rise of various smart applications, such as Virtual Reality (VR) and the Internet of Everything (IoE) <cit.>. These services demand enhanced communication efficiency to manage the massive inflow of data traffic. In this context, semantic communications have emerged as a new paradigm,  attracting significant attention. Unlike traditional transmission systems that rely on separate source and channel coding, semantic communications focus on accurately transmitting the underlying semantic information of digital data. This approach integrates source and channel coding for joint optimization, a technique known as joint source-channel coding (JSCC).
	
	Recently, the integration of deep learning into wireless communication system designs has gained traction, driven by the exceptional information processing capabilities of various deep learning models. In this context, deep learning-based JSCC (DJSCC) has stimulated significant interest, particularly through the use of autoencoders (AEs) and their variants, such as variational AEs (VAEs). DJSCC maps input images into low-dimensional symbol vectors for transmission. A pioneering work in this area is DeepJSCC<cit.>, which is built on an AE-based framework that allows for joint optimization and mitigates the cliff-edge effect commonly observed in traditional separation-based schemes. Moreover, the authors in <cit.> introduced a Transformer-based framework that enhances image fidelity by incorporating channel feedback to adaptively reconstruct images under varying wireless conditions. Inspired by entropy model-based source coding <cit.>, the authors in <cit.> proposed a rate-adaptive JSCC system, where the number of transmitted symbols is determined by the entropy estimated by the entropy models. This work was further evolved in <cit.> into a digital version, enabling practical application in modern digital systems.
	
	
	
	Despite the superior performance of current DJSCC transmission systems, they are typically optimized for mean-square error (MSE) distortion metrics such as peak signal-to-noise ratio (PSNR), which assesses pixel-level similarity between reconstructed and source images. However, it is increasingly recognized that high pixel-level similarity does not necessarily indicate high perceptual quality, which reflects how humans perceive the image's realism and visual appeal <cit.>. To improve the perceptual quality of transmitted images, the authors in <cit.> introduced adversarial and perceptual losses into DJSCC. Moreover, <cit.> developed a dual-path framework to extract both pixel-level information and perceptual information. Alternatively,  <cit.> combined DJSCC with the denoising diffusion probabilistic model (DDPM) <cit.>, a powerful generative model that creates images from  Gaussian noise. In <cit.>, the reconstructed image is split into two components: range-space and null-space. The range-space that captures the primary structure is transmitted via DJSCC, while the null-space, refining details, is generated at the receiver using diffusion models.
	Another approach, presented in <cit.>, leveraged diffusion models by first performing an initial reconstruction with DJSCC, followed by a neural network to guide the denoising process. Although <cit.> and <cit.> have achieved performance improvements, the decoding process at the receiver is time-consuming, as denoising requires hundreds of steps. Additionally, the training is performed module by module without joint optimization, potentially leading to performance loss. Besides, they typically support fixed-rate transmission, lacking the ability to determine optimal bandwidth for each image.
	
	
	
	In this letter, we introduce a novel framework called conditional diffusion models-based generative DJSCC (CDM-JSCC) for wireless image transmission.
	Unlike previous methods that rely on DDPM noise prediction (referred to as ϵ-prediction)<cit.>, our method utilizes 𝒳-prediction <cit.> to directly predict the source image. 
	
	𝒳-prediction offers comparable performance to ϵ-prediction but with significantly fewer steps, as its objective function resembles an autoencoder loss, enabling data reconstruction in a single iteration.
	Moreover, we optimize transmission bandwidth by utilizing entropy models to estimate the entropy of transmitted symbols. At the receiver, the transmitted symbols are utilized as conditional information to guide the conditional diffusion decoder to reconstruct images.
	Recently, the mamba-like linear attention (MLLA)  technique has demonstrated superior performance in image processing tasks, benefiting from parallel computation and fast inference <cit.>. Building on this, we design our JSCC encoder and decoder with MLLA as the core architecture.
	Furthermore, we propose a multi-stage training strategy to achieve joint optimization, resulting in substantial enhancements in perceptual quality.
	
	
	
	
	
	
	
	

§ PROPOSED CDM-JSCC FRAMEWORK
 
	In this section, we propose the framework of the proposed CDM-JSCC to realize rate-adaptive image transmission. To provide a comprehensive understanding, we begin by giving	an overview of CDM-JSCC and then proceed to detail the conditional diffusion models-based decoder.
	
	
	

 §.§ Rate-Adaptive JSCC

	As illustrated in Fig. <ref>, the source image, represented
	by a pixel intensity vector x_0 ∈ℝ^n, is first processed by a nonlinear transform coding encoder g_e, producing a low-dimensional latent representation z = g_e(x; θ_g), where θ_g encompasses the trainable parameters.
	Inspired by <cit.>, we estimate the distribution of z̃, i.e., the quantized version of z, by utilizing hyperprior entropy models. Specifically, we introduce an additional latent feature y = h_e(z; θ_h), serving as side information to capture the dependencies among the elements in z̃, where h_e represents the parametric analysis transform, and θ_h denotes its trainable parameters. 
	Each z̃_i in z̃ is variationally modeled as a Gaussian distribution with the standard deviation σ_i and mean μ_i predicted based on the quantized ỹ as:
	
    p_z̃ | ỹ(z̃ | ỹ)= ∏_i (𝒩(μ̃_i, σ̃_i^2) * 𝒰(-1/2, 1/2)) (z̃_i),

	where (μ̃, σ̃)=h_s(z̃; ϕ_h), h_s is the parametric synthesis transform, ϕ_h denotes its trainable parameters, and * represents the convolutional operation. Moreover, we convolve each element with a standard uniform density 𝒰(-1/2, 1/2) to enable a better match of the prior to the distribution of z̃ <cit.>.
	
	
	In this way, we are able to control the channel bandwidth ratio (CBR) based on the estimated entropy of z̃. Specifically, z comprises multiple embedding vectors z_i of length C. The JSCC encoder f_e adjusts the length of each vector to k_i=Q(-βlog p_z̃_i | ỹ(z̃_i | ỹ)), where Q denotes the quantization operation, β controls the relation between the prior p_ẑ_i | ỹ(ẑ_i | ỹ), and the expected length of the corresponding symbol vector s_i. 	In particular, when the entropy of z̃_i is high, a higher CBR is adopted, resulting in a larger k_i. This process is expressed as s = f_e(z, p_z̃ | ỹ(z̃ | ỹ); θ_f) ∈ℂ^k, where s denotes the channel input symbols, k is the number of transmitted symbols, and θ_f denotes the trainable parameters.  Additionally, to meet the energy constraints of real-world communication systems, we ensure that s satisfies an average power constraint before transmission. 
	
	
	Then, s is transmitted through a noisy wireless channel, which is denoted by η. As the additive white Gaussian noise (AWGN) channel is adopted in our work, this process follows ŝ≜η(s) =s + n, where n∼𝒞𝒩(0,σ^2 I_k × k) is a complex Gaussian vector with variance σ^2. 
	At the receiver, ŝ is first decoded to obtain the reconstructed latent ẑ = f_d(ŝ; ϕ_f), where  f_d represents the JSCC decoder, and ϕ_f denotes its trainable parameters.
	The receiver utilizes ẑ as an additional “content” latent to guide the conditional diffusion decoder g_d to reconstruct images x̂_0 = g_d(x̂_N, ẑ; ϕ_g), where x̂_N is the randomly sampled Gaussian source and ϕ_g represents the trainable parameters, as shown in Fig. <ref>. The details of the conditional diffusion decoder will be provided in the following section.
	
	
	
	

 §.§ Conditional Diffusion Decoder
 	
	We build our decoder on conditional diffusion models for their significant success in generative tasks. 
	
	The core idea of diffusion models is to transform an image x_0 into a Gaussian distribution by progressively adding noise to it, referred to as the forward process q, resulting in a sequence of increasingly noisy versions x_1, x_2,..., x_N. Then, the reverse process p_θ generates high-quality samples by reversing this process.
	The two Markov processes at step n can be respectively described as follows:
	
    q(x_n|x_n-1) = 𝒩(x_n; √(1-β_n)x_n-1, β_n I),

	
    p_θ(x_n-1 | x_n) = 𝒩(x_n-1; μ_θ(x_n,n), Σ_θ(x_n, n)),

	where the variance β_n is held constant as hyperparameters, the reverse process mean μ_θ(x_n, n) is parameterized by a neural network, and the variance Σ_θ(x_n, n) is always set to β_n I.
	
	The diffusion models are typically trained to predict the accumulated noise ϵ that perturbs  x_0 to x_n, a process known as ϵ-prediction, with the loss function being:
	
    ℒ(θ, x_0)=𝔼_x_0, n,ϵ||ϵ - ϵ_θ(x_n, n)||^2,

	where n ∼Uniform(1, ..., N), ϵ∼𝒩(0, I),  x_n = √(α̅_n)x_0 + √(1 - α̅_n)ϵ, and α̅_n=∏_i=1^n (1-β_i). 
	
	
	In our conditional diffusion decoder, ẑ is taken as the condition for the reverse process. Consequently, the forward process remains unchanged as (<ref>), and the reverse process is replaced by:
	
    p_θ(x_n-1 | x_n, ẑ) =  N(x_n-1; μ_θ(x_n, ẑ, n), β_n I).

	
	Analogous to (<ref>), we modify the training objective as follows:
	
    ℒ(θ, x_0) 
    			= 𝔼_n, ϵ||x_0 - 𝒳_θ(x_n, ẑ, n/N)||^2.
	
	First, we employ 𝒳-prediction <cit.> to predict the source image instead of using ϵ-prediction. 𝒳-prediction requires only a few denoising steps yet achieves performance comparable to ϵ-prediction, which involves hundreds of denoising steps. This is because the optimization objective resembles an autoencoder loss, allowing the model to reconstruct the source image in a single iteration. Moreover, we replace the step n with a normalized step n/N, enabling the use of a smaller N during testing compared to training, thereby accelerating the decoding process.
	
	During inference, images are generated using ancestral sampling with Langevin dynamics as follows:
	
    x_n-1    = √(α̅_n)𝒳_θ(x_n, ẑ, n/N) - √(1-α̅_n)ϵ_θ(x_n,n,n/N), 
       ϵ_θ(x_n, ẑ, n/N)=x_n-√(α̅_n)𝒳_θ(x_n, ẑ, n/N)/√(1-α̅_n).

	
	

§ TRAINING STRATEGY AND MODEL ARCHITECTURE
 
	In this section, we first present the architecture of CDM-JSCC. Then, we analyze the proposed loss function. Finally, we propose a multi-stage training algorithm to ensure the stability and improve the overall performance.
	

 §.§ Model Architecture

	As shown in Fig. <ref>, the proposed model consists of a transmitter, a wireless channel, and a receiver. We adopt the AWGN channel in our work, which is incorporated as a non-trainable layer in the architecture.
	

  §.§.§ Transmitter

	The transmitter consists of an entropy encoder and a rate-adaptive JSCC encoder. The entropy encoder comprises a hyperprior model to effectively capture
	spatial dependencies in the latent representation <cit.>. 
	The encoder g_e consists of 4 residual blocks followed by 4 down-sampling convolution layers connected sequentially. The hyperprior encoder is composed of a down-sampling convolution layer followed by a ReLU activation function, while the hyperprior decoder is composed of an up-sampling convolution layer and a ReLU activation function.
	Then, z is input to the JSCC encoder, where it is compressed into a variable-length latent representation based on the estimated entropy p(z̃|ỹ). Specifically, elements are selected and masked sequentially in a one-dimensional checkerboard pattern, which captures global features more effectively than random selection or linear sequencing. In addition, our proposed model is built on the advanced MLLA skeleton, known for its exceptional performance in image processing tasks and fast inference speed.
	The JSCC encoder consists of 3 MLLA blocks and multiple multi-layer perceptrons (MLPs).
	
	
	

  §.§.§ Receiver

	
	The receiver consists of a rate-adaptive JSCC decoder and a conditional diffusion decoder. The JSCC decoder inverts the operations performed by the JSCC encoder. Then, ẑ is input to the diffusion decoder as conditional information to guide the denoising process. The diffusion denoising unit is built on a U-Net architecture, featuring skip connections that allow the output from each down-sampling unit to be directly passed to the corresponding up-sampling unit.  The architecture includes 6 down-sampling units and 6 up-sampling units, with a mid-block in between. Each down-sampling unit, up-sampling unit, and mid-block is composed of residual blocks and linear attention mechanisms. In addition, ẑ is encoded into features of different scales by 4 layers of hierarchical encoders. These features are concatenated with the output of the preceding down-sampling unit and embed time information as input for the next down-sampling unit. During testing, the output x̂_n-1 is passed through the diffusion denoising unit iteratively until x̂_0 is produced.
	
	

 §.§ Loss Function

	In a typical end-to-end JSCC system, the optimization focuses sorely on the MSE between the source image and the reconstructed image, with the loss function defined as:
	
    ℒ_D = 
    		𝔼_x_0||x_0-x̂_0||^2,
	
	which aligns with the training objective of our conditional diffusion models as defined in (<ref>), with the exception that step n and noise ϵ are randomly sampled.
	
	
	
	Moreover, for an entropy model-based rate-adaptive JSCC system, the loss function is typically defined as:
	
    ℒ_RD=   𝒟+λℛ
    			= 𝔼_x_0||x_0-x̂_0||^2_JSCC distortion + 𝔼_x_0||x_0-x̅_0||^2_compression distortion
    
    			+   λ𝔼_x_0[- log p(z̃|ỹ)- log p(ỹ)]_rate,
	
	where 𝔼_x_0||x_0-x̅_0||^2 is the compression distortion between the source image x_0 and the compressed image x̅_0, 𝔼_x_0[- log p(z̃|ỹ)- log p(ỹ)] represents the compression rate estimated by entropy models, and 𝔼_x_0||x_0-x̂_0||^2 denotes the typical JSCC loss term.
	
	To improve the perceptual quality of reconstructed images, we introduce an additional perceptual loss, expressed as follows: 
	
    ℒ_P = 
    		𝔼_x_0[d_p(x_0 , x̂_0)] + 𝔼_x_0[d_p(x_0,x̅_0)],
	
	where d_p(·, ·) represents the perceptual loss term. Here we adopt widely-used learned perceptual image patch similarity (LPIPS) loss based on VGGNet <cit.>. Therefore, the objective for our model can be expressed as:
	
    ℒ =    
    			(1 - η)(𝔼_x_0||x_0-x̂_0||^2_JSCC distortion + 𝔼_x_0||x_0-x̅_0||^2_compression distortion) 
     
    			+    η (𝔼_x_0[d_p(x_0,x̂_0)]_JSCC perceptual loss + 𝔼_x_0[d_p(x_0,x̅_0)]_compression perceptual loss)	
    
    			+    λ𝔼_x_0[- log p(z̃|ỹ)- log p(ỹ)]_rate,
	
	where η∈ [0,1] balances the trade-off between MSE distortion and the perceptual loss term, while λ adjusts the trade-off between the image quality and transmission rate. 
	

 §.§ Training Strategy

	To ensure training stability and enhance overall performance, we propose a multi-stage training strategy. Initially, we train each module individually to reduce complexity and facilitate convergence. Once all modules have been trained separately, we fine-tune the entire model to optimize performance. The complete multi-stage training process is detailed in Algorithm <ref>.
	
	
	
	
	
	
	
	
	
	
	
	
	
	
		
		

§ SIMULATION RESULTS
 
		In this section, we perform simulations to evaluate the performance of the proposed CDM-JSCC.
		

 §.§ Simulation Settings

		

  §.§.§ Basic Settings

		We evaluate our CDM-JSCC on the Kodak image dataset, which consists of 24 images, each with a resolution of 512 × 768 × 3. During testing, we only employ 30 denoising steps at the receiver, ensuring an efficient and fast decoding process. All experiments are conducted using Pytorch. For the training process, we use the Adam optimizer for stochastic gradient descent, starting with a learning rate of 1 × 10^-4 and planning to reduce it after several epochs. The training dataset comprises 50,000 randomly sampled images from the ImageNet dataset, which are randomly cropped to 256 × 256. The batch size is set to 4. We set λ=0.0512 and η=0.5 for the experiments.
		
		
		
		
		

  §.§.§ Benchmarks
 
		For benchmarks, we consider both deep learning-based methods and traditional separation-based schemes, incorporating optimization strategies for both MSE distortion and perceptual loss. The benchmarks are detailed as follows:
		
			
  * “BPG+LDPC": This method employs BPG for source coding and LDPC for channel coding, followed by quadrature amplitude modulation (QAM).
			
  * “BPG+capacity": In this approach, we use an ideal capacity-achieving channel coding in conjunction with BPG and QAM. 
			
  * “DeepJSCC": As a representative learning-based method, we include the classic DeepJSCC optimized for MSE distortion <cit.>.
			
  * “GAN-JSCC": This method optimizes both MSE distortion and perceptual loss <cit.>.
		 
		
		
		
		
		

  §.§.§ Considered Metrics

		We evaluate the performance of our method along with benchmarks using typical metric PSNR as well as perceptual metrics LPIPS<cit.>, and Fréchet Inception Distance (FID). PSNR calculates pixel-wise MSE distortion. LPIPS measures the l2 distance between two latent embeddings extracted by a pre-trained network <cit.>. FID calculates the distribution of source images and reconstructed images in the feature space using the Fréchet distance <cit.>.
		
		
		
		

 §.§ Performance Comparison

		
		We evaluate the performance of our proposed model on the AWGN channel, a representative and widely adopted method in this research community. In all subsequent experiments, the model is trained on a single SNR value and tested on the same SNR. In addition, by carefully adjusting the hyper-parameters, we control the CBR =1/48. 
		
		Fig. <ref> compares CDM-JSCC with benchmarks across various SNRs, illustrating a comparative analysis of our proposed CDM-JSCC against the benchmarks. Notably, our model outperforms others across all perceptual quality metrics, including LPIPS and FID. This enhancement is attributed to the proposed loss function and the powerful generative diffusion models. However, it is important to note that our scheme shows some performance degradation in the PSNR metric compared to the MSE-optimized DeepJSCC and traditional separation-based schemes. This is because generative models tend to generate realistic and clear images and may overlook pixel-wise fidelity sometimes. Despite this, our model still outperforms GAN-JSCC in the PSNR metric. Unlike GAN-JSCC, our model does not plateau as SNR increases, maintaining the potential to transmit clear, realistic, and high-fidelity images at higher SNR levels.	
		
		
		

§ CONCLUSION
 
		In this letter, we proposed a framework of conditional diffusion models-based generative DJSCC system for image transmission. Moreover, we employed 𝒳-prediction with a few denoising steps to accelerate the decoding process. Furthermore, we effectively managed the transmission bandwidth based on the estimated entropy of transmitted symbols. Besides, we proposed a multi-state training strategy to ensure the stability of the training process. Simulation results demonstrated that the proposed method can significantly surpass existing methods in terms of perceptual quality. 
		
		
		IEEEtran
		
		
	