| --- |
| 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. |
|
|
