⚡ electrical engineering

MIDI-RAE-JEPA: Hierarchical Representation Learning and Generation for Symbolic Music

This paper introduces MIDI-RAE-JEPA, a self-supervised framework that combines a Swin Transformer V2 encoder with pitch- and time-shift equivariance objectives to learn hierarchical, semantically rich representations of symbolic music that excel in downstream tasks like emotion classification and enable high-quality generative modeling.

Original authors: Scott H. Hawley

Published 2026-07-17

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Scott H. Hawley

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to understand music, not just by listening to the sound waves, but by reading the sheet music itself. In the world of artificial intelligence, this is like teaching a computer to read a map of musical notes. Usually, AI learns by looking at millions of photos of cats or dogs, learning that a "cat" is a cat whether it's zoomed in or zoomed out. But music is different. If you zoom in on a piece of sheet music, you don't just see a bigger note; you see a different rhythm or a different pitch entirely. The rules that work for photos break down for music.

To solve this, scientists use a clever trick called "self-supervised learning." Instead of hiring a human teacher to label every single note, they let the AI play a game of "fill in the blanks." They hide parts of the music and ask the AI to guess what's missing. If the AI can guess correctly, it means it has learned the hidden rules of how music is built. The big question this paper tackles is: Can we teach an AI to understand the structure of music—how a small melody fits into a big song—without ever showing it a single labeled example? And if we do, can that AI then write its own music that sounds good?

Meet MIDI-RAE-JEPA, a new AI system designed to learn the secret language of symbolic music (music written as digital instructions, like a piano roll). Think of a piano roll as a giant, binary grid: 128 squares wide for every possible piano key, and 128 squares tall for time. A black square means a note is playing; a white square means silence. The goal of this project was to teach an AI to look at these grids and understand the deep, hierarchical relationships of music—how a single note builds a phrase, how phrases build a section, and how sections build a whole song.

The researchers built a system that acts like a musical detective. They used a special type of AI architecture called a Swin Transformer V2, which is like a set of magnifying glasses that look at the music at different levels of detail, from tiny single notes up to the whole song structure. But the real magic is in how they taught it. Instead of just guessing missing notes, they taught the AI a game of "musical translation."

Imagine you have a melody. If you shift that melody up by a few notes (pitch shift) or move it forward in time (time shift), the structure of the music stays the same, even though the notes are different. The AI was trained to recognize that these shifted versions are related. If you shift a song up by one note, the AI's internal "fingerprint" for that song should move just a little bit. If you shift it up by an entire octave, the fingerprint should move a lot further away. This is called equivariance. It's like teaching a child that a red ball is still a red ball whether it's on the floor or on the table; the AI learns that a musical idea is still the same idea, even if it's in a different key or at a different time.

To make sure the AI didn't just memorize the answers or get confused, the researchers added a few safety nets. They used a "masked embedding predictor," which is like covering up parts of the sheet music and asking the AI to predict what's underneath based on the rest of the song. They also added a special rule called SIGReg to stop the AI from getting lazy and giving the same answer for every song (a problem known as "collapse").

The results were surprisingly successful. When they froze the AI's brain (stopped it from learning anything new) and attached a simple decoder to it, the system could reconstruct the original piano roll with incredible accuracy, scoring an F1 of 0.995. That's almost perfect. It means the AI had truly captured the essence of the music in its internal memory.

But the real test was generation. The researchers took these internal "fingerprint" representations and fed them into a generative model (a type of AI that creates new things). When they gave the model a specific musical excerpt to "condition" on, the AI generated new music that matched the pitch register (how high or low the notes are) and the rhythmic density (how busy or sparse the notes are) of the original. If they gave it a mismatched condition, the AI still made music that sounded plausible, but it was unrelated to the prompt. This suggests the AI isn't just copying; it's understanding the style and structure of the input.

