Learning Bayes-Optimal Channel Estimation for Holographic MIMO in Unknown EM Environments

Our target is a Bayes-optimal channel estimator, $\mathbf{\hat{h}}=D(\mathbf{y})$, that minimizes the mean-square-error (MSE), i.e.,

where the expectation above is taken with respect to (w.r.t.) the unknown channel $\mathbf{h}$, while $p(\mathbf{h}|\mathbf{y})$ is the posterior density. Taking a derivative of the above equation w.r.t. $\hat{\mathbf{h}}$ and nulling it, we reach the Bayes-optimal, i.e., minimum MSE (MMSE), channel estimator, given by

where $p(\mathbf{h},\mathbf{y})$ is the joint density, and $p(\mathbf{y})$ is the measurement density obtained via marginalization, i.e.,

The last equality holds since the likelihood function $p(\mathbf{y}|\mathbf{h})\sim\mathcal{N}(\mathbf{h},\frac{1}{\rho}\mathbf{I})$ expresses $p(\mathbf{y})$ as a convolution between the prior distribution $p(\mathbf{h})$ and the i.i.d. Gaussian noise. Taking the derivative of both sides of (11) w.r.t. $\mathbf{y}$ gives

Dividing both sides of (12) w.r.t. $p(\mathbf{y})$ results in the following:

where the second equality holds owing to the Bayes’ theorem, and the last equality holds due to (10) and $\int p(\mathbf{h}|\mathbf{y})d\mathbf{h}=1$. By rearranging the terms and plugging in $\frac{\nabla_{\mathbf{y}}p(\mathbf{y})}{p(\mathbf{y})}=\nabla_{\mathbf{y}}\log p(\mathbf{y})$, we reach the foundation of the proposed algorithm:

where $\nabla_{\mathbf{y}}\log p(\mathbf{y})$ is called the Stein’s score function in statistics [12]. From (14), we notice that the Bayes-optimal MMSE channel estimator can be achieved solely based on the received pilot signals $\mathbf{y}$, without access to the prior distribution $p(\mathbf{h})$ or a supervised dataset of the ground-truth channels $\mathbf{h}$, which are unavailable in practice. With an efficient estimator of the score function $\nabla_{\mathbf{y}}\log p(\mathbf{y})$, the Bayes-optimal MMSE channel estimator can be computed in a closed form with extremely low complexity. In addition, since (14) holds regardless of the distribution of the HMIMO channel $\mathbf{h}$, one can construct the MMSE estimator in arbitrary EM environments without any assumptions on the scatterers or the array geometry.

It is noticed that two terms that should be obtained in (14), i.e., the received SNR $\rho$ and the measurement score function $S(\mathbf{y})$. In the following, we discuss on how to utilize machine learning tools to obtain an accurate estimation of them based solely on the measurement $\mathbf{y}$.

We discuss how to get the score function $\nabla_{\mathbf{y}}\log p(\mathbf{y})$. Given that a closed-form expression is intractable to acquire, we instead aim to achieve a parameterized function with a neural network, and discuss how to train it based on score matching. We first introduce the denoising auto-encoder (DAE) [12], the core of the training process, and explain how to utilize it to approximate the score function.

To obtain the score function $\nabla_{\mathbf{y}}\log p(\mathbf{y})$, the measurement $\mathbf{y}$ is treated as the target signal that the DAE should denoise. The general idea is to obtain the score function $\nabla_{\mathbf{y}}\log p(\mathbf{y})$ based on its analytical relationship with the DAE of $\mathbf{y}$, which will be established later in Theorem 1. We first construct a noisy version of the target signal $\mathbf{y}$ by manually adding some additive white Gaussian noise, $\sigma\mathbf{u}$, where $\mathbf{u}\sim\mathcal{N}(\mathbf{0},\mathbf{I}_{2N})$ and $\sigma$ controls the noise level⁴⁴4Note that the extra noise is only added during the training process. , and then train a DAE to denoise the manually added noise. The DAE, denoted by $R_{\bm{\theta}}(\cdot;\cdot)$, is trained by the $\ell_{2}$-loss function, i.e.,

