Multi-Stage Music Source Restoration with BandSplit-RoFormer Separation and HiFi++ GAN

This paper presents the CP-JKU team's two-stage system for the ICASSP 2025 Music Source Restoration Challenge, which combines a curriculum-trained BandSplit-RoFormer model for separating eight stems and a specialized HiFi++ GAN for restoring instrument-specific waveforms from mastered audio.

The Gist

Imagine you have a delicious, complex stew that someone has already cooked, seasoned, and served in a bowl. Now, imagine you want to take that stew apart and get back the original, raw ingredients (the carrots, the beef, the herbs) exactly as they were before they were cooked.

That is essentially what Music Source Restoration (MSR) tries to do. But instead of a stew, it's a finished song. And instead of just separating ingredients, it has to undo all the "cooking" (mixing, mastering, compression, and digital glitches) that happened to the song.

Here is how the team at Johannes Kepler University (CP-JKU) solved this puzzle for the 2025 challenge, explained simply:

The Big Problem: The "Over-Processed" Song

In the music industry, songs aren't just a simple mix of instruments playing together. They go through a "production pipeline" that adds effects like echo, compression (making loud parts quieter and quiet parts louder), and distortion. They also get compressed for streaming (like MP3s), which adds digital noise.

Trying to separate these instruments is like trying to un-bake a cake to get the raw eggs and flour back. Standard methods fail because the ingredients have been chemically changed by the "cooking" process.

The Solution: A Two-Stage Assembly Line

The team built a robot factory with two distinct stations to handle this messy job.

Stage 1: The "Rough Slicer" (Separation)

First, they need to cut the finished song into eight rough piles: Vocals, Guitar, Drums, Bass, Keyboards, Synthesizers, Percussion, and Orchestra.

• The Tool: They used a smart AI called BandSplit-RoFormer. Think of this AI as a master chef who looks at the frequency "flavors" of the song and tries to guess which ingredient belongs to which pile.
• The Training Trick (The Curriculum): You can't just throw the AI into the deep end. They taught it in three steps:
  1. Warm-up: They started with a simple version of the AI that only knew how to separate 4 things (Vocals, Drums, Bass, and "Everything else") using clean, perfect recordings.
  2. Practice: They then taught it to handle "messy" recordings that had been processed and degraded, just like real songs.
  3. The Expansion: Finally, they gave the AI a "head transplant." They kept the brain it had learned and added new "heads" (specialized outputs) to handle the other 4 instruments (Guitar, Keys, etc.). This allowed them to teach the AI to handle 8 instruments without forgetting how to handle the first 4.

Result: The AI now produces 8 rough, slightly imperfect stems. They are separated, but they still sound a bit "muddy" or damaged because the AI made some mistakes.

Stage 2: The "Polishing Station" (Restoration)

Now that the AI has 8 rough piles, they need to clean them up and make them sound like the original, pristine recordings.

• The Tool: They used a HiFi++ GAN. Imagine a high-end photo editor that doesn't just fix a blurry photo, but imagines what the sharp details should look like and paints them in.
• The "Generalist to Expert" Strategy:
  1. The Generalist: First, they trained one super-AI to clean up any instrument. It learned to remove noise, fix distortion, and sharpen the sound.
  2. The Specialists: Then, they took that Generalist and created 8 "Experts." One Expert only knows how to fix Vocals, another only Drums, etc.
  3. The Secret Sauce: To make these Experts perfect, they didn't train them on clean songs. They trained them on the rough, messy outputs from Stage 1. This is like training a doctor to fix a broken leg by showing them a cast that is already slightly cracked, so they learn exactly how to fix the specific kind of damage the "Rough Slicer" creates.

The Results

When they tested this system, it worked surprisingly well.

• The Score: It achieved high marks in separating the instruments and making them sound natural.
• The Human Test: When humans listened to the results, they rated the quality very highly (a "MOS" of 3.55 out of 5, which is very good for AI).

The Limitations (Where it still struggles)

Even with this fancy two-stage factory, the system isn't perfect yet.

• Noisy Recordings: If the original song is very old, recorded in a live concert, or has a lot of background noise, the "Rough Slicer" gets confused. If the first stage fails to separate the ingredients correctly, the "Polishing Station" can't fix it.
• The "Dry" Target: Sometimes, the original song had effects like reverb (echo) or delay. The AI is asked to remove "production effects," but it's hard to know if an echo was a mistake or part of the artist's vision. The AI sometimes accidentally removes cool artistic effects, making the song sound too "dry" or flat.

Summary

The CP-JKU team solved the problem by breaking it down:

1. Separate first: Use a smart AI to split the song into 8 messy piles.
2. Clean second: Use specialized AI experts to polish each pile, trained specifically to fix the mistakes the first AI makes.

It's a bit like hiring a team to first sort a pile of mixed laundry into 8 baskets, and then hiring 8 different ironing experts to fix the wrinkles in each specific type of fabric, knowing exactly how the sorting machine messed them up.

Technical Summary

Here is a detailed technical summary of the paper "Multi-Stage Music Source Restoration with BandSplit-RoFormer Separation and HiFi++ GAN" by the CP-JKU team.

1. Problem Definition

Music Source Restoration (MSR) aims to recover original, unprocessed instrument stems from fully mixed and mastered audio tracks. Unlike traditional Music Source Separation (MSS), which assumes a linear mixture of clean sources, MSR addresses a more complex reality:

