Audio Source Separation Architecture, in depth: masking, the phase problem, Conv-TasNet, Demucs, PIT training, evaluation and streaming inference

The problem, stated precisely

Model the recording as x(t) = s_1(t) + s_2(t) + ... + s_N(t), possibly with each source filtered by the room or the mix. Three tasks share this equation. Speech enhancement has two sources, speech and everything else, and keeps one. Speech separation has several speakers of the same kind, so there is no natural label that says which output is speaker 1. Music demixing has a fixed set of named stems; the common research setup has four (vocals, drums, bass, other), defined by the MUSDB18 dataset of 150 professionally mixed tracks split into 100 for training and 50 for testing.

The distinction matters for training. Named stems give each output a fixed target. Interchangeable speakers do not, and the model would be punished for producing the right voices in the 'wrong' order unless the loss ignores order, which is what permutation-invariant training does.

Time-frequency masking and the phase problem

The short-time Fourier transform (STFT) slices the waveform into overlapping windows, for example 4,096 samples with a hop of 1,024 at 44.1 kHz, and turns each window into complex frequency bins. In that representation different sources often occupy different bins at a given moment, so you can separate them by multiplying the mixture spectrogram by a mask per source. The ideal binary mask assigns each bin to whichever source is loudest there; the ideal ratio mask gives each source the fraction |S_i| / sum_j |S_j|. A network learns to predict the mask from the mixture magnitude.

The catch is phase. Most magnitude-masking models reuse the mixture's phase for every output, because phase is hard to predict: it looks like noise and wraps around. Where sources overlap in a bin the mixture phase is wrong for each of them, and the result is smearing and the warbling artefact called musical noise. Remedies include complex ratio masks that scale and rotate each bin, multichannel Wiener filtering as a post-step (Open-Unmix uses one to refine its per-target magnitude estimates), and models that work on the waveform directly. Before training anything, measure the ceiling your representation allows with oracle masks computed from the true stems:

import numpy as np, librosa

def oracle_irm(stems, n_fft=4096, hop=1024):
    """Separate the mixture with ideal ratio masks built from the true stems."""
    mix = sum(stems)
    X = librosa.stft(mix, n_fft=n_fft, hop_length=hop)
    S = [np.abs(librosa.stft(s, n_fft=n_fft, hop_length=hop)) for s in stems]
    total = sum(S) + 1e-8
    return [librosa.istft((Si / total) * X, hop_length=hop, length=len(mix)) for Si in S]

def si_sdr(est, ref):
    ref = ref - ref.mean(); est = est - est.mean()
    target = (est @ ref) / (ref @ ref + 1e-8) * ref
    return 10 * np.log10((target @ target) / ((est - target) @ (est - target) + 1e-8))

If oracle masks with mixture phase already leave audible artefacts on your material, a better mask network will not remove them; you need complex masks or a time-domain model.

Learned-basis and waveform models

Conv-TasNet replaced the STFT with a learned encoder: a 1-D convolution with short kernels (a few milliseconds) and hundreds of filters, producing a non-negative representation. A separator built from stacked dilated temporal convolution blocks predicts one mask per source on that representation, and a transposed-convolution decoder turns each masked representation back into a waveform. Because the basis is learned and the decoder is trained end to end, there is no fixed phase to reuse. Short kernels give low latency, which is why Conv-TasNet-style models suit speech separation and real-time use; they were designed and mostly evaluated on clean speech mixtures such as WSJ0-2mix.

Demucs took the music route: a U-Net over the raw waveform, with strided convolutions down, a recurrent or attention bottleneck, and transposed convolutions up, skip connections joining matching levels. Later versions became hybrids. Hybrid Transformer Demucs runs a waveform branch and a spectrogram branch side by side, joins them with cross-domain transformer layers, and sums the decoded outputs, taking advantage of the spectrogram's frequency resolution and the waveform's phase accuracy. Band-split models take the spectrogram route instead: they split frequency into sub-bands of different widths, narrow where vocals and bass carry detail and wide at the top, model each band, and then model across time and across bands. Spleeter, a spectrogram U-Net with pretrained 2-, 4- and 5-stem models, remains a fast baseline.