The theorem below explains the relationship between the score function and the trained DAE.

The above theorem indicates that, for a sufficiently small $\sigma$, we can approximate the score function based on the DAE by $\nabla_{\mathbf{y}}\log p(\mathbf{y})\thickapprox\frac{R_{\bm{\theta}}(\mathbf{y};\sigma)-\mathbf{y}}{\sigma^{2}}$, assuming that parameter of the DAE, $\bm{\theta}$, is near-optimal, i.e., $\bm{\theta}\thickapprox\bm{\theta}^{*}$. Nevertheless, the approximation can be numerically unstable as the denominator, $\sigma^{2}$, is close to zero. To alleviate the problem, we improve the structure of the DAE and rescale the original loss function.

First, we consider a residual form of the DAE with a scaling factor. Specifically, let $R_{\bm{\theta}}(\mathbf{y};\sigma)=\sigma^{2}S_{\bm{\theta}}(\mathbf{y};\sigma)+\mathbf{y}$. Plugging it into (16), the score function is approximately equal to

when $\sigma\rightarrow 0$ holds. This reparameterization enables $S_{\bm{\theta}}(\mathbf{y};\sigma)$ to approximate the score function directly, thereby circumventing the need for division that may lead to numerical instability. Also, the residual link significantly enhances the capability of the DAE, since it can easily learn an identity mapping [13].

Second, since the variance $\sigma^{2}$ of the manually added noise is small, the gradient of the DAE loss function (15) can easily vanish to zero and may lead to difficulties in training. Hence, we rescale the loss function by a factor of $\frac{1}{\sigma}$ to safeguard the vanishing gradient problem, i.e.,

where (17) is plugged into the loss function.

We are interested in the region where $\sigma$ is sufficiently close to zero, in which case $S_{\bm{\theta}}(\mathbf{y};0)$ can be deemed to be equal to the score function $\nabla_{\mathbf{y}}\log p(\mathbf{y})$ according to (17). Nevertheless, directly training the network using a very small $\sigma$ is difficult since the SNR of the gradient signal decreases in a linear rate $\mathcal{O}(\sigma)$ with respect to $\sigma$, which introduces difficulty for the stochastic gradient descent [14]. To exploit the asymptotic optimality of the score function approximation when $\sigma\rightarrow 0$, we propose to simultaneously train the network $S_{\bm{\theta}}(\mathbf{y};\sigma)$ with varying $\sigma$ values, to handle various $\sigma$ levels and then naturally generalize to the desired region, i.e., $S_{\bm{\theta}}(\mathbf{y};0)$. To achieve the goal, we control the manually added noise by letting $\sigma$ follow a zero-mean Gaussian distribution $\sigma\sim\mathcal{N}(0,\text{$\xi^{2}$})$ and gradually anneal $\xi\in[\sigma_{\text{min}},\sigma_{\text{max}}]$ from a large value $\sigma_{\text{max}}$ to a small one $\sigma_{\text{min}}\thickapprox 0$ in each iteration. That is, we condition $S_{\bm{\theta}}(\mathbf{y};\sigma)$ on the manually added noise level $\sigma$ during training.

The proposed algorithm is shown in Algorithm 1. The DAE is trained using stochastic gradient descent for $Q$ epochs. In each epoch, we draw a random vector $\mathbf{u}$ and anneal $\xi$ in $\sigma\sim\mathcal{N}(0,\text{$\xi^{2}$})$ to control the extra noise level according to the current number of iterations $q$. Then, the DAE loss function $\mathcal{L}_{\text{DAE}}$ in (18) is minimized by stochastic optimization. Note that in the training process, nothing but a dataset of the received pilot signals $\mathbf{y}$ is necessary, which is readily available in practice. In the inference stage, one can apply formula (14) to compute the score-based MMSE estimator, in which the score function can be approximated by using $S_{\bm{\theta}}(\mathbf{y};0)$, i.e., setting $\sigma$ as zero, and the received SNR $\hat{\rho}$ can be estimated by the PCA-based algorithm in the next subsection.

