{'fraction_non_alphanumeric': 0.03283759162747341, 'fraction_numerical': 0.008177619592390153, 'mean_word_length': 4.6491144088080425, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Deep Diffusion Models for Robust Channel Estimation

Channel estimation is a critical task in digital communications that greatly impacts end-to-end system performance. In this work, we introduce a novel approach for multiple-input multiple-output (MIMO) channel estimation using deep diffusion models. Our method uses a deep neural network that is trained to estimate the gradient of the log-likelihood of wireless channels at any point in high-dimensional space, and leverages this model to solve channel estimation via posterior sampling. We train a deep diffusion model on channel realizations from the CDL-D model for two antenna spacings and show that the approach leads to competitive inand out-of-distribution performance when compared to generative adversarial network (GAN) and compressed sensing (CS) methods. When tested on CDL-C channels which are never seen during training or finetuned on, our approach leads to end-to-end coded performance gains of up to 3 dB compared to CS methods and losses of only 0.5 dB compared to ideal channel knowledge. To encourage open and reproducible research, our source code is available at https://github.com/utcsilab/diffusion-channels.


I. INTRODUCTION
Channel estimation [1] represents the task of recovering a communications channel tensor using a number of measurements. In the vast of majority of cases, these measurements come from pilot symbols transmitted at specific intervals, which are known ahead of time by the receiver. This task is critical for the end-to-end performance, since the channel estimate is used in downstream tasks, such as symbol equalization and soft demodulation of the data transmissions, and as is thus an active area of research. With the adoption of 5G and development of 6G standards [2], [3], channel estimation via deep learning has been identified as a promising solution, especially in high-dimensional cases such as MIMO millimeter-wave communications [4], [5].
Our main motivation in this work is to develop robust, datadriven MIMO channel estimation algorithms: we want our method to work reasonably well under distributional shifts compared to the training conditions, with zero adaptation. Although online adaptation methods are possible [6], they still require a set of ground-truth channels from the test distribution to be known, which may incur a significant pilot overhead. This work was supported by ONR grant N00014- 19-1-2590. Thus, algorithms that generalize and retain their performance in completely novel test-time distributions without fine-tuning are an important component toward real-world utilization.
In this paper, we use deep diffusion models [7] for robust, high-dimensional channel estimation. Inspired by recent research that shows the potential of this method for magnetic resonance imaging (MRI) reconstruction [8], we introduce a training and hyper-parameter tuning approach for wireless channel estimation in a single-carrier MIMO communication scenario. For training, a database of known channels is used to train a deep diffusion model in an unsupervised manner, that is independent of the pilot symbols. During inference, we tackle the main challenges that arise when applying this method to channel estimation as the need to operate in a very wide (e.g., up to 30 dB) signal-to-noise ratio (SNR) range and out-ofdistribution generalization. Our experiments show that with the proposed choice of hyper-parameters, test-time performance is retained even in simulated test channels that come from a completely different distribution than during training, and gains of up to 3 dB are achieved against compressed sensing based methods in end-to-end system simulations.

A. Related Work
Previous work on data-driven deep learning methods for channel estimation largely falls in the category of supervised methods. In [9], the matrix of pilots (with time and frequency dimensions) is passed through a deep convolutional neural network that is trained to output the ground-truth channel that generates the pilots. The work in [10] follows the same approach and shows that outside of the training distribution, a severe performance loss (tens of dB) occurs. This is not the case for diffusion models, which are shown to be theoretically robust outside of the training set under certain assumptions, as well as in practice [8]. The work in [5] takes a different approach than supervised learning by applying the compressed sensing with generative models (CGSM) to channel estimation by training a deep generative adversarial network to serve as an implicit prior. Their results show excellent performance in the low SNR and few-measurements regime, but saturated performance in the high-SNR or high-measurements regimes [11]. This serves as a strong baseline for our approach, since both methods are unsupervised.
Deep diffusion (or score-based) models are originally developed in [7], where the authors also introduce annealed Langevin dynamics to efficiently sample from a distribution of interest using a learned diffusion model. The work in [12] introduces practical techniques for improving the training of diffusion models, which we fully leverage. The work in [13] theoretically proves the in-distribution optimality of posterior sampling and links this to their fully-dimensional support. Finally, the work in [8] investigates the robustness of posterior sampling in the setting of MRI reconstruction, where it is shown that these models are competitive with state-of-the-art deep learning methods in-distribution and significantly more robust out-of-distribution.

B. Contributions
Our contributions in this work are the following: 1) We formulate channel estimation via posterior sampling: given measurements y obtained from pilot transmissions, we estimate the underlying channel h by sampling from the posterior distribution p H (h|y). 2) We successfully train a noise-conditional deep diffusion model for millimeter-wave wireless channel matrices from the CDL-D channel model and show that this can be used to estimate channels even when a reduced number of pilots makes the inverse problem underdetermined. 3) We develop a hyper-parameter selection algorithm for posterior sampling with diffusion models that optimizes the average relative channel estimation error over a wide range of operating SNR values. We test our scheme with and without knowledge of the exact operating SNR and show superior estimation performance both in-and outof-distribution (on CDL-C channels) with zero online adaptation or fine-tuning.

