FLASepformer: Efficient Neural Speech Separation

• FLASepformer is an efficient speech separation model that replaces quadratic attention with Focused Linear Attention to achieve O(N) complexity.
• It integrates 1D depthwise convolution and a gated MLP for enhanced local and global context, significantly boosting runtime performance.
• The model demonstrates competitive accuracy and reduced GPU memory usage, making it ideal for real-time and embedded speech applications.

FLASepformer is an efficient neural speech separation model designed to address the computational bottlenecks inherent in transformer-based architectures for long sequence processing. It builds upon the limitations of quadratic-complexity self-attention by introducing the Focused Linear Attention (FLA) mechanism, and further augments feature aggregation through a novel gating module. The model is implemented in two backbone variants, FLA-SepReformer and FLA-TFLocoformer, each demonstrating linear computational complexity and notable improvements in runtime and memory usage across diverse speech separation benchmarks.

1. Model Architecture

FLASepformer replaces the standard quadratic attention mechanisms found in many transformer-based speech separation models with Focused Linear Attention (FLA), resulting in a fundamental shift from $O(N^2)$ to $O(N)$ complexity, where $N$ denotes input sequence length. The architecture is composed primarily of three elements:

• Focused Linear Attention (FLA):

FLA generalizes vanilla linear attention (VLA), which computes similarity as $\mathrm{Sim}(Q,K) = \phi(Q)\phi(K)^\top$, by introducing the "Focused Function" $\phi_p$. In the vanilla formulation, simple kernels (e.g., $\mathrm{ReLU}$) lead to smooth weights and reduced attention rank, especially for $d \ll N$. FLA addresses this by defining:

$$\phi_p(x) = f_p(\mathrm{ReLU}(x)), \quad f_p(x) = \left(\frac{\|x\|}{\|x^{(**p)}\|}\right) x^{(**p)}$$

where $x^{(**p)}$ denotes element-wise power with focus factor $p$. This formulation sharpens attention weights, contrasting similar and dissimilar query–key pairs, and approximates softmax behavior.

• 1D Depthwise Convolution (DWC1d):

Post-FLA computation, DWC1d with kernel size $k$ enriches local feature diversity and compensates for reduced rank in $\phi(Q)\phi(K)^\top$. The overall output is:

$$O = \phi(Q) \, [\phi(K)^\top V] + \mathrm{DWC1d}(V)$$

which is particularly important for modeling sequential local context in speech.

• Gated Module:

A Gated MultiLayer Perceptron (MLP), incorporating a normalization (LayerNorm or RMSGroupNorm per variant), linear projection, and activation, modulates the FLA output. The resultant gating signal regulates feature flow per token, enhancing global feature aggregation.

This Gated FLA module is systematically embedded into transformer backbones. In SepReformer, it replaces all global attention components, obviating the need for downsampling and securing linear-time operations. In TF-Locoformer, it supplants multi-head self-attention in the temporal modeling block, offering analogous complexity reductions while retaining global dependency capture.

2. Variants

FLASepformer is instantiated in two principal backbone configurations, inheriting structural elements (separation encoder, reconstruction decoder, and local blocks) from prior models:

Variant	Backbone	Attention Replacement	Scales Explored
FLA-SepReformer	SepReformer	Global Linear Attention → Gated FLA	T, B, L
FLA-TFLocoformer	TF-Locoformer	Temporal Attention → Gated FLA	S, M, L

• FLA-SepReformer:

Removes subsampled multi-head global attention, substituting with Gated FLA throughout. This achieves $O(N)$ complexity and avoids the memory spike associated with $O(N^2)$ attention modules. Evaluated in three scales (Tiny, Base, Large), it demonstrates speedup and reduction in GPU memory expenditure.

• FLA-TFLocoformer:

Within the TF-Locoformer framework, the standard multi-head self-attention for temporal modeling is replaced with Gated FLA. Although TF-Locoformer uses STFT to preprocess and reduce sequence length, quadratic attention still limits scalability; FLA circumvents this, enhancing efficiency.

Both variants retain key functional blocks and can be extended for general speech separation tasks.

3. Performance Metrics

FLASepformer demonstrates competitive separation performance while substantially exceeding previous approaches in resource efficiency:

• Accuracy:

On WSJ0-2Mix, a standard clean speech separation benchmark, FLA-SepReformer-B achieves SI-SNRi close to 23.5 dB, matching SepReformer and TF-Locoformer baselines.

• Speed and Memory Usage:
  • Tiny (T): 2.29× faster inference, ~15.8% GPU memory usage
  • Base (B): 1.91× faster, ~20.9%
  • Large (L): 1.49× faster, ~31.9%
  • These figures result directly from replacing quadratic attention with linear FLA.
• Datasets and Benchmarks:
  • WSJ0-2Mix (clean mixture separation)
  • WHAM! and WHAMR! (noisy, reverberant speech)
  • Libri2Mix (broader generalization)
  • Robust performance across all datasets evidences resilience to noise, reverberation, and extended signal duration.

4. Applications

The combination of linear scalability, effective global context modeling, and reduced resource overhead renders FLASepformer applicable across a spectrum of scenarios:

• Real-Time/Embedded Speech Separation:

Linear complexity and low memory use make it suitable for real-time deployment in hearing aids, mobile devices, and embedded acoustic IoT systems.

• Voice Communication Systems:

Separation of overlapping talkers ("cocktail party problem") enhances teleconferencing, voice-over-IP, and automated speech recognition in multi-speaker settings.

• Noisy and Reverberant Environments:

Performance on WHAM! and WHAMR! indicates robustness, enabling use in public spaces, urban environments, and telecommunications with significant background interference.

• Cloud-Based Batch Processing:

Efficient handling of long sequences (broadcast, meetings, call centers) improves scalability and lowers operational cost for cloud transcription or virtual assistants.

5. Future Directions

Several avenues are identified for further investigation:

• Scaling to Multi-Speaker Mixtures:

Extension to more than two simultaneous speakers necessitates evaluation of FLA’s behavior and rank properties under rising source count and greater signal complexity.

• Attentional Mechanism Refinements:

Investigation into alternate kernel functions or dynamically adaptive focus factors may refine the balance between expressivity and efficiency. Likewise, optimization of the DWC1d parameters and exploration of alternative convolutional mechanisms may further advance local and global feature modeling.

• Efficiency–Richness Trade-Off:

While the model closes much of the gap between linear and quadratic attention, future work will target residual disparities in context modeling, focusing on ultralong sequences and highly challenging acoustic scenarios.

• Multimodal Integration:

A plausible implication is potential extension to multimodal speech separation, utilizing visual or contextual signals alongside audio for environments where acoustic-only separation reaches fundamental limits.

6. Summary

FLASepformer presents a substantial advance in the field of speech separation, combining a mathematically principled Focused Linear Attention mechanism with a gating module to realize $O(N)$ complexity end-to-end. The two primary variants, FLA-SepReformer and FLA-TFLocoformer, illustrate the utility of this approach across backbone architectures, yielding competitive accuracy, dramatic acceleration, and memory reduction on standard benchmarks. The robust performance, scalability, and deployment suitability suggest broad applicability; future research is likely to extend FLASepformer’s capabilities to higher-order mixtures and multimodal scenarios.