TUSS: Task-Aware Unified Source Separation

• TUSS is a conditional prompt-based audio source separation framework that uses a variable-length set of learnable prompts to define the separation task.
• It redefines source separation as task-aware rather than task-agnostic, enabling one model to handle diverse tasks such as speech enhancement, music demixing, and sound event separation.
• Empirical evaluations and variants like FasTUSS and SUNAC demonstrate TUSS's capability to improve separation performance, even as it highlights ongoing challenges related to prompt design and computational efficiency.

Searching arXiv for the core TUSS paper and closely related work to ground the article in current literature. arXiv search query: "Task-Aware Unified Source Separation" Task-Aware Unified Source Separation (TUSS) is a conditional, prompt-based audio source separation framework in which a single model, with a single parameter set, performs multiple source separation tasks by changing a variable-length set of learnable prompts at the input. In its canonical formulation, TUSS addresses the fact that source separation is not defined solely by the waveform mixture but also by the intended decomposition: the same acoustic mixture may require per-speaker separation, speech-versus-noise enhancement, four-stem music demixing, or three-stem cinematic grouping. TUSS encodes that intended decomposition directly in the prompt set, allowing one model to handle speech enhancement, speech separation, sound event separation, music source separation, and cinematic audio source separation, including cases where the tasks are logically contradictory (Saijo et al., 2024).

1. Problem formulation and conceptual scope

The motivating claim of TUSS is that “general” or “universal” source separation is intrinsically under-specified unless the model is told what counts as a source. In music source separation, instruments are separated into distinct stems; in cinematic audio source separation, those same instruments may need to be grouped into a single music stem; in speech enhancement, non-speech content is treated as a residual “noise” mixture; in universal sound separation, those residual components may instead be the explicit targets. TUSS therefore treats source separation as a task-aware problem rather than a task-agnostic one (Saijo et al., 2024).

This position distinguishes TUSS from large-scale task-agnostic unified separators such as GASS. GASS trains a fixed four-output model across speech, music, and sound events without explicit task metadata or prompt conditioning, and its results show both the feasibility and the limitation of that approach: strong in-distribution performance and useful transfer to FUSS and Libri2Mix, but weak zero-shot generalization to DnR and MUSDB when the downstream definition of a “source” differs from the upstream training definition (Pons et al., 2023). TUSS directly targets that mismatch by making the desired grouping structure part of the input specification rather than leaving it implicit.

A common misconception is that TUSS is merely a multi-domain separator. In the original formulation, the essential novelty is not only that one model is trained on multiple domains, but that the model’s behavior changes as a function of the prompt set. Another misconception is that TUSS requires separate output heads per task. In fact, the architecture uses shared modules and treats the prompt configuration itself as a task specification language (Saijo et al., 2024).

2. Prompt vocabulary and core architecture

In the original TUSS model, the mixture waveform $\mathbf{x} \in \mathbb{R}^{L}$ is transformed to a time-frequency representation $\mathbf{X} \in \mathbb{R}^{2 \times T \times F}$, and then to a band-split latent $\mathbf{Z} \in \mathbb{R}^{D \times T \times F}$ by an encoder derived from BSRNN. TUSS defines eight learnable prompts, each parameterized as a learned tensor $\mathbf{P}_n \in \mathbb{R}^{D \times 1 \times 1}$: <Speech>, <SFX>, <SFX-mix>, <Drums>, <Bass>, <Vocals>, <Other inst.>, and <Music-mix> (Saijo et al., 2024).

The prompts are tiled along frequency and prepended to the mixture representation along the time axis, producing a joint sequence

$$\mathbf{Z}' = [\mathbf{P}_1^{\mathrm{tile}}, \ldots, \mathbf{P}_N^{\mathrm{tile}}, \mathbf{Z}] \in \mathbb{R}^{D \times (N+T) \times F}.$$