II. CHANNEL ESTIMATION AS AN INVERSE PROBLEM
We consider MIMO channel estimation in a narrowband digital communication scenario. Given a complex-valued channel matrix H ∈ C Nr×Nt and the matrix of complexvalued pilots P ∈ C Nt×Np the measurements seen at the receiver are expressed as [14,Eq. 7.55]: where N is complex additive white Gaussian noise, N t and N r are the number of transmitter and receiver antennas, respectively, and N p is the number of pilot transmissions. In practice, the transmitter will form the matrix P as the product between a (potentially varying across pilots) beamforming matrix W t and a matrix of pilot symbols S. The beamforming matrix itself could be further factorized in its digital and analog parts [4]. In this work, we do not assume any kind of factorized structure for P. This does not impact our method or the baselines, since we use unsupervised algorithms that are independent of the pilot matrix used. For the remainder of the paper, we also assume that the channel is constant across the pilot transmissions. In practice, this is a reasonable assumption if the pilots are staggered in time and the coherence time of the channel is much larger than the pilot sequence [15]. The model in (1) is also of practical interest in OFDM scenarios, since it corresponds to a single pilot subcarrier.
Eq. (1) can be rewritten by introducing the flattened channel and received symbol matrices h and y, respectively, as: where A ∈ C NpNr×NtNr is the consolidated measurement matrix such that A = P T ⊗ I Nt , and n is the flattened noise vector. If the distribution of h is Gaussian, then the maximum a posteriori estimator is linear and optimal with respect to the squared error [16]. Channels of practical interest may severely deviate from the i.i.d. Gaussian assumption, especially in millimeter-wave scenarios [4], and, more importantly, their underlying prior distribution p H (h) may be unknown and intractable. The importance of a well-chosen prior is increased when the number of pilots makes the inverse problem under-determined. Denoting α = N p /N t as the undersampling ratio, a lower α leads to fewer pilots being used for channel estimation. When α < 1, compressed sensing methods aim to estimate h by solving the following regularized optimization problem: where R is a sparsity-promoting regularization function and λ ≥ 0 is a hyper-parameter that depends on the SNR and the chosen R. When dealing with complex signal distributions -such as wireless channels -this parameter is chosen in a data-driven way by using a set of validation channels and a search strategy (e.g., line search [17]). A commonly used prior for wireless channel estimation is sparsity in the 2D Fourier (DFT) domain [18]. The approach of compressed sensing with generative models [5], [11] has been applied to solve (3) by using deep generative adversarial networks as implicit priors. The algorithm recovers h by first training a GAN G that takes as input a lowdimensional latent vector z and outputs a sample h from p H (h). With measurements y, the channel is then estimated by solving the optimization problem: and λ z enforces the Gaussian assumption placed on z. To solve (4) in practice, a gradient descent approach is used to find the solution z , since the deep neural network used is differentiable.

