Iterative Audio Separation with Mixture Consistency via MIMO Model Extension

01. Concept

An Iterative Refinement Paradigm Across Arbitrary Methods

In deep-learning-based audio source separation, conventional research has primarily focused on architectural design improvements (e.g., Conv-TasNet, Band-Split RNN, BS-RoFormer, SCNet). These models operate under single-step inference (single-input single-output or single-input multi-output configurations), where a single forward pass estimates target sources directly from the mixture input.

Iterative refinement paradigm across arbitrary methods.

While architectural advancements have brought impressive gains, single-step inference leaves the power of progressive refinement untapped. Conversely, generative iterative models (such as diffusion-based methods) achieve high perceptual quality through multi-step sampling, but struggle to enforce strict mixture consistency—an essential constraint for high-fidelity audio stem separation.

The Core Concept of This Work:

Architecture-Agnostic Framework: Instead of proposing yet another custom network backbone, we propose a general framework that converts any existing single-step separation model into an iterative refinement model.
Mutual Refinement via MIMO: By extending the backbone to a Multi-Input Multi-Output (MIMO) configuration, the model receives all estimated source signals from the previous iteration step as inputs, enabling joint, mutual refinement across all target stems.
Strict Mixture Consistency: Preserves exact mixture reconstruction ($\sum \hat{s}_k = x$) after every iteration step, stabilizing iterative prediction and maintaining phase and timbral fidelity.

02. Algorithm

Algorithmic Explanation

Naively feeding back a single output into a SISO/SIMO model causes error accumulation and over-separation because the model cannot correct errors across all stems simultaneously. Our MIMO framework resolves this through a simple, effective algorithm:

Initial Input Setup (Iteration 0): To satisfy mixture consistency from the outset, the input source estimates are initialized by dividing the input mixture equally among target stems: s_k[0] = mixture / K.
Joint MIMO Forward Pass (Iteration i): At step i, all estimated source signals from step i-1 are fed back into the MIMO model alongside an iteration step condition i. The model updates estimates by modeling inter-stem interactions and removing cross-stem leakage.
Mixture Consistency Projection: After each iteration step i, a projection is applied so that the estimated stems sum back exactly to the original mixture: s_k[i] ← s_k[i] + (mixture - Σ s_j[i]) / K.
Discriminator & Generative Extension (Optional): Stem-wise discriminators enhance fine high-frequency details. Optionally, a zero-sum noise perturbation scheme enables stochastic generative refinement while maintaining strict mixture constraints.

Overview of the proposed MIMO iterative separation framework

Overview of the proposed framework for iterative audio separation.

03. Experiments

Experimental Results

Evaluated on the MUSDB18-HQ dataset across 2-stem and 4-stem separation tasks.

1. Vocal-Accompaniment Separation (BS-RoFormer Backbone)

Applying our MIMO iterative framework to BS-RoFormer achieves consistent improvements of separation for both vocals and accompaniment over strong single-step baselines.

Model	Iter	Params	Vocal			Accompaniment
Model	Iter	Params	SDR	SIR	SAR	SDR	SIR	SAR
BS-RoFormer (Baseline)	N/A	72.2 M	11.92	21.75	13.04	24.54	40.15	25.47
MIMO BS-RoFormer	3	71.4 M	12.29	23.93	13.71	26.01	39.85	26.73

Vocal-Accompaniment separation evaluation on MUSDB18-HQ.

2. 4-Stem Separation (SCNet Backbone)

To demonstrate general applicability across architectures and stem counts, we evaluated 4-stem separation (vocals, bass, drums, other) using SCNet. The iterative MIMO SCNet outperforms single-step SCNet across all four stems simultaneously.

Model	Iter	Params	Vocals			Bass			Drums			Other
Model	Iter	Params	SDR	SIR	SAR	SDR	SIR	SAR	SDR	SIR	SAR	SDR	SIR	SAR
SCNet (Baseline)	N/A	10.6 M	10.39	18.15	11.77	9.10	15.64	12.47	11.32	19.58	12.72	7.21	12.81	9.58
MIMO SCNet	2	10.6 M	10.57	18.62	11.92	9.80	16.98	13.65	11.73	20.39	12.94	7.63	13.61	9.75

4-Stem separation (Vocals, Bass, Drums, Other) evaluation on MUSDB18-HQ.

04. Implementation Guide

Applying the MIMO Framework to Your Method

Because the proposed framework is general and architecture-agnostic, you can apply it to your custom source separation models by following these 3 steps:

Develop your base model: Start with your standard PyTorch single-step source separation backbone (e.g., SISO or SIMO model architecture).
Extend the model architecture to MIMO: Widen the input channel dimension to receive $K \times C$ multi-source audio tensors (where $K$ is the number of target stems and $C$ is audio channels), output $K$ stems (or $K \times K$ cross-stem spectral masks), and accept a time-step embedding input `t`. Additionally, if your model is spectral masking-based, you can adopt mask-weighted prediction for the MIMO extension by referring to Section 3.1 of the paper.
Wrap with MIMOBase: Wrap your MIMO backbone using the provided MIMOBase class (`src/model/mimo_base.py`), which handles multi-step iterative refinement and mixture consistency projection during training and inference.

import torch
import torch.nn as nn

# Step 1: Develop your base separation backbone (SISO / SIMO)
class MyBaseSeparator(nn.Module):
    def __init__(self, in_channels=2, out_channels=2):
        super().__init__()
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1)
        
    def forward(self, x):
        # Single-step prediction: [B, C, T] -> [B, C, T]
        return self.conv(x)

import torch
import torch.nn as nn

# Step 2: Extend your model to a MIMO configuration with step conditioning
class MyMIMOSeparator(nn.Module):
    def __init__(self, base_model, num_sources=2, in_channels=2):
        super().__init__()
        self.num_sources = num_sources
        self.in_channels = in_channels
        
        # Input layer receives K * C channels (all estimated stems as multi-channel inputs)
        self.input_layer = nn.Conv1d(num_sources * in_channels, in_channels, kernel_size=1)
        self.base_model = base_model
        
        # Time step condition embedding for iteration step i
        self.time_emb = nn.Embedding(num_embeddings=10, embedding_dim=in_channels)

    def forward(self, x, t=None):
        # x: Multi-source input tensor [B, K, C, T]
        B, K, C, T = x.shape
        x_flat = x.view(B, K * C, T)  # Flatten stems along channel dimension
        
        h = self.input_layer(x_flat)
        if t is not None:
            h = h + self.time_emb(t).unsqueeze(-1)
            
        out = self.base_model(h)  # Output tensor [B, K*C, T]
        return out.view(B, K, C, T)  # Reshape to [B, K, C, T]

from src.model.mimo_base import MIMOBase

# Step 3: Wrap your MIMO backbone with MIMOBase
mimo_backbone = MyMIMOSeparator(base_model=MyBaseSeparator(), num_sources=2, in_channels=2)

model = MIMOBase(
    num_channels=2,
    num_sources=2,
    backbone_model=mimo_backbone,
    max_iter=3,                     # Number of iterative refinement steps
    model_output_style="direct",      # "direct" or "diff"
    mixture_consistency=True,         # Enables mixture consistency projection
    use_time_emb=True
)

# The model automatically performs I-step iterative refinement
# and enforces mixture consistency: sum(pred_stems) == mixture
outputs = model(x_multi_sources, t)

Iterative Audio Separation with Mixture Consistencyvia MIMO Model Extension