Audio-JEPA: pretrained encoder

Weights for the encoder trained in Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning (Tuncay, Labbé, Benetos, Pellegrini, ICME 2025). Source code: LudovicTuncay/Audio-JEPA.

Files

File	Purpose
JEPA.ckpt	PyTorch Lightning checkpoint (state_dict + trainer metadata).
config.json	Standalone documentation of the architecture and audio pipeline. Machine-readable.
README.md	This document.

Model summary

• Encoder: VisionTransformer (ViT-Base, 12 layers, 768-dim, 12 heads).
• Input: log-mel spectrogram of shape (target_time_bins=256, n_mels=128) from a 10 s mono waveform at 32 kHz.
• Patchification: (16, 16), giving a grid 8 x 16 = 128 patches.
• Each patch spans ~625 ms of audio × 16 mel bins.
• Output: (128, 768) embeddings per 10 s clip.
• Effective temporal resolution: 1.6 positions/s (the 12.8 tokens/s counts 8 temporal × 16 frequency patches per second).

Minimal inference (CPU, no flash-attn install required)

The training repo depends on flash-attn, which requires CUDA to build. For inference-only use, flash_attn.modules.mha.MHA can be substituted with a torch-native equivalent that matches the checkpoint's parameter names (qkv, proj). See inference_example.py in this repository for a ~150-line standalone script.

# 1. Clone the source code (needed for the ViT class)
git clone --depth 1 https://github.com/LudovicTuncay/Audio-JEPA.git

# 2. Install a lean set of deps
pip install torch torchaudio numpy huggingface_hub

# 3. Run the example (downloads JEPA.ckpt on first run)
python inference_example.py --audio-jepa-src ./Audio-JEPA

The script prints the loading diagnostics (should be 0 missing, 0 unexpected) and the embedding shape.

Domain fit: where the model excels vs where it doesn't

Audio-JEPA is designed to learn a generic audio representation. The 16×16 patch shape is a compromise across speech, music, and environmental sounds. Consequences:

• Strong on: music, environmental sounds, audio captioning, general audio tagging.
• Weaker on: speech-only downstream tasks (see the paper's X-ARES tables). Speech-specific SSL models such as wav2vec 2.0, HuBERT and the Whisper encoder currently outperform Audio-JEPA on speech-centric benchmarks.
• A follow-up from the same author, BEST-RQ-2 (Tuncay, Labbé, Pellegrini, Interspeech 2026, arXiv 2606.30700), combines the encoder-predictor decomposition of Audio-JEPA with BEST-RQ discrete targets. It yields substantially better cross-domain results while sharing the exact same encoder inference speed.

Citation

@inproceedings{tuncay2025audio,
  title  = {Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning},
  author = {Tuncay, Ludovic and Labb{\'e}, Etienne and Benetos, Emmanouil and Pellegrini, Thomas},
  booktitle = {ICME 2025},
  address = {Nantes, France},
  year   = {2025},
  url    = {https://hal.science/hal-05128180}
}

License

MIT.