For the neural architecture of the DAE $S_{\bm{\theta}}(\cdot;\cdot)$, we adopt a simplified UNet architecture [15]. Note that depending on the complexity budget, many other prevailing neural architectures could also be applied [6]. Other details of the training process are deferred to Section IV.

We propose a low-complexity PCA-based algorithm to estimate the received SNR $\rho$ in (14) based on a single instance of the pilot signals $\mathbf{y}$. Before further discussion, we stress that the received SNR estimation is executed only at the inference stage, not at the training stage, as shown in Algorithm 1.

The basic idea behind the PCA-based algorithm is the low-rankness of the spatial correlation matrix $\mathbf{R}$ of HMIMO due to the dense deployment of the antenna elements. Specifically, for isotropic scattering environments, the rank of the correlation matrix $\mathbf{R}_{\text{iso}}$ is approximately $\text{rank}(\mathbf{R}_{\text{iso}})\thickapprox\pi Nd_{a}^{2}/\lambda_{c}^{2}$ [17]. It decreases with the shrink of antenna spacing and the increase of the carrier frequency. For example, when $d_{a}=\lambda_{c}/4$, around $80\%$ of the eigenvalues of $\mathbf{R}_{\text{iso}}$ shrinks towards zero. The rank deficiency of $\mathbf{R}$ tends to be even more prominent in the case of non-isotropic scattering environments [4].

Similar to Fig. 2 in our previous work [16], we decompose multiple virtual subarray channels (VSCs) from the HMIMO channel by a sliding window. Specifically, we reshape the real-valued HMIMO channel $\mathbf{h}\in\mathbb{R}^{2N\times 1}$ into a tensor form $\mathbf{H}\in\mathbb{R}^{\sqrt{N}\times\sqrt{N}\times 2}$. We then decompose $\mathbf{H}$ into $s=(\sqrt{N}-\sqrt{d}+1)^{2}$ VSC tensors using a sliding window of size $\mathbb{R}^{\sqrt{d}\times\sqrt{d}\times 2}$, and then reshape them back into the vector form to obtain a set of VSCs, denoted by $\{\mathbf{h}_{t}\in\mathbb{R}^{2d\times 1}\}_{t=1}^{s}$. Similarly, the received pilot signals and the noise could also be decomposed as $\{\mathbf{y}_{t}\in\mathbb{R}^{2d\times 1}\}_{t=1}^{s}$ and $\{\mathbf{n}_{t}\in\mathbb{R}^{2d\times 1}\}_{t=1}^{s}$, and should satisfy

Due to the low-rankness of the spatial correlation matrix $\mathbf{R}$, the decomposed VSCs $\mathbf{h}_{t}$ should also lie in a low-dimensional subspace. In Fig. 2, we plot the eigenvalues of the covariance matrices of the decomposed VSCs $\{\mathbf{h}_{t}\in\mathbb{R}^{2d\times 1}\}_{t=1}^{s}$ from an HMIMO channel in isotropic scattering environments and their corresponding received pilots $\{\mathbf{y}_{t}\in\mathbb{R}^{2d\times 1}\}_{t=1}^{s}$, respectively, with a reference line marking the inverse of the received SNR $\frac{1}{\rho}$. According to the figure, the eigenvalues of the covariance of $\mathbf{h}_{t}$ quickly shrinks to zero with about 30 principal dimensions. The zero eigenvalues correspond to the redundant dimensions. In contrast, we observe that the eigenvalues of the covariance of $\mathbf{y}_{t}$ are concentrated around $\frac{1}{\rho}$ in the redundant dimensions. This example suggests that it is indeed possible to estimate the received SNR based on the redundant eigenvalues. Rigorously, these eigenvalues could be proved to follow a Gaussian distribution $\mathcal{N}(\frac{1}{\rho},\frac{2}{s\rho^{2}})$ [16]. Plus, the principal and redundant eigenvalues could be separated by an iterative process. Hence, we could leverage [16, Algorithm 1] to accurately estimate the received SNR. Detailed setups are discussed in Section IV.