This sequence is processed by a cross-prompt module built from TF-Locoformer blocks. The same self-attention mechanism updates both prompt tokens and mixture tokens, so prompts attend to each other and to the mixture, while mixture frames attend to prompts and to other frames. The output can be partitioned into updated prompt features $\tilde{\mathbf{P}}_n$ and an updated mixture representation $\tilde{\mathbf{Z}}$ (Saijo et al., 2024).

For each prompt, TUSS forms a prompt-conditioned representation by broadcasted element-wise multiplication,

$$\tilde{\mathbf{Z}}_n = \tilde{\mathbf{Z}} \odot \tilde{\mathbf{P}}_n,$$

and sends each $\tilde{\mathbf{Z}}_n$ through a shared conditional target sound extraction module, again implemented with TF-Locoformer blocks. A shared decoder reconstructs one waveform per prompt, yielding

$$(\hat{\mathbf{s}}_1,\dots,\hat{\mathbf{s}}_N)=f_\theta(\mathbf{x};\mathbf{P}_1,\dots,\mathbf{P}_N).$$

Because prompts are part of the sequence and identical prompts can appear multiple times, repeated categories such as multiple <Speech> prompts can be used for multi-instance separation (Saijo et al., 2024).

Separation task	Prompt set
Speech enhancement	<Speech>, <SFX-mix>
Noisy speech separation	<Speech> \times N, (<SFX-mix>)
Universal sound separation	<SFX> \times N
Music source separation	<Drums>, <Bass>, <Vocals>, <Other inst.>
Cinematic audio source separation	<Speech>, <SFX-mix>, <Music-mix>

The original paper emphasizes that prompts in TUSS are not text-derived tokens. They are learned embeddings trained jointly with the network, and the prompt set determines the number of outputs, the grouping relation among outputs, and the decomposition criterion itself (Saijo et al., 2024).

3. Training procedure, data construction, and optimization

TUSS is trained on dynamically synthesized six-second mixtures at 48 kHz. For each example, the number of prompts is sampled from $\mathbf{X} \in \mathbb{R}^{2 \times T \times F}$0, a valid prompt combination is chosen under explicit constraints, and one source is sampled for each prompt category. <SFX> and <SFX-mix> cannot co-occur; <Music-mix> cannot co-occur with individual instrument prompts; <Speech> and <SFX> may repeat; other prompts occur at most once. Sources are drawn from VCTK, WSJ0, LibriVox, FSD50K, WHAM!, DEMAND, MUSDB-HQ, MOISESDB, and FMA, RMS-normalized, randomly gain-scaled within category-specific ranges, and summed into the mixture (Saijo et al., 2024).

The objective is a negative SNR loss, with permutation-invariant training applied per category when multiple outputs of the same category are present. This category-wise PIT is necessary for repeated prompts such as multiple <Speech> or multiple <SFX> prompts, because the task is to separate multiple same-type instances without fixing an ordering in advance (Saijo et al., 2024).

The original paper also introduces prompt dropout to support partial separation at inference. Starting from epoch 124 of the Large unified TUSS model, the network is fine-tuned for 26 additional epochs; in 25% of training steps, a uniformly sampled number of prompts is randomly dropped, with the proviso that same-category repeated prompts are not partially dropped. The intended effect is to teach the model that the absence of a prompt means the corresponding source should not be separated, so that inference can request only a subset of the stems present in the mixture (Saijo et al., 2024).

An important architectural finding is that prompt–mixture coupling is essential, but prompt updates need not depend strongly on mixture content. In the original ablations, decoupling prompts and mixture entirely produced a large performance drop, whereas keeping mixture conditioned on prompts while making prompt updates independent of the mixture produced only a small degradation. This result later became central to the causal formulation of TUSS (Paissan et al., 15 Jul 2025).

4. Empirical behavior across the five major separation tasks

