Diffusion Inpainting MIMO-OFDM Channels with Limited Noisy Observations

Abstract: Acquiring the channel state information from limited and noisy observations at pilot positions is critical for wireless multiple-input multiple-output (MIMO)-orthogonal frequency division multiplexing (OFDM) systems. In this paper, we view this process as a conditional generative task in which the partial noisy channel estimates at the pilots are utilized as a prompt'' to guide the diffusioninpainting'' of the underlying channel. To this end, we resort to a general Conditional Diffusion Transformer (CDiT) framework with a well-designed network architecture and update rule. In particular, we design a dedicated embedding strategy to encode and adapt to different pilot patterns and noise levels, and utilize a special cross-attention mechanism to align the partial raw channel observations with the denoised channel at each time step of the generation process. This architecture effectively anchors the diffusion process, enabling the model to accurately recover full channel details from limited noisy observations. Comprehensive experimental results show that, the proposed approach achieves a performance gain of over 5 dB compared to the baselines under varying noise conditions, and provides robust channel acquisition even under a sparse pilot density of 1/32 without significant performance loss compared to the denser pilot cases. Moreover, it is capable of generating high-quality channel matrices within just 10 inference steps, effectively balancing estimation accuracy with computational efficiency and inference speed. Ablation studies demonstrate the rationality of the model design and the necessity of its modules.

• The paper presents a conditional diffusion transformer (CDiT) that significantly improves NMSE by over 5 dB compared to traditional methods.
• It integrates weighted mask embedding, patchify tokenization, and cross-attention fusion to robustly inpaint channels from sparse and noisy pilot measurements.
• The approach enables rapid inference with as few as 4–10 steps, offering computational efficiency for real-time, high-dimensional channel estimation.

Conditional Diffusion Inpainting for MIMO-OFDM Channel Estimation with Noisy Observations

Introduction and Motivation

Accurate channel state information (CSI) acquisition remains a core challenge for high-dimensional MIMO-OFDM systems, particularly under sparse and noisy pilot observations. Traditional linear estimators (e.g., LS, LMMSE) either exhibit marked performance degradation in low pilot regimes or face excessive computational overhead due to matrix inversions, especially for large antenna or subcarrier arrays. The rise of data-driven approaches and, more recently, diffusion models (DMs) opens new opportunities for robust channel estimation under challenging conditions. Diffusion models, especially those incorporating conditional architectures, have demonstrated outstanding performance in high-noise, data-imputation, and inverse problem settings.

This paper introduces the Conditional Diffusion Transformer (CDiT), a diffusion-based transformer architecture designed for MIMO-OFDM channel inpainting from limited, noisy pilot measurements. The work positions channel estimation as a conditional generative task, wherein raw, noisy pilot-based channel estimates act as a "prompt" to guide the inpainting of the full channel tensor. CDiT's design emphasizes robustness to pilot pattern and noise level variability, adaptability to different system configurations, and computational efficiency.

Figure 1: The proposed Conditional Diffusion Transformer (CDiT) architecture integrates weighted masking, multi-level conditional embeddings, patchify tokenization, and cross-attention fusion for robust channel inpainting.

---

System Model and Problem Formulation

The MIMO-OFDM model considered consists of a high-dimensional receive antenna array and a single-user uplink. Pilots are sparsely inserted in frequency, with raw pilot-based channel estimates being both partial and perturbed by additive noise. These observations are structured as:

$\widetilde{H} = (H + \widetilde{N}) \otimes M$

where $H$ is the full channel matrix, $\widetilde{N}$ is i.i.d. Gaussian noise, and $M$ encodes the pilot pattern. The challenge is to estimate $H$ globally from these partial, noisy measurements given arbitrary pilot masks and system noise levels.

---

Conditional Diffusion Transformer (CDiT) Architecture

Embedding Conditional Information

To handle arbitrary pilot positions and adapt to changing noise regimes, CDiT embeds several conditional signals:

• Weighted Mask Embedding: The pilot mask is weighted by estimated noise variance and fused via $1 \times 1$ convolutions, conveying both spatial support and observation reliability.
• Conditional Embeddings: Scalar conditions (diffusion time step $t$, SNR, pilot interval class) are embedded using sinusoidal spatial encoding and MLPs, appended to feature tokens.
• Classifier-Free Guidance: Pilot interval is treated as a class label, supporting both unconditional and class-conditional score guidance for enhanced control during inference.

Tokenization and Patchify

A 2D convolutional patchify module projects noisy channel matrices and conditional masks into token sequences. This enables scalable, patch-wise processing suitable for transformer-based architectures, ensuring positional alignment via consistent sinusoidal positional encodings.

Figure 2: Patchify operation converts frequency-antenna channel arrays into embedded token sequences for transformer processing.

Cross-Attention Fusion

In each CDiT block, latent channel tokens interact with conditional tokens through cross-attention. Cross-attention allows conditional features—such as pilot-derived prompt information, weighted masks, and noise descriptors—to guide the restoration of missing channel entries at each diffusion timestep. Adaptive layer normalization further integrates conditional embeddings during both self- and cross-attention operations.

