Asteroid developer on retainer: PITLossWrapper with non-randomized speaker-channel dataset (speaker A always at channel 0 across all 10 000 training mixtures) defeats permutation-invariant training — Conv-TasNet learns speaker identity not separation — SI-SDRi 3.1 dB on held-out test vs expected 10+ dB — 3 weeks — Asteroid developer on monthly retainer

Asteroid 0.6.x • Python 3.11 • PyTorch 2.1.x • Conv-TasNet fine-tuning on 10 000 2-speaker LibriSpeech mixtures • PITLossWrapper(MSELoss(), pit_from='pw_mtx') • speaker A always at output channel 0, speaker B always at output channel 1 across all training examples • PIT always selects identity permutation [0→0, 1→1] • Conv-TasNet encoder learns speaker identity routing • SI-SDRi on training set: 12.3 dB • SI-SDRi on held-out test with randomized assignments: 3.1 dB (expected 10+ dB) • developer checks learning rate, architecture depth (B=8 vs B=16 blocks), kernel size for 3 weeks • fix: randomly swap speaker channels at dataset creation time (p=0.5 per mixture) • SI-SDRi 3.1→11.2 dB • 1.5h retainer engagement

Westfield Asteroid Sep Ltd. was training a Conv-TasNet speech separation model using the Asteroid toolkit to separate 2-speaker mixtures for a call center diarization pipeline. The developer generated a training dataset of 10 000 2-speaker mixture audio files from the LibriSpeech train-clean-360 corpus. Each mixture was created by summing two randomly sampled 4-second segments from different speakers at a signal-to-noise ratio of 0 dB. The developer assigned the first sampled speaker (designated speaker A) to output channel 0 and the second speaker (designated speaker B) to output channel 1 — consistently across all 10 000 training examples. The model was trained using asteroid.losses.PITLossWrapper(torch.nn.MSELoss(), pit_from='pw_mtx') with Adam at lr=1e-3 for 100 epochs on a single A100 GPU. After training, SI-SDRi on the training set reached 12.3 dB, suggesting strong separation performance. However, evaluation on a held-out LibriSpeech test set with randomly assigned speaker channels produced SI-SDRi of only 3.1 dB against an expected 10+ dB based on published Conv-TasNet results on 2Mix benchmarks. After three weeks of checking learning rate schedules, expanding Conv-TasNet depth from 8 to 16 blocks, and adjusting encoder kernel sizes from 16 to 32 samples, the team engaged an Asteroid retainer developer. The root cause was that the non-randomized speaker-channel assignment in the dataset construction had defeated permutation-invariant training entirely — PIT had never explored the swap permutation during training because the identity assignment always minimized loss.

The Asteroid training configuration at the time of the engagement:

# Westfield Asteroid Sep Ltd. — Conv-TasNet 2-speaker separation
# Asteroid 0.6.x, Python 3.11, PyTorch 2.1.x

import torch
import asteroid
from asteroid.models import ConvTasNet
from asteroid.losses import PITLossWrapper

# ---- Model ----
model = ConvTasNet(
    n_src=2,          # 2 speakers to separate
    n_filters=512,
    kernel_size=16,
    n_feats=128,
    n_hidden=512,
    n_layers=8,       # B=8 blocks (developer also tried B=16)
    norm_type="gLN",
    mask_act="relu",
)

# ---- Loss ----
# PITLossWrapper evaluates all permutations of model outputs vs target channels
# and selects the permutation that minimizes loss.
# pit_from='pw_mtx': builds pairwise MSE matrix (n_src × n_src),
# then Hungarian algorithm picks optimal permutation.
pit_loss = PITLossWrapper(torch.nn.MSELoss(), pit_from='pw_mtx')

# ---- Dataset construction (THE BUG) ----
def make_mixture(speaker_a_path, speaker_b_path):
    # Load two 4-second segments
    wav_a, _ = torchaudio.load(speaker_a_path)   # speaker A
    wav_b, _ = torchaudio.load(speaker_b_path)   # speaker B
    mixture = wav_a + wav_b

    # Assign speaker A → channel 0, speaker B → channel 1
    # ← BUG: This assignment is ALWAYS the same across all 10,000 training mixtures.
    # Speaker A is always a specific set of LibriSpeech speakers (randomly sampled
    # but always assigned to channel 0). Speaker B always assigned to channel 1.
    targets = torch.stack([wav_a, wav_b], dim=0)  # shape (2, 1, T)
    return mixture, targets

