Lab 8.3 — Learned Denoising / Noise Suppression
← Course 2 syllabus · Bonus Module 8 · Prev: « Lab 8.2 · Next: Lab 8.4 »
Goal
Build a small learned noise-suppression network — a denoising autoencoder or, better, a spectral-masking network operating on the STFT — that cleans up live noisy audio in real time on the edge, and compare it head to head against the classical Wiener filter (Hayes) on the same noisy input. The measured axes are SNR improvement, audible artifacts (musical noise), and latency on the Pi 5 versus the Jetson. This is the lab where the classical/statistical DSP you built all course meets the learned approach on equal footing: the Wiener filter is the optimal linear estimator under a stationarity assumption; the network is a nonlinear estimator that learns the noise structure from data. Seeing where each wins is the whole lesson.
Recommended reading
- Hayes Ch. 7 — Wiener filtering: the orthogonality principle and the optimal linear estimator. This is the classical baseline you must beat. → ../books/hayes-statistical-dsp/index.qmd
- Hayes Ch. 9 — adaptive filtering (LMS/RLS): the “learns online” middle ground between a fixed Wiener filter and an offline-trained net.
- Grewal (Kalman Filtering) — the recursive optimal estimator: a steady-state Kalman filter is the Wiener filter computed sample-by-sample. The model-based counterpart to this learned denoiser is the real-time Kalman filter of Lab 6.6.
- Course 1 Weeks 1–3 — linear algebra and least squares: the Wiener–Hopf equations are a linear least-squares problem, and the autoencoder’s linear layers are the same normal-equations machinery generalized. → ../course1/index.qmd#week-1
- Course 1 Weeks 7–8 — convex optimization: the mask-prediction training objective and why the reconstruction loss is minimized by gradient descent. → ../course1/index.qmd#week-7
Equipment & parts
- Host Mac (PyTorch) for training; ONNX export.
- Jetson Orin Nano (TensorRT/PyTorch/CuPy) and Raspberry Pi 5 (onnxruntime/tflite).
- Audio input (USB mic / I²S / ADS1115) as in Labs 8.1–8.2.
- A repeatable noise source for controlled tests: a fan/AC hum, or synthesized noise played from the MCP4725/STM32 or a second speaker, so you can create matched clean/noisy pairs and compute true SNR.
Safety & don’t-break-it
- Same analog-input discipline: within-rails, mid-rail-biased, common-ground — a DC-offset or clipped input makes the noise estimate garbage and can saturate the ADC.
- Latency is a safety-of-the-experiment issue here. A denoiser adds algorithmic latency (its STFT window + hop + any lookahead). If you close the loop through a speaker you can get feedback howl; keep input and output acoustically separated, or process offline files first and only then go live.
- Don’t clip the output. A mask can boost as well as attenuate bands; clamp the reconstructed waveform to full scale before it hits a DAC/output, or you’ll add your own distortion on top of the noise you removed.
- Thermals/clocks on the Jetson as before — a throttled run corrupts the latency comparison.
Background
Signal model. Observed audio is clean speech plus additive noise, \(y[n] = x[n] + d[n]\). In the STFT domain, \(Y[m,k] = X[m,k] + D[m,k]\). Both the Wiener filter and the mask network estimate a per-bin gain \(G[m,k] \in [0,1]\) and reconstruct \(\hat{X} = G \odot Y\), then inverse-STFT (overlap-add) back to a waveform.
The classical Wiener gain. Under the (locally) stationary, zero-mean, uncorrelated-signal-and-noise assumption, the gain that minimizes the mean-squared error \(\mathbb{E}\,|X - \hat X|^2\) is the ratio of clean-signal power to observed power:
\[ G_{\text{W}}[m,k] \;=\; \frac{S_{xx}[m,k]}{S_{xx}[m,k] + S_{dd}[m,k]} \;=\; \frac{\xi[m,k]}{1 + \xi[m,k]}, \qquad \xi \;=\; \frac{S_{xx}}{S_{dd}} \;=\; \text{a priori SNR}. \]
Where the SNR \(\xi\) is high the gain \(\to 1\) (pass the bin); where it’s low the gain \(\to 0\) (suppress it). In practice you estimate \(S_{dd}\) from noise-only segments (a VAD-gated running average) and \(S_{xx}\) by subtraction — this is the orthogonality-principle result from Hayes Ch. 7. Its failure modes are exactly what you’ll hear: nonstationary noise breaks the \(S_{dd}\) estimate, and per-bin thresholding produces musical noise (random isolated surviving bins that warble).
The learned mask. Train a small network \(G_\theta = f_\theta(|Y|)\) that predicts the gain mask directly from the noisy magnitude spectrogram (optionally several context frames). Train on matched pairs by minimizing a reconstruction loss, e.g. mean-squared error between the masked magnitude and the clean magnitude,
\[ \mathcal{L}(\theta) \;=\; \frac{1}{MT}\sum_{m,k} \big(\, G_\theta[m,k]\,|Y[m,k]| - |X[m,k]| \,\big)^2 , \]
with a sigmoid output so \(G_\theta \in [0,1]\). A denoising-autoencoder variant instead reconstructs the clean spectrogram directly through a bottleneck. Because the mask is nonlinear and sees spectral context, it can track nonstationary noise and avoid musical noise better than the per-bin Wiener rule — at the cost of needing training data and compute.
Measuring the result. Report SNR improvement,
\[ \text{SNR}_{\text{dB}} = 10\log_{10}\frac{\sum_n x[n]^2}{\sum_n (\hat x[n]-x[n])^2}, \qquad \Delta\text{SNR} = \text{SNR}_{\text{out}} - \text{SNR}_{\text{in}}, \]
on held-out clean/noisy pairs, and note artifacts qualitatively (musical noise, speech distortion). Latency = STFT window + hop + inference + inverse-STFT.
Procedure
Part A — Make matched pairs.
- Record or synthesize clean speech/tones, and a set of noise clips (fan, hum, white/pink). Mix at known SNRs (e.g. 0, 5, 10 dB) so every noisy clip has its clean reference and true input SNR. Hold out a test set with unseen noise.
Part B — Classical Wiener baseline.
- Implement the Wiener filter: STFT the noisy signal, estimate \(S_{dd}\) from noise-only frames (or a running minimum), compute \(G_{\text{W}}\) from the formula, mask, inverse-STFT (overlap-add). Measure \(\Delta\)SNR and listen for musical noise. This runs fine on either board’s CPU — it’s cheap.
Part C — Train the mask network on the host.
- Define a small masking net (a few conv layers or a tiny U-Net over the spectrogram, sigmoid output). Train on the matched pairs with the reconstruction loss above.
# illustrative — owner writes/trains the real network
mask = torch.sigmoid(net(mag_noisy)) # (B,1,M,T) in [0,1]
est = mask * mag_noisy
loss = F.mse_loss(est, mag_clean) # + optional phase-aware / SI-SNR term- Reconstruct with the noisy phase (standard) or add a phase-aware / SI-SNR loss if you want to push quality. Export to ONNX.
Part D — Deploy and measure on-device.
- Jetson: run the mask net in PyTorch or as an fp16/int8 TensorRT engine, streaming via the Lab 8.1 ring-buffer pipeline (STFT → mask → iSTFT). Pi 5: run the ONNX in onnxruntime (or tflite).
- On each board, process the held-out test set through the board’s audio path. Record \(\Delta\)SNR and end-to-end latency (feature + inference + reconstruction, p95). Do the same for the Wiener baseline on each board.
Deliverable & expected results
- Cleaned audio files (Wiener and learned) for the test set, with a spectrogram triptych (noisy / Wiener / learned) per example.
- A results table comparing \(\Delta\)SNR, artifact notes, and latency across methods and boards.
| Quantity | Predicted / expectation | Measured |
|---|---|---|
| Wiener \(\Delta\)SNR, stationary noise | good (few–several dB) | … |
| Wiener \(\Delta\)SNR, nonstationary noise | degrades; musical noise | … |
| Learned \(\Delta\)SNR, unseen noise | ≥ Wiener, fewer artifacts | … |
| Wiener latency (Pi 5) | small (FFT-bound) | … |
| Learned latency, Pi 5 (onnx) | (ms, p95) | … |
| Learned latency, Jetson (TensorRT fp16) | (ms, p95) | … |
| Algorithmic latency (window + hop) | \(N/f_s + H/f_s\) | … |
Analysis & reconciliation
Show why the Wiener filter does well on stationary noise (it is provably the optimal linear MMSE estimator when its assumptions hold — the Hayes Ch. 7 result) and where it breaks: nonstationary noise defeats the \(S_{dd}\) estimate, and its per-bin decisions produce musical noise you can point to in the spectrogram. Then show the learned mask holding up on the same clips because it uses spectral context and a nonlinear decision — but be honest about its costs: it needs matched training data, it may overfit to the noise types it saw, and its latency is higher. Reconcile the latency numbers: the Jetson fp16 engine should give the lowest inference time but the mask net is small enough that on the Pi 5 the STFT/iSTFT overlap-add may dominate — meaning the classical Wiener and the learned mask have similar real-world latency on the Pi, and the accelerator only matters as the model grows. Feed these latency and \(\Delta\)SNR numbers into Lab 8.5.
Going further
- Add the adaptive LMS/RLS filter (Hayes Ch. 9) as a third baseline — the online-learning middle ground — and compare its tracking of nonstationary noise to both the fixed Wiener and the offline-trained net.
- Try an SI-SNR or perceptual loss instead of magnitude MSE and A/B the audible quality.
- Feed the denoised audio into the Lab 8.2 keyword classifier and measure whether denoising improves downstream accuracy in noise — the real reason to denoise in a pipeline.
- Quantize the mask net to int8 and measure the quality/latency trade, as in Lab 8.2.