SlowGuess/NeurIPS_2023_OCR huggingface dataset:
3.3 Relation to Other Works
BASIS separation. Jayaram and Thickstun [5] introduced the BASIS separation algorithm that leverages the score from a generative model to perform source separation using annealed Langevin dynamics. Unlike our method, the BASIS algorithm relies on a specially tuned noise schedule for separation that is distinct from the diffusion model training noise schedule. We circumvent the challenges associated with tuning such a schedule and instead re-use the pre-determined training noise schedule in a randomized fashion. As such, the only parameter that we tune is the learning rate. Our results (§5.3) show that we outperform BASIS in the context of RF source separation.
Score Distillation Sampling. Our work can be viewed as an extension of the recently proposed Dreamfusion architecture [15], which uses pre-trained image diffusion models as critics to guide the generation of novel 3D visual content. Given a generated sample,  $g(\pmb{\theta})$ , Score Distillation Sampling (SDS) updates the 3D object realization using gradient descent with a gradient given by,
$$
\nabla_ {\theta} \mathcal {L} _ {\mathrm {S D S}} (\phi , \tilde {\mathbf {s}} = g (\boldsymbol {\theta})) = \mathbb {E} _ {\mathrm {t}, \mathrm {z}} [ w (t) \left(r _ {\phi} \left(\sqrt {\alpha_ {t}} g (\boldsymbol {\theta}) + \sqrt {1 - \alpha_ {t}} \mathbf {z}, t\right) - \mathbf {z}\right) ], \tag {13}
$$
where  $w(t)$  is a scaling related to the noise variance. Our updates contain the same gradient terms in (13) and hence it can be viewed as a multi-source extension of SDS, shedding light on its use in applications to problems beyond sampling with interactions between numerous individual priors.
Diffusion models as priors. More recently diffusion priors have been used in inverse problems such as MRI reconstruction [29, 30], image restoration/colorization [23, 31-34] and for developing general inversion methods [32, 35]. A majority of these works solve a MAP problem with domain-specific likelihood constraints, thus resembling our formulation. However, they either use annealed Langevin dynamics or a noise contracting reverse diffusion process to sample from the posterior distribution, necessitating the tuning of a special annealing schedule. In our work, we demonstrate convergence to the same modes using a simpler scheme based on randomized levels of Gaussian smoothing (see §4). The  $\alpha$ -posterior parameter can be interpreted as a weighting of the regularization term, thus shedding light on the potential applicability of our method to general inverse problems in the future.
4 Characterization of  $\alpha$ -RGS
In this section, we characterize the behavior of our objective function through a simple yet intuitive example. We first present a sufficient condition for perfect signal separation in the context of discrete sources that follow the signal generation model in (1).
Proposition 1. Let  $s(t)$  and  $b(t)$  be two sources following (1) with underlying symbols  $c_p^s$  and  $c_p^b$  respectively. Assume that the symbols are obtained as,
$$
c _ {p} ^ {s} = f \left(\left\{u _ {i} \right\} _ {i = p} ^ {L + p}\right) \quad a n d \quad c _ {p} ^ {b} = h \left(\left\{v _ {i} \right\} _ {i = p} ^ {L + p}\right), \tag {14}
$$
where  $f: \mathcal{U}^L \to \mathbb{C}$  and  $h: \mathcal{V}^L \to \mathbb{C}$  are mappings from a sequence of length  $L$  over the discrete alphabets  $\mathcal{U}$  and  $\mathcal{V}$  respectively. If the mapping between the discrete representation and the symbol representation of the sources is unique, perfect recovery is possible.
<sup>3</sup>In digital communications, the discrete alphabet is typically the binary alphabet, representing the underlying bits that need to be communicated, and the symbols refer to the digital constellation (see §5.1)
Figure 2: The single-carrier digital modulation pipeline with an intelligent decoder that performs a pre-processing stage of source separation. Illustrated is an example with a QPSK constellation and a root-raised cosine (RRC) pulse shaping function.
We now argue that under such conditions, our method in Algorithm 1 asymptotically approaches a local extremum corresponding to the following loss function,
$$
\mathcal {L} (\boldsymbol {\theta}) := - \mathbb {E} _ {\mathbf {t}, \mathbf {z} _ {s}} \left[ \log p _ {\tilde {\mathbf {s}} _ {t}} \left(\tilde {\mathbf {s}} _ {t} (\boldsymbol {\theta})\right) \right] - \omega \mathbb {E} _ {\mathbf {u}, \mathbf {z} _ {u}} \left[ \log p _ {\tilde {\mathbf {b}} _ {u}} \left(\tilde {\mathbf {b}} _ {u} (\boldsymbol {\theta}, \mathbf {y})\right) \right]. \tag {15}
$$
To see this, let  $s$  and  $b$  be two discrete sources with equiprobable modes at  $\{-1, +1\}$  and  $\{\pm 6/\sqrt{20}, \pm 2/\sqrt{20}\}$  respectively. Figure 1 shows the contours of the negative log joint probability, with the sources augmented with an infinitesimally small amount of Gaussian noise (see §3.2.1). An observed mixture  $y$ , adds a linear constraint in the  $(s, b)$  plane. The goal of Algorithm 1 is to pick out the constraint-satisfying mode at  $(1, -2/\sqrt{20})$ . If  $\omega = 1$ , as shown in the middle plot, and given a poor initialization  $\theta^{(0)}$  closer to  $-1$ , the gradients in this region may prevent the estimate from escaping the suboptimal local minimum. If instead, we use  $\omega = \kappa^2 > 1$  ( $\kappa = 15.85$ ), as shown in rightmost plot, the optimization landscape is better conditioned in the same (absolute) neighborhood around  $\theta = -1$  with the mode at  $+1$  much more accentuated. Specifically, if the algorithm is now initialized at  $\theta^{(0)} = -1$ , after enough iterations, gradient steps at larger noise levels would lead the solution towards the direction of  $\theta = 1$ . Thus, on average (black curve), after enough iterations, the solution will approach  $+1$ . In constraint, Langevin-dynamics-based approaches without the  $\alpha$ -posterior weighting [5], if initialized poorly, could potentially get stuck at local extrema due to sharpening of the optimizing landscape as the level of noise decreases.
The estimate of the discrete source  $\mathbf{s}$  can be obtained by mapping the solution returned by Algorithm 1 to the closest point (in the Euclidean sense) in the discrete alphabet  $S$ . This relies on the extremum  $\theta^{*}$  of (15) being sufficiently close to the desired mode of  $p_{\mathrm{s|y}}(\mathbf{s|\mathbf{y}};\omega)$ , so that  $\theta^{*}$  can be mapped to the correct point in  $S$  with high probability. In the context of the above example, a key observation in achieving this desired behavior is that the mode at  $(1, -2/\sqrt{20})$  is still prominent at large noise levels. Thus, randomizing across different noise levels helps balance the exploration between modes and the resolution of the estimate. This mainly requires a suitable noise level range (i.e., a lower bound on  $\alpha_{T}$ ), to ensure that the modes are sufficiently resolved. In our experiments, we show that no additional tuning is required and that the training noise levels from the pre-trained diffusion models can be re-used, provided that the models have learned the source's structure sufficiently well.
A more detailed characterization of our algorithm's convergence to the modes can be found in Appendix B. We underscore that in this work we focus on general signal types and mixtures, where perfect separability is not guaranteed and hence performance bounds are not analytically tractable.
5 Applications and Experiments
5.1 Digital Communication Signals
Digital communication pipeline. At a high level, digital communications deals with the transmission of bits by modulating a so-called "carrier signal". Groups of bits, from which the underlying discreteness of these sources originate, are first mapped to symbols  $c_{p} \in \mathbb{C}$  via the digital constellation—a mapping between groups of bits and a finite set of complex-valued symbols. These symbols are subsequently aggregated via (1) to form a complex-valued continuous waveform. In the RF domain,  $g(\cdot)$  is known as the pulse shaping filter and helps limit the signal's bandwidth [36, Sec 4.4.3]. The constellation is chosen (among other considerations) by the number of bits modulated simultaneously. Common schemes include modulating two bits at a time (Quadrature Phase Shift Keying, or QPSK), or one bit at a time (Binary Phase Shift Keying, or BPSK). Figure 2 illustrates a representative modulation pipeline. To recover the bits at the receiver, one may adopt matched filtering (MF) [18, Sec 5.8] before the estimation of the underlying symbols, and thereafter decode them back to bits. For commonly used pulse shaping functions, such as the root-raised cosine
(RRC) shown in Figure 2, the matched filter and pulse shaping filter coincide. We refer readers to [18, 36, 37] for a more thorough exposition of the topic.
Interference mitigation. Mitigating co-channel interference is a challenging problem, especially in heterogeneous networks [9-11]. If the system parameters and signal generation model are known ahead of time, one can leverage this knowledge to devise hand-crafted priors. However, one often deals with interference from more complicated wireless sources, for which the signal model is unknown. In such scenarios, data-driven methods that learn priors from background recordings can be useful. Our results demonstrate that  $\alpha$ -RGS, which leverages these priors, sets a new state-of-the-art in SCSS for RF mixtures. We envision that our SCSS solution could help mitigate such interference prior to decoding the signal, as shown in the right hand (receiver) side of Figure 2.
Deep learning for RF systems. Recently deep learning methods have demonstrated the potential to reap significant gains over handcrafted model-based methods in RF applications [38, 39]. Some works have studied the problem of symbol detection [40, 41], where they assume that the channel is stationary. Other works, such as DeepSIC [42], use deep learning for interference cancellation in the multi-user setting within the same channel. In contrast, we deal with the more challenging setting of non-stationary interference, thereby requiring efficient exploitation of intricate temporal structures. While the latter works consider the superposition of independent and identically distributed sources (same technology), we assume unknown additive interference (cross technology), a hard problem to solve with naive decoding methods in the absence of explicit prior knowledge about the interference. Our problem formulation is closer to recent work in [43]. However, they learn an end-to-end estimate of the signal from paired data samples. As motivated in §1, we instead assume restricted access to joint data, with a focus on capturing properties of the components through independent priors.
We now detail the setup for training diffusion models on RF signals. We subsequently explain how these models are used in the implementation of our method at inference time (i.e, for separation).
RF SCSS formulation. We are interested in recovering  $\mathbf{s}$ , the signal of interest (SOI), from a mixture  $\mathbf{y} = \mathbf{s} + \kappa \mathbf{b}$ , where  $\mathbf{b}$  is assumed to be a co-channel interference with unknown system parameters. We evaluate performance using two metrics—i) the MSE, that measures the distortion between the estimated SOI and the ground truth; and ii) the BER of the decoded discrete representation, which is obtained from the estimated SOI by extracting the underlying bits (see Appendix C). The latter measure is particular to digital communication signals as it captures the fidelity of the estimated representation that is only partially reflected in the MSE criterion.
Datasets. We train diffusion models on different RF datasets —i) synthetic QPSK signals with RRC pulse shaping (see §5.1), ii) synthetic OFDM signals (BPSK and QPSK) with structure similar to IEEE 802.11 WiFi signals; and iii) signals corresponding to "CommSignal2" from the RF Challenge [44], which contains datasets of over-the-air recorded signals. All synthetic datasets were created using the NVIDIA Sionna toolkit [45]. All datasets contain between  $150\mathrm{k}$  -  $500\mathrm{k}$  samples and we use a 90-10 train-validation split during training. Additional details are in Appendix C.
Diffusion model training. We adopt the Diffwave [24] architecture for our experiments, with a minor changes (see Appendix D) to accommodate the complex-valued nature of our signals. Our models are trained in the waveform domain on inputs of length 2560 with the real and imaginary components concatenated in the channel dimension. We train unconditional diffusion models and assume no access to knowledge about the signal generation model. We use noise variance levels in the range  $(1\mathrm{e} - 4,0.72)$  discretized into 50 levels. We train for  $500\mathrm{k}$  steps with early stopping on  $2\times$  NVIDIA 3090 GPUs. Detailed hyperparameters for our training setup are provided in Appendix D.
Source separation setup. We consider three different mixtures in our experiments, all using an RRC-QPSK signal as the SOI s. The interference signal b is one of OFDM (BPSK), OFDM (QPSK) or a windowed recording from the CommSignal2 dataset. Our proposed algorithm uses  $\omega = \kappa^2$  and is initialized with the MF solution given the mixture y (see §5.1). Note that  $\kappa$  can be equivalently described as the signal to interference ratio ( $SIR := 1 / \kappa^2 := \mathbb{E}_s\left[\|s\|_2^2\right] / \mathbb{E}_b\left[\|\kappa b\|_2^2\right]$ ).
Figure 3: Left: Source separation results for a mixture with RRC-QPSK SOI and OFDM (QPSK) interference. All curves are obtained by averaging 400 different mixtures per SIR and using  $\omega = \kappa^2$ . Our proposed method significantly outperforms traditional and learning-based baselines (BASIS) in terms of BER and MSE across all noise levels. Middle: Similar comparisons only in BER for mixtures with OFDM (BPSK) and CommSignal2 interference, respectively. BER is slightly higher for the CommSignal2 mixture since it contains background noise that is amplified for large  $\kappa$ . The black dotted line in the bottom figure denotes the (presumed) BER lower bound assuming the background noise is an additive white Gaussian noise. Right: BER and MSE versus  $\omega / \kappa^2$  for different SIR levels. Clearly a good choice in this setting, in the sense of minimum BER and MSE, is  $\omega = \kappa^2$ .
We assume that  $\kappa$  is known<sup>4</sup> and use  $N = 20,000$  with a cosine annealing learning rate schedule [46]. The OFDM mixtures use  $(\eta_{\mathrm{max}},\eta_{\mathrm{min}}) = (5\mathrm{e} - 3,1\mathrm{e} - 6)$  and the CommSignal2 mixture uses  $(\eta_{\mathrm{max}},\eta_{\mathrm{min}}) = (2\mathrm{e} - 3,1\mathrm{e} - 6)$ . Importantly, we re-use the training noise levels from the diffusion models and randomize over all but the smallest noise level resulting in a noise variance range of  $(1 - \alpha_{1} = 1.2\mathrm{e} - 3,1 - \alpha_{T} = 0.72)$  discretized into  $T = 49$  levels. We test performance across SIR levels ranging from  $-24$  dB to  $-3$  dB ("strong interference" regime), by changing the value of  $\kappa$  in the mixture. Each set of separation experiments was conducted on a single NVIDIA V100 GPU.
Baselines. We compare our proposed method against baselines that also leverage independent statistical or structural priors over the sources. The simplest baseline, which nevertheless is still commonly used in most communication systems, is the matched filtering solution, which treats the interference as white Gaussian noise. The linear minimum mean square error (LMMSE) solution, a commonly used technique for signal estimation, is another baseline that leverages (up to) second-order statistics of the underlying source distributions. More details on these methods are in Appendix E.
We also compare with the BASIS separation algorithm (see §3.3), which is the closest learning-based method that resembles our problem formulation. Applying their method as is yielded poor results, and hence we modified the original hyperparameters by tuning the annealing schedule to the best of our abilities for a fair comparison. For more details on this, refer to Appendix F.
To study the fidelity of our learned score models, we derive the analytical score function of the QPSK SOI in the symbol domain (i.e., before pulse shaping). We use this analytical score as another comparison to demonstrate the performance of our method if the score was known perfectly. This formulation is closer to a learning-based interference mitigation setting, where we assume perfect knowledge about the SOI model, and rely on a learned interference model.
5.3 Source Separation Results
Figure 3 shows the source separation results for the three different mixture types in §5.2, obtained by averaging 100 independent trials per SIR. Our proposed method (analytical or trained SOI score) significantly outperforms MF, LMMSE and BASIS in terms of both BER and MSE. As expected, our best results are obtained by leveraging the prior knowledge, in the form of the analytical score
4Many communication systems have power constraints and equalization capabilities, and with the endowment of such knowledge it is possible to estimate the signal to interference ratio (SIR) within reasonable margin.
for the SOI. Nevertheless, our learned SOI score can nearly mimic this performance, and despite the slight degradation still outperforms all baselines in terms of BER as well. It should be noted that CommSignal2 contains small amount of background noise, which is amplified at low SIR (high  $\kappa$ ). This noise constrains the minimal achievable BER even under the assumption of only having the residual AWGN present, illustrated by the black dotted line at the bottom middle plot in Figure 3. We also outperform the baselines in terms of MSE across all SIR levels, an example of which is shown at the bottom left of Figure 3. More detailed results are included in Appendix G.
Characterizing the choice of  $\omega$ . We numerically verify that  $\omega = \kappa^2$  is a good choice for the  $\alpha$ -posterior term. To this end, we find a suitable order of magnitude for  $\omega$ , by varying  $\omega / \kappa^2$  between  $10^{-2}$  and  $10^2$  and computing the resulting BER and MSE for four representative SIR levels using a held-out hyperparameter validation set of 400 mixtures. The right panel of Figure 3 shows that  $\omega = \kappa^2$  achieves the minimum BER and MSE on average. More characterizations are in Appendix G.
In this work, we propose  $\alpha$ -RGS, a method that extends MAP estimation with an  $\alpha$ -posterior across randomized levels of Gaussian smoothing, which stems from a new objective function, whose extrema points correspond to the modes of the underlying discrete distribution of interest. Our method relies only on pre-trained diffusion models as priors via a simple randomized algorithm that does not require cumbersome tuning of a special annealing schedule, as is done in existing Langevin-dynamics-based works. Through simple analytical illustrations, we demonstrate the favorable mode-preserving nature of our objective, and show that  $\alpha$ -RGS is a generalized extension of the recently proposed SDS loss. Experiments on RF sources demonstrate superior separation performance, with gains of up to  $95\%$  in terms of both BER and MSE over classical and existing score-based SCSS methods.
Limitations and future work. We believe there is a potential to use our algorithm to develop a general toolkit of algorithms for source separation and inverse problems. However, we do list some limitations of our method that will be a focus of future work: 1) inference time ( $\sim$  5 min per mixture) requires speed-up in real-time systems; 2) studying settings with more than two mixture components or other mixture models; 3) developing a systematic approach to choose  $\omega$  for practitioners; and 4) addressing the robustness of our method for applications to signals with new properties, e.g., gravitational waves. We further elaborate on the broader impact of our work in Appendix A.
Research was supported, in part, by the United States Air Force Research Laboratory and the United States Air Force Artificial Intelligence Accelerator under Cooperative Agreement Number FA8750-19-2-1000. The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the official policies, either expressed or implied, of the United States Air Force or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation herein.
The authors acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing HPC resources that have contributed to the research results reported within this paper.
Alejandro Lancho has received funding from the European Union's Horizon 2020 Research and Innovation Programme under the Marie Sklodowska-Curie Grant Agreement No. 101024432. This work is also supported, in part, by the National Science Foundation (NSF) under Grant No. CCF-2131115.
[38] Yonina C Eldar, Andrea Goldsmith, Deniz Gündüz, and H Vincent Poor. Machine learning and wireless communications. Cambridge University Press, 2022.
[39] Taiwo Oyedare, Vijay K. Shah, Daniel J. Jakubisin, and Jeffrey H. Reed. Interference suppression using deep learning: Current approaches and open challenges. IEEE Access, 10: 66238-66266, 2022. doi: 10.1109/ACCESS.2022.3185124.
[40] Neev Samuel, Tzvi Diskin, and Ami Wiesel. Learning to detect. IEEE Transactions on Signal Processing, 67(10):2554-2564, 2019. doi: 10.1109/TSP.2019.2899805.
[41] Nir Shlezinger, Nariman Farsad, Yonina C. Eldar, and Andrea J. Goldsmith. Viterbinet: A deep learning based viterbi algorithm for symbol detection. IEEE Transactions on Wireless Communications, 19(5):3319-3331, 2020. doi: 10.1109/TWC.2020.2972352.
[42] Nir Shlezinger, Rong Fu, and Yonina C. Eldar. Deep soft interference cancellation for mimo detection. In 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 8881-8885, 2020. doi: 10.1109/ICASSP40776.2020.9054732.
[43] Gary CF Lee, Amir Weiss, Alejandro Lancho, Jennifer Tang, Yuheng Bu, Yury Polyanskiy, and Gregory W Wornell. Exploiting temporal structures of cyclostationary signals for data-driven single-channel source separation. In 2022 IEEE 32nd International Workshop on Machine Learning for Signal Processing (MLSP), pages 1-6. IEEE, 2022.
[44] MIT RLE. Single-Channel RF Challenge. URL https://rfchallenge.mit.edu/challenge-1/. Accessed 2022-11-07.
[45] Jakob Hoydis, Sebastian Cammerer, Fayçal Aït Aoudia, Avinash Vem, Nikolaus Binder, Guillermo Marcus, and Alexander Keller. Sionna: An open-source library for next-generation physical layer research. arXiv preprint, Mar. 2022.
[46] Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017.
[47] Herbert E. Robbins. An empirical bayes approach to statistics. 1956.
[48] Paul Bromiley. Products and convolutions of Gaussian probability density functions. Tina-Vision Memo, 3(4):1, 2003.
[49] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. GANs trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[50] David Tse and Pramod Viswanath. Fundamentals of wireless communication. Cambridge University Press, 2005.
[51] Vladimir Kulikov, Shahar Yadin, Matan Kleiner, and Tomer Michaeli. SinDDM: A single image denoising diffusion model. arXiv preprint arXiv:2211.16582, 2022.
A.1 Regarding the Proposed  $\alpha$ -RGS Method