A. Diffusion Models
The key component of our approach is the class of models termed denoising diffusion (score-based) models, which we briefly term diffusion models, introduced in [7]. These models are implicit likelihood models that are trained to estimate the score (the gradient of the log-likelihood) of the underlying distribution p X (x) that governs a high-dimensional signal x. Given a data point x sampled from p X (x), the goal of a diffusion model s θ with learnable weights θ is to output the gradient of the log-likelihood log p X (x) with respect to x: To train these models, an unsupervised approach is used, since the exact gradients of log p X (x) are intractable and unavailable for high-dimensional real-world signals, e.g., for wireless channels. To solve this issue and generate paired training samples, i.i.d. Gaussian noise with zero mean and standard deviation σ is added to x, yielding a perturbed samplẽ x, for which the gradient is estimated via the finite difference approximation:∇ The approximate gradients are then used to train the diffusion model using the simple mean squared error loss L(x; θ, σ) = s θ (x) − (x −x)/σ 2 2 . Importantly, during training, the value of σ is randomly chosen from a broad interval at each training step, leading to diffusion models learning a gradient at every point in space, including outside of the support of p X (x). This is a key aspect that aids the robustness of these models, since it effectively allows them to extrapolate and be well-behaved beyond the training dataset. In the sequel, we introduce channel estimation through posterior sampling, which is different than all previously considered channel estimation approaches.

B. Channel Estimation via Posterior Sampling
Our solution for channel estimation comes in the form of posterior sampling. An overview of this approach, as well as a proof of optimality (given access to p H (h|y)) are given in [13]. A block diagram of the entire approach is summarized in Fig. 1. Fundamentally, we solve the estimation problem by sampling from the posterior, that is:h ∼ p H (h|y). In this work, we use annealed Langevin dynamics [13] to sample from the posterior by running noisy gradient ascent, but other strategies are possible. At inference step t + 1, the current channel estimate is updated using the rule: where η t is the step size, β is a hyper-parameter that controls the injected noise power, and ζ t is additive white Gaussian noise added at each step. Since the distribution p Y (y|h t ) is Gaussian, we can apply Bayes rule to p H (h t |y) and omit the p Y (y) term since its gradient is zero with respect to h t . Furthermore, by using that: , we obtain the final update rule as: where inference starts from a pure noise channel h 0 ∼ CN (0, I) and is carried out for N steps. Finally, the returned channel estimate ish = h N .
The three terms involved in (8) each contribute to solving the inverse the problem differently: • The first term takes a gradient step that maximizes the log-likelihood of the current channel estimate, i.e., makes the channel more naturally structured under its governing distribution p H (h). • The second term takes a negative gradient step in the direction of the measurement loss, i.e., makes the channel more consistent with the observed pilots y. • The third term adds noise proportional to the step size at each step in order to prevent the approach from converging to the MAP estimate, which is sub-optimal when the posterior is multi-modal [13].

