Analysing Adversarial Priors for Data-driven Unsupervised Speech Enhancement

Speech enhancement performed inside a neural audio codec. A noisy waveform is encoded once; a two-branch RoFormer bottleneck splits the latent into a clean-speech latent and a noise latent; both are decoded through a shared decoder and recombined. Training is adversarial against three independent discriminators, and works unsupervised — the clean and noise corpora are sampled independently of the noisy input, so no example is ever a matched pair.

                                 ┌──────────────┐   clean_z   ┌─────────┐  clean_audio
 noisy ──▶ Encoder ──▶ latent ──▶│  RoFormer    │────────────▶│         │─────────────▶
                    │            │  clean branch│             │ Decoder │
                    │            ├──────────────┤   noise_z   │(shared) │  noise_audio
                    │            │  RoFormer    │────────────▶│         │─────────────▶
                    │            │  noise branch│             └─────────┘
                    │            └──────────────┘                   ▲
                    └────────────── skip connections ───────────────┘

This is a rewrite of an earlier research codebase. It keeps the model and the training behaviour, and drops everything that was not on the live path: WavLM, WavTokenizer, Wave-U-Net, FSQ, a standalone clean decoder, four alternative separator backbones, the audiotools dependency, and roughly 250 experiment configs.

Quick start

pip install -r requirements.txt

# Train the default (ICASSP) recipe
python scripts/train.py

# Swap objective, model or hardware
python scripts/train.py losses=msc_covariance
python scripts/train.py model=one_branch_roformer
python scripts/train.py trainer=ddp data.batch_size=4

# Enhance audio with a trained checkpoint
python scripts/enhance.py --checkpoint runs/.../last.ckpt \
    --input /data/test_noisy --output-dir /tmp/enhanced

# Run the test suite (parity against the original implementation)
python -m pytest tests -q

# A 6-step end-to-end run on synthetic audio, useful for checking a new machine
python scripts/train.py --config-name smoke

Layout

Path	Contents
src/dsp/	STFT, mel filterbank, ITU-R BS.1770 loudness. Everything else depends on these.
src/data/	Manifests, cropping, SNR mixing, the Dataset and the LightningDataModule.
src/models/	DAC encoder/decoder, minimal RoFormer, two-branch separator, RVQ, discriminators.
src/losses/	Reconstruction, adversarial, regularisers, and the loss-term registry.
src/training/	The LightningModule, schedules, metrics, checkpoint import.
src/inference/	Batch enhancement and the least-squares branch mix.
configs/	Hydra config groups.
tests/	Parity tests against the original, plus integration tests.

Losses, regularisers, mixing and DSP are plain functions on tensors; nn.Module appears only where parameters live. There is no AudioSignal type — everything is a tensor plus an explicit sample rate.

Configuration

Hydra, composed through defaults: in configs/train.yaml. Any key can be overridden on the command line, and unknown keys are rejected — the schema in src/config.py is authoritative. Machine-specific paths live in configs/paths/local.yaml; everything else refers to them by interpolation.

The loss objective is a plain weights mapping, validated at startup against src/losses/registry.py. A key that is not a real loss term raises rather than being silently dropped, and a term that the model configuration cannot produce (a VQ loss with use_vq: false, say) is reported.

Verifying against the original

tests/ compares this implementation against the original, which is expected at ./neural_codec_enhancement. Those tests skip if it is absent.

Test	What it pins down
test_parity_dsp.py	STFT, mel and loudness match audiotools to <1e-5.
test_parity_losses.py	Every loss and regulariser matches the original.
test_parity_data.py	Crops match the original AudioLoader bit for bit.
test_parity_step.py	Every loss term and the weighted total match, for both live objectives.
test_smoke_train.py	Training runs, all four optimisers step, checkpoints round-trip.

test_parity_step.py is the one that certifies "behaves like the old code": per-term parity is not enough, because the objective also depends on which signal is paired with which.