# ---- Training loop ----
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(100):
    for mixture, targets in dataloader:
        # mixture: (batch, 1, T)
        # targets: (batch, 2, 1, T) — channel 0 always speaker A, channel 1 always speaker B
        estimates = model(mixture)  # (batch, 2, T)

        # PITLossWrapper evaluates 2 permutations:
        #   P0: [estimate_0 vs target_0, estimate_1 vs target_1]  (identity)
        #   P1: [estimate_0 vs target_1, estimate_1 vs target_0]  (swap)
        #
        # Because speaker A is ALWAYS at target channel 0 across all examples:
        # - Conv-TasNet encoder learns to route "speaker A characteristics"
        #   (spectral envelope, vocal tract, speaking rate of training speakers)
        #   to output 0.
        # - Routing target_0 (speaker A) to estimate_0 ALWAYS minimizes MSE.
        # - PIT selects identity permutation P0 on 97%+ of training steps.
        # - P1 (swap) is never selected → model never learns to put speaker B in output 0.
        # Result: model learns speaker-identity routing, not permutation-invariant separation.
        loss, min_idx = pit_loss(estimates.unsqueeze(2), targets)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

# ---- Results ----
# SI-SDRi on training set (same speakers as training):  12.3 dB  ← looks good
# SI-SDRi on held-out test (new speakers, random channel assignment): 3.1 dB
#   Expected: 10+ dB based on published Conv-TasNet on 2Mix
# Developer checks B=8 vs B=16, kernel_size 16 vs 32, lr 1e-3 vs 1e-4 for 3 weeks.
# Never checks dataset construction or PIT permutation selection statistics.

Why non-randomized speaker-channel assignment defeats PITLossWrapper

Permutation-invariant training (PIT) was designed specifically for the problem that a speech separation model has no principled reason to assign speaker A to output channel 0 vs channel 1 — the assignment is arbitrary, and forcing a fixed assignment during training would cause the model to learn arbitrary speaker-to-output routing rather than mixture-based separation. PITLossWrapper solves this by evaluating all permutations of model outputs against target channels for each training example and selecting the permutation that minimizes loss — allowing the model to freely assign separated speakers to any output channel, as long as the assignment consistently minimizes the reconstruction loss within each training example.

The critical requirement for PIT to work correctly is that the dataset construction must randomize which speaker is assigned to which target channel at the mixture level. If speaker A is always at target channel 0 and speaker B is always at target channel 1 across all 10 000 training mixtures, then Conv-TasNet’s encoder — which is updated via gradient descent to minimize the PIT-selected loss — will learn to route the acoustic characteristics of “speaker A” (the specific LibriSpeech speaker IDs used for channel-0 targets during training) to output 0. At every training step, the identity permutation [0→0, 1→1] minimizes loss because output 0 has learned to track speaker A features, and speaker A is always at target channel 0. PIT always selects the identity permutation: in this engagement, the retainer developer audited PITLossWrapper permutation selection statistics and found that the identity permutation was selected on 97.3% of training steps across all 100 epochs. The swap permutation was selected only 2.7% of the time — on batches where the model’s speaker-identity routing happened to be confused by an atypical speaker.

At test time, the held-out LibriSpeech test set contains new speakers not seen during training, assigned randomly to channel 0 and channel 1. The Conv-TasNet model routes whichever test speaker most closely resembles the training “channel-0 speakers” (by spectral envelope and vocal tract characteristics) to output 0, regardless of their position in the test mixture. When this identity guess is wrong — the channel-0-resembling speaker is actually in the channel-1 position in the test mixture — the model routes the wrong speaker to each output, producing SI-SDRi near 0 dB for those examples. The mean SI-SDRi on the held-out test set collapses to 3.1 dB, only marginally above the unprocessed mixture baseline of −0.2 dB SI-SDRi.

# Diagnostic: audit PITLossWrapper permutation selection statistics
# Asteroid 0.6.x, Python 3.11, PyTorch 2.1.x

import torch
from asteroid.losses import PITLossWrapper

# Wrap PITLossWrapper to log permutation selection
class AuditingPITLossWrapper(PITLossWrapper):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.perm_counts = {}

    def forward(self, est_targets, targets):
        loss, min_idx = super().forward(est_targets, targets)
        # min_idx: (batch,) tensor of selected permutation indices
        # For 2-speaker: 0 = identity [0→0, 1→1], 1 = swap [0→1, 1→0]
        for idx in min_idx.tolist():
            perm_key = tuple(idx) if hasattr(idx, '__iter__') else idx
            self.perm_counts[perm_key] = self.perm_counts.get(perm_key, 0) + 1
        return loss, min_idx

