Residual Diffusion Models for Variable-Rate Joint Source Channel Coding of MIMO CSI

1. Problem Statement

The paper addresses the critical challenge of Channel State Information (CSI) compression for massive Multiple-Input Multiple-Output (MIMO) systems in Frequency Division Duplexing (FDD) networks.

• The Bottleneck: Accurate downlink precoding requires the User Equipment (UE) to feedback CSI to the Base Station (BS). As antenna counts ($N_t$) increase, the feedback overhead becomes prohibitive.
• Limitations of Current Methods:
  • Separate Source/Channel Coding: Suffers from the "cliff effect," where performance collapses under poor channel conditions.
  • Standard Autoencoders (AE): While effective for simple channels, they struggle to model complex, non-linear channel distributions (e.g., COST2100 outdoor). They often exhibit poor scalability and fail to capture high-dimensional statistical nuances.
  • Existing Diffusion Models: While powerful, standard generative diffusion models are computationally expensive and typically initialized from random noise, which is suboptimal for reconstruction tasks where a coarse estimate already exists.
  • Fixed-Rate Constraints: Most deep learning solutions require separate models for different compression rates, increasing UE storage and switching overhead.

2. Methodology: RD-JSCC

The authors propose Residual Diffusion Joint Source-Channel Coding (RD-JSCC), a hybrid framework that combines a lightweight autoencoder with a residual diffusion module.

A. System Architecture

The framework operates in two stages at the receiver (BS):

1. Stage 1: Coarse Reconstruction (Autoencoder):

   • A lightweight Convolutional Neural Network (CNN) encoder at the UE compresses the high-dimensional CSI matrix into a continuous-valued latent vector $s$.
   • The vector is transmitted over a noisy uplink channel.
   • A CNN decoder at the BS produces a coarse estimate $\hat{H}_d$.
   • Key Feature: Uses continuous-valued latents rather than vector quantization (VQ) to avoid the "cliff effect" and enable graceful degradation under noise.
   • SNR Adaptation: Both encoder and decoder include modules that dynamically scale features based on the uplink Signal-to-Noise Ratio (SNR).

2. Stage 2: Residual Refinement (Diffusion):

   • Instead of generating CSI from random noise, the diffusion model treats the coarse estimate $\hat{H}_d$ as a prior.
   • It learns to predict the residual error ($r = H_d - \hat{H}_d$) using a U-Net architecture.
   • The process iteratively refines the coarse estimate to produce the final high-fidelity CSI $H'_d$.
   • Adaptive Decoding: The system can "early exit" after Stage 1 if channel conditions are favorable, or invoke Stage 2 for complex channels, balancing latency and performance.

B. Key Technical Innovations

• Residual Diffusion Formulation: The forward diffusion process is modified to start from the coarse estimate $\hat{H}_d$ rather than pure Gaussian noise. The model learns to denoise the residual between the ground truth and the coarse estimate, which is sparser and easier to model.
• $\chi$-Prediction: Instead of predicting added noise (standard diffusion), the model predicts the clean signal component directly. This enables low-latency inference (e.g., 2-step or even 1-step) with minimal performance loss.
• Variable-Rate Compression (Matryoshka Representation Learning - MRL):
  • The model learns a single hierarchical latent space.
  • By truncating the latent vector to different lengths ($k$), the system supports multiple compression rates (e.g., $k \in \{8, 16, 24, 32\}$) within a single unified model, eliminating the need for retraining or model switching.
• Preprocessing: CSI is transformed from the spatial-frequency domain to the angular-delay domain via 2D IFFT. This exploits channel sparsity, allowing the matrix to be cropped (e.g., to $32 \times 32$) without significant information loss.

3. Key Contributions

1. Novel Framework (RD-JSCC): A flexible two-stage architecture that integrates a lightweight autoencoder with a residual diffusion module, enabling adaptive switching between low-complexity decoding and high-fidelity refinement.
2. Variable-Rate Capability: Implementation of MRL to support a continuous range of compression rates (25 distinct rates) using a single model, facilitating dynamic bandwidth adaptation.
3. Robustness to Practical Constraints:
   • Demonstrated robustness to imperfect channel estimation by training with additive noise models.
   • Validated fixed-bit quantization (down to 4 bits per dimension) with negligible performance degradation.
4. Comprehensive Evaluation: Extensive ablation studies quantifying the trade-offs between encoder complexity, SNR adaptation, and diffusion steps.

4. Experimental Results

The model was evaluated on the COST2100 outdoor (complex) and 3GPP indoor (simpler) datasets.

• Performance Gains:
  • On the COST2100 dataset, RD-JSCC achieved an order-of-magnitude improvement in Normalized Mean Squared Error (NMSE) compared to state-of-the-art non-diffusion baselines (e.g., ADJSCC, CSINet+).
  • At 5 dB uplink SNR, RD-JSCC achieved an NMSE of $\approx -12.2$ dB, compared to $\approx -9.2$ dB for U-Net-based JSCC without diffusion.
• Block Error Rate (BLER):
  • The NMSE improvements translated to significant reliability gains. RD-JSCC achieved a BLER of $10^{-6}$ at 5 dB SNR, whereas non-diffusion methods saturated at $10^{-4}$.
• Efficiency:
  • 2-Step Inference: Using $\chi$-prediction, the model achieves near-optimal performance with only 2 diffusion steps, reducing inference latency significantly compared to 20-step diffusion.
  • Complexity vs. Channel: The performance gap between diffusion and autoencoders narrows for simple channels (3GPP indoor), validating the need for adaptive decoding (early exit) based on channel complexity.
• Quantization: The latent vector can be quantized to 4 bits with negligible NMSE degradation, making it suitable for bandwidth-constrained feedback.

5. Significance

This paper represents a significant advancement in physical layer communications by bridging the gap between generative AI and wireless engineering.

• Overcoming the Cliff Effect: By moving away from discrete vector quantization to continuous latent spaces and residual diffusion, the system handles channel noise gracefully.
• Practical Deployment: The framework addresses real-world constraints such as variable bandwidth, imperfect CSI, and computational limits through adaptive decoding and quantization.
• Paradigm Shift: It demonstrates that for complex channel distributions, diffusion-based refinement is superior to traditional supervised learning, provided the computational cost is managed via residual initialization and low-step inference.

In summary, RD-JSCC offers a robust, scalable, and high-fidelity solution for MIMO CSI feedback, particularly in challenging wireless environments where traditional compression methods fail.