Frequency-Aware Self-Supervised Music Representation Learning

This is the official implementation of the paper "Frequency-Aware Self-Supervised Music Representation Learning". The main intuition of this work can be simply explained by the figure below:

In particular, the top panel displays a multitrack project in a Digital Audio Workstation (DAW), while the bottom panel shows its corresponding 2D spectrogram. As illustrated on the right, an experienced producer, Pupu (``bunny'' in Finnish), can intuitively and explicitly map individual tracks to their distinct time-frequency patterns in the spectrogram, as shown by the colored bounding boxes. Based on this information, music information retrieval (MIR) tasks can be easily performed by music producers with minimal observation of spectrograms.

Therefore, we adapted the SOTA Masked Modeling Duo (M2D) Architecture in Audio Self-Supervised Learning (SSL) to the music domain and thus propose PupuM2D. We release training code and high-quality pre-trained checkpoints in the repo to facilitate Music SSL research:

This script provides an example of how to set up the environment and train your own models with custom datasets.

NOTE: You need to run every command of this recipe in the PupuM2D root path:

cd PupuM2D

📀 Installation

git clone https://github.com/sizigi/PupuM2D.git
cd PupuM2D

# Install Python Environment
conda create --name pupum2d python=3.12
conda activate pupum2d

# Install Python Packages Dependencies
pip install -r requirements.txt

🐍 Usage in Python

1. Data Preparation

PupuM2D trains directly from local .flac files. The current dataloader scans one directory level under the dataset root:

{DATA_ROOT}/
  subset_a/
    track_000001.flac
    track_000002.flac
  subset_b/
    track_000003.flac

The statistics and training scripts currently contain dataset placeholders. First, set root_dir in calc_stat.py to your dataset root:

root_dir = "{DATA_ROOT}"

Run calc_stat.py before training to compute the dataset normalization statistics:

python calc_stat.py --config exp_config_pupum2d_base.json

The script prints two numbers: the mean and standard deviation. Replace the normalization constants in train_pupum2d.py with those values:

mean = {MEAN_FROM_CALC_STAT}
std = {STD_FROM_CALC_STAT}

If you computed variance separately, use its square root for std.

Then set the same dataset root in train_pupum2d.py:

train_dataset = StreamingAudioWebDataset(
    "{DATA_ROOT}", cfg
)

2. Training

Choose one of the provided model-size configs:

exp_config_pupum2d_tiny.json
exp_config_pupum2d_small.json
exp_config_pupum2d_base.json
exp_config_pupum2d_large.json
exp_config_pupum2d_huge.json
exp_config_pupum2d_giant.json

Before launching, set log_dir in the config to the directory where checkpoints and TensorBoard logs should be written:

{
  "log_dir": "{LOR_DIR}"
}

Then start training:

accelerate launch train_pupum2d.py \
  --config exp_config_pupum2d_base.json \
  --exp_name pupum2d_base

TensorBoard logs and checkpoints are saved under:

{log_dir}/{exp_name}/
  args.json
  log/
  checkpoint/

For multi-GPU training, use Accelerate in the usual way, for example:

accelerate launch --multi_gpu train_pupum2d.py \
  --config exp_config_pupum2d_base.json \
  --exp_name pupum2d_base

To resume training from an Accelerate checkpoint directory:

accelerate launch train_pupum2d.py \
  --config exp_config_pupum2d_base.json \
  --exp_name pupum2d_base \
  --resume_type resume \
  --checkpoint experiments/pupum2d_base/checkpoint/step-0000000_loss-0.000000

NOTE: The default configs target large GPUs. If you run out of memory, reduce train.batch_size, increase train.gradient_accumulation_step, or use a smaller model config.

🙏 Acknowledgement

• M2D for codebase reference.
• MATPAC++ for hyperparameter reference.
• AudioMAE++ for hyperparameter reference.

©️ License

Our project is under the MIT License. It is free for both research and commercial use cases.

📚 Citation

@article{pupum2d,
  title        = {Frequency-Aware Self-Supervised Music Representation Learning},
  author       = {Yicheng Gu and Junan Zhang and Jerry Li and Zhizheng Wu and Lauri Juvela},
  year         = {2026},
  journal      = {arXiv:2606.25713},
}