The inference complexity of the proposed algorithm consists of the computation of $S_{\bm{\theta}}(\mathbf{y};0)$ and the PCA-based estimation of $\rho$. The former depends upon the specific neural architecture of $S_{\bm{\theta}}(\cdot,\cdot)$, and costs a constant complexity, denoted by $p$, once the network is trained. The computational complexity for the latter, as analyzed in [16], is given by $\text{$\mathcal{O}$}(Nd^{2}+d^{3})$, where $d$, the size of the sliding window, is usually quite small. Hence, the overall complexity is $\text{$\mathcal{O}$}(Nd^{2}+d^{3}+p)$, which scales only linearly with respect to the number of antennas $N$.

In practice, the fluctuation of the received SNR may not be frequent. In such a case, it is not a necessity to estimate the received SNR for every instance of the received pilot signals $\mathbf{y}$. Therefore, the actual complexity of the proposed score-based algorithm is within the range of $\mathcal{O}(p)$ and $\text{$\mathcal{O}$}(Nd^{2}+d^{3}+p)$, which is extremely efficient given near-optimal performance⁵⁵5Most previous works on channel estimation assume that the received SNR is perfectly known. In this case, the complexity reduces to $\mathcal{O}(p)$. . By sharp contrast, the (oracle) MMSE estimator requires time-consuming matrix inversion, which is as complex as $\mathcal{O}(N^{3})$. In Section IV, we provide a running time complexity to offer a straightforward comparison.

We consider a typical HMIMO system setup with $N=1024$ and $d_{a}=\lambda_{c}/4$. In the training stage, the hyper-parameters of the proposed score-based algorithm are chosen as $\gamma=0.001$, $\sigma^{\text{min}}=0.001$, $\sigma^{\text{max}}=0.1$, $Q=100$, and $M=10\text{,}000$. Also, the learning rate $\gamma$ is decayed by half after every 25 epochs. In the inference stage, the size of the sliding window is set as $7\times 7$. The performance discussed below is all averaged over a held-out dataset consisting of $L=10,000$ testing samples.

For isotropic scattering, the covariance $\mathbf{R}_{\text{iso}}$ is given via (7). In the non-isotropic case, we follow [4] and also construct the covariance $\mathbf{R}$ by truncating the leading $\nicefrac{{\text{rank}(\mathbf{R}_{\text{iso}})}}{{8}}$ eigenvalues of $\mathbf{R}_{\text{iso}}$. Mathematically, the non-isotropic covariance could be expressed as $\mathbf{R}=\mathbf{V}\boldsymbol{\Lambda}_{\text{trunc}}\mathbf{V}^{H}$, where $\mathbf{V}$ is a matrix consisting of the eigenvectors of $\mathbf{R}_{\text{iso}}$, while $\boldsymbol{\Lambda}_{\text{trunc}}$ is a diagonal matrix that contains the eigenvalues of $\mathbf{R}_{\text{iso}}$ arranged in descending order. To truncate the matrix, we select only the first $\nicefrac{{\text{rank}(\mathbf{R}_{\text{iso}})}}{{8}}$ eigenvalues and eigenvectors of $\mathbf{R}_{\text{iso}}$.