A. Training a Diffusion Model for MIMO Channels
We train an NCSNv2 model [12] on complex-valued channel matrices H with N t = 64 and N r = 16 as a proof-ofconcept for our method. The backbone network we use is a RefineNet [19] with a total of 5.2 million parameters. To handle complex valued-inputs, we treat the real and imaginary parts of H as two input channels. The model is trained on a total of 20000 training channel realizations drawn from the CDL-D channel model available in MATLAB, equally split between antenna spacings λ/10 and λ/2. Training for 800 epochs takes approximately six hours on an NVIDIA RTX 2080Ti GPU. Complete details about the learning architecture are available in the repository linked in the abstract.
A part of our contribution that aids the practical use of the proposed approach is to provide a robust method for tuning the hyper-parameters for the annealed Langevin dynamics in (8). Importantly, we do not tune the step size schedule η twhich is chosen according to [12] -but instead only tune two scalar hyper-parameters: the noise amplification factor β and the number of iterations N . We chose these hyper-parameters since they affect estimation performance the most: too much added noise causes degrades performance in the low SNR regime, while running inference for too many steps is costly and overfits the measurements in the under-sampled regime. To measure raw estimation quality we use the normalized mean square error metric, defined as NMSE = H −H 2 F / H 2 F . All test results are for a pilot matrix P containing random unit power QPSK symbols, but the methodology holds for any pilot sequence known ahead of time. Given a fixed number of pilots, a known antenna spacing at the transmitter and receiver, a set of K validation SNR values {ρ i } i=1...K , and the resulting validation NMSE metrics n i (N, β) as a function of the tunable hyper-parameters and the SNR point, selecting the best hyperparameters is done by maximizing: This criterion represents the average, normalized log-scale gain across a range of SNR values. Since optimal recovery algorithms for sparse signals (given a sufficient number of measurements) exhibit a linear gain with unit slope in the loglog domain [20], the above metric is constant across all SNR points. This serves as a principled heuristic for choosing hyperparameters that perform well across the entire range of values, since it promotes low errors at high SNR values, favouring solutions that do not have an error floor.
We search for β among the set of values {1, 0.1, 0.01, 0.001} and for the stopping point N among 6933 inference steps. The best pair (β, N ) is found and stored for each antenna spacing and α value. If the SNR is assumed known, we store a separate pair for each SNR value and we find that, as the SNR increases, the optimal β decreases. Intuitively, this corresponds to a higher importance placed on the first two terms, since low SNR scenarios are implicitly regularized by the measurement noise. In the blind (unknown SNR) setting we use (9) to select a single pair of hyper-parameters that maximizes the relative gain in the entire SNR range. Exact values are given in the source repository. Finally, we note that only the two scalar hyper-parameters are tuned as a function of the scenario parameters, with the weights θ being completely frozen in the remainder of the experiments.

B. In-Distribution Performance
We compare our results with two baselines. The Lasso baseline [21] uses the 1 -norm of the spatial frequency representation of the channel as a regularizer in (3). We implement this using the SigPy library [22], a well-optimized Python package designed for inverse problems. For the known SNR setting, we tune a separate λ value for each SNR value, otherwise we use the criterion in (9) to find a single value in the blind setting. The optimal values for λ vary as a function of α, and are generally in the [0.03, 0.3] range.
The second baseline is the method in [5] -which we term Wasserstein GAN (WGAN). We implement and successfully reproduce the authors' results. To further boost the performance of this baseline we tune the size of the latent representation z, and we increase the number of layers by two, as well as the number of hidden channels to match the number of learnable weights. Complete details are available in the repository linked in the abstract. Fig. 2 shows the test NMSE on in-distribution CDL-D channels in the blind SNR setting with α = 0.4, making the inverse problem in (3) under-determined. From the top plot it can be seen that for a very low antenna spacing, the WGAN approach captures some of the structure of the channel, but the performance quickly saturates in the high SNR regime, and the same is true for Lasso. This effect is worsened when the antenna spacing is λ/2 and there is less structure, indicating that neither baseline uses a suitable prior. In contrast, the proposed approach exhibits almost a linear decay of the NMSE, thus being near-optimal and inline with the theoretical results in [8], all while never explicitly learning a prior. At the SNR value of 15 dB the NMSE is more than 12 dB lower than both baselines, but as results in the sequel Fig. 3. Estimation NMSE on 16 × 64 5G-NR CDL-C channels at λ/2 antenna spacing in the blind (left plot) and known (right plot) settings, respectively. Our proposed approach is never trained or tuned on any channel realization from the distribution of CDL-C channels and retains its usefulness as a channel estimator and outperforms Lasso when α < 1 in both cases.

