msaidov/audioseal-robust-sgmse-16bits

AudioSeal generator — robustness fine-tuned against SGMSE (16 bits)

A fine-tuned AudioSeal watermark generator hardened against SGMSE, a score-based diffusion speech-enhancement model used here as a watermark-removal attack. It is a drop-in replacement for the stock facebook/audioseal audioseal_wm_16bits generator and stays compatible with the unmodified audioseal_detector_16bits detector.

Sibling model: msaidov/audioseal-robust-audioldm-16bits, the same experiment run against a structurally different diffusion attack (AudioLDM latent-diffusion resynthesis).

• Training code: martysai/psiml11-audio-with-diffusion (src/audioseal_robust)
• Built at: PSIML 11 — Practical Seminar in Machine Learning, a machine learning summer school in Serbia
• Checkpoint file: generator.pth, repackaged for stock AudioSeal (see Checkpoint format). The raw training checkpoint is kept alongside as generator_train_ckpt.pth
• Exact training config: training_config.yaml, read back out of the checkpoint's own embedded xp.cfg

What was actually trained

Only the generator is trained. The AudioSeal detector is loaded pretrained, frozen (requires_grad_(False), never taken out of eval()), and used purely as a differentiable objective. Nothing about the detector, the architecture, the 16-bit payload, or the 16 kHz sample rate changes — so existing AudioSeal detection code keeps working unchanged.

The architecture is the stock AudioSeal SEANet encoder/decoder with the 16-bit message processor (nbits: 16, dimension: 128, n_filters: 32, ratios: [8, 5, 4, 2], 2 LSTM layers), initialised from the pretrained audioseal_wm_16bits weights — this is a fine-tune, not a from-scratch train.

The attack it was trained against

SGMSE (sp-uhh/sgmse) — an OU-VE SDE score-based generative model for speech enhancement, vendored under src/sgmse/ and driven from the authors' VoiceBank-DEMAND pretrained checkpoint (sgmse_vb_pretrained.ckpt).

Used as an attack, it is a reconstruction threat: an adversary who wants to strip a watermark can run the audio through a diffusion enhancer, which regenerates the signal and discards perturbations that do not survive the score model's prior.

The attack is run fully differentiably inside the training loop — its 30 predictor-corrector reverse steps are not wrapped in no_grad, so gradients propagate all the way back through the whole reverse-diffusion chain into the generator. The corruption depth t* is sampled per example, uniformly over the full [0, 1] range, so the generator sees everything from a barely perturbed signal to a full re-synthesis.

Per training step, one attack branch is sampled from the sgmse_mixed recipe:

branch	weight	meaning
identity	0.5	unattacked — anchors bit accuracy on the easy case
sgmse	0.5	full differentiable SGMSE reverse-diffusion attack

The identity half is deliberate: training against the attack alone gives the generator no easy steps to anchor bit accuracy on, and doubles its exposure to the Langevin corrector's occasional gradient-norm spikes.

AudioLDM was held out of this run entirely (attack.weights.audioldm: 0.0) so it can serve as a generalization probe — i.e. "does robustness learned against one diffusion attack transfer to a structurally different one?"

Training data

LibriSpeech (16 kHz, English read speech):

• Train: the first 10 h of train-clean-100 by cumulative duration, taken in sorted file order for determinism (2.9 k utterances).
• Validation: ~60 min of dev-clean — a split never seen in training.
• Segments are 1.0 s random crops (longer files cropped, shorter ones tiled, never zero-padded).

Training procedure

Each step embeds a fresh random 16-bit message:

x_wm = x + scale · G(x, m)

where scale is set per example so the watermark lands at a target SNR drawn uniformly from [24, 36] dB relative to the host signal, rather than using whatever raw amplitude the generator happens to emit. The attack is then applied, and the frozen detector scores the result.

Objective

L = λ_det · (BCE(presence) + λ_bit · BCE(message bits)) + λ_perc · L_mel

L_mel is a psychoacoustic mel loss: L1 between log-mel spectrograms of x and x_wm, with each mel bin weighted by a Terhardt absolute-threshold-of-hearing curve so perturbation energy in less audible bands is penalised less.

hyperparameter	value
λ_det	1.0
λ_perc	0.0 — perceptual loss disabled for this run
λ_bit	2.0 — up-weights the hard sub-problem (bit decoding) over the easy one (mere presence)
optimizer	Adam, lr 5e-5, betas (0.5, 0.9), weight decay 0.0
gradient clipping	max_norm = 3.0 (clip only, no gradient normalization)
activation-gradient clamp	max_x_wm_grad_norm = 1000.0 at the generator/attack boundary
precision	bf16 autocast on the forward pass; BCE and SGMSE's log/exp forced back to fp32
batch size	8
mel loss	n_fft 1024, hop 256, win 1024, 80 mels, f_min 20 Hz
seed	1234

Note on λ_perc = 0.0. Perceptual loss was switched off for this run to isolate the detection objective while diagnosing a bit-accuracy plateau. The watermark is still amplitude-constrained by the 24–36 dB SNR scaling above, so it is not unbounded — but this checkpoint was not optimized for perceptual transparency. Measure SI-SNR/PESQ yourself before assuming imperceptibility.

Schedule. epochs: 100 and updates_per_epoch: 1000 are caps, not targets — the inner loop also ends when the dataloader is exhausted, whichever comes first. On this ~10 h subset at batch size 8 that is ≈365 optimizer steps per epoch, so the cap never bound. This checkpoint is saved at the end of epoch index 3, i.e. after 4 completed passes over the subset (≈1.5 k optimizer steps). Trained on a single A100.

Checkpoint format

generator.pth is a plain torch.save dict, in the same shape AudioSeal's own generator_base.pth uses:

{"model": <generator state_dict>, "xp.cfg": <architecture config, plain dicts>}

Three properties are what make the one-line load above work:

• xp.cfg describes the architecture, not the training run. It is the stock audioseal_wm_16bits config (nbits, seanet, decoder) as plain dicts and lists — exactly what AudioSeal.parse_config reads. It references no OmegaConf or project-specific classes, so torch.load needs nothing but torch, and nbits is picked up automatically.
• The conv keys use the flat naming (....conv.conv.weight). AudioSeal picks its SEANet by interpreter version — AudioCraft's (flat) below Python 3.10, Moshi's (an extra inner_conv level) at or above it — and its loader only converts flat → inner_conv. Flat is therefore the only naming that loads on both sides of that split, which is why upstream publishes it and why this checkpoint does too.
• The training config is still in the file, converted to plain containers, under an "audioseal_robust" key the loader ignores — together with the source checkpoint name and the exporting commit. training_config.yaml is that same data as YAML.

generator_train_ckpt.pth is the unmodified file the training run wrote (generator_epoch3.pth), kept for provenance. Its tensors already use the flat naming, because it was written by a Python < 3.10 process. It is still not usable on its own: its xp.cfg pickles audioseal_robust.config.TrainConfig and thirteen sibling dataclasses by reference, so torch.load fails with ModuleNotFoundError: No module named 'audioseal_robust' unless the training repo is importable — and even then AudioSeal.load_generator rejects it, because a TrainConfig has no seanet block. Reach for it only if you are reproducing the run.

The repackaging is done by audioseal_robust.export_checkpoint, which reloads its own output through AudioSeal.load_generator and compares it tensor-by-tensor against the source before writing the file.

Usage

pip install audioseal huggingface_hub

import torch
from audioseal import AudioSeal
from huggingface_hub import hf_hub_download

generator = AudioSeal.load_generator(
    hf_hub_download("msaidov/audioseal-robust-sgmse-16bits", "generator.pth")
)

# Watermark exactly as with stock AudioSeal.
wav, sr = ..., 16000                       # (batch, channels, samples), 16 kHz
msg = torch.randint(0, 2, (wav.shape[0], 16))
watermarked = wav + generator.get_watermark(wav, sr, message=msg)

# The stock detector is unchanged and still applies.
detector = AudioSeal.load_detector("audioseal_detector_16bits")
prob, decoded = detector.detect_watermark(watermarked, sr)

There is no nbits= to pass and no state-dict reconciliation to do: both come out of the checkpoint itself (see Checkpoint format).

If you would rather not add huggingface_hub, AudioSeal will fetch the URL itself through torch.hub:

generator = AudioSeal.load_generator(
    "https://huggingface.co/msaidov/audioseal-robust-sgmse-16bits/resolve/main/generator.pth"
)

To reproduce the training-time embedding exactly, scale the watermark to a target SNR in [24, 36] dB instead of adding it raw — see embed_watermark in src/audioseal_robust/train.py.

Evaluation

No benchmark numbers are published with this checkpoint. The evaluation harness lives in the training repo and reports bit accuracy, TPR@FPR (with explicit FPR-resolution checks), ROC-AUC, confusion counts, SI-SNR and PESQ, per attack branch and along a t* robustness curve:

PYTHONPATH=src python -m audioseal_robust.evaluate \
  generator_checkpoint=generator.pth \
  eval_dir=/path/to/LibriSpeech/test-clean \
  recipe=after_sgmse_training \
  attack.sgmse.checkpoint=/path/to/sgmse_vb_pretrained.ckpt

The after_sgmse_training eval recipe reports identity, bigvgan, dac and sgmse, and holds audioldm and mbd out as unseen-attack generalization probes.

Intended use and limitations

• Intended use: research on the robustness of audio watermarking to diffusion-based removal attacks, and as a baseline for reproducing this experiment.
• Not production-ready. It is a short fine-tune (≈1.5 k steps) on ~10 h of clean read English speech, with perceptual loss disabled. Expect degradation on other languages, noisy/far-field audio, music, and non-speech.
• Robustness is attack-specific and unverified here. Robustness learned against SGMSE does not automatically transfer to structurally different attacks; in this project's own observations a sgmse_mixed-trained generator measured considerably worse against held-out AudioLDM than against SGMSE itself. Do not treat "robust" in the model name as a security guarantee — measure it on your own threat model.
• Detection requires the stock detector. These weights only change the embedder.

Acknowledgements

This checkpoint was produced as a student project at PSIML 11 — Practical Seminar in Machine Learning, a machine learning summer school in Serbia.

Thanks to the PSIML organisers and mentors for the compute and for the guidance on the experiment design.

License and provenance

MIT, inherited from facebook/audioseal (Meta Platforms), from which these weights are derived. The SGMSE (sp-uhh/sgmse, MIT) model is used only as a training-time attack; its weights are not redistributed here.

Citation

If you use this checkpoint, please cite AudioSeal and SGMSE:

@article{sanroman2024proactive,
  title   = {Proactive Detection of Voice Cloning with Localized Watermarking},
  author  = {San Roman, Robin and Fernandez, Pierre and Elsahar, Hady and
             D{\'e}fossez, Alexandre and Furon, Teddy and Tran, Tuan},
  journal = {ICML},
  year    = {2024}
}

@article{richter2023speech,
  title   = {Speech Enhancement and Dereverberation with Diffusion-based
             Generative Models},
  author  = {Richter, Julius and Welker, Simon and Lemercier, Jean-Marie and
             Lay, Bunlong and Gerkmann, Timo},
  journal = {IEEE/ACM Transactions on Audio, Speech, and Language Processing},
  volume  = {31},
  pages   = {2351--2364},
  year    = {2023},
  doi     = {10.1109/TASLP.2023.3285241}
}