One-sentence Summary

Researchers from New York University, McGill University, MBZUAI, and other institutions propose MUSIC-JEPA, a JEPA-based world model for piano sound that treats audio as state and pianoroll as action, predicting future audio states from current audio and actions, enabling downstream tasks like beat tracking, composer identification, key estimation, and piano transcription via planning.

Key Contributions

• Music-JEPA is a world model that learns action-conditioned temporal dynamics of piano audio by treating the audio as state and pianoroll as action, using paired data to predict future latent representations.
• The learned representations outperform an audio-only JEPA baseline and achieve performance comparable to MERT on downstream music information retrieval tasks, including beat tracking, composer identification, and key estimation, while using only 7% of MERT's parameters.
• The learned dynamics enable piano transcription via amortized planning, where an inverse predictor infers actions and a decoder generates outputs, producing coherent transcriptions that effectively capture continuous pedal variations.

Introduction

The authors address music understanding from an active engagement perspective, where humans learn to anticipate musical outcomes by interacting with instruments. This motivates a music world model that learns representations by capturing temporal dynamics between performance actions and resulting audio states. Existing self-supervised audio models, including those based on Joint Embedding Predictive Architecture (JEPA), treat audio passively and do not explicitly model how music unfolds over time, leaving their sensitivity to music-specific concepts like key, chord, and beat largely unexplored. The authors propose Music-JEPA, a JEPA-based framework that learns action-conditioned latent dynamics by predicting future audio states from past states and actions, without reconstructing raw audio. The learned representations outperform audio-only JEPA baselines on piano music information retrieval tasks, match much larger models despite using only 7% of the parameters, and enable planning, yielding meaningful piano transcription through amortized inference.

Method

The authors propose Music-JEPA, which models music as a dynamical system where audio evolves over time under the influence of performance actions. The framework is instantiated on piano music using paired audio and control sequences $(x_{1:T}, y_{1:T})$(x1:T​,y1:T​), where $x_t$xt​ represents a 2-second audio segment and $y_t$yt​ denotes the corresponding pianoroll and pedal signals. Audio is represented as log-mel spectrograms, while the action consists of a pianoroll encoding note velocity and a sustain pedal signal aligned to the same frame rate.

As shown in the figure below:

The model learns latent state and action representations $s_t = \mathcal{E}_{\mathrm{state}}(x_t)$st​=Estate​(xt​) and $a_t = \mathcal{E}_{\mathrm{action}}(y_t)$at​=Eaction​(yt​), modeling the dynamics as:

Here, $f(\cdot, \cdot)$f(⋅,⋅) predicts the next state, and $g(\cdot)$g(⋅) models the temporal structure of actions. Modeling the evolution of high-dimensional piano control provides a structured prior over the action space.

The architecture follows the Joint Embedding Predictive Architecture (JEPA) framework, utilizing a Vision Transformer (ViT) backbone. The spectrogram and pianoroll are treated as image-like inputs and mapped into patch-based representations. The spectrogram is divided into patches along time and frequency, projected into a D-dimensional embedding via a convolutional layer. Similarly, the pianoroll is partitioned into patches along time and pitch. The sustain pedal signal is treated as a one-dimensional temporal sequence, encoded with a 1D convolution, and fused with the pianoroll embeddings. After patch embedding, transformer encoders process the state and action tokens to produce latent representations $s \in \mathbb{R}^{K_s \times D}$s∈RKs​×D and $a \in \mathbb{R}^{K_a \times D}$a∈RKa​×D.

The state and action predictors are also ViT-based. The state predictor $f(s, a)$f(s,a) operates on a partially masked state sequence, taking $s_t$st​ and a masked version of $s_{t+1}$st+1​ as input, and predicts the missing regions via self-attention. Cross-attention to the action tokens $a_{t+1}$at+1​ is applied at each layer to condition the prediction. In contrast, the action predictor $g(a)$g(a) is a standard Transformer without cross-attention, trained to reconstruct $a_{t+1}$at+1​ from $a_t$at​ and masked tokens.

Given inputs $(x_{t-1}, y_{t-1})$(xt−1​,yt−1​) and $(x_t, y_t)$(xt​,yt​), the training objective is formulated as:

where $\lambda$λ controls the strength of the action prior. To mitigate representation collapse, the authors adopt an exponential moving average (EMA) strategy, maintaining stop-gradient teacher encoders with parameters $\bar{\theta}$θˉ updated as:

where $\tau$τ is a momentum coefficient. Layer normalization is applied to both student and teacher encoder outputs to control representation scale.

For downstream tasks, the model is frozen, and the state encoder serves as a feature extractor. Beyond representation learning, the learned dynamics enable piano transcription via planning in latent space. Given a known sequence of latent states $s_{1:T}$s1:T​, the problem is formulated as finding actions $a_{1:T}^*$a1:T∗​ that minimize the prediction error. To handle the high-dimensional action space, the authors adopt an amortized optimization approach by training an inverse predictor $a_{t+1} = h(s_t, s_{t+1}, a_t)$at+1​=h(st​,st+1​,at​) to approximate the solution. The predicted latent actions are then mapped back to the observation space using a separately trained action decoder.

Experiment

Music-JEPA is trained on classical piano music (MAESTRO) and evaluated against an audio-only JEPA and MERT. Latent dynamics analysis demonstrates strong sensitivity to state and action perturbations, validating the benefit of action-conditioned modeling over the audio-only baseline. On downstream MIR tasks, the model outperforms the audio-only JEPA and remains competitive with MERT, while transcription via planning yields strong continuous pedal estimation but lags behind fully supervised note-level methods.

Music-JEPA demonstrates robust sensitivity to mismatched state transitions, with high win rates under both temporal and random perturbations. AO-JEPA lags notably behind, especially on temporal perturbations of the target state, where its win rate drops substantially. Music-JEPA's win rate for temporal input perturbations is higher than AO-JEPA's, indicating better detection of input timing mismatches. For target state temporal shifts, Music-JEPA maintains a near-perfect win rate, while AO-JEPA's falls to 0.576, highlighting a large gap in temporal sensitivity. Both models approach perfect win rates under random perturbations, but Music-JEPA consistently outperforms AO-JEPA across all conditions.

The lightweight Music-JEPA model (6M parameters) outperforms the similarly sized AO-JEPA on all three tasks. It achieves the best beat tracking F1 scores and, together with AO-JEPA, surpasses the 95M-parameter MERT model by roughly 10 points on key recognition. For composer identification, Music-JEPA yields the highest top-3 and top-5 accuracy, while MERT leads in top-1 accuracy and weighted F1. Music-JEPA achieves the best F1 scores for beat tracking at both 70 ms and 100 ms tolerance thresholds. Despite having only 6 million parameters, Music-JEPA surpasses the 95-million-parameter MERT on key recognition and beat tracking. Both JEPA-based models outperform MERT by about 10 points across all key recognition metrics. Music-JEPA attains the highest top-3 and top-5 composer identification accuracy, but MERT achieves higher weighted F1 and top-1 accuracy. AO-JEPA shows lower performance than Music-JEPA on every task, with the largest gap in composer identification weighted F1.

The proposed model trails fully supervised methods on note-level transcription, with lower F1 scores across frame, onset, onset+offset, and onset+offset+velocity, but achieves the best continuous sustain pedal estimation, recording the lowest mean absolute error among all compared systems. The model's note-level transcription F1 scores are consistently lower than those of Yan et al. and Kong et al. across all evaluation stages. For sustain pedal estimation, the model attains the lowest MAE (0.1222), outperforming the supervised pedal model by Zhang et al. and the other transcription systems.

Music-JEPA is evaluated across multiple tasks, including sensitivity to state transitions, lightweight downstream music understanding, and note-level transcription with pedal estimation. It exhibits robust detection of temporal mismatches, far outperforming AO-JEPA, and as a 6M-parameter model it surpasses the 95M-parameter MERT on key recognition and beat tracking while lagging behind in composer identification top-1 accuracy. On transcription, the model trails fully supervised approaches but uniquely achieves the lowest mean absolute error for continuous sustain pedal estimation, highlighting its efficient capture of musically relevant structure.