• Non-Linear Transformations: Professional production pipelines apply equalization, compression, reverb, saturation, stereo widening, and limiting.
• Artifacts: Distribution codecs and historical degradations (e.g., vinyl noise) further corrupt the signal.
• Challenge: The target sources are not just mixed but systematically transformed. Direct separation under "clean-stem" assumptions is insufficient because the model must simultaneously de-mix, de-master, and de-artifact the audio.

The specific goal of the ICASSP 2025 MSR Challenge is to recover eight instrument classes (vocals, guitar, keyboard, synthesizer, bass, drums, percussion, orchestra) from such degraded mixtures.

2. Methodology

The CP-JKU team proposes a modular, two-stage pipeline that explicitly decouples the de-mixing problem from the restoration problem.

Stage 1: Source Separation (De-mixing)

• Architecture: A BandSplit-RoFormer (BS-RoFormer) model.
  • It uses a band-split front-end to process frequency regions separately.
  • It employs RoFormer blocks to model temporal and cross-band dependencies for mask estimation.
  • A single model predicts 9 stems: 8 target instruments + 1 auxiliary "other" stem.
• Training Strategy (Three-Stage Curriculum):
  1. Warm-up (4 Stems): Fine-tuning on clean mixtures (vocals, drums, bass, other) using LoRA (Low-Rank Adaptation) on transformer layers.
  2. Degradation Adaptation (4 Stems): Continued fine-tuning on mixtures created by an online degradation pipeline (simulating mastering and artifacts). The target here is the degraded stems, teaching the model to handle production effects.
  3. Extension (8 Stems): The model is expanded from 4 to 8 stems. The backbone and original heads are frozen (initialized from Stage 2), while new heads are randomly initialized and trained. This allows the model to learn new instrument classes without catastrophic forgetting.

Stage 2: Restoration (De-mastering & De-artifacting)

• Architecture: A HiFi++ GAN bundle comprising:
  • SpectralUNet front-end.
  • Upsampling stage.
  • WaveUNet refinement network.
  • SpectralMaskNet for residual spectral correction.
• Training Strategy (Five-Stage Curriculum):
  1. Generalist Warm-up: Training on original music content.
  2. GAN Introduction: Introducing feature matching.
  3. Perceptual Metrics: Incorporating music-specific perceptual losses.
  4. Noise Suppression: Augmenting data with gramophone noise to handle historical artifacts.
  5. Expert Specialization: Fine-tuning eight instrument-specific experts. Crucially, the training inputs for these experts are generated by running the trained separator on synthetic mixtures. This ensures the restorers are exposed to realistic separation errors (mismatch between training and test conditions).

3. Key Contributions

1. Curriculum-Based BS-RoFormer Adaptation: A novel approach to scaling a separation model from 4 to 8 stems using parameter-efficient LoRA fine-tuning and head expansion, avoiding full retraining.
2. Instrument-Specific Restoration Experts: Instead of a single general restorer, the system trains eight specialized GAN models.
3. Realistic Train-Test Alignment: By generating restoration training data using the separator's own outputs (rather than clean ground truth), the system explicitly learns to correct the specific types of errors introduced by the separation stage, significantly improving robustness.
4. Comprehensive Degradation Modeling: The pipeline simulates complex production chains (EQ, compression, limiting) and historical noise (gramophone) during training.

4. Results

The system was evaluated on the MSR Challenge 2025 test set and the MSRBench dataset.

• Official Challenge Performance (Test Set):
  • MMSNR (Multi-Mel-SNR): 0.8329 (Average)
  • Zimtohrli (Perceptual Metric): 0.0189
  • FAD (Fréchet Audio Distance): 0.6223 (Note: The abstract mentions 0.3814 in the text, but Table 1 lists 0.6223 as the average; the text likely refers to a specific iteration or subset).
  • System-Level MOS (Mean Opinion Score): 3.5510.
• MSRBench Performance:
  • FAD: 0.638
  • MMSNR: 2.338
• Per-Stem Performance: The system showed varying performance across instruments, with Vocals achieving the best FAD (0.3092) and Percussions showing the lowest MMSNR (0.0994), indicating that complex, transient-heavy instruments remain challenging.

5. Limitations and Future Work

• Sensitivity to Noise: The pipeline struggles with highly noisy mixtures (e.g., live recordings or severe historical degradation), where the separator fails to produce faithful estimates, limiting the downstream restoration.
• Data Mismatches: Discrepancies in large-scale datasets (e.g., RawStems) can bias training, leading to residual artifacts.
• Ambiguous Targets: Time-varying effects (reverb, chorus) present in the reference "clean" stems make it difficult to define a consistently "dry" target, as the model may inadvertently remove desirable artistic effects.
• Future Directions: The authors plan to prioritize data curation, alignment verification, and "effect-aware" conditioning to better handle domain shifts and ambiguous targets.

6. Significance

This paper presents a state-of-the-art approach to the difficult problem of MSR by acknowledging that separation and restoration are coupled but distinct tasks. The use of a curriculum learning strategy for both the separator (scaling from 4 to 8 stems) and the restorer (generalist to expert) demonstrates a scalable and effective methodology. The specific innovation of training restorers on separator-generated inputs addresses the common "domain gap" in multi-stage pipelines, making the system more robust to real-world inference conditions. This work sets a strong baseline for future research in recovering high-fidelity audio from complex, mastered recordings.