SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
p _ {\mathbf {s}} (s) = \sum_ {i = 1} ^ {K} \lambda_ {i} \mathcal {N} \left(s; \mu_ {i}, \sigma_ {i} ^ {2}\right), \quad \sum_ {i = 1} ^ {K} \lambda_ {i} = 1. \tag {54}
$$
For the Gaussian smoothing model in (17), the score at timestep  $t$  is,
$$
\nabla_ {\tilde {s} _ {t} (s)} \log p _ {\tilde {s} _ {t}} (\tilde {s} _ {t} (s)) = \frac {1}{Z (\tilde {s} _ {t} (s))} \sum_ {i = 1} ^ {K} \lambda_ {i} c _ {i} (\tilde {s} _ {t} (s)) \left(\frac {\gamma_ {t} \mu_ {i} - \tilde {s} _ {t} (s)}{\gamma_ {t} ^ {2} \sigma_ {i} ^ {2} + \sigma_ {t} ^ {2}}\right), \tag {55}
$$
$$
c _ {i} \left(\tilde {s} _ {t} (s)\right) := \mathcal {N} \left(\tilde {s} _ {t} (s); \gamma_ {t} \mu_ {i}, \gamma_ {t} ^ {2} \sigma_ {i} ^ {2} + \sigma_ {t} ^ {2}\right), \tag {56}
$$
$$
Z \left(\tilde {s} _ {t} (s)\right) := \sum_ {i = 1} ^ {K} \lambda_ {i} c _ {i} \left(\tilde {s} _ {t} (s)\right). \tag {57}
$$
Similar to the BPSK case, since a closed-form expression for the solution to (18) cannot be obtained, we simulate (16) for a four-component GMM with two large equiprobable modes and two smaller modes. As shown in the colored curves on the right in Figure 5, we again notice similar behavior where the modes are sharper at lower noise levels, with the sharpness decreasing at large noise levels. Hence, gradient-based methods can leverage the randomly chosen larger noise levels to move between modes and use the smaller noise levels to resolve the estimate. On average, the minima of the loss approach the modes of the original GMM source, as shown by the black curve on the right in Figure 5.
In this section, we particularize the signal models introduced in §2.1 to RF signals. We first consider a single-carrier signal modulating, e.g., QPSK symbols, which can be expressed as,
$$
s (t) = \sum_ {p = - \infty} ^ {\infty} c _ {p, \ell} g \left(t - p T _ {s}\right). \tag {58}
$$
The second form corresponds to multi-carrier signals, such as OFDM signals, where multiple symbols are modulated in parallel,
$$
b (t) = \sum_ {p = - \infty} ^ {\infty} \sum_ {\ell = 0} ^ {L - 1} c _ {p, \ell} \exp \left\{j 2 \pi \ell t / L \right\}. \tag {59}
$$
Both types of signals can be grouped in a single family of signals represented as,
$$
u (t) = \sum_ {p = - \infty} ^ {\infty} \sum_ {\ell = 0} ^ {L - 1} c _ {p, \ell} g (t - p T _ {s}, \ell) \exp \{j 2 \pi \ell t / L \}. \tag {60}
$$
We use  $s(t)$  to represent the signal of interest (SOI) and  $\mathbf{s}$  as the vector representation for a collection of  $N$  consecutive samples thereof. Similarly, the interference is represented by  $b(t)$  or  $\mathbf{b}$ .
Before we describe the construction of our datasets, we will provide some terminology that we use in our description. This terminology is meant to supplement the details already provided in §5.1 in the main paper.
Symbol constellation: A digital constellation describes the underlying scheme for mapping bits to symbols. Typically used digital constellations (the mapping between bits and symbols) include binary phase shift keying (BPSK) and quadrature phase shift keying (QPSK). Specific BPSK and QPSK examples are depicted in Figure 6.
Pulse/Symbol shaping function: This corresponds to  $g(\cdot, \cdot)$  in (60) and is meant to help limit the bandwidth of the signal. Each symbol  $c_{p,l}$  is multiplied with a pulse shaping function with different time offsets and a main lobe width  $T_s$ , helping smoothen the signal and remove high frequencies. In this work, we will make use of the root-raised cosine (RRC) pulse shaping for our single-carrier QPSK SOI. On the other hand, OFDM signals do not use an explicit pulse shaping function (as shown in (59)), i.e.,  $g(t, \cdot)$  corresponds to a rectangular function in (60).
(a) BPSK constellation.
(b) QPSK constellation.
Figure 6: Discrete constellations for a BPSK and QPSK signal. Application of the RRC filter in the time domain leads to interpolation between constellation point as shown in (c). Phase offsets in the waveform domain can be corrected by looking at the rotated constellation as in (d).
(c) Effect of RRC filtering.
(d) Effect of  $45^{\circ}$  phase.
Total number of subcarriers: This corresponds to  $L$  in (60) and specifically to OFDM signals of the form in (59). The subcarriers are complex exponentials of the form  $\exp \{j2\pi \ell /L\} ,\ell \in \{0,\ldots ,L - 1\}$  with spacing  $2\pi /L$  in the (angular) frequency domain, and are orthogonal to each
other. Note that the weighted summation with complex exponentials can be interpreted as computing the inverse discrete Fourier transform (IDFT) of the pulse-shaped symbols, for which the inverse fast Fourier transform (IFFT) is used during implementation. Thus, in an interference-free setting, decoding the symbols from  $b(t)$  can be performed by computing the FFT of  $b(t)$ .
Starting point: When decoding an RF waveform to extract the underlying symbols, we need them to be time-synchronized to a specified frame of reference known as the starting point,  $\tau$ . If synchronized correctly, this indicates the number of time samples to stream in before bit decoding can begin. If unsynchronized, this corresponds to a shift in time that results in "non-natural" looking decoded constellations at all other time shifts except for the correct time-shift  $\tau$ . Here, decoding of  $u(t)$  can be performed either via an FFT for OFDM, or sampling after matched filtering for a single-carrier signal.
C.2 Synthetic and Over-the-air RF Datasets
Below we describe the details relating to the RF signals used in our single-channel source separation (SCSS) experiments. To simulate representative types of RF signals, we generated single-carrier QPSK signals with RRC pulse shaping and OFDM signals using the NVIDIA Sionna library [45]. We also used over-the-air recordings corresponding to "CommSignal2" from the dataset in the RF challenge. We do not have explicit knowledge on CommSignal2 source model as it is not provided in the challenge. Hence, we are unable to explicitly describe the pulse shaping function of such a signal type—leaving learning of such properties to our data-driven diffusion models instead.
RRC-QPSK. We use QPSK signals with root raised cosine (RRC) pulse shaping [18, Sec 11.3] as the SOI across all our experiments. These signals are used in a variety of different protocols ranging from low-power radios to cellular signals and even for satellite communication. For our experiments, we generated a dataset of 100,000 RRC-QPSK waveforms with the following parameters:
- Modulation scheme: QPSK (2 bits per symbol)
- Pulse shaping function: RRC, roll-off factor 0.5, spanning  $\sim 8$  symbols
- Default starting point:  $\tau_{s} = 8$
- QPSK-specific details:
Oversampling rate: 16 [samples/symbol], using zero-padding between symbols.
OFDM (QPSK/BPSK). For the OFDM signals (specifically, cyclic prefix OFDM), we generated waveforms with the following parameters (so as to simulate an 802.11n WiFi waveform):
- Modulation scheme: Either QPSK (2 bits per symbol) or BPSK (1 bit per symbol)
- Total number of subcarriers (FFT size): 64
- Symbol shaping: None as per (59) or rectangular window as per (60)
- Other notes: Random time offset and random phase rotation for each time series realization
OFDM-specific details:
$\circ$  Number of non-zero subcarriers: 56 (null at DC subcarrier  $(\ell = 0)$ )
Length of cyclic prefix: 16 (25% of FFT size)
Number of non-zero symbols: 56 symbols per 80 samples (1792 symbols in a 2,560 window)
$\circ$  Symbol length: 80 (64 subcarriers + cyclic prefix of length 16)
In our experiments, we choose to use OFDM signals as one class of interference signals. This means that, given a new realization, we assume that the starting point of a new OFDM symbol (i.e.,  $\tau_{b}$ ) is unknown. Further, we assume no access to knowledge about the parameters outlined above, such as the number of subcarriers, the FFT size and the length of cyclic prefix. Hence, we seek a data-driven approach where characteristics of this signal should be learned through data.
Table 1: Hyperparameters used for training our diffusion models.
QPSK	OFDM (BPSK)	OFDM (QPSK)	CommSignal2
Number of residual blocks (R)	30	30	30	30
Dilation cycle length (m)	10	10	10	10
Residual channels (C)	64	128	256	128
Random shift + phase rotation	X	✓	✓	✓
Batch size	128	128	64	128
Learning rate	5e-4	5e-4	5e-4	1e-4
Early stopping iteration	360,000	220,000	340,000	90,000
CommSignal2. CommSignal2 is a dataset comprising over-the-air recordings of type of communication waveforms, provided in [44]. This dataset contains 150 frames each of length 43, 560 samples. For purposes of this work, we use the first 100 frames in creating our training set for training the diffusion model, and we extract windows of relevant lengths (2560 samples) from it, to create a dataset of size  $\sim 160,000$  samples. The last 50 frames are reserved for creating our test set. Similar to the OFDM case, we also perform random time offsets and phase rotations for each time series realization.
D Diffusion Models for RF Signals
We train diffusion models to learn the statistical structure inherent to digital RF signals. As detailed in Appendix C, these signals exhibit both continuous and discrete properties. To this end, we adopt an architecture based on an open-sourced implementation of the DiffWave [24] diffusion model, initially developed for speech synthesis.
Broadly speaking, as shown in [24, Figure 2, Appendix C], DiffWave employs  $R$  residual blocks with dilated convolutions, where the output of block  $i - 1$  serves as an input to block  $i$ ,  $i \in \{0, \dots, R - 1\}$ . The dilated convolutions assist in the learning of long range temporal and periodic structures. The dilations first start small and successively get larger, such that the dilation at block  $i$  is given by  $2^{i \bmod m}$ , where  $m$  is the dilation cycle length. For example, if the dilation periodicity is  $m = 10$ , then in block  $i = 9$  the dilation is 512 and in block 10 the dilation is reset to 1. This allows the network to efficiently tradeoff between learning local and global temporal structures. All residual blocks use the same number of channels,  $C$ .
Modifications. Several changes have been made in order to facilitate training with RF signals. First, since we are dealing with complex-valued continuous waveforms, we train on two channel signals where the real and imaginary components of the RF signals are concatenated in the channel dimension. Second, while the open-sourced implementation uses an  $\ell_1$  loss for training, we train with an MSE (squared  $\ell_2$ ) loss to match the training objective in [21]. We monitor the validation MSE loss, and once the loss stops decreasing substantially, we choose to stop training early. Lastly, we had to increase the channel dimension  $C$  to learn more complicated RF signals, e.g., OFDM (QPSK). We train all our models on  $2\times$  NVIDIA 3090 GPUs. Additionally, during data loading, we perform random time shifts and phase rotations on the OFDM (BPSK), OFDM (QPSK) and CommSignal2 signals. Physically, these simulate transmission impairments in RF systems. The hyperparameters that we use for training each model are shown in Table 1.
Metrics. In the image domain, metrics such as the Fréchet Inception Distance (FID) [49] can be used to assess the quality of generative models. However, such perceptual metrics are less relevant to the digital communications domain. In the RF domain, the fidelity is measured by the ability to extract the underlying transmitted bits. Hence, in the context of our synthetic datasets, we probe generated samples and compare the estimated received symbols with the underlying constellation. However, for real world signals such as CommSignal2, for which we do not know the system parameters, we have no other means to assess the fidelity apart from looking at time-domain structure. This is another
$^{8}$  Namely, multiplication of the signal with a complex exponential with a random (imaginary) exponent.
<sup>9</sup>https://github.com/lmnt-com/diffwave
Figure 7: Left: A ground truth RRC-QPSK time-domain waveform. Right: A sample generated by our trained RRC-QPSK diffusion model. Evidently, the generated waveform resembles the true one.
motivation for studying RF source separation, as it provides us with a framework to assess the quality of the learned statistical priors.
D.1 RRC-QPSK
We train a diffusion model on the RRC-QPSK dataset (see Appendix C.2). While a QPSK signal has four distinct constellation points, application of the RRC filter results in interpolation between points as shown in Figure 6c. Figure 7 shows an example of a ground truth RRC-QPSK waveform on the left. On the right we show a RRC-QPSK sample generated by our diffusion model. While it is hard to judge whether it has learned the underlying discrete structure, visually, the waveform seems to have characteristics similar to the ground truth.
Figure 8: Two samples generated from the diffusion model trained on RRC-QPSK samples. Within each box, the image on the left is underlying constellation of the realization generated by the diffusion model. Notice that the RRC filtering results in interpolation between the QPSK constellation points. After applying MF, the effects of pulse shaping are reversed and the original QPSK constellation is recovered.
In Figure 8, we show that by probing generated samples for the underlying (interpolated) symbols, we do indeed recover the constellation for an RRC-QPSK signal. Note that we can do such probing since we know the parameters and signal generation model for an RRC-QPSK signal. Furthermore, by performing MF and removing the effects of pulse shaping we are able to recover the original QPSK constellation, thus demonstrating that the diffusion model has indeed successfully learned the discrete constellation along with the pulse shaping function from data samples.
D.2 OFDM (BPSK and QPSK)
We train two OFDM diffusion models, on the OFDM (BPSK) and the OFDM (QPSK) dataset respectively (see Appendix C.2). Figure 9, shows an example of an OFDM (BPSK) signal from our dataset. On the left we plot the real and imaginary components of the waveform. These components visually look like Gaussian noise and it is not evident that there is actually inherent structure to these signals. As mentioned in Appendix C.2, the OFDM signals can be demodulated using oracle knowledge about the FFT size and the cyclic prefix. Additionally, since the signal undergoes a random time shift and phase rotation, the underlying (rotated) BPSK constellation will only be visible when the time shift has been compensated for, i.e., the signal has been time-synchronized, as shown on the right in Figure 9.
Figure 9: Left: A complex-valued OFDM (BPSK) source augmented with a time shift of 8 samples and with a random phase rotation plotted in the time domain. In the time domain, the discrete structure is not discernible and the time-domain waveforms visually look like Gaussian noise. Right: Extracting the underlying (rotated) symbols from the waveform on the left by demodulating the OFDM signal using oracle knowledge about the FFT size, cyclic prefix, and time shift.
Figure 10: Top: Probing six generated samples from the OFDM (BPSK) diffusion model recovers the underlying BPSK constellation after time synchronization. Bottom: Probing six generated samples from the OFDM (QPSK) diffusion model recovers the underlying BPSK constellation after time synchronization.
The top row of Figure 10 shows the recovered constellation for six generated OFDM (BPSK) symbols and the bottom row shows the same for six generated OFDM (QPSK) signals. Note that since there are some null symbols, i.e., unused subcarriers in the frequency domain, there is an additional "constellation point" at the origin. In general the recovered constellations are clean, except for a few symbols that lie slightly off the constellation. When converting back to bits the symbols are mapped to the nearest constellation point.
D.3 CommSignal2
The CommSignal2 dataset was recorded over-the-air and hence, unlike the synthetic datasets, we cannot probe it without knowledge of the true system parameters and signal generation model, which is not provided in [44]. These complications serve as one of our motivations to develop data-driven source separation models. In the real-world, CommSignal2 could interfere with an SOI such as a QPSK signal for which the signal generation model is known. A learned prior for CommSignal2 could help mitigate such interference.
As shown in Figure 11, we observe that the generated CommSignal2 waveforms generally display similar global temporal structure to the ground truth waveforms. To further demonstrate that our diffusion model has truly learned the statistical structure of the signal, we carry out source separation experiments and demonstrate that leveraging this diffusion model outperforms conventional and learning-based source separators.
Figure 11: Left: Ground truth CommSignal2 waveform from the dataset. Right: A generated CommSignal2 waveform from the learned diffusion model.
Figure 12: Block diagram of the matched filtering demodulation pipeline used in our work. (First symbol at  $\tau_{s} = 8$ , symbol period  $T_{s} = 16$ , as detailed in Appendix C
E Classical RF Interference Mitigation Techniques
E.1 Matched Filtering
Matched filtering (MF) is a signal processing technique commonly used in the demodulation of RF waveforms. It exploits knowledge about the signal waveform to enhance the detection and recovery of the transmitted symbols/bits, and is optimal in the maximum-SNR sense for signals contaminated with additive Gaussian noise.
The basic principle involves filtering (or similarly viewed as correlating) the received sampled RF waveform with a known reference waveform called the "matched filter". The goal is to maximize the signal-to-interference-plus-noise ratio (SINR) at the filtered output, which consequently minimizes the error probability in the subsequent symbol detection when the noise is Gaussian. In the following segment, we develop the theory of the matched filter, with particular focus to our problem formulation.
Consider the baseband RRC-QPSK signal as described in Appendix C. Suppose, for the purposes of this exposition, that we adopt a simple additive white Gaussian noise (AWGN) channel model, thereby representing our received signal as
$$
\begin{array}{l} y (t) = \sum_ {p} c _ {p} g _ {\mathrm {t x}} (t - p T _ {s}) + w (t) (61) \\ = g _ {\mathrm {t x}} (t) * \sum_ {p} c _ {p} \delta \left(t - p T _ {s}\right) + w (t), (62) \\ \end{array}
$$
where  $c_{p}$  are the symbols from a QPSK constellation (see Figure 6b), * denotes the convolution operator,  $\delta(\cdot)$  is the dirac delta function, and  $w(t) \sim \mathcal{N}(0, \sigma_{\mathrm{AWGN}}^2)$  is the additive noise in the observed signal, statistically independent of all  $\{c_{p}\}$ . Of particular interest in this formulation is the transmit pulse shaping function  $g_{\mathrm{tx}}(t)$ , where we chose to use the RRC function. (This is due its favorable properties in practice; however, the optimal transmit filter choices and waveform designs
are outside the scope of this work, and we defer readers to related resources for more information on this topic [18, 36].)
At the receiver, we seek a receiver filter,  $g_{\mathrm{rx}}(t)$ , such that the filtered and sampled output
$$
y _ {\text {f i l t}} (t) = \underbrace {g _ {\mathrm {r x}} (t) * g _ {\mathrm {t x}} (t)} _ {:= g (t)} * \sum_ {p} c _ {p} \delta \left(t - p T _ {s}\right) + g _ {\mathrm {r x}} (t) * w (t) \tag {63}
$$
$$
\begin{array}{l} y [ n ] = y _ {\text {f i l t}} (n T _ {s}) = \sum_ {p} c _ {p} g ((n - p) T _ {s}) + \underbrace {\int w (\tau) g _ {\mathrm {r x}} (n T _ {s} - \tau) \mathrm {d} \tau} _ {:= v [ n ]} (64) \\ = \underbrace {c _ {n} g (0)} _ {:= y _ {s} [ n ]} + \underbrace {\sum_ {p \neq n} c _ {n} g ((n - p) T _ {s}) + v [ n ]} _ {:= y _ {v} [ n ]} (65) \\ \end{array}
$$
would maximize the output SINR. In other words, we are looking to maximize
$$
\operatorname {S I N R} = \frac {\mathbb {E} \left[ | y _ {s} [ n ] | ^ {2} \right]}{\mathbb {E} \left[ | y _ {v} [ n ] | ^ {2} \right]} = \frac {\mathbb {E} \left[ | c _ {n} | ^ {2} \right] | g (0) | ^ {2}}{\mathbb {E} \left[ | c _ {n} | ^ {2} \right] \sum_ {p \neq n} | g (p T _ {s}) | ^ {2} + \sigma_ {\mathrm {A W G N}} ^ {2} \int | G _ {\mathrm {r x}} (f) | ^ {2} \mathrm {d} f} \tag {66}
$$
(where  $G_{\mathrm{rx}}(f)$  is the Fourier transform of  $g_{\mathrm{rx}}(t)$ ) via an appropriate choice of  $g(t)$  and thereby,  $g_{\mathrm{rx}}(t)$ . This can be done by finding an upper bound on the SINR that reaches equality for the appropriate filter choices. Ultimately, one such choice is  $g_{\mathrm{rx}}(t) = g_{\mathrm{tx}}^{*}(-t)$  —termed as the matched filter—that leads to a maximized SINR. In the case of an RRC pulse shaping function (which is real and symmetric), the matched filter is also the same RRC function.
In this work, we refer to matched filtering more broadly to also include the detector/decoding step, as shown in Figure 12. As part of the demodulation pipeline, the filtered output is sampled (as in (65)), and then mapped to the closest symbol in a predefined constellation (in the Euclidean distance sense). Finally, we can map these complex-valued symbols back to their corresponding bits to recover the underlying information. We use this as a standard demodulation/detection pipeline for our RRC-QPSK SOI following the corresponding signal separation/interference mitigation step (where applicable).
Demodulation with matched filtering is optimal for waveforms in the presence of additive Gaussian noise. However, in our signal separation problem, we consider the presence of an additive interference, which is not necessarily Gaussian. Thus, exploiting the non-Gaussian characteristics of the interference would likely lead to enhanced decoding performance.
E.2 LMMSE Estimation
Recall that our observation model is
$$
\mathbf {y} = \mathbf {s} + \kappa \mathbf {b},
$$
where we assume  $\mathbf{s}$  and  $\mathbf{b}$  are zero-mean and that they are statistically independent. The linear minimum mean square error (LMMSE) estimator is the estimator  $\widehat{\mathbf{s}} = \mathbf{W}_{\mathrm{LMMSE}}\mathbf{y}$ , such that
$$
\boldsymbol {W} _ {\text {L M M S E}} = \underset {\boldsymbol {W} \in \mathbb {C} ^ {T \times T}} {\arg \min } \mathbb {E} \left[ \| \mathbf {s} - \boldsymbol {W} \mathbf {y} \| _ {2} ^ {2} \right]. \tag {67}
$$
In this case, the optimal linear transformation (in the sense of (67)) can be written as
$$
\boldsymbol {W} _ {\text {L M M S E}} = \boldsymbol {C} _ {s y} \boldsymbol {C} _ {y y} ^ {- 1} = \boldsymbol {C} _ {s s} \left(\boldsymbol {C} _ {s s} + \kappa^ {2} \boldsymbol {C} _ {b b}\right) ^ {- 1}
$$
where  $C_{sy} \coloneqq \mathbb{E}[\mathbf{sy}^{\mathrm{H}}]$  corresponds to the cross-covariance between s and y,  $C_{yy}, C_{ss}, C_{bb}$  are the auto-covariance of y, s and b respectively. The second equality is obtained by statistical independence, thereby  $C_{sy} = C_{ss}, C_{yy} = C_{ss} + \kappa^{2}C_{bb}$ .
In relevant literature, this may also be referred to as linear MMSE receivers [50, Sec 8.3.3]. Note that we aim to mitigate the effects of an additive interference channel. For the purposes of this work, we use LMMSE as one of the baselines as a signal separation (interference mitigation) method. Thereafter, we assess performance based on the squared error between  $\hat{\mathbf{s}}$  with the ground truth  $\mathbf{s}$ . To obtain the underlying bits, we perform a standard matched filtering operation on the estimator  $\hat{\mathbf{s}}$ .