Music-JEPA: Learning a World Model of Sound from Action

Music-JEPA: Learning a World Model of Sound from Action

Paper: Music-JEPA: Learning a World Model of Sound from Action

Links: arXiv:2607.22000 · Demo Page

---

One-Sentence Summary

Music-JEPA models music as an action-conditioned system — audio as state, pianoroll as action — using a Joint Embedding Predictive Architecture (JEPA) to predict future audio states conditioned on actions in latent space, with EMA training to prevent representation collapse, yielding representations that support downstream tasks including beat tracking, composer identification, and key estimation, and enabling piano transcription via latent-space planning.

---

Background and Motivation

Human understanding of music typically arises from active engagement: learning piano involves interacting with the instrument and understanding how keyboard actions produce acoustic outcomes. Cognitive studies suggest humans develop internal representations enabling them to anticipate and mentally simulate musical outcomes without direct auditory input. This learning process naturally suggests a music world model — a self-supervised framework learning music representation by modeling temporal dynamics between actions (e.g., playing an instrument) and resulting states (e.g., sound).

Joint Embedding Predictive Architectures (JEPA) are a paradigm for learning world models by predicting latent representations rather than reconstructing raw observations, avoiding the burden of modeling unnecessary details. However, existing music-domain JEPA work primarily treats it as a passive self-supervised approach for learning general audio representations, not modeling how music evolves through time, and rarely evaluated on music-specific concepts like key, chord, or beat. Whether JEPA-based representations can capture temporal dynamics required for music understanding and planning remains unclear.

Music-JEPA's core insight is to model music as an action-conditioned system: a 2-second piano audio segment as state, the corresponding pianoroll and sustain pedal signals as action. The model predicts the next state from the previous state and current action directly in latent space, without reconstructing raw audio.

Figure 1: Model diagram of Music-JEPA. $\mathcal{E}_s$, $\mathcal{E}_a$, $f$, $g$ denote the state encoder, action encoder, state predictor, and action predictor. The shaded region is latent temporal dynamics.

---

Method

1. Data Representation

The state is a 2-second piano audio segment converted to a spectrogram and split into $K_s = 16 \times 16 = 256$ patches of $25 \times 15$. The action is the corresponding pianoroll plus sustain pedal signal, split into $K_a = 16 \times 15 = 240$ patches of $25 \times 6$. With embedding dimension $D = 256$, state and action representations are $(256, 256)$ and $(240, 256)$ respectively.

2. Model Architecture

The state encoder uses 12 Transformer layers, the action encoder 8 layers, and state and action predictors 6 layers each. All modules use $d_{\text{model}}=256$, $d_{\text{ff}}=512$, 4 attention heads. Total ~19M parameters.

3. Training Objective and Regularization

The state predictor predicts the next state from the current state and next action; the action predictor predicts the next action from the current action. An EMA training scheme prevents representation collapse — teacher network parameters are an exponential moving average of the student's:

$$ \theta_{\text{tea}} \leftarrow \tau \cdot \theta_{\text{tea}} + (1 - \tau) \cdot \theta_{\text{stu}}, \quad \tau = 0.95 $$

The training loss combines state prediction loss and action prediction loss:

$$ \mathcal{L} = \text{MSE}\bigl(f(s_t, a_{t+1}),\, \mathrm{sg}(\mathcal{E}_s^{\text{tea}}(x_{t+1}))\bigr) + \lambda \cdot \text{MSE}\bigl(g(a_t),\, \mathrm{sg}(\mathcal{E}_a^{\text{tea}}(y_{t+1}))\bigr) $$

where $\mathrm{sg}$ is the stop-gradient operator, $\lambda = 0.5$. Teacher targets use stop-gradient to prevent gradient backflow. The specific teacher targets are:

$$ s_{t+1}^{\text{tgt}} = \mathrm{sg}\bigl(\text{LN}(\mathcal{E}_s^{\text{tea}}(x_{t+1}))\bigr), \quad a_{t+1}^{\text{tgt}} = \mathrm{sg}\bigl(\text{LN}(\mathcal{E}_a^{\text{tea}}(y_{t+1}))\bigr) $$

where $\text{LN}$ is layer normalization. Student representations are also layer-normalized: $s_{\text{curr}} = \text{LN}(\mathcal{E}_s^{\text{stu}}(x_{\text{curr}}))$. Prediction occurs in the normalized latent space, avoiding the burden of reconstructing raw spectrograms. The model is trained on MAESTRO v3.0.0 (~200 hours of piano recordings with time-aligned MIDI) using Adam, learning rate $6 \times 10^{-4}$, weight decay $10^{-4}$, batch size 128, on a single A100 for 15-25 epochs.

4. Downstream Tasks and Planning

After training, the model is frozen and the state encoder serves as a feature extractor. For audio of length $T$ seconds, representations are concatenated along the temporal dimension to $(\frac{T}{2} \times K_s, D)$. The learned dynamics also enable piano transcription via planning in latent space — searching for actions that best explain a target sound:

$$ a_{1:T}^{*}=\operatorname*{argmin}_{a_{1:T}}\sum_{t=1}^{T-1}\Big(\|f(s_{t},a_{t+1})-s_{t+1}\|^{2}+\|g(a_{t})-a_{t+1}\|^{2}\Big) $$

Since the action space is high-dimensional, an amortized optimization approach is adopted — training an inverse predictor $a_{t+1} = h(s_t, s_{t+1}, a_t)$ to approximate the solution. Predicted latent actions are mapped back to observation space via a separately trained action decoder.

flowchart TD
    A["Piano audio x_t
spectrogram 256 patches"] --> B["State encoder Es_stu
12-layer Transformer"]
    B --> C["State rep s_t
(256, 256)"]
    D["Pianoroll+pedal y_t
240 patches"] --> E["Action encoder Ea_stu
8-layer Transformer"]
    E --> F["Action rep a_t
(240, 256)"]
    C --> G["State predictor f
6-layer Transformer"]
    F2["Next action a_{t+1}"] --> G
    G --> H["Predicted state s_hat_{t+1}"]
    H --> I{"MSE vs
teacher target sg(Es_tea(x_{t+1}))"}
    F --> J["Action predictor g"]
    J --> K["Predicted action a_hat_{t+1}"]
    K --> L{"MSE vs
teacher target sg(Ea_tea(y_{t+1}))"}
    I --> M["Total loss L
= MSE_state + λ·MSE_action"]
    L --> M
    M --> N["EMA update teacher
τ=0.95"]
    style B fill:#e1f5fe
    style G fill:#fff3e0
    style N fill:#e8f5e9

---

Experiments

Latent Dynamics Evaluation

Figure 2: Perturbation sensitivity measured by loss difference $\Delta\mathcal{L}$ under input state, target state, and action perturbations.

The model is evaluated on whether it captures action-conditioned temporal dynamics. Given input state $s_t$ and action $a_{t+1}$, the model predicts $\hat{s}_{t+1} = f(s_t, a_{t+1})$, and error against ground truth $s_{t+1}^{\text{gt}}$ is measured:

$$ \mathcal{L}(s_{t},a_{t+1},s_{t+1}^{\text{gt}})=\|f(s_{t},a_{t+1})-s_{t+1}^{\text{gt}}\|^{2} $$

A well-formed dynamics model assigns lower error to correct state-action-target triplets and higher error when any component is perturbed. The win rate is the fraction of cases with $\Delta\mathcal{L} > 0$.

Table 1: Win rate under temporal and random perturbations — Music-JEPA significantly outperforms AO-JEPA on temporal perturbations

Model	Input State-Temporal	Input State-Random	Target State-Temporal	Target State-Random
Music-JEPA	0.929	0.999	0.991	0.992
AO-JEPA	0.787	0.986	0.576	0.984

Music-JEPA significantly outperforms AO-JEPA on temporal perturbations (input state 0.929 vs 0.787, target state 0.991 vs 0.576), showing that action-conditioned modeling substantially improves temporal dynamics capture. Pitch shifts have the strongest effect, particularly at tritone and other dissonant intervals.

Downstream MIR Tasks

Table 2: Comparison of pretrained music representations on three downstream MIR tasks

Model	Params	Beat F1@70ms	Composer wF1	Key wF1
Music-JEPA	6M	0.6208	0.5520	0.7617
AO-JEPA	6M	0.6013	0.4606	0.7615
MERT	95M	0.5780	0.6690	0.6469

Music-JEPA with only 6M parameters surpasses 95M-parameter MERT on beat tracking and key recognition, and is second to MERT but better than AO-JEPA on composer identification. The temporal dynamics understanding from action-conditioned modeling directly translates to downstream music understanding improvements.

Figure 3: Four action-conditioned synthesis examples. Each group: left=ground truth, middle=model prediction, right=counterfactual with different action.

Figure 4: Loss difference distribution under input state perturbation.

---

Limitations

1. Piano only: The model is trained and evaluated only on piano recordings from the MAESTRO dataset, without verifying generalization to other instruments or multi-instrument music.
2. Offline training limitation: Although modeled in an action-conditioned manner, training is fully offline using paired audio-pianoroll data without environment interaction — still a gap from a truly interactive world model.
3. High-dimensional planning challenge: Piano transcription planning requires optimization in a high-dimensional latent action space where gradient-based planning remains challenging, relying on an amortized inverse predictor approximation rather than direct optimization.

---

Conclusion and Outlook

Music-JEPA learns a world model of piano sound by modeling music as an action-conditioned system within the JEPA framework. The state predictor predicts the next audio state from the current audio state and pianoroll action, with EMA training preventing collapse. The learned representations outperform passive JEPA and the larger MERT on beat tracking, composer identification, and key estimation, and enable piano transcription via latent-space planning. This demonstrates that action-conditioned modeling is crucial for musical temporal dynamics understanding.

Golden quote: "Modeling music as an action-conditioned system — audio as state, performance actions as input — transforms JEPA from passive representation learning into an active world model that captures the causal dynamics of how actions produce sound."