Stem-JEPA: A Joint-Embedding Predictive Architecture for Musical Stem Compatibility Estimation

TL;DR

We address automatic stem compatibility estimation by learning to predict embeddings of a missing stem from a context mix using Stem-JEPA, a joint-embedding predictive architecture. Stem-JEPA comprises an encoder and a predictor trained in a self-supervised fashion on a 20k-track multi-track dataset, operating on 8-second Log Mel spectrogram chunks with instrument-label conditioning. Embeddings capture timbre, rhythm, and harmony and support stem retrieval, temporal alignment analysis, and several downstream MIR tasks; optimization minimizes the mean-squared error between normalized predicted and target embeddings, $\mathcal{L}(\tilde{\boldsymbol{z}}, \boldsymbol{\bar{z}})$, with the target encoder updated by EMA. Results on MUSDB18 and a user study demonstrate competitive retrieval performance and meaningful musical structure in the learned representations, suggesting potential applications in stem generation and automatic arrangement.

Abstract

This paper explores the automated process of determining stem compatibility by identifying audio recordings of single instruments that blend well with a given musical context. To tackle this challenge, we present Stem-JEPA, a novel Joint-Embedding Predictive Architecture (JEPA) trained on a multi-track dataset using a self-supervised learning approach. Our model comprises two networks: an encoder and a predictor, which are jointly trained to predict the embeddings of compatible stems from the embeddings of a given context, typically a mix of several instruments. Training a model in this manner allows its use in estimating stem compatibility - retrieving, aligning, or generating a stem to match a given mix - or for downstream tasks such as genre or key estimation, as the training paradigm requires the model to learn information related to timbre, harmony, and rhythm. We evaluate our model's performance on a retrieval task on the MUSDB18 dataset, testing its ability to find the missing stem from a mix and through a subjective user study. We also show that the learned embeddings capture temporal alignment information and, finally, evaluate the representations learned by our model on several downstream tasks, highlighting that they effectively capture meaningful musical features.

Figures (6)

• Figure 1: Overview of the Stem-JEPA framework. From an audio clip composed of 4 stems, we first crop a chunk of 8 seconds, then sample the target $\bold{\bar{x}}$ (one of the stems) and the context $\bold{x}$ (a mix of some of the remaining stems) as described in section \ref{['sec:sampling']}. They are then converted into Log Mel Spectrograms and passed through the context and target encoders, respectively. Finally, the predictor (conditioned on the target instrument label) is trained so that each of its outputs individually predicts each target representation.
• Figure 2: Analysis of the closest embedding $\bold{z}^*$ for all queries $\bold{q}$ from the MUSDB18 dataset MUSDB18. Left: Categories of failures for each instrument (same song but wrong instrument, the opposite, or both wrong). Right: confusion matrix between conditioning instruments and retrieved instruments.
• Figure 3: Box plot of the listening test for the different instrument classes. The $\times$ represents the mean of the data.
• Figure 4: Average pairwise cosine similarity between embeddings and predictions across various temporal shifts. Each curve corresponds to a different track.
• Figure 5: Key/Chord co-occurrence matrix between segments within the same clusters.