msaidov/audioseal-robust-audioldm-16bits

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

A fine-tuned AudioSeal watermark generator hardened against AudioLDM latent-diffusion resynthesis, 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-sgmse-16bits, the same experiment run against a structurally different diffusion attack (SGMSE's OU-VE SDE speech enhancement).

• 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

AudioLDM — a text-to-audio latent diffusion model (haoheliu/AudioLDM-training-finetuning, vendored under src/audioldm_train/), driven from the audioldm-s-full pretrained checkpoint with the audioldm_original.yaml config and run fully unconditionally.

As an attack it implements partial-noise-then-regenerate: mel-spectrogram → VAE latent → forward-diffuse to timestep t* → reverse-diffuse back with the pretrained UNet → decode → re-vocode to a waveform with HiFi-GAN. Because the vocoder resynthesises phase from scratch, the output is essentially uncorrelated with the input at the sample level — a far more destructive threat than a codec or an enhancer.

The whole chain — VAE encoder, the reverse UNet loop, and the vocoder — is run differentiably inside the training loop, so gradients propagate back through it into the generator.

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

branch	weight	meaning
identity	0.5	unattacked — anchors bit accuracy on the easy case
audioldm	0.5	full differentiable AudioLDM noise-and-regenerate attack

The identity half is not cosmetic. At 100 % AudioLDM the generator never sees an unattacked step, and the attack destroys the watermark outright rather than merely degrading it: on an earlier run the bit loss sat at 0.698 against ln 2 = 0.6931 — chance, i.e. zero recoverable bits — flat across hundreds of steps, with no gradient signal to learn from. The identity half is what gives the generator something to anchor on.

Attack strength. attack.audioldm.strength_max: 0.02, lowered from the project default of 0.08. t* is drawn uniformly from [0, 0.02], i.e. up to ~20 of the DDPM's 1000 timesteps. This ceiling exists because backprop through a deeper reverse loop is not memory-tractable — and it caps the strength this checkpoint was actually hardened against.

SGMSE was held out of this run entirely (attack.weights.sgmse: 0.0) so it can serve as a generalization probe for a structurally different diffusion attack.

Gradient normalization

This run enables optim.normalize_grad: True — the generator's gradient is rescaled to exactly max_norm every step rather than only clipped when it exceeds it. This is load-bearing for a mixed recipe, not a cosmetic tweak.

Backprop through the diffusion chain amplifies the gradient by roughly 250× at the attack boundary, so on a measured earlier run the two branches produced median parameter-gradient norms of 1.69 (identity, under the 3.0 clip on 26/30 steps) versus 42.94 (AudioLDM, over the clip on 33/33 steps, scaled down ~14×). Plain clipping therefore let identity steps through at full strength while shrinking every AudioLDM step, leaving the optimizer on a ~93 % identity diet — a task the pretrained checkpoint has already solved. Over 1550 steps the AudioLDM branch did not move at all under that regime. Adam does not rescue this: it is scale-invariant to a uniform rescale, but its moment estimates are shared across steps, so a consistently down-scaled branch contributes proportionally less to the running direction.

Training data

LibriSpeech (16 kHz, English read speech):

• Train: the full train-clean-100 split (~100 h, ~28.5 k utterances).
• Validation: the full dev-clean split — never seen in training.
• Segments are 2.0 s random crops (longer files cropped, shorter ones tiled, never zero-padded). The AudioLDM attack tiles each segment up to its native 10.24 s window internally, so it sees repeated real audio rather than silence padding.

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	1.0 — perceptual loss active
λ_bit	1.0
optimizer	Adam, lr 5e-5, betas (0.5, 0.9), weight decay 0.0
gradient handling	normalize_grad: True, max_norm = 3.0, floor 1e-6
activation-gradient clamp	max_x_wm_grad_norm = 1000.0 at the generator/attack boundary
precision	bf16 autocast on the forward pass; BCE and the VAE/UNet's sensitive ops forced back to fp32
batch size	4 per rank (config batch size is per-GPU under DDP)
mel loss	n_fft 1024, hop 256, win 1024, 80 mels, f_min 20 Hz
seed	1234

Schedule. epochs: 20 × updates_per_epoch: 250. On the full ~100 h split the dataloader is far from exhausted at 250 batches, so the cap binds and every epoch is exactly 250 optimizer steps. This checkpoint is saved at the end of epoch index 18, i.e. after 19 completed epochs ≈ 4 750 optimizer steps.

Trained on A100 GPUs under Azure ML (run train-audioldm-gradnorm-0814-171739), with MLflow tracking.

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_epoch18.pth), kept for provenance. Its tensors use the inner_conv-wrapped naming, because it was written by a Python ≥ 3.10 process. It is 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 original file was distributed as generator_epoch18.pth.zip. That is not a zip wrapper — a torch.save file already is a zip archive, and the suffix was simply appended. generator_train_ckpt.pth is that file, byte-identical.

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-audioldm-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-audioldm-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-fixed \
  segment_duration=10.24 \
  recipe=after_audioldm_training \
  attack.audioldm.checkpoint=/path/to/audioldm-s-full \
  attack.audioldm.config=src/audioldm_train/config/2023_08_23_reproduce_audioldm/audioldm_original.yaml

Two details matter when evaluating against AudioLDM: use a fixed-duration 10.24 s eval set and set segment_duration=10.24, matching AudioLDM's native window. Otherwise the attack pads every shorter segment, which is a structurally different input than anything it saw in pretraining, and the resulting numbers are not comparable. The after_audioldm_training recipe reports identity, bigvgan, dac and audioldm, holding sgmse and mbd out as unseen-attack 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. ≈4.75 k optimizer steps on clean read English speech. Expect degradation on other languages, noisy/far-field audio, music, and non-speech.
• Only hardened at low attack strength. Training capped t* at 0.02 (~20 of 1000 diffusion steps). This says nothing about robustness to a full-strength AudioLDM regeneration, which discards the sample-level signal almost entirely.
• Robustness is attack-specific and unverified here. 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. AudioLDM (haoheliu/AudioLDM-training-finetuning, MIT) is used only as a training-time attack; its weights are not redistributed here.

Citation

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

@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{liu2023audioldm,
  title   = {{AudioLDM}: Text-to-Audio Generation with Latent Diffusion Models},
  author  = {Liu, Haohe and Chen, Zehua and Yuan, Yi and Mei, Xinhao and
             Liu, Xubo and Mandic, Danilo and Wang, Wenwu and Plumbley, Mark D.},
  journal = {Proceedings of the International Conference on Machine Learning},
  pages   = {21450--21474},
  year    = {2023}
}