The paper also tested if these learned representations could help with other tasks. They tried using the AI's internal fingerprints to guess the emotion of a song (like "happy" or "sad"). The results suggest that this method works better than older, simpler techniques (like a Haar scattering transform), proving that the AI learned something musically meaningful without needing to be explicitly taught emotions.

However, the authors are careful not to claim they have solved everything. They note that while the "finer" levels of their AI (which look at small details) are very good at capturing specific notes and rhythms, the "coarser" levels (which look at the big picture) haven't quite learned to capture obvious, interpretable musical abstractions yet. The system is a strong proof-of-concept, showing that combining equivariance (understanding shifts) with hierarchical learning (understanding scales of detail) is a viable path forward. It suggests that we can build AI assistants that listen to a musical idea and understand its structure well enough to help refine a composition, but we are still working on getting the AI to fully grasp the grand, abstract architecture of a symphony.

In short, this paper shows that by teaching an AI to respect the rules of musical shifts and scales, we can create a system that understands music deeply enough to reconstruct it perfectly and generate new, plausible music based on specific musical moods. It's a step toward an AI that doesn't just mimic music, but truly understands its grammar.

Technical Summary: MIDI-RAE-JEPA

Problem Statement

The paper addresses the lack of rich, self-supervised internal representations for symbolic music that capture its hierarchical, multiscale nature. While large multimodal language models exist, they perform poorly on music understanding tasks. Furthermore, standard self-supervised learning (SSL) objectives designed for natural images (e.g., DINOv2's crop-and-resize strategy) are semantically problematic for piano roll representations. "Zooming in" on a piano roll merges adjacent semitones and distorts rhythmic structure, unlike the smooth scale invariance found in photographs. The goal is to develop a domain-specific approach that learns hierarchical representations of symbolic music (encoded as piano roll images) capable of supporting both discriminative tasks and high-quality generation without end-to-end fine-tuning.

Methodology

Data Representation

The system operates on the POP909 dataset, rendering 8-bar segments as 128 × 128 binary piano roll images. The horizontal axis represents time (64 eighth-note steps), and the vertical axis represents pitch (all 128 MIDI pitches). The training pipeline utilizes random cropping with pitch and time transposition augmentation to expand the effective training variety.

Architecture

The core of the system is a Hierarchical Swin Transformer V2 encoder with six stages. It processes the input using shifted windows for self-attention, downsampling by 2× at each stage while doubling the embedding dimension. The architecture produces embeddings at multiple resolutions, from fine-grained local details (Level 5) to coarse global context (Level 0).

Training Objectives

The encoder is trained purely on self-supervised objectives, combining three main components:

Equivariance Loss ( $L_{equiv}$ ): Instead of forcing all shifted views to identical embeddings (which causes collapse), the model enforces a target embedding distance proportional to the shift magnitude.
- Mechanism: Two views ( $x_1$ and $x_2$ ) are created by shifting the original crop in pitch and time. The loss minimizes the squared difference between the Euclidean distance of the embeddings and a target distance scaled by the shift magnitude ( $\alpha \sqrt{d} \|\hat{\delta}\|$ ).
- Effect: This encourages the model to internalize temporal and pitch relationships, ensuring that larger shifts result in proportionally more distant embeddings.
Masked Embedding Predictor (MEP) & LeJEPA: Following the LeJEPA framework, a lightweight predictor network attempts to reconstruct masked patch embeddings from the student's full-context representation, using stable targets generated by an Exponential Moving Average (EMA) teacher. This ensures the model learns predictive representations of musical context across the hierarchy.
Regularization and Factorization:
- SIGReg: Chunked Sketched Isotropic Gaussian Regularization is applied to prevent representation collapse and enforce an isotropic Gaussian prior, theoretically optimal for downstream tasks.
- Soft Factorization Loss ( $L_{fact}$ ): A novel loss term encourages pitch-shift and time-shift directions to be geometrically orthogonal in the latent space. It uses cosine similarity on augmentation difference vectors to enforce parallel, anti-parallel, or orthogonal relationships depending on the type of shift, without prescribing specific directions.

Decoder and Generative Model

Decoder: A mirrored Swin hierarchy with a Feature Pyramid Network (FPN) reconstructs the piano roll from the frozen encoder embeddings. It is trained using binary cross-entropy.
Generative Model: A Flow Matching model operates in pixel space, conditioned on the encoder embeddings. To manage dimensionality, embeddings undergo PCA (retaining 95% variance), with the mean pitch explicitly appended to preserve register information. Training employs a Multi-Level Dropout scheme to enable unconditional, partially conditioned, and fully conditioned generation.

Key Contributions

First Hierarchical JEPA for Symbolic Music: The paper presents the first hierarchical JEPA-based approach for symbolic music, combining equivariance objectives, LeJEPA, and a Swin V2 encoder.
Complete RAE Pipeline: It demonstrates a full Representation Autoencoder (RAE) pipeline for music, showing that frozen encoder embeddings support high-fidelity reconstruction, conditional generation, and downstream classification.
Measurable Equivariance: The authors provide empirical validation that embedding distances increase monotonically with pitch and time shift magnitude, confirming measurable equivariance.
Novel Loss Functions: The introduction of a smooth, delta-scaled equivariance loss and a soft factorization loss that encourages geometric orthogonality between pitch and time directions in the latent space.
Feasibility on Consumer Hardware: All experiments were conducted on consumer-grade GPUs (e.g., RTX 4090), demonstrating that such complex SSL training does not require large-scale compute clusters.

Results

Equivariance and Factorization

Equivariance: Embedding distances increase monotonically with shift magnitude, extending beyond the nominal patch size of each hierarchy level.
Factorization: The soft factorization loss successfully establishes geometric relationships where cross-type shifts (pitch vs. time) are nearly orthogonal ( $|\cos| < 0.02$ ), and same-type shifts are parallel or anti-parallel.

Downstream Tasks

Emotion Classification: On the EMOPIA dataset, linear probes trained on the frozen embeddings (specifically at Level 3) outperformed both a Haar scattering transform baseline and a DINOv2 baseline across all metrics (4-class, Arousal, and Valence).
Structure Probes: Fine levels (L4/L5) strongly encoded note density ( $R^2 > 0.98$ ), while coarser levels captured broader context. Cross-song similarity decreased monotonically from coarse to fine levels, indicating increasing discriminability.

Reconstruction and Generation

Reconstruction: The decoder achieved an F1 score of 0.995 on held-out test data, confirming that the SSL objective forces the encoder to retain sufficient information for accurate reconstruction.
Conditional Generation: The flow matching model produced generations that closely matched the pitch register and rhythmic density of the conditioning excerpt. Mismatched or unconditional conditioning yielded musically plausible but unrelated outputs.
Conditioning Validity: Analysis of the conditioning projection norm showed that the model learns to ignore shuffled (uninformative) conditioning signals while maintaining stable weights for valid conditioning.

Significance and Claims

The paper claims that equivariance-based SSL objectives, when combined with sufficient fine-level encoder capacity, provide a viable path toward semantically rich, generatively useful representations of symbolic music.

Key takeaways include:

The paradigm of using frozen pretrained representations for both discrimination and generation (RAE) extends successfully to the symbolic music domain.
The learned representations capture musically meaningful structures (pitch intervals, rhythmic periods) without requiring fine-tuning for downstream tasks.
The approach avoids the semantic pitfalls of image-based SSL on piano rolls by utilizing translation equivariance rather than scale invariance.
While the finer levels of the hierarchy capture local details effectively, the paper notes that capturing "obviously interpretable musical abstractions" at the coarser levels remains an open problem for future work.

The authors conclude that this method offers a foundation for future work in extending these approaches to audio spectrograms, scaling to larger datasets, and exploring richer downstream tasks like arrangement or accompaniment generation.