---

Training and Inference

Variance-Preserving Diffusion with Conditional Prompts

CDiT leverages a modified Denoising Diffusion Probabilistic Model (DDPM) framework, using $T=1000$ diffusion steps during training but only a handful ($S \ll T$, e.g., $S=10$) during inference. Training incorporates random pilot patterns and SNRs to ensure model robustness over a wide operational envelope.

Accelerated Inference

At inference, a linear schedule for time-step sub-sampling allows for rapid denoising from Gaussian noise to channel estimate with minimal loss relative to full-step diffusion. The paper empirically demonstrates that CDiT achieves near-optimal performance using as few as 4–10 inference steps, substantially reducing latency compared to both DDPMs with more steps and deterministic DDIM variants.

Figure 3: Schematic of the diffusion forward (corruption) and reverse (generative sampling) process. Most inference occurs in a fraction of the training steps.

---

Experimental Results and Ablations

Simulation Settings

CDiT is evaluated on 3.5 GHz urban MIMO-OFDM channels via the sionna ray-tracing simulator. Experiments analyze channel restoration accuracy (NMSE, cosine similarity) under varying SNRs, pilot densities (down to $H$0), and pilot patterns. Competing methods include linear interpolation, LMMSE with sample covariances, and CMixer (a large MLP-based network).

Comparative Performance

• Channel Recovery: CDiT achieves $H$1 dB NMSE improvement across SNRs compared to baselines. It maintains robust performance at $H$2 pilot interval (pilot density $H$3), with negligible degradation versus denser pilot regimes.
• Cumulative Error Distribution: Performance remains tightly concentrated near optimal even as SNR and pilot sparsity vary.
• Efficiency: Optimal performance is retained with as few as 4 inference steps, a marked efficiency advantage over alternatives.

  Figure 4: NMSE comparison across SNRs ($H$4, $H$5) demonstrates superior denoising and generalization by CDiT.

  

  Figure 5: Cumulative probability distribution of errors (NMSE) indicates consistently low error rates for CDiT.

  

  Figure 6: Robustness of CDiT to pilot interval variation.

  

  Figure 7: Channel estimation over time: CDiT rapidly denoises from noise to high-fidelity channel matrices over 10 steps.

Impact of Inference Steps and Modeling Stochasticity

CDiT’s stochastic denoising process (DDPM) outperforms deterministic DDIM procedures, especially as inference steps are reduced or noise levels increase—a crucial regime for real-time communications.

Figure 8: NMSE vs. inference steps $H$6 shows outstanding performance for $H$7; further reduction increases error modestly only at high SNR.

Figure 9: Comparison of DDIM and DDPM reveals the necessity of stochasticity for robust channel inpainting, especially at low SNR or sparse pilots.

Model Scaling, Patch Size, and Ablations

• Model Size: Moderate reductions in CDiT depth ($H$8) or width ($H$9) yield minimal performance loss, supporting computational tradeoffs for resource-constrained deployment (Table I).
• Patch Size: Smaller patch sizes enhance representation but increase computation; careful tuning aligns cost and accuracy.
• Pilot Pattern Generalization: CDiT generalizes to unseen pilot patterns with marginal loss, especially under low-SNR training, highlighting its adaptability.
• Ablation Studies: Weighted masking, noise embedding, class embedding, and independent patchify modules all contribute to model robustness, particularly in smaller models.

  Figure 10: Influence of training epochs and patch size on performance: smaller patches and longer training yield increased NMSE improvements.

---

Implications and Future Directions

The CDiT framework advances the principled integration of conditional generative models with transformer architectures for wireless channel estimation:

• Theoretical Impact: Positions channel estimation as a general conditional inpainting problem; provides a general recipe for leveraging diffusion-based generative modeling for structured sensing and restoration tasks beyond wireless communications.
• Practical Relevance: Achieves high estimation fidelity from minimal/noisy pilots with limited computational burden, supporting deployment in dense, massive MIMO and rapidly-varying OFDM systems where pilot resources are scarce.
• Modularity: The framework's ability to embed arbitrary structural, environmental, or protocol information suggests extensibility to other multi-modal inference tasks or joint channel-data problems.
• Future Research: Directions include multi-modal data fusion (user/environment context), model compression (distillation, quantization), and further acceleration of inference (few-step generation), as well as exploration in broader inverse problems and physical-layer AI.

---

Conclusion

CDiT demonstrates that conditional diffusion-based transformers with principled conditional fusion and cross-attention can achieve robust, scalable MIMO-OFDM channel estimation under challenging scenarios involving limited and noisy observations. The architecture strikes a favorable balance between generalization, estimation accuracy, and inference efficiency, setting a new benchmark for diffusion-powered wireless inference. The approach's architectural and algorithmic innovations offer a promising foundation for the next generation of robust, adaptive wireless AI systems.

---

References:

Zhou, W., Zhang, Z., Yang, Y., Yan, S., Kong, Z., & Debbah, M. "Diffusion Inpainting MIMO-OFDM Channels with Limited Noisy Observations" (2604.09039)