Soft Bit Estimation
Soft LDPC Decoding Channel Estimation Fig. 4. Block diagram with the components of the simulated end-to-end scenario. x and P represent the data and pilot transmissions, respectively. The red block marks channel estimation, where our proposed approach is used and further feeds into the symbol equalization and maximum-likelihood soft bit estimation blocks. The coded bit error rate is measured between b andb. This implementation is part of our released source code.
show, this gap is reduced in end-to-end systems by channel coding. Fig. 3 shows test performance of our proposed approach and Lasso when tested on channels drawn from the CDL-C channel model at antenna spacing λ/2, in the blind and known SNR settings, respectively. Compared to the CDL-D channels used for training, CDL-C are significantly less sparse and structured differently. We omit WGAN from these results, since CDL-C channels are not in the range of the trained generative model and the high estimation NMSE makes them unusable out-of-distribution. In both cases, our method outperforms Lasso in the low SNR and few measurements regime, however performance still plateaus in the unknown SNR regime.

C. Generalization Performance
Finally, we evaluate coded end-to-end performance in a proof-of-concept system shown in Fig. 4 in order to evaluate the practical impact of different channel estimation algorithms. Our transmitter performs LDPC encoding with a rate of 1/2 and codeword size 648 bits, followed by a four stream QPSK modulation and random Gaussian precoding with unit gain to 64 antennas. The receives estimates the channels from random QPSK pilot transmissions P with a transmit power boosted by 20 dB relative to the data signal power, performs symbol equalization (including the known precoding matrix) to recover the four data streams, followed by soft bit estimation and LDPC decoding until convergence. We measure performance in coded bit error rate. Fig. 5 shows end-to-end performance of our method compared to Lasso. Note that this experiment is still using CDL-C channels with α = 1, thus is completely out-of-distribution for the proposed approach. We can see that, in the case of blind SNR, Lasso is highly sub-optimal in the low SNR regime of up to 0 dB, exactly where the channel coding cannot correct all errors in the transmission, thus leading to a 4 dB loss in end-toend performance. Thus, even though Lasso has a lower NMSE at high SNR values, the loss at low SNR is a determining factor for this scenario configuration and drives the performance. In contrast, the proposed approach is robust to unknown SNR and exhibits a 0.5 dB loss against ideally known channels. Even in the known SNR regime, our method still outperforms Lasso by 0.5 dB, even though no CDL-C channels were used for training or fine-tuning either approach.  5. End-to-end coded bit error performance in a comprehensive digital communication chain simulation of the proposed method, compared with ideal channel knowledge, in 5G-NR CDL-C channels. Our method is never trained or tuned on any channel realization from the distribution of CDL-C channels and offers excellent and robust end-to-end performance under realistic impairments.

V. DISCUSSION AND CONCLUSION
In this paper, we have introduced a deep diffusion model approach for MIMO wireless channel estimation. Our method performs posterior sampling using a pretrained diffusion model and includes a careful tuning of the inference hyperparameters. When tested on out-of-distribution channels, the proposed approach retains its approach One current disadvantage of the proposed approach is its expensive inference time. Posterior sampling for a single channel of size 16 × 64 takes approximately one minute on a graphical processing unit. On the other hand, Lasso is considerably faster, taking as little as 3 seconds per channel. However, both approaches are still too slow for real-world deployment, where channels are re-estimated as fast as every millisecond. Speeding up inference with diffusion models is an active area of research, with very recent work showing speed-ups of up to three orders of magnitude [23]. This is a promising future research direction for channel estimation, given the latency-sensitive nature of the task. Finally, in this paper we have only trained and tested diffusion models on channels of fixed size. Due to the fully-convolutional nature of the RefineNet, a future promising direction is to train on a mixture of sizes along with antenna spacings and fading distributions, to obtain a near-universal estimation approach.