docs/layer_extension_utils.md

Sortformer Output Layer Extension: Script Architecture

Background

Extending the Sortformer from N to M speakers involves three sequential operations:

Stage	Operation	Input	Output
1	Extend weights (SVD orthogonal init) + create split head	N-spk unified model	M-spk split-head model
2	Fix config (recursively update all num_spks)	M-spk model with stale config fields	M-spk model with consistent config
3	Merge split head back to unified head	Post-finetune split-head model	Standard unified inference model

Architecture Decision

extend_output_layer.py     →  Stages 1 + 2  (extension + config fix)
merge_split_head.py        →  Stage 3        (post-finetune merge)

Why merge Stages 1+2

• Stage 1 and Stage 2 are always executed together after extension. There is no valid use case for extending weights without fixing the config — doing so produces a checkpoint with stale num_spks fields in train_ds/validation_ds/test_ds, which causes hungarian_get_perm IndexError during training or inference.
• Running both in a single script avoids intermediate .nemo file I/O (load + save ~10-15s for a 500MB checkpoint).
• Eliminates the "forgot to run step 2" class of bugs entirely.

Why Stage 3 is separate

• Stage 1+2 runs before finetuning. Stage 3 runs after finetuning.
• The two workflows are temporally separated (days/weeks apart in practice) and have completely different inputs: Stage 1+2 takes a pretrained N-spk checkpoint, Stage 3 takes a finetuned split-head checkpoint.
• Keeping Stage 3 as its own script makes it discoverable and self-documenting — a user looking at src/finetune_pipeline/scripts/ sees merge_split_head.py and immediately understands its purpose.

Usage

# --- Pre-finetune: extend 4spk → 10spk (Stages 1+2) ---
python src/finetune_pipeline/scripts/extend_output_layer.py \
    --src checkpoint_4spk.nemo \
    --dst-spk 10 \
    --out checkpoint_10spk_extended.nemo

# --- Finetune with split head (differential LR) ---
bash train.sh \
    --init_nemo_path checkpoint_10spk_extended.nemo ...

# --- Post-finetune: merge split head → unified inference model ---
python src/finetune_pipeline/scripts/merge_split_head.py \
    --src exp/.../checkpoints/model--val_loss=xxx.nemo \
    --out checkpoint_10spk_inference.nemo

Implementation Details

extend_output_layer.py (Stages 1+2)

1. Stage 1 (weight extension):

   • Load source model, extract unified weight matrix from either unified or already-split head via get_unified_output_weights()
   • Generate new rows via orthogonal_extend_weight() (SVD-based: uses right singular vectors as directions for new rows, falls back to random-magnitude noise when Vh is exhausted)
   • Extend bias similarly with random values centered at source bias statistics
   • Build target model with split head: single_hidden_to_spks_base (n_base_spks=N) + single_hidden_to_spks_new (n_new=M-N)
   • Copy matching-shape weights from source; inject extended base/new weights

2. Stage 2 (config fix):

   • _update_all_num_spks(cfg, n_dst) recursively traverses the OmegaConf tree, replacing every max_num_of_spks and num_spks with the target value
   • This covers train_ds, validation_ds, test_ds, and any nested config sections
   • Sets sortformer_modules.n_base_spks = n_src to enable differential LR during subsequent finetuning

merge_split_head.py (Stage 3)

1. Load finetuned split-head model; verify n_base_spks > 0 and split keys exist
2. Concatenate base.weight with new.weight (and biases) along dim=0
3. Build target config with n_base_spks=0 → model uses unified single_hidden_to_spks path
4. Copy all state dict entries from finetuned model, excluding the 4 split head keys, then assign the concatenated weights to single_hidden_to_spks.{weight,bias}
5. Save unified inference checkpoint

Shared Constants

Both scripts share the same state dict key constants:

SK_BASE_W = "sortformer_modules.single_hidden_to_spks_base.weight"
SK_BASE_B = "sortformer_modules.single_hidden_to_spks_base.bias"
SK_NEW_W  = "sortformer_modules.single_hidden_to_spks_new.weight"
SK_NEW_B  = "sortformer_modules.single_hidden_to_spks_new.bias"
SK_UNI_W  = "sortformer_modules.single_hidden_to_spks.weight"
SK_UNI_B  = "sortformer_modules.single_hidden_to_spks.bias"

And _update_all_num_spks():

def _update_all_num_spks(cfg, target: int):
    """Replace every num_spks / max_num_of_spks occurrence in the config tree."""
    if OmegaConf.is_list(cfg):
        return
    if "max_num_of_spks" in cfg:
        cfg.max_num_of_spks = target
    if "num_spks" in cfg:
        cfg.num_spks = target
    for key in list(cfg.keys()):
        val = cfg[key]
        if OmegaConf.is_dict(val) or OmegaConf.is_config(val):
            _update_all_num_spks(val, target)

Existing Script Disposition

Script	Disposition
extend_output_layer.py	Rewritten — now includes Stage 2 config fix
convert_extended_sortformer.py	Deleted — logic absorbed into extend_output_layer.py
merge_split_head.py	New — standalone Stage 3 script