Frequency-Aware Self-Supervised Music Representation Learning
Abstract
Self-supervised learning (SSL) has emerged as an essential paradigm for music information retrieval (MIR). While current SSL models achieve state-of-the-art performance across various MIR tasks, they typically treat audio as 1D sequences, either operating on time-domain waveforms or on flattened time-frequency-domain spectrograms. This discards the rich spatial and structural information in time-frequency representations and overlooks a fundamental intuition in music production. In particular, music is naturally represented as time-frequency grids in MIDI-based workflows, a structure that tightly corresponds to 2D spectrograms and inherently makes many MIR tasks trivial. Motivated by this intuition, we propose PupuJEPA, a visual Joint-Embedding Predictive Architecture (JEPA) that is trained directly on 2D spectrograms. Instead of applying masked language modeling (MLM) to 1D sequences, PupuJEPA learns robust representations by predicting the latent embeddings of masked 2D spectrogram patches from unmasked contexts. To optimally adapt such a visual framework to music signals, we also apply domain-specific modifications to model architecture, training scheme, and inference paradigm, with comprehensive ablation studies showing their effectiveness. Evaluations on the MARBLE benchmark show that PupuJEPA outperforms the 1D sequence-based SSL models across multiple MIR tasks in linear probing. Additionally, case studies of the attention maps also confirm that PupuJEPA captures musically meaningful patterns within the 2D time-frequency domain. Codes and checkpoints are available at: https://www.yichenggu.com/PupuJEPA/.
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 presents PupuJEPA, a 2D spectrogram-based JEPA model for music representation learning that achieves state-of-the-art results on the MARBLE benchmark by introducing domain-specific architectural modifications and inference strategies. The work is a solid contribution to audio SSL, effectively bridging the gap between visual JEPA successes and music information retrieval needs, although the novelty is somewhat incremental given the existing landscape of 2D audio models.
Methodology
The paper proposes PupuJEPA, a Joint-Embedding Predictive Architecture (JEPA) adapted for music information retrieval (MIR) by operating directly on 2D Mel-spectrograms. The core methodological contribution lies in adapting the visual JEPA framework to the audio domain through specific architectural and training modifications. Key innovations include: 1) Using an asymmetric patch size (4x16) to maintain high temporal resolution suitable for MIR tasks; 2) Implementing a restricted target encoder that only processes masked patches to prevent shortcut learning, diverging from standard JEPA implementations; 3) Introducing domain-specific masking strategies (blockwise and time-frequency masking) alongside random masking, with a curriculum-based scheduling mechanism; 4) Proposing novel inference paradigms for 2D models, including weighted layer fusion and structure-aware patch aggregation (Time-, Frequency-, and Block-Partitioned) to replace standard Global Average Pooling (GAP). The authors also identify that standard ViT components like DropPath and LayerScale cause representation collapse in this specific audio-SSL context, recommending their removal.
Results & evaluation
The evaluation is conducted on the MARBLE benchmark, covering a wide range of MIR tasks including emotion recognition, key detection, genre classification, beat tracking, structure analysis, and music tagging. PupuJEPA-Large achieves state-of-the-art (SOTA) or near-SOTA performance across most tasks compared to 1D sequence-based models (MERT, MusicFM, MuQ) and 2D audio models (AudioMAE++, A-JEPA). The ablation studies are comprehensive, validating the necessity of SwiGLU, QK-Norm, the smoothed L1 loss, and the specific masking/inference strategies. The paper demonstrates that 2D modeling preserves structural information beneficial for both global and local tasks. However, the performance gain over strong baselines like A-JEPA (which also uses 2D spectrograms) is modest in some metrics, suggesting that the specific JEPA adaptation provides incremental rather than revolutionary gains over existing 2D audio SSL approaches.
Reproducibility
The paper provides detailed implementation details, including hyperparameters, dataset preprocessing (24kHz mono, 10.24s crops), and model configurations. The code and checkpoints are made publicly available. The training setup (500k steps, 32 B200 GPUs) is clearly described. The reproducibility is high, although the reliance on a large in-house dataset (100k hours) for pre-training might limit independent verification of the pre-training phase, though the downstream evaluation is on standard benchmarks.
Limitations
The paper notes that scaling beyond the Large variant yields diminishing returns, likely due to the limitations of linear probing on highly complex representations. The performance on HookTheory structure analysis is only on par with baselines, indicating that 2D pooling strategies may still struggle with fine-grained local temporal dependencies compared to 1D sequence models in some contexts. The claim that music is "naturally" represented as 2D grids is a strong intuition but may not hold for all musical styles or production techniques, potentially limiting generalizability. Additionally, the comparison with some baselines involves retraining them on the authors' in-house dataset, which introduces a potential bias if the dataset distribution differs from the original training data of the baselines.
Broader impact
This work advances the field of self-supervised learning for audio by demonstrating the efficacy of 2D visual architectures (JEPA) for music processing. It challenges the dominance of 1D sequence models in MIR and provides a robust framework for learning rich musical representations. The findings could influence future model architectures for audio understanding, potentially leading to more efficient and accurate MIR systems. The open-source release contributes to the community by providing a strong baseline and codebase for future research in audio SSL. The paper presents PupuJEPA, a 2D spectrogram-based JEPA model for music representation learning that achieves state-of-the-art results on the MARBLE benchmark by introducing domain-specific architectural modifications and inference strategies. The work is a solid contribution to audio SSL, effectively bridging the gap between visual JEPA successes and music information retrieval needs, although the novelty is somewhat incremental given the existing landscape of 2D audio models.
Author of this paper? Send a correction.