The primary empirical comparison in the original TUSS paper is between a prompting unified TUSS model and a conventional unified fixed-output baseline. On Medium models, the unified TUSS model improves VCTK-DEMAND speech from 19.2 to 19.4 dB and SFX-mix from 8.7 to 10.2 dB, WHAM! speech from 3.4 to 7.0 dB and SFX-mix from 10.2 to 11.2 dB, FUSS SFX from 8.1 to 9.6 dB, and DnR speech from 12.0 to 14.5 dB, music-mix from 3.8 to 5.7 dB, and SFX-mix from 2.1 to 7.1 dB. The same trend appears on Large models. These results show that explicit task-aware prompting is materially more effective than training one unconditional model on all datasets with a fixed output structure (Saijo et al., 2024).

Against task specialists, the picture is more differentiated. Large unified TUSS exceeds the Large prompting task specialist on FUSS, where the specialist reaches 10.0 dB and unified TUSS reaches 12.2 dB, and is essentially matched on DnR, where both reach 15.1 dB on speech and 7.0 dB on music-mix, with unified TUSS slightly higher on SFX-mix at 8.2 dB versus 7.9 dB. On MUSDB-HQ, however, unified TUSS still trails the task specialist: for example, vocals are 10.4 dB for the specialist and 8.6 dB for unified TUSS, with similar gaps for bass, drums, and other. The original paper explicitly presents this as a remaining specialization gap rather than as a failure of the unified framework (Saijo et al., 2024).

Qualitative evaluations in the original study show that the same DnR mixture can be decomposed differently by changing only the prompt set. A standard cinematic decomposition [<Speech>, <Music-mix>, <SFX-mix>] can be replaced by [<Speech>, <Music-mix>, <SFX>, <SFX>] to split the sound-effects stem into multiple SFX outputs, or by replacing <Music-mix> with instrument prompts to obtain a hybrid cinematic-plus-music-separation behavior. This suggests that, within the learned prompt vocabulary, TUSS functions as a user-programmable separator rather than as a fixed demixer (Saijo et al., 2024).

Prompt dropout specifically addresses a failure mode of the base unified model: without dropout training, using only a single category of prompts at inference causes large performance degradation, whereas the dropout-fine-tuned model preserves much more of its full-prompt performance under partial prompt sets. The underlying issue is that the model must learn not only what to extract when a prompt is present, but also what not to represent when a prompt is absent (Saijo et al., 2024).

5. Efficiency, causality, and later realizations of the TUSS idea

The computational cost of the original TUSS architecture motivated "FasTUSS: Faster Task-Aware Unified Source Separation" (Paissan et al., 15 Jul 2025). FasTUSS retains the prompt-driven conditional formulation but modifies the TF-Locoformer blocks by increasing convolution stride, removing one Conv-SwiGLU block, and, in the most aggressive variant, using grouped convolutions. The resulting FasTUSS-11.7G and FasTUSS-8.3G reduce the original model’s operations by 73% and 81%, with average performance drops of 0.4 dB and 1.2 dB, respectively. The same study also derives a causal TUSS variant by constraining the cross-prompt attention mask so that prompt tokens can condition the mixture while mixture frames attend only to current and past frames; this causal version incurs an approximately 1.8 dB drop (Paissan et al., 15 Jul 2025).

A distinct but conceptually direct realization of TUSS appears in "SUNAC: Source-aware Unified Neural Audio Codec" (Aihara et al., 20 Nov 2025). SUNAC instantiates the TUSS idea inside a neural audio codec rather than in a standalone separator: a mixture encoder produces latent features, prompts such as <Speech>, <Music>, <SFX>, and <Mix> condition latent-space source extraction, and a shared RVQ quantizer and decoder reconstruct either the mixture or selected sources. In this setting, TUSS becomes a joint separation-and-coding framework. SUNAC supports repeated same-type prompts, uses feature-space PIT constrained within source types, and at 6 kbps achieves lower computational cost than cascade systems such as TUSS-DAC or FasTUSS-DACT while remaining competitive in perceptual quality; in two-speaker mixtures it reaches approximately 11.8 dB SI-SDR and 4.12 ViSQOL, whereas SDCodec variants effectively do not separate repeated speech prompts (Aihara et al., 20 Nov 2025).