Note the tests pin torch.set_num_threads(1). The decoder's convolution stack is not reproducible with multiple CPU threads — repeated forwards of the same module on the same input differ by ~3e-3. Individual convolutions are deterministic; the accumulation order across the stack is not.

Deviations from the original

Behaviour is reproduced exactly except where listed. Each of these is a deliberate choice, and the ones that change numbers are behind a flag.

Change	Why	Revert with
Regularisers computed in float32	In bfloat16 reg/cs_zero_mean has a noise floor ~1500× the value it is minimising	training.force_fp32_regularizers=false
No gradient scaler	The original paired a float16 GradScaler with bfloat16 autocast, which does nothing but silently skip steps	—
SNR mixing in float32	The original round-tripped through 16-bit PCM, so noisy − noise ≠ clean (off by ~3e-5), and clipped loud clean signals	—
Noise crop offset is random	The original passed random_mix_offset=True to lhotse but rebuilt the mixer each call with a fixed seed. Measured during analysis as returning the same region of a given noise file on every call; not re-asserted here, since lhotse is no longer a dependency	—
All sampling seeded per index	The original drew noise type, SNR, noise file and RIR from the global RNG, so those were not reproducible	—
Noise files are resampled	The original ignored the sample-rate argument when loading a noise path and mixed at the wrong rate	—
--extend_pe extends both branches	The original extended only the clean branch, then ran a forward that used the noise branch too	—
Peak normalisation has a floor	The original divided by the peak unguarded — a division by zero on silent output	—
Unknown loss keys raise	The original dropped them silently from the objective	—

Findings in the original worth knowing

These are properties of the original code that carry over into how the existing results should be read.

1. scale_two_branches_audio never affected training. The code did out['audio'].audio_data = scale_audio(...) where out['audio'] is a plain tensor, so it attached a dead attribute and the next line re-wrapped the unscaled sum. Only sample logging applied it. So the two branches were never constrained to be least-squares optimal during training, yet inference applies exactly that transform. Here it is off for losses by default (training.apply_branch_scaling_to_losses) and on for inference, matching the original — but it is now a visible switch.
2. train_csd_with_noisy_data: true was inert in the live configs: the term sits in the branch that sample_train_clean_independently never takes.
3. max_grad_norm: 5e8 disabled the step-rejection guard. Gradients are always clipped to max_clip_grad_norm (10.0) and the pre-clip norm was compared against 5e8, so only the NaN check ever rejected anything.
4. Several mel scales contribute nothing. n_mels=5 over a 32-sample window, and n_mels=320 over 1025 bins, give all-zero librosa filters: log10(clamp_eps) is a constant with zero gradient.
5. reg/sum_energy is mostly a constant. Its clamp(min=0.1) is a floor, so the gradient vanishes wherever the residual is already small — which at −16 LUFS is most bins.
6. RoFormer dropout was silently on at 0.1 for both hidden and attention. Nothing in the original set it. It is now explicit in configs/model/.
7. The rotary table caps sequence length at 1536 frames = 30.7 s at the 50 Hz latent rate. Longer inference needs the table extended.
8. WANDB_API_KEY was a literal in scripts/train.py and a SLURM script. It is read from the environment here; that key is in git history and should be rotated.
9. The pretrained codec declares 12 codebooks against a configured 8. That mismatch is harmless with use_vq: false, but the original's non-strict load would equally have tolerated an unloaded encoder — so load_dac_pretrained now hard-errors if the encoder or decoder receives no weights.

Environment

scipy>=1.16 requires numpy>=1.26.4. A numpy downgrade to 1.23.5 will break scipy.signal and pyloudnorm, and therefore the STFT and loudness paths of both this codebase and the original.

Batch size under DDP

data.batch_size is per device. The original divided a global batch size by the world size inside its own accelerator, so a config reading batch_size: 16 on 4 GPUs actually trained at 4 per device. Divide accordingly when porting a config, or the effective batch — and the meaning of the LR schedule — changes.