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

Ziyu Wang, Kun Fang, Yann LeCun

New York University  •  MBZUAI  •  McGill University  •  CIRMMT  •  AMI Labs

This is the demo page for the paper Music-JEPA: Learning a World Model of Piano Sound from Action. We introduce a Joint Embedding Predictive Architecture (JEPA) that learns to predict future piano sound representations conditioned on performance actions. This page provides additional examples to supplement the paper and is organized into three parts:

Part 1: Action-Conditioned Prediction

To illustrate the action-conditioned predictions of Music-JEPA, we provide samples showing the transition from the current audio state $s_t$ (spectrogram) and an action $a_{t+1}$ (pianoroll and pedal) to the resulting audio state $s_{t+1}$. The ground-truth and prediction rows use the same action, while the counterfactual row uses a modified action: an unrelated action randomly sampled from the dataset (Samples 1–2), or a constructed action such as arpeggios (Sample 3) or scales (Sample 4).

Predictions are performed in the Music-JEPA latent space and then decoded into spectrograms using a separately trained decoder, with Music-JEPA kept frozen during decoder training. The decoder is a 12-layer ViT trained to reconstruct normalized log-mel spectrograms using a combination of Pearson correlation and L1 reconstruction losses. The corresponding pianorolls and spectrograms are shown below. The audio below each spectrogram is synthesized from the reconstructed spectrogram via the Griffin–Lim algorithm, while the audio below each pianoroll is rendered from MIDI using a standard soundfont. These reconstructions are not intended for high-fidelity audio reproduction, but rather to illustrate the information captured by the learned latent representations and their dependence on the provided actions. As phase is not explicitly modeled, the resulting audio may exhibit perceptual artifacts.

Click the tabs to switch among samples.

Part 2: Piano transcription via Amortized Action Inference

Music-JEPA supports piano transcription by inferring performance actions from consecutive audio states. We train an inverse predictor to estimate the latent representation of the next action, followed by a ViT-based denoising autoencoder (DAE) that refines the estimate toward the learned action representation space. A separate ViT decoder then maps the refined representation to a pianoroll and continuous pedal values. These modules are trained separately, with Music-JEPA kept frozen.

Below are selected audio examples taken from the test split of MAESTRO [1]; use the tabs to switch among samples. In each tab, the first section presents the original audio and the synthesized ground-truth MIDI. The second section presents the Music-JEPA transcription, two variants with an additional action-prediction loss using either an EMA or student target, and a supervised end-to-end baseline that uses the same transcription pipeline but is trained from scratch directly on the transcription objective. The third section presents results from two supervised transcription systems, Yan et al. [2] and Kong et al. [3]. All transcribed MIDI files are rendered to audio using PianoTeq. Note that these examples are provided for qualitative comparison; differences in transcription metrics do not necessarily translate directly to differences in perceived audio quality.

Part 3: Transcription on Out-of-Distribution (OOD) Samples

In addition to standard piano transcription, we present results on out-of-distribution (OOD) audio from the Piano Concerto Dataset (PCD) [4], which contains both piano and orchestral instruments. We use the same transcription methods and audio rendering setup as above, but ground-truth MIDI is not available for these samples. The examples therefore serve as a qualitative comparison of how different transcription methods respond to audio beyond solo piano.

In these examples, the supervised transcription methods of Yan et al. [2] and Kong et al. [3] tend to transcribe primarily the piano component, while Music-JEPA may also produce piano actions corresponding to parts of the surrounding orchestral texture. This listening-based observation does not imply that either behavior is necessarily preferable or ideal.

References

  1. C. Hawthorne, A. Stasyuk, A. Roberts, I. Simon, C.-Z. A. Huang, S. Dieleman, E. Elsen, J. Engel, and D. Eck, “ Enabling factorized piano music modeling and generation with the MAESTRO dataset ,” in International Conference on Learning Representations (ICLR), 2019.
  2. Y. Yan and Z. Duan, “ Scoring time intervals using non-hierarchical transformer for automatic piano transcription ,” in Proceedings of the 25th International Society for Music Information Retrieval Conference (ISMIR), 2024.
  3. Q. Kong, B. Li, X. Song, Y. Wan, and Y. Wang, “ High-resolution piano transcription with pedals by regressing onset and offset times ,” IEEE/ACM Transactions on Audio, Speech, and Language Processing (TASLP), 2021.
  4. Y. Özer, S. Schwär, V. Arifi-Müller, J. Lawrence, E. Sen, and M. Müller, “ Piano Concerto Dataset (PCD): A Multitrack Dataset of Piano Concertos ,” Transactions of the International Society for Music Information Retrieval (TISMIR), 2023.