The practical rule: masking models over magnitude are cheap and good enough for enhancement; hybrid and band-split models lead on music quality and cost much more compute; learned-basis time-domain models with short frames are the natural choice when latency must be tens of milliseconds.

Training: loss, permutation invariance and data

The most common loss for speech is negative scale-invariant SNR (SI-SNR): project the estimate onto the reference to find the part that is just a scaled copy of the target, and compare that part's energy with the energy of everything else. Scale invariance stops the model from gaming the loss by changing volume. Music models often train with L1 on the waveform or on spectrogram magnitudes, sometimes both.

For interchangeable speakers, permutation-invariant training (PIT) computes the loss under every assignment of outputs to references and backpropagates the best one. With two or three speakers that is two or six permutations; for many sources the Hungarian algorithm finds the best assignment without enumerating them.

import itertools, torch

def si_snr(est, ref, eps=1e-8):            # est, ref: (batch, time)
    est = est - est.mean(-1, keepdim=True); ref = ref - ref.mean(-1, keepdim=True)
    s = (est * ref).sum(-1, keepdim=True) / (ref.pow(2).sum(-1, keepdim=True) + eps) * ref
    n = est - s
    return 10 * torch.log10(s.pow(2).sum(-1) / (n.pow(2).sum(-1) + eps) + eps)

def pit_loss(est, ref):                     # (batch, n_src, time)
    n = est.shape[1]
    scores = torch.stack([
        torch.stack([si_snr(est[:, i], ref[:, p]) for i, p in enumerate(perm)]).mean(0)
        for perm in itertools.permutations(range(n))], dim=1)   # (batch, n_perm)
    return -scores.max(dim=1).values.mean()

for mix, ref in loader:                    # training step
    est = model(mix)                       # (batch, n_src, time)
    loss = pit_loss(est, ref)
    opt.zero_grad(); loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
    opt.step()

Data is the limiting factor for music. Public multitrack datasets are small, so augmentation does real work: remix stems from different songs into new mixtures, apply random gains, swap stereo channels, flip polarity, and shift pitch or tempo. For speech, synthesise mixtures from clean utterances at random overlaps and signal-to-noise ratios, then add reverberation and noise if you will deploy in real rooms. A model trained on anechoic mixtures falls apart on a conference-room recording.

Measuring separation

The BSS Eval family decomposes each estimate into the target, interference from other sources, and artefacts, giving SDR (overall), SIR (leakage from other sources) and SAR (artefacts). Music papers usually report SDR computed on one-second windows with the museval toolkit and summarised by a median; speech papers usually report SI-SNR or SI-SDR improvement over the unprocessed mixture. Numbers from different protocols are not comparable, so fix one protocol for your own evaluations.

Metrics and ears disagree. A model can gain SDR while adding a metallic artefact listeners hate, or lose SDR by leaving a little bleed that nobody notices. Keep a fixed listening set of difficult clips (dense mixes, reverb, solo instruments that resemble voices) and listen to every candidate before it ships.

Inference on long files: chunking and overlap-add

Models are trained on segments of a few seconds and memory grows with input length, so a four-minute track is processed in chunks. Cutting chunks naively produces clicks at the boundaries, because a model behaves differently near the edges of its input. Process overlapping chunks, weight each with a window that fades in and out, add them, and divide by the summed weights.

import torch

@torch.no_grad()
def separate_long(model, mix, seg=10 * 44100, overlap=0.25):
    """mix: (channels, time) -> (n_src, channels, time) with crossfaded chunks."""
    hop = int(seg * (1 - overlap))
    win = torch.hann_window(seg, periodic=False).clamp_min(1e-3)
    n_src = model.n_src
    out = torch.zeros(n_src, *mix.shape)
    norm = torch.zeros(mix.shape[-1])
    for start in range(0, mix.shape[-1], hop):
        chunk = mix[:, start:start + seg]
        length = chunk.shape[-1]
        if length < seg:
            chunk = torch.nn.functional.pad(chunk, (0, seg - length))
        est = model(chunk.unsqueeze(0))[0, ..., :length]  # (n_src, ch, length)
        out[..., start:start + length] += est * win[:length]
        norm[start:start + length] += win[:length]
    return out / norm