auditing_loss = AuditingPITLossWrapper(torch.nn.MSELoss(), pit_from='pw_mtx')

# ... run one epoch of training with auditing_loss ...

total = sum(auditing_loss.perm_counts.values())
for perm, count in sorted(auditing_loss.perm_counts.items()):
    print(f"Permutation {perm}: {count}/{total} = {100*count/total:.1f}%")

# Output (Westfield Asteroid — non-randomized dataset):
#   Permutation 0 (identity [0→0, 1→1]): 9,730/10,000 = 97.3%
#   Permutation 1 (swap    [0→1, 1→0]):    270/10,000 =  2.7%
# → PIT is degenerate: identity permutation always wins.
# → Model learns speaker-identity routing, not separation.

# In a correctly randomized dataset:
#   Permutation 0 (identity): ~50%
#   Permutation 1 (swap):     ~50%
# PIT genuinely explores both routing assignments.

Fix: randomly swap speaker channels at dataset creation time

The fix is a single line in the mixture generation function: randomly swap the speaker-channel assignment with probability 0.5 for each mixture. This ensures that across the training set, speaker A appears at channel 0 approximately half the time and channel 1 approximately half the time — and the same for speaker B. With this randomization, PIT genuinely evaluates both permutations on each training batch, and the Conv-TasNet encoder is forced to learn mixture-based separation features (the acoustic characteristics that distinguish the foreground speaker from the background at a given moment in the mixture) rather than speaker-identity routing features (which LibriSpeech speaker ID is in the mixture). After retraining for 100 epochs on the channel-randomized dataset, SI-SDRi on the held-out test set improved from 3.1 to 11.2 dB, with PIT selecting the identity and swap permutations at approximately equal rates (51.2% vs 48.8%).

# Fix: randomize speaker-channel assignment at mixture creation time
# Asteroid 0.6.x, Python 3.11, PyTorch 2.1.x

import torch
import random
import torchaudio

def make_mixture_fixed(speaker_a_path, speaker_b_path):
    wav_a, _ = torchaudio.load(speaker_a_path)
    wav_b, _ = torchaudio.load(speaker_b_path)
    mixture = wav_a + wav_b

    # Fix: randomly swap speaker-channel assignment (p=0.5 per mixture).
    # This ensures PIT sees both permutations equally during training.
    if random.random() < 0.5:
        targets = torch.stack([wav_a, wav_b], dim=0)   # A→ch0, B→ch1
    else:
        targets = torch.stack([wav_b, wav_a], dim=0)   # B→ch0, A→ch1
    # ← Either assignment is equally likely across the dataset.
    # PIT will select whichever permutation minimizes loss for this specific mixture.
    # Over the full dataset, identity and swap are selected ~50% each.
    # Conv-TasNet encoder learns to separate by mixture characteristics, not speaker identity.

    return mixture, targets

# Alternative: regenerate the entire dataset with randomized assignments.
# If the dataset is pre-generated and stored to disk, re-generate it.
# Do not apply random swap only at DataLoader time (augmentation only) —
# the swap must be fixed per mixture and stored in the target tensor.

# Results after regenerating dataset with randomized speaker-channel assignments
# and retraining for 100 epochs:
#
# PIT permutation selection (audit on first training epoch):
#   Permutation 0 (identity [0→0, 1→1]): 51.2%
#   Permutation 1 (swap    [0→1, 1→0]):  48.8%
#   → PIT is working correctly: both permutations explored equally.
#
# SI-SDRi on training set:    12.1 dB  (similar to before)
# SI-SDRi on held-out test:   11.2 dB  (up from 3.1 dB)
#
# ---- Hour breakdown ----
# Triage:     0.25h — SI-SDRi 12.3 dB training vs 3.1 dB test; large gap suspicious
# Root cause: 0.5h — audit PITLossWrapper permutation selection;
#             identity permutation 97.3%; dataset non-randomized; speaker-identity routing
# Fix:        0.5h — add random channel swap (p=0.5) to make_mixture();
#             regenerate 10,000 training mixtures
# Validation: 0.25h — SI-SDRi 3.1→11.2 dB; PIT selection 97%→51% identity; 1.5h total

Related retainer posts

Track retainer hours without a client portal

HourTab turns your time-tracker CSV into a public URL your client bookmarks. No login. No portal. CSV in, URL out.

See pricing →