A broader multichannel and multi-task generalization is provided by DeepASA, which treats source separation as one component of unified auditory scene analysis. DeepASA uses object-centric representations shared across MIMO separation, dereverberation, sound event detection, classification, and direction-of-arrival estimation, and refines them through a chain-of-inference mechanism with temporal coherence matching. In this formulation, TUSS is extended from prompt-conditioned source grouping to joint optimization with downstream ASA tasks. On ASA2, the full model with chain-of-inference reaches SI-SDRi 11.2 dB, ER 25.0%, F1 74.1%, LE 17.0°, LR 78.1%, and SELD 0.206; an ablation shows that adding full CoI improves SI-SDRi from 11.0 to 12.0 dB and SELD from 0.230 to 0.206 (Lee et al., 21 Sep 2025).

6. Limitations, adjacent lines of research, and open directions

The original TUSS paper identifies several limitations. Unified TUSS still underperforms task specialists on instrument-level MSS, it does not estimate how many prompts are required for a given mixture, its prompts are learned vectors rather than human-readable semantic tokens, and the TF-Locoformer backbone remains computationally heavy despite a relatively small parameter count (Saijo et al., 2024). FasTUSS mitigates the compute issue but does not eliminate it, and its causal variant still pays a nontrivial quality penalty (Paissan et al., 15 Jul 2025). SUNAC adds source-aware coding but operates at a fixed 6 kbps without explicit bitrate control, and its training regime only covers mixtures up to three prompts with limited repetition patterns (Aihara et al., 20 Nov 2025).

Several adjacent research directions clarify what TUSS is not yet. Query-by-example music models show that unification can be driven by source-conditioned representations rather than a fixed discrete prompt vocabulary. In "A Unified Model for Zero-shot Music Source Separation, Transcription and Synthesis," a single model uses a query embedding, a shared latent $\mathbf{X} \in \mathbb{R}^{2 \times T \times F}$1, a transcriptor, and pitch–timbre disentanglement to couple separation, transcription, and synthesis; MSI improves overall SDR to $\mathbf{X} \in \mathbb{R}^{2 \times T \times F}$2 dB while maintaining transcription precision of $\mathbf{X} \in \mathbb{R}^{2 \times T \times F}$3 (Lin et al., 2021). This suggests that future TUSS systems need not restrict prompts to discrete categories.

Another line of work shows that auxiliary tasks can affect separation not only through training-time regularization but also through inference-time control. "Multitask learning for instrument activation aware music source separation" couples separation with instrument activation detection and then gates the separated output by the predicted activation sequence, improving SDR on a six-instrument MedleyDB-plus-Mixing-Secrets setup for all six instruments, including vocals from 3.68 to 4.78 dB and bass from 4.04 to 5.26 dB (Hung et al., 2020). A plausible implication is that TUSS-like systems may benefit from treating side tasks as explicit control signals rather than only as auxiliary losses.

A different open direction concerns computational adaptivity. TISDiSS is not task-conditioned, but it demonstrates a unified separation backbone with training-time and inference-time scalable depth, shared-parameter iterative refinement, and early-split multi-loss supervision. On WSJ0-2mix, an 8.0M-parameter TISDiSS configuration reaches 25.1 dB SI-SNRi and can be run at different inference depths with the same weights (Feng et al., 19 Sep 2025). This suggests that future TUSS architectures could make not only the decomposition criterion but also the computation budget task-aware.

Taken together, the literature defines TUSS less as a single architecture than as a principle: source separation should be conditioned on the intended decomposition, and that conditioning should be strong enough to alter the internal representation rather than merely select an output head. The original prompt-driven TUSS model provides the canonical formulation; FasTUSS shows that the formulation can be made substantially cheaper; SUNAC demonstrates that it can be pushed into codec latents; and DeepASA shows that separation can be jointly optimized with broader scene-analysis objectives. The remaining open problems are largely structural: richer prompt semantics, automatic source-count inference, better handling of music-specific decompositions, tighter integration with downstream tasks, and stronger compute-quality trade-offs under causal and low-latency constraints.