Normalise input loudness to the range the model was trained on, run in float32 or carefully tested half precision, and check that the stems sum back to something close to the mixture. A large residual is a quick sign that a model or a chunking bug is losing energy. Some music models also average predictions over a few small random time shifts, which trades extra compute for slightly cleaner output.

Real-time and streaming separation

In a call or a hearing device, the budget is algorithmic latency, the audio the model must buffer before it can emit output, plus compute time per frame. For an STFT model that is roughly the window length plus any future frames the model looks at; a 4,096-sample window at 44.1 kHz is already 93 ms before compute, which is why streaming speech models use windows of 20 to 32 ms or learned-basis encoders with kernels of a few milliseconds. Bidirectional recurrent layers, non-causal convolutions and attention over the whole clip all need future input and must be replaced by causal versions, carrying state between frames. Expect a quality drop compared with the offline model, and measure it on your own audio.

Compute must also fit: every frame must finish in less than its hop duration on the slowest target device, with headroom for the rest of the audio pipeline. Profile on the device, not on a workstation GPU.

Worked example: a vocal-removal service

A karaoke product removes vocals from uploaded songs. Requests are offline, so latency is minutes, not milliseconds, and quality dominates: choose a hybrid or band-split music model and run it on GPUs. A four-minute stereo song at 44.1 kHz is 10.6 million samples per channel. With 10-second chunks and 25% overlap the hop is 7.5 seconds, so the song needs 32 chunks. Batch the chunks on the GPU, measure seconds of GPU time per minute of audio for your chosen model, and size the fleet from expected minutes per hour at peak plus a queue.

Output only the stem you need: the instrumental is the mixture minus the vocal estimate, or the sum of the other three stems; listen to both, because they fail differently. Store the input's loudness and restore it, encode the result at a bit rate that does not add its own artefacts, and cache by audio fingerprint so repeat uploads of popular songs cost nothing. Before launch, check the licensing of the model weights for commercial use and the rights you need to process users' uploads.

Failure modes and operational guidance

• Bleed. Hi-hats in the vocal stem, vocal reverb tails in the instrumental. Expected to a degree; track it with SIR on a test set and with listening.
• Musical noise and phasiness. Magnitude masks with mixture phase. Use complex masks, Wiener post-filtering or a waveform model.
• Out-of-domain material. Genres, languages or instruments absent from training: a sustained synth lead treated as vocals, a choir split between stems. Build an evaluation set from your real traffic.
• Sample-rate mismatch. A 44.1 kHz model fed 16 kHz audio, or the reverse. Resample explicitly at the boundary and assert the rate.
• Clipping and loudness. Heavily limited masters and loudness outside the training range produce distortion. Normalise in, restore out, and use float audio internally.
• Chunk seams. Clicks every few seconds mean missing overlap or a mismatched window normalisation.
• Speaker count mismatch. A two-speaker model given three speakers merges two of them. Detect speaker count first, or use models trained for a variable number of speakers.

Related reading on this site

• Noise suppression, the real-time two-source special case
• Music audio architecture
• Dereverberation, often needed before or after separation
• Resampling, and getting sample rates right at system boundaries
• Real-time audio pipelines and latency budgets

What to do next

1. Decide which task you have: enhancement, speaker separation or named-stem demixing, and offline or streaming.
2. Collect 30 to 50 representative clips with references, or synthesise mixtures from clean stems that match your domain.
3. Compute oracle-mask SI-SDR on those clips to learn the ceiling of a magnitude-masking approach.
4. Evaluate two pretrained models of different families with one fixed metric protocol and a listening session.
5. Implement chunked overlap-add inference, check that the stems sum back to the mixture, and profile seconds of compute per minute of audio.
6. Fine-tune on your domain with remixing augmentation if the gap matters, and monitor bleed and artefacts on a fixed test set with every model change.