Music-JEPA: Learning a World Model of Sound from Action
Abstract
Joint Embedding Predictive Architectures (JEPA) have recently emerged as a paradigm for learning world models by predicting latent representations, offering a promising direction for self-supervised learning. While initial attempts have applied JEPA to the music domain, it remains unclear how such frameworks can naturally support the formation of a world model for music. In this work, we propose to learn a world model of piano sound using JEPA by framing music as an action-conditioned system: the audio is treated as the state, and the pianoroll as the instrument action. Given a current audio state and an action, the model predicts the resulting future audio state, mirroring how humans learn musical sound through interaction. The model is trained in a fully offline setting using paired audio-pianoroll data, without environment interaction. Experiments show that the learned model captures the relationships between musical actions and their resulting sound. The resulting representations support downstream tasks, including beat tracking, composer identification, and key estimation, and enable piano transcription via planning, by searching for actions that best explain a target sound.
The following analysis is generated by AI and may contain errors. It is separate from the authors’ abstract and is not peer review.
Research summary
The paper introduces Music-JEPA, an action-conditioned world model that learns latent dynamics of piano music by predicting future audio states from current states and pianoroll actions. This work makes a significant contribution to self-supervised learning in audio by moving beyond passive representation learning to active, action-conditioned modeling. The technical execution is solid, with clear demonstrations of learned dynamics and competitive performance on downstream tasks. While the novelty is incremental (applying known JEPA techniques to a new domain with a specific conditioning mechanism), the significance lies in the rigorous evaluation of action-conditioned latent spaces for music, providing a strong baseline for future research in controllable music generation and understanding. The limitation of requiring paired data is a notable constraint, but the results validate the potential of this paradigm.
Methodology
The paper proposes Music-JEPA, an action-conditioned Joint Embedding Predictive Architecture for piano music. The core innovation lies in framing music as a dynamical system where audio spectrograms are states and pianoroll/pedal signals are actions. The model predicts future latent states given current states and actions, leveraging the EMA-based training scheme to prevent collapse. The methodology is technically sound, adapting recent vision JEPA techniques to the audio domain with a specific focus on action-conditioning. The use of ViT backbones for both spectrogram and pianoroll inputs is standard but effective. The introduction of an inverse predictor for transcription via planning is a notable conceptual extension, although the implementation relies on amortized optimization rather than direct latent-space planning, which is a pragmatic choice given the high dimensionality of the action space.
Results & evaluation
The experimental evaluation is comprehensive, covering temporal dynamics analysis, downstream MIR tasks (beat tracking, composer identification, key recognition), and transcription via planning. The results show that Music-JEPA outperforms a passive audio-only JEPA baseline and remains competitive with MERT, a much larger model, on several tasks. The ablation studies on perturbation sensitivity provide strong evidence that the model learns meaningful action-conditioned dynamics rather than just temporal proximity. The transcription results are mixed; while pedal estimation is strong, note-level transcription lags behind supervised methods, which is expected for a generative/planning approach but highlights current limitations. The comparison with MERT is fair and demonstrates the efficiency of the proposed approach.
Reproducibility
The paper provides detailed hyperparameters, dataset splits (MAESTRO v3.0.0), and architectural details. The training algorithm is clearly outlined. However, the code and pretrained models are stated to be "released soon," which currently hinders immediate reproducibility. The specific patch sizes and embedding dimensions are provided, allowing for architectural replication.
Limitations
The primary limitation is the reliance on paired audio-MIDI data, which restricts the model to domains with precise action annotations (like piano) and prevents direct application to unannotated audio. The transcription performance, particularly for notes, is inferior to supervised methods, indicating that the latent space may not be perfectly invertible or that the planning objective is not yet optimized for high-fidelity reconstruction. The model's ability to generalize to other instruments or genres is untested. Additionally, the "planning" aspect is implemented via an inverse model (amortized planning) rather than search-based planning, which limits the interpretability of the "world model" aspect in terms of explicit trajectory optimization.
Broader impact
This work contributes to the growing field of world models for audio and music information retrieval. By demonstrating that action-conditioned latent dynamics can capture musical structure, it opens avenues for more interactive and controllable music AI systems. The approach could inspire similar frameworks for other domains where actions influence observations, such as robotics or speech synthesis. The focus on self-supervised learning with minimal labeled data aligns with broader trends in efficient AI. The paper introduces Music-JEPA, an action-conditioned world model that learns latent dynamics of piano music by predicting future audio states from current states and pianoroll actions. This work makes a significant contribution to self-supervised learning in audio by moving beyond passive representation learning to active, action-conditioned modeling. The technical execution is solid, with clear demonstrations of learned dynamics and competitive performance on downstream tasks. While the novelty is incremental (applying known JEPA techniques to a new domain with a specific conditioning mechanism), the significance lies in the rigorous evaluation of action-conditioned latent spaces for music, providing a strong baseline for future research in controllable music generation and understanding. The limitation of requiring paired data is a notable constraint, but the results validate the potential of this paradigm.
Author of this paper? Send a correction.