The accuracy of the received SNR estimation can influence the performance of the proposed score-based estimator. Hence, different from previous works that assume perfect knowledge of the received SNR, we propose a practical PCA-based means to estimate it and examine its performance. In Table I, we list the estimation accuracy under different SNRs. We present the performance in terms of the inverse of the SNR, i.e., $\frac{1}{\rho}$, since it is used in (14). The bias, the standard deviation (std), and the root MSE (RMSE) are given by $\mathbb{E}[|\frac{1}{\rho}-\mathbb{E}[(\frac{1}{\hat{\rho}})]|]$, $\sqrt{\mathbb{E}[((\frac{1}{\hat{\rho}})-\mathbb{E}[(\frac{1}{\hat{\rho}})])^{2}]}$, and $\sqrt{\mathbb{E}[(\frac{1}{\rho}-(\frac{1}{\hat{\rho}}))^{2}]}$, respectively. The estimator’s accuracy and robustness are reflected with the bias and std, while the RMSE provides an assessment of its overall performance. From Table I, we observe that the performance of the proposed PCA-based method is both highly accurate and robust, and outperforms the sparsity-based median absolute deviation (MAD) estimator in [18], as HMIMO channels are not exactly sparse. Also, the performance is close to the oracle bound which assumes perfect knowledge of the channel $\mathbf{h}$, and estimate $\rho$ directly from $\mathbf{y-\mathbf{h}}$ [16]. Later, in Section IV-D, we will illustrate that the proposed score-based channel estimator is robust to SNR estimation errors.

We compare the proposed score-based estimator with three benchmarks, including LS, sample MMSE, and oracle MMSE. Here, oracle MMSE refers to the MMSE estimator with perfect knowledge of both the covariance $\mathbf{R}$ and the received SNR $\rho$, which is the Bayesian performance bound, given by

This is difficult, if not impossible, to acquire in practice since $\mathbf{R}$ contains $N^{2}$ entries and is prohibitive to estimate when $N$ is particularly large in an HMIMO system. The sample MMSE method utilizes the same equation as (20), but replaces the true covariance $\mathbf{R}$ with an estimated one $\mathbf{R}_{\text{sample}}$ based upon the testing samples, given by $\mathbf{R}_{\text{sample}}\triangleq\frac{1}{L}\sum_{l=1}^{L}\mathbf{y}_{l}\mathbf{y}_{l}^{H}-\frac{1}{\rho}\mathbf{I}$ [8], where $L$ is the number of testing samples and $\mathbf{y}_{l}$ denotes the $l$-th sample of the received pilot signals in the testing dataset. We also utilize the perfect received SNR $\rho$ in sample MMSE.

In Fig. 3(a), we present the NMSE as a function of the SNR $\rho$ in isotropic scattering environments. It is illustrated that the proposed score-based algorithm significantly outperforms the LS and the sample MMSE estimator, and achieves almost the same NMSE as the oracle MMSE bound at every SNR level. Note that the oracle MMSE method utilizes the true covariance and received SNR, but the proposed method requires neither.

In Fig. 3(b), we compare the NMSE in non-isotropic environments, in which the rank of the covariance matrix $\mathbf{R}$ is further reduced. As a result, the performance gap between the LS and the oracle MMSE estimator is enlarged as the rank deficiency becomes more significant. Nevertheless, similar to the isotropic case, the proposed score-based estimator exhibits almost the same performance as the oracle bound and significantly outperforms the sample MMSE method, illustrating its effectiveness in different scattering environments.

In Fig. 3(c), we provide discussions on how the accuracy for the received SNR estimation will influence the performance of the proposed score-based estimator. In the simulations, the true value of the received SNR is set as 10 dB, i.e., $\frac{1}{\rho}=0.1$. We vary the adopted values of $(\frac{1}{\hat{\rho}})$ in (14) within $(\frac{1}{\hat{\rho}})\in[0,0.24]$, and plot the corresponding NMSE performance curve in blue. Particularly, we utilize red and green dots to denote the NMSE achieved with the estimated and true SNR values, respectively. The LS and the oracle MMSE algorithms are presented as the performance upper and lower bounds. It is observed that even when an inexact received SNR is adopted, the performance of the score-based algorithm is still quite robust and significantly outperforms the LS method. Also, the estimated received SNR by the PCA-based method is accurate enough to offer a near-optimal performance, even in unknown EM environments.

We introduce the CPU running time of the proposed algorithm. For the considered setups, the proposed score-based estimator takes as low as 3 ms on an Intel Core i7-9750H CPU, which is much shorter than that of the oracle MMSE method involving high-dimensional matrix inverse (requiring around 70 ms). The high efficiency and the Bayes-optimal performance in unknown EM environments thus make the proposed score-based algorithm an ideal candidate in practice.