Fix speaker_encoder silent decay under AdamW weight decay in SFT

[MR-AI-HuiMin]

Summary

• Freeze speaker_encoder and exclude it from AdamW in finetuning/sft_12hz.py.
• After locking the first speaker embedding, reuse that same vector for training conditioning, so train-time speaker identity matches the row written into codec_embedding (and used by CustomVoice inference).

Why this is risky (detailed)

What the current script does

1. Every step extracts speaker_embedding = model.speaker_encoder(ref_mels).detach().
2. Optimizer is AdamW(model.parameters(), weight_decay=0.01), which includes speaker_encoder.
3. Checkpoint export drops speaker_encoder.* and writes a fixed vector into talker.model.codec_embedding.weight[spk_id].
4. Inference for CustomVoice uses that fixed codec_embedding row — not a live speaker_encoder.

Root cause

.detach() means the TTS loss does not produce useful gradients for speaker_encoder.
AdamW still applies weight decay to those parameters every optimizer step (θ ← θ - lr * (… + wd * θ)).
Net effect: speaker_encoder is slowly driven toward zero / smaller norms even though it is “not being trained” in the usual sense.

Impact on the official single-speaker script

Official code already locks the first embedding for saving, so checkpoint norms do not follow a pure 1/(epoch+1) average collapse.

However there is still a train / infer mismatch risk:

• During training, later epochs keep injecting online encoder outputs. If the encoder is being weight-decayed, those conditioning vectors shrink / drift over time.
• At inference, CustomVoice uses the first-locked vector stored in codec_embedding.
• Talker / MTP therefore learn under a moving speaker condition that is not what deployment uses. Longer SFT (many epochs, small lr, non-trivial weight_decay) makes this worse.

Higher risk for community / multi-speaker forks

Several community trainers (e.g. multi-speaker ports) keep the same AdamW(model.parameters(), weight_decay=0.01) pattern but average speaker embeddings across the run when saving. In that setting, saved norms empirically collapse roughly as:

[
|e_{\text{avg}}| \approx |e_0| / (epoch + 1)
]

when encoder outputs shrink while being accumulated. This directly hurts speaker similarity after SFT.

Related fix already proposed for a multi-speaker trainer: vspeech/Qwen3-TTS-Train#6

Even without averaging, any recipe that:

• keeps speaker_encoder in AdamW with nonzero weight_decay, and
• relies on detached encoder outputs for identity,

is exposing users to silent speaker-identity degradation.

Why freezing is the right fix

• SFT does not need to update speaker_encoder (outputs are detached; weights are discarded from the CustomVoice export anyway).
• Freezing + excluding from AdamW stops the silent decay.
• Reusing the locked embedding for conditioning aligns training with CustomVoice inference.

Changes

1. speaker_encoder.requires_grad_(False) and build AdamW only on remaining trainable params.
2. Lock target_speaker_embedding from the first batch sample; expand/reuse it for input_codec_embedding[:, 6, :] on all subsequent steps.

Test plan

• Run a short single-speaker SFT (e.g. 2–3 epochs) with and without this patch.
• Log speaker_encoder parameter norms (or a probe forward embedding norm) across epochs — without the patch they trend down under weight decay; with the patch they stay flat.
• Confirm exported model.safetensors still has no speaker_encoder.* keys and still writes codec_embedding.weight[3000].
• Smoke generate_custom_voice on the exported checkpoint and check speaker similarity is stable vs an early epoch (no progressive timbre wash-out).

Notes / non-goals

• This PR does not add multi-speaker training to the official script.
• If a future recipe intentionally fine-tunes speaker_encoder with real gradients, it should remove .detach(), keep the encoder in the optimizer intentionally, and not drop it from the export — that is a different design.

Made with Cursor