Task-Aware Unified Source Separation (2410.23987v1)

Abstract: Several attempts have been made to handle multiple source separation tasks such as speech enhancement, speech separation, sound event separation, music source separation (MSS), or cinematic audio source separation (CASS) with a single model. These models are trained on large-scale data including speech, instruments, or sound events and can often successfully separate a wide range of sources. However, it is still challenging for such models to cover all separation tasks because some of them are contradictory (e.g., musical instruments are separated in MSS while they have to be grouped in CASS). To overcome this issue and support all the major separation tasks, we propose a task-aware unified source separation (TUSS) model. The model uses a variable number of learnable prompts to specify which source to separate, and changes its behavior depending on the given prompts, enabling it to handle all the major separation tasks including contradictory ones. Experimental results demonstrate that the proposed TUSS model successfully handles the five major separation tasks mentioned earlier. We also provide some audio examples, including both synthetic mixtures and real recordings, to demonstrate how flexibly the TUSS model changes its behavior at inference depending on the prompts.

• The paper introduces Task-Aware Unified Source Separation (TUSS), a single model that uses learnable prompts to handle diverse and contradictory audio separation tasks, such as Speech Enhancement and Music Source Separation.
• The TUSS model utilizes a Transformer-based architecture with learnable prompts, a cross-prompt module, and a prompt dropout mechanism for flexible and robust source extraction.
• Experimental results show TUSS performs competitively or better than specialized models, demonstrating potential for unified systems in automatic audio editing and real-time analysis.

Task-Aware Unified Source Separation: A Comprehensive Analysis

This paper presents a novel approach to addressing multiple audio source separation tasks using a single model, coined Task-Aware Unified Source Separation (TUSS). The fundamental innovation of this research is the integration of learnable prompts that guide the model to conditionally change its behavior based on the specific task, thereby overcoming the limitations presented by previous models that could not handle contradictory tasks simultaneously.

Core Contributions

The paper illustrates the application of TUSS across five principal source separation tasks: Speech Enhancement (SE), Speech Separation (SS), Universal Sound Separation (USS), Music Source Separation (MSS), and Cinematic Audio Source Separation (CASS). The proposal extends the capacity of current models by allowing them to efficiently respond to variable input conditions through dynamic prompts. The prompts serve as flexible, trainable parameters that identify which sources or groups of sources to output.

1. Innovative Use of Prompts: The TUSS model leverages a variable number of prompts, each specifying a different combination of audio elements to separate, utilizing a model architecture based on Transformers. This flexible, prompt-based control permits the handling of a wide array of separation tasks even when they require opposing decisions (e.g., grouping vs. separating musical instruments in CASS versus MSS).
2. Structured Model Architecture: The model is structured with an encoder, learnable prompts, and a cross-prompt module that aligns source-specific information with the encoded audio. The architecture supports the simultaneous extraction of multiple sources, advancing the model's capability beyond traditional methods.
3. Prompt Dropout Mechanism: In order to maintain adaptability and performance when fewer than all specified sources are to be extracted, the paper introduces a prompt dropout mechanism. This form of dropout improves the model's robustness when tasked with separating a subset of input sources.

Experimental Results

The researchers present comprehensive experimental results validating TUSS's efficacy across distinct datasets, each representing one of the primary tasks for source separation:

• Performance Analysis: The TUSS outperformed several data-specific and task-specific baseline models, particularly showing strength in scenarios that involve multiple, contradictory source separation tasks. The performance metrics, including SI-SNR and SNR, are extensively discussed, highlighting the competitive nature of the TUSS models against specialized models.
• Model Scalability: A significant insight is the indication that scaling the model and training data further could potentially outperform specialized models, suggesting a pathway for further enhancements.

Implications and Future Directions

The work presents meaningful implications for advancing audio processing systems capable of multitasking in dynamic audio environments. The introduction of a unified model capable of performing complex arrays of audio separations induces speculation about potential applications, notably in developing systems for automatic audio editing, assistive listening devices, and real-time audio analysis.

Future research directions identified include the integration of speaker ID and text embeddings as prompts, which could further broaden the applicability of the TUSS framework. Including these functionalities represents an exciting trajectory that could encompass more intricate tasks, such as specific voice extraction amidst complex and varied background audio landscapes.

Conclusion

The Task-Aware Unified Source Separation model demonstrates a remarkable evolution in handling multifaceted audio separation tasks, employing innovative prompt-based conditioning and a robust architectural framework. This research underscores an essential step towards more sophisticated, versatile, and adaptable audio processing models. The foundational work laid here opens numerous avenues for expanding AI capabilities in auditory scenarios, propelling systems toward unprecedented levels of task comprehension and execution.