Stem-JEPA: Self-Supervised Stem Compatibility

• Stem-JEPA is a self-supervised model that estimates stem compatibility in musical mixes using joint-embedding predictive and contrastive strategies.
• It processes log-Mel spectrogram patches through a ViT-based context encoder and an MLP predictor to capture both global and local audio features.
• The framework excels in retrieval and alignment tasks on MIR benchmarks and supports zero-shot conditioning for flexible instrument queries.

Stem-JEPA is a self-supervised neural architecture designed to estimate compatibility among musical stems—single-instrument audio tracks—relative to a given musical mix. It applies a Joint-Embedding Predictive Architecture (JEPA) trained on multi-track datasets, enabling retrieval, alignment, or conditioning for compatible stems, even under zero-shot instrument label constraints. Through a combination of contrastive pretraining and predictive modeling, Stem-JEPA learns patchwise and global audio representations sensitive to timbre, harmony, and rhythm, supporting both standard music information retrieval (MIR) tasks and novel generative scenarios (Riou et al., 2024, Riou et al., 2024).

1. Model Architecture and Data Pipeline

Stem-JEPA comprises three core components: a context encoder $f_\theta$, a target encoder $f_{\bar\theta}$, and a predictor $g_\phi$. The approach ingests log-Mel spectrograms of audio mixes and stems, split into non-overlapping $16 \times 16$ patches (5 in frequency × 50 in time, for 8 s crops under standard settings), each processed independently (Riou et al., 2024, Riou et al., 2024).

• Context Encoder: $f_\theta$, instantiated as a Vision Transformer (ViT-Base), maps an audio mix spectrogram to a sequence of $K = 250$ patch embeddings $z_1, \ldots, z_K \in \mathbb{R}^d$ ($d=768$).
• Target Encoder: $f_{\bar\theta}$, a parameter-averaged copy of $f_\theta$ (EMA), embeds the missing stem patch sequence $z^*_k$ for supervision.
• Predictor: $g_\phi$ is an MLP (depth $L=6$, hidden $1024$) or, optionally, a transformer. It consumes each patch $z_k$ concatenated or FiLM-conditioned with an instrument label embedding, predicting $\hat{z}_k$ as an estimate of the corresponding stem patch embedding.

Label conditioning is realized either via a learned lookup for categorical instrument class (Riou et al., 2024) or through CLAP text embedding vectors (Riou et al., 2024), generalizing to arbitrary instrument names.

Data Preparation

Training employs multi-track corpora:

• Log-Mel spectrogram extraction: $80$ Mel bins, $25\,\mathrm{ms}$ windows, $10\,\mathrm{ms}$ hop ($800$ frames/8 s).
• Patchification into $5\times50$ tokens.
• Tracks grouped by four coarse labels (bass, drums, vocals, other) (Riou et al., 2024) or up to $38$ fine-grained labels (Riou et al., 2024).

Chunks are constructed by masking non-silent, randomly sampled stems as targets and mixing subsets of remaining stems as context.

2. Self-Supervised Objectives and Training Scheme

Stem-JEPA applies the joint-embedding predictive paradigm, explicitly avoiding audio reconstruction in the waveform or spectrogram space.

• Contrastive Pretraining (optional, phase 1) (Riou et al., 2024):

$$\mathcal{L}_{\text{contrast}}(s_x,s_y) = -\log\frac{\exp(sim(s_x, s_y)/\tau)}{\sum_{s' \in B \setminus \{s_x\}} \exp(sim(s_x, s')/\tau)}$$

where $sim(u, v) = u^\top v / (\|u\| \|v\|)$ encourages the global average embeddings of context-target pairs to be near, repelling others.

• Predictive (JEPA) Objective (phase 2):

$$\mathcal{L}_{\text{pred}}(\{\hat{z}_k\}, \{z^*_k\}) = \frac{1}{K} \sum_{k=1}^K \left\| \frac{\hat{z}_k}{\|\hat{z}_k\|} - \frac{z^*_k}{\|z^*_k\|} \right\|^2$$

This MSE is between the $\ell_2$-normalized predictor output and target patch embeddings.

Combined loss: $$\mathcal{L} = \lambda_{\text{contrast}}\,\mathcal{L}_{\text{contrast}} + \lambda_{\text{pred}}\,\mathcal{L}_{\text{pred}}$$

Target encoder parameters are updated via an exponential moving average: $$\bar{\theta}_t = \tau_t \bar{\theta}_{t-1} + (1 - \tau_t) \theta_t$$ with $\tau_t$ increasing across training to encourage target encoder stability.

Negative sampling is unnecessary; gradients do not flow through $f_{\bar\theta}$, avoiding representational collapse.

3. Stem Compatibility Estimation and Retrieval Methodology

The central application is predicting, aligning, or retrieving stems compatible with a context mix.

• At inference, given a context mix and (possibly zero-shot) instrument label $c$, context spectrogram patches are encoded via $f_\theta$, and predictions are made for each patch via $g_\phi(z_k, c)$. The mean predicted embedding forms a query vector.
• Retrieval is performed by nearest neighbor search in a reference bank containing time-averaged real stem embeddings, identifying stems likely to be compatible in timbre, rhythm, and style.

Alignment analysis is facilitated by patch-based embeddings at 160 ms resolution, enabling temporal similarity as a function of relative shift. The framework also supports prospective use of predicted stem embeddings as conditioning for waveform generation models (e.g., diffusion or autoregressive), broadening generative MIR pipelines (Riou et al., 2024).

Zero-shot conditioning is enabled by representing free-form instrument descriptors as CLAP embeddings; the FiLM-style conditioning in $g_\phi$ allows direct generalization to new timbres and instrument names at test time (Riou et al., 2024).

4. Experimental Results and Observations

Comprehensive experiments were conducted on MUSDB18 and MoisesDB.

Retrieval Performance

Model Variant	R@1 (MUSDB18)	R@5	R@10	mean NR (%)	med NR (%)
MLP w/ cond. (Riou et al., 2024)	33.0	63.2	76.2	2.0	0.5
+ CLAP + FiLM (Riou et al., 2024)	38.8	89.7	95.0	0.7	0.3
Transformer predictor (Riou et al., 2024)	5.2	17.5	25.7	12.1	6.0
AutoMashUpper (Riou et al., 2024)	1.0	8.8	15.5	29.1	19.5

• CLAP+FiLM conditioning enables robust zero-shot retrieval: performance drops only slightly when using fine-grained or unseen instrument name queries (Riou et al., 2024).
• Ablations show that contrastive pretraining notably improves recall@5/10 and normalized rank on both coarse and fine-grained instrument splits, but can hurt R@1 unless effective conditioning is used.
• Most retrieval errors are "correct track, wrong instrument of same category" or "correct instrument, wrong track," especially for underrepresented timbres.

User Study and Subjective Evaluation

Participants (20 with ≥10 years’ musical experience) consistently rated the top model retrieval nearly as compatible as the ground truth, and far above random, on criteria of genre, timbre, and style, but independent of pitch or key. Variance was highest for drums; “Other” category showed largest model-ground-truth gaps, reflecting broad and ambiguous label scope (Riou et al., 2024).

Embedding Structure and MIR Utility

• Patchwise embeddings exhibit strong local sensitivity: temporal alignment metrics display sharp peaks at zero shift, with side-peaks aligning with musical beats, bars, or absolute chunk boundaries.
• Clustering analysis on Beatles’ tracks shows embedding clusters align with harmonic relations (key/chord co-occurrences), indicating the representations encode tonal proximity.
• Downstream performance on key detection, genre classification, auto-tagging, and instrument recognition benchmarks is competitive with larger systems (MULE, Jukebox-5B), despite using two orders-of-magnitude less data. The choice of predictor has only minor effect on downstream classification; key detection was highest for transformer predictors (Riou et al., 2024).

Beat tracking with linear probes on patch-embeddings achieves F1, AMLt, and CMLt scores near strong contrastive SSL baselines, confirming retention of temporal signal features (Riou et al., 2024).

5. Limitations, Ablations, and Prospective Directions

• Predictor structure: MLPs with instrument conditioning yield higher retrieval accuracy than transformer predictors or unconditioned variants.
• Label taxonomy: The use of only four instrument labels, with a very broad "Other" class, limits granularity; CLAP-based zero-shot conditioning on fine-grained labels improves flexibility and adapts to open-vocabulary instrument retrieval (Riou et al., 2024).
• Dataset scale: Current reliance on a proprietary 20k-song multi-track dataset and sparse stem-separated corpora constrains coverage. Larger or more diverse separated datasets (possibly via source separation advances) are anticipated to benefit performance.
• Absolute positional encodings: Periodic alignment peaks tied to chunk/window boundaries suggest that absolute positional codes leak timing information; relative or no positional encoding is a potential avenue for improvement, especially for generative applications.
• Generative integration: Use of predicted embeddings as conditioning vectors for generative stem models (e.g., latent diffusion, autoregressive audio), or as plug-ins within end-to-end generative architectures, represents an active direction. The predictor’s use at inference is a paradigm that could generalize to multi-modal compatibility estimation (Riou et al., 2024).

6. Comparative Analysis and Significance

Stem-JEPA advances the state of the art in musical stem compatibility estimation by unifying contrastive and predictive embedding-based self-supervision. Its architecture and training design induce representations sensitive to both global (e.g., timbre, harmony, style) and local (temporal, rhythmic) musical structure.

Key distinguishing features relative to prior work include:

• Application of a JEPA predictive objective at patch level for context-to-target stem embedding mapping, instead of solely contrastive objectives or signal reconstruction.
• Zero-shot retrieval via CLAP/FILM conditioning, substantially extending the generality to arbitrary or unseen instrument/timbre queries.
• Rich latent representations validated via both subjective audition and MIR benchmarks, with performance robust to ablations and dataset perturbations (Riou et al., 2024, Riou et al., 2024).

A plausible implication is the emergence of a new class of compatibility-predictive self-supervised learning frameworks applicable beyond music, potentially to multi-modal and creative AI domains where flexible alignment and conditioning on semantic attributes are central.