MOLM-Audio: SPDMark-Style Segment-Wise Audio Watermarking

LoRA-routing audio watermarks for three generators (HiFi-GAN, VibeVoice acoustic decoder, DiffWave). Per SPDMark (https://arxiv.org/abs/2512.12090), each audio is split into S=8 segments; each segment carries an HMAC-derived M-bit message embedded via a parallel LoRA "basis dictionary." The verifier matches recovered per-segment bits against the expected HMAC sequence using Hungarian assignment + Binomial hypothesis test.

Repo layout

checkpoints/
  diffwave_spdmark_spec/final/{lora_weights.pt, extractor.pt, diffwave_full.pt}
  diffwave_v5_2step/final/{lora_weights.pt, extractor.pt, diffwave_full.pt}
  hifigan_spdmark_spec/final/{lora_weights.pt, extractor.pt, hifigan_full.pt}
  vibevoice_spdmark_spec/final/{lora_weights.pt, extractor.pt, model_full.pt}
  vibevoice_14bit_v2/final/{lora_weights.pt, extractor.pt, model_full.pt}
eval/<run>_<regime>/results.json
README.md

lora_weights.pt + the base model from torchaudio/pretrained is sufficient for inference. *_full.pt is provided for one-step loading where the base model isn't available locally.

Training/test wavs live in MOLM-Audio/molm-audio-data under data/.

Checkpoints

Current SPDMark runs (spec-trained, with Hungarian verifier)

Run	Backbone	Routing	Paths	Bits/seg	Train attack
`hifigan_spdmark_spec`	HiFi-GAN V3 (torchaudio LJSpeech)	`0,1,2,3,4,5,6`	4	14	`nvlceqr`
`vibevoice_spdmark_spec`	VibeVoice acoustic decoder	`0,1,2,3,4,5,6` (×2 slots)	2	14	`nvlceqr`
`diffwave_spdmark_spec`	DiffWave (LJSpeech)	`0,4,8,12,16,20,24`	4	14	`nvlceqr`

Legacy SPDMark runs (older eval pipeline; included for reference)

Run	Backbone	Routing	Paths	Bits/seg	Train attack	Notes
`vibevoice_14bit_v2`	VibeVoice acoustic decoder	`0,1,2,3,4,5,6` (×2 slots)	2	14	`nvlceq`	Pre-SPDMark-temporal-attacks training; eval JSON has no Hungarian verify block.
`diffwave_v5_2step`	DiffWave (LJSpeech)	`0,4,8,12,16,20,24`	2	7	`nvlceq`	2-step diffusion (`--diffusion_steps 2`), `lambda_perc 0.1` → louder watermark, faster inference.

Spectral attack codes: n noise, v gain, l lowpass, c crop, e erase, q quantize, r resample. SPDMark temporal codes: d segment-drop, s segment-swap, i segment-insert.

Inference

All three generators share a common eval interface; --verify enables the Hungarian + Binomial verifier (auto-sets --message_scheme hmac).

HiFi-GAN

python generate_molm_hifigan.py \
    --lora_weights checkpoints/hifigan_spdmark_spec/final/lora_weights.pt \
    --extractor_weights checkpoints/hifigan_spdmark_spec/final/extractor.pt \
    --routing_blocks 0,1,2,3,4,5,6 --num_paths 4 --lora_rank 64 \
    --test_dir <path-to-wavs> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/hifigan_spec_spectral --device cuda

DiffWave (current `spdmark_spec`)

python generate_molm_audio.py \
    --diffwave_checkpoint pretrained/diffwave-ljspeech.pt \
    --lora_weights checkpoints/diffwave_spdmark_spec/final/lora_weights.pt \
    --extractor_weights checkpoints/diffwave_spdmark_spec/final/extractor.pt \
    --routing_layers 0,4,8,12,16,20,24 --num_paths 4 --lora_rank 64 \
    --diffusion_steps 4 \
    --test_dir <path-to-wavs> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/diffwave_spec_spectral --device cuda

DiffWave (legacy `v5_2step`)

python generate_molm_audio.py \
    --diffwave_checkpoint pretrained/diffwave-ljspeech.pt \
    --lora_weights checkpoints/diffwave_v5_2step/final/lora_weights.pt \
    --extractor_weights checkpoints/diffwave_v5_2step/final/extractor.pt \
    --routing_layers 0,4,8,12,16,20,24 --num_paths 2 --lora_rank 64 \
    --diffusion_steps 2 \
    --test_dir <path-to-wavs> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/diffwave_v5_2step_spectral --device cuda

Note: lambda_perc 0.1 was used in training, so the watermark is louder (higher bit accuracy, lower SNR) than _spec.

VibeVoice (current `spdmark_spec`)

python generate_molm_vibevoice.py \
    --components_dir pretrained/vibevoice \
    --lora_weights checkpoints/vibevoice_spdmark_spec/final/lora_weights.pt \
    --extractor_weights checkpoints/vibevoice_spdmark_spec/final/extractor.pt \
    --routing_blocks 0,1,2,3,4,5,6 --num_paths 2 --lora_rank 16 --lora_alpha 8 \
    --num_frames 8 \
    --data_dir <path-to-precomputed-latents> --num_samples 20 \
    --chunked_generation --num_chunks 8 \
    --verify --gamma_f 0.01 --gamma_v 0.01 --attacks nvlceqr \
    --output_dir eval_out/vibevoice_spec_spectral --device cuda

VibeVoice (legacy `14bit_v2`)

Same flags as spdmark_spec but point at checkpoints/vibevoice_14bit_v2/final/. Note: this run was trained without SPDMark temporal attacks (dsi), so its robustness profile under --attacks d/s/i reflects the verifier alone.

Eval snapshots

Each eval/<run>_<regime>/results.json is the literal output of the corresponding generate_molm_*.py --output_dir call. Schemas differ between the "current" runs (which include a spdmark_verify block with tau_f, tau_v, detection_rate_*, avg_|Q|_*, avg_bit_accuracy_on_Q_*) and the "legacy" runs (older schema, just raw bit accuracy and quality metrics).

DiffWave numbers under eval/diffwave_spdmark_spec_step7000_* are from an undertrained snapshot (step 7000 of a 40k-step run) — re-evaluate against checkpoints/diffwave_spdmark_spec/final/ for the converged numbers.

License / citation

LJSpeech audio is CC0. HiFi-GAN/DiffWave/VibeVoice base models retain upstream licenses. Watermarking technique: Fares et al., SPDMark, arXiv 2512.12090.