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