Source Separation by Flow Matching (2505.16119v1)

Abstract: We consider the problem of single-channel audio source separation with the goal of reconstructing $K$ sources from their mixture. We address this ill-posed problem with FLOSS (FLOw matching for Source Separation), a constrained generation method based on flow matching, ensuring strict mixture consistency. Flow matching is a general methodology that, when given samples from two probability distributions defined on the same space, learns an ordinary differential equation to output a sample from one of the distributions when provided with a sample from the other. In our context, we have access to samples from the joint distribution of $K$ sources and so the corresponding samples from the lower-dimensional distribution of their mixture. To apply flow matching, we augment these mixture samples with artificial noise components to ensure the resulting "augmented" distribution matches the dimensionality of the $K$ source distribution. Additionally, as any permutation of the sources yields the same mixture, we adopt an equivariant formulation of flow matching which relies on a suitable custom-designed neural network architecture. We demonstrate the performance of the method for the separation of overlapping speech.

• The paper introduces Flowsep, framing source separation as a constrained generative problem that maps noisy mixtures to clean speech signals using flow matching.
• It employs a permutation equivariant training loss that identifies the optimal source permutation at the outset to maintain mixture consistency throughout ODE integration.
• Experiments demonstrate that Flowsep attains state-of-the-art SI-SDR, speech intelligibility, and perceptual quality metrics, balancing effective performance with inference efficiency.

This paper (2505.16119) introduces Flowsep, a novel method for single-channel audio source separation (specifically, speech) based on the concept of flow matching. The goal is to reconstruct $K$ individual source signals ($\vs_1, ..., \vs_K$) given only their linear mixture ($\bar{\vs} = K^{-1} \sum_{k=1}^K \vs_k$). This is an ill-posed inverse problem because multiple combinations of sources can produce the same mixture, and the permutation of the sources is ambiguous. Flowsep tackles this by framing source separation as a constrained generative problem, ensuring the separated sources sum up correctly to the original mixture.

The core idea is to use flow matching to learn a deterministic mapping (an Ordinary Differential Equation or ODE) that transforms samples from a simple initial distribution, conditioned on the mixture $\bar{\vs}$, to samples from the complex target distribution of separated sources, also conditioned on $\bar{\vs}$ ($p(\mS|\bar{\vs})$).

Problem Formulation and Flow Matching Setup

The sources $\vs_{1:K}$ are represented as a $K \times L$ matrix $\mS$, where $L$ is the length of the audio signals. The mixture is $\bar{\vs}$, and its $K$ stacked copies form $\bar{\mS}$. The mixture constraint is $\bar{\mS} = \mP \mS$, where $\mP = K^{-1}\vone\vone^\top$ is a projection matrix. The problem involves recovering the components of $\mS$ that lie in the subspace orthogonal to $\vone$, spanned by $\mP^\perp = \mI_K - \mP$. Figure 1 in the paper illustrates this geometry, showing the mixture living in a 1D subspace and the missing information (the differences between sources) residing in the orthogonal subspace.

Flow matching learns the drift $v^\theta(t, \vx, \vc)$ of an ODE $\frac{d\vx_t}{dt} = v^\theta(t, \vx_t, \vc)$ that transports samples from an initial distribution $p_0(\vx_0|\vc)$ to a target distribution $p_1(\vx_1|\vc)$. The training objective is to minimize the difference between the learned drift and the true drift of a simple linear interpolation between samples from $p_0$ and $p_1$.

For source separation, the conditioning variable $\vc$ is the mixture $\bar{\vs}$.

1. Target Distribution ($p_1|\bar{\vs}$): This is the conditional distribution of the separated sources given the mixture, $p(\mS|\bar{\vs})$.
2. Initial Distribution ($p_0|\bar{\vs}$): A sample $\vx_0$ is constructed by taking the mixture stacked $K$ times ($\bar{\mS}$) and adding noise $\mZ \sim \mathcal{N}(\vzero, \mI_{K \times L})$ projected onto the difference subspace ($\mP^\perp$). Specifically, $\vx_0 = \bar{\mS} + \mP^{\perp} \sigma(\bar{\vs}) \mZ$. The scaling factor $\sigma(\bar{\vs})$ is crucial; the paper explores scaling by the active power $g_{\operatorname{pwr}}(\bar{\vs})$ or the square root of the energy envelope $g_{\operator{env}}(\bar{\vs})$, the latter being preferred as it adds less noise in inactive regions. This setup ensures $\mP \vx_0 = \bar{\mS}$.

The linear interpolation used in flow matching is $\vx_t = t \vx_1 + (1-t) \vx_0$. With $\vx_1=\mS$, the objective for the network $v^\theta$ is to approximate the time-dependent vector field $(\mS - \vx_0)$, projected onto the $\mP^\perp$ subspace. The network output is parameterized as $v^\theta(t, \vx_t, \bar{\vs})=\mP^\perp \tilde{v}^\theta(t,\mP^\perp\vx_t,\bar{\vs})$. This structural choice ensures that the learned flow maintains the mixture consistency $\mP\vx_t = \bar{\mS}$ throughout the ODE integration, a key advantage.

Handling Permutation Equivariance

A major challenge in source separation is that $p(\mS|\bar{\vs}) = p(\pi \mS|\bar{\vs})$ for any permutation $\pi$ of the sources. The standard flow matching loss, based on the squared Euclidean distance $||\vx_1 - \vx_0||^2$, implicitly favors one arbitrary permutation during training, which doesn't generalize well.

The paper proposes a Permutation Equivariant Training (PET) loss (2505.16119) that leverages the concept of Permutation Invariant Training (PIT) [kolbaekMultitalkerSpeechSeparation2017]. The core idea is to identify the optimal source permutation at the start of the flow ($t=0$) based on the minimal distance objective, and then fix this permutation for training the flow network for all time steps $t \in (0, 1]$. The loss for a given $\vx_0, \vx_1$ pair and permutation $\pi$ is defined as $\calL(\pi,t) = \| v^{\theta}(t, \vx_t(\vx_0,\pi\vx_1),\bar{\vs}) - (\pi\vx_1 - \vx_0)\|^2$. The PET loss finds the permutation $\pi^{\text{PIT}}$ that minimizes $\calL(\pi, 0)$ and then minimizes $\calL(\pi^{\text{PIT}}, 0) + \calL(\pi^{\text{PIT}}, t)$ averaged over $t$ and data samples. This ensures the network learns to map $\vx_0$ to the correct permuted $\vx_1$ initially and maintains that mapping throughout the trajectory.

For practical audio separation, the mean squared error is often insufficient. The authors found significant improvements using a normalized and decibel-valued loss based on the negative Signal-to-Noise Ratio (SNR).

$\calL^{\operatorname{N}}(\pi, t) = \frac{\| v^{\theta}(t, \vx_t(\vx_0,\pi\vx_1 ),\bar{\vs}) - (\pi\vx_1 - \vx_0)\|^2}{\|\pi\vx_1 - \vx_0\|^2}$

The actual training loss is $10\log_{10} \calL^{\operatorname{N}}(\pi^{\text{PIT}}, t)$, averaged with the $t=0$ term. This objective is less sensitive to signal amplitude, encouraging better separation across different energy levels.

Permutation Equivariant Network Architecture

To implement the drift network $v^\theta(t, \vx, \bar{\vs})$ in a permutation-equivariant manner, a specialized architecture called PEM-Loco is introduced (Figures 3 and 4 in the paper). The network processes audio signals in the time-frequency domain using a Short-Time Fourier Transform (STFT) and its inverse (iSTFT).

Key architectural choices:

• Time-Frequency Processing: STFT/iSTFT transforms the waveform into magnitude and phase. The network processes the complex STFT coefficients (real and imaginary parts concatenated). Log-magnitude compression can be applied.
• Mel-band Splitting: The frequency axis is split into mel-bands, mapping irregularly shaped bands to fixed-size embeddings. This is common in speech processing networks.
• Source Handling: The network takes $K$ sources and the mixture as inputs, treating them as separate channels or entities.
• Permutation Equivariance: Interactions between sources (and the mixture) are handled by multi-head self-attention (MHSA) modules. Critically, no positional encoding is used across sources within these attention layers, allowing the network's output for a permuted input to be the same permutation of the original output. Operations applied within each source (like convolutions along time or frequency bands) are independent across sources.
• Attention Modules: Two types of attention blocks are used alternately: Band-Source Joint Attention (BSJA) and Time-Source Parallel Attention (TSPA). BSJA allows interaction between sources within each frequency band, while TSPA allows interaction between sources across time frames. Both use convolutional MLPs and RMSGroupNorm for normalization.
• Mixture Conditioning: The mixture $\bar{\vs}$ is included as an additional input to the network and is differentiated from the sources using a special positional encoding in the attention blocks. This allows the network to condition its separation on the observed mixture while maintaining permutation equivariance among the true sources.
• Diffusion Time Conditioning: The time parameter $t$ is incorporated into the network, similar to Diffusion Transformer (DiT) architectures.

Implementation and Results

The model is trained on a large dataset of 5-second speech samples from LibriVox, mixed dynamically with varying SNR. Training uses the Adam optimizer with a cosine learning rate schedule.

For inference, the ODE is solved numerically using an Euler-Maruyama sampler. The number of steps is a trade-off between separation quality and computation time. The paper finds good performance with 25 steps and even reports improved results with a custom 5-step schedule, suggesting potential for faster inference compared to diffusion models that may require hundreds or thousands of steps.

Experiments comparing Flowsep (PEM-Loco network, dB loss, Envelope noise, PET) against standard predictive models (Conv-TasNet, MB-Locoformer) and a diffusion baseline (Diffsep) show that Flowsep achieves state-of-the-art performance across various objective metrics including SI-SDR, speech intelligibility (ESTOI), and perceptual quality (PESQ, POLQA, DNSMOS). The ablation paper confirms the proposed PET loss and the envelope-based noise shaping contribute significantly to performance.

Practical Implementation Considerations:

• Computational Cost: Training Flowsep with the proposed PEM-Loco architecture on a large dataset is computationally expensive, requiring significant GPU resources and training time. Inference, while faster than typical diffusion models, still involves solving an ODE with multiple steps, making it slower than single-forward-pass predictive models.
• Data Requirements: Flowsep requires a dataset of clean, separated sources to train the generative process, as it learns to map from noisy/mixed versions back to the clean sources.
• Architecture Complexity: Implementing the PEM-Loco architecture involves integrating several components (STFT/iSTFT, mel-band processing, custom attention modules, normalization, time conditioning), requiring careful coding.
• Loss Function: The use of a normalized, decibel-valued loss (negative SNR in dB) is a key practical detail for training on audio data and suggests that simple MSE is not optimal.
• Noise Conditioning: The dynamic, envelope-based noise injection strategy is important for effectively modeling the initial distribution and guiding the separation process, especially in silence regions.
• Permutation Handling: The PET loss is a crucial implementation strategy to address the inherent permutation ambiguity in source separation. Implementing the PIT step at $t=0$ correctly is essential.
• Sampling Steps: The trade-off between sampling speed and quality needs to be managed by selecting an appropriate number of ODE steps and potentially optimizing the sampling schedule. The results with 5 steps are promising for deployment scenarios where low latency is required.

Flowsep represents a powerful generative approach to audio source separation that enforces mixture consistency and handles permutation equivariance effectively through tailored architectural and training innovations. Its strong performance on a high-quality speech separation task demonstrates the practical viability of flow matching for this problem, although the computational cost remains a consideration for real-time or resource-constrained applications. Future work is noted to extend the method to more challenging scenarios like noisy and distorted mixtures.