File size: 3,593 Bytes
9358c80 bdd5c2a 9358c80 bdd5c2a 9358c80 bdd5c2a 9358c80 bdd5c2a 9358c80 bdd5c2a 9358c80 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | ---
license: mit
tags:
- audio-jepa
- jepa
- audio-representation-learning
- self-supervised
paper:
title: "Audio-JEPA: Joint-Embedding Predictive Architecture for Audio Representation Learning"
authors:
- Ludovic Tuncay
- Etienne Labbé
- Emmanouil Benetos
- Thomas Pellegrini
venue: "ICME 2025"
arxiv: "2507.02915"
hal: "hal-05128180"
---
# 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`](https://github.com/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.
```bash
# 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`](https://huggingface.co/ltuncay/BEST-RQ-2) (Tuncay, Labbé, Pellegrini, Interspeech 2026, [arXiv 2606.30700](https://arxiv.org/abs/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
```bibtex
@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.
|