Audio-JEPA-base

Audio-JEPA-base

Recommended: BEST-RQ-2.2-base has the strongest reported X-ARES results in this comparison.

Audio-JEPA-base is a self-supervised audio encoder trained on AudioSet for 200,000 optimizer steps. It produces 768-dimensional clip and frame embeddings from mono 16 kHz waveforms, and supports downstream fine-tuning.

Audio-JEPA learns by predicting the latent representations of masked spectrogram patches from an exponential-moving-average teacher. This export contains the student encoder, matching the research model's inference path; the pretraining predictor and teacher are omitted.

This is an updated training run, not an exact reproduction of the ICME 2025 model.

The author reports better X-ARES results for this newer checkpoint. The Audio-JEPA scores below are for this release and should not be presented as the original ICME results. The original checkpoint remains available in its own repository.

This repository contains the trained encoder, its preprocessing configuration, and the custom Transformers implementation. No installation of the research repository is needed.

X-ARES results

Scores on X-ARES (0–100, higher is better). Audio-JEPA, BEST-RQ (Conformer), BEST-RQ (ViT), and all BEST-RQ-2 variants reported below are trained on the same AudioSet split for 200,000 steps. The pretrained baselines are shown for comparison.

Global Mean averages all benchmark task scores. Mean of Means gives equal weight to the Speech, Music, and Environment category means.

Audio-JEPA scores were supplied by the author for run jp6l70l6. The remaining scores are reported in the project README. These are reported research results, not a new benchmark run of the Transformers exports.

The Audio-JEPA row refers to the newer 16 kHz, 200,000-step run, not the original ICME model (32 kHz, 100,000 steps). The author reports better results for this newer checkpoint.

Model and training

Training recipe

The matching experiment is configs/experiment/audio_jepa/audioset/default.yaml. The run used the earlier name audio_jepa/baseline. Its saved model settings, optimizer, EMA schedule, data selection, batch size, seed and 200,000-step budget match the current recipe. Newly explicit settings retain the defaults: white masking noise, compilation disabled, and automatic RoPE mode (unused with the saved sinusoidal positional embeddings). See recipe_match.json for the comparison.

In the research repository, reproduce the configuration with:

uv run src/train.py experiment=audio_jepa/audioset/default
uv run slurm_scripts/submit.py audio_jepa/audioset/default --gpus 1 --time 20:00:00 --dry-run

The Slurm command previews submission; remove --dry-run to submit on a configured cluster. Dataset paths and worker counts depend on your environment. The recipe matches the saved hyperparameters; exact numerical reproduction also depends on the training-code version and runtime.

Load the model

Install the runtime dependencies in your Python environment:

pip install "torch>=2.9.1" "torchaudio>=2.9.1" "timm>=0.9" "einops>=0.7" "transformers>=4.57,<6"

Use matching PyTorch and torchaudio versions. GPU installations may require the appropriate PyTorch build for your CUDA version.

import torch
from transformers import AutoFeatureExtractor, AutoModel

model_name = "ltuncay/Audio-JEPA-base"
processor = AutoFeatureExtractor.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True).eval()
audio = torch.zeros(processor.sampling_rate)  # Replace with real mono audio.
inputs = processor(audio, sampling_rate=processor.sampling_rate, return_tensors="pt")
with torch.inference_mode():
    outputs = model(**inputs)
clip_embeddings = outputs.pooler_output
frame_embeddings = outputs.last_hidden_state
output_dim = model.config.encoder_kwargs["embed_dim"]

Resample to processor.sampling_rate and downmix stereo before preprocessing. The extractor performs padding only; spectrogram/convolution features are computed inside the model. Pass the returned sample attention mask for variable durations. Outputs include a frame attention mask and timestamps in milliseconds (-1 for padding). pooler_output uses the saved HEAR extraction preset, including phase-balanced pooling. Frame features average frequency patches; they are not the raw frequency-time ViT grid.

Fine-tuning

For fine-tuning call model.train(), attach a task head and optimize its parameters alongside the model. Save with model.save_pretrained(path) and processor.save_pretrained(path). This encoder export excludes pretraining predictors, quantizers, teachers and optimizer state. Continue self-supervised research training with the original Lightning code and checkpoints.

Reproducibility and provenance

revision is optional; pin both loaders to the same full commit hash for reproducibility. Custom Python code is included in this repository and requires trust_remote_code=True. The weights use safetensors. See export_manifest.json for source and validation details. Source run: jp6l70l6 (access may require permission). The export used checkpoints/last.ckpt and the matching saved Hydra configuration. The Lightning checkpoint was converted to safetensors before encoder extraction. The manifest records the converted checkpoint and configuration hashes.

The exported model matched the pre-export encoder exactly on the exporter's test batch. Loading with AutoFeatureExtractor and AutoModel was also checked in a fresh Python process on unequal-duration inputs. These are integration checks, not a downstream benchmark evaluation of this release.

Intended uses and limitations

Use these embeddings as features for audio research, classification, retrieval, or downstream fine-tuning. This is an encoder: it does not generate transcripts, audio, or class labels without a downstream model.

AudioSet training does not establish performance on every language, acoustic domain, demographic group, or downstream task. Evaluate the model on your target data. No new downstream benchmark scores are claimed for this export.

The feature extractor pads waveforms and checks the sample rate; it does not resample or downmix. Unequal-length clips are processed individually after padding is removed. The output embeddings follow the saved windowing policy rather than exposing the raw spectrogram token grid.

Browse the BEST-RQ-2 family collection.

Source and citation

Research code: audio-embeddings. If you use Audio-JEPA, cite:

@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},
  year={2025}
}

License

The model weights and bundled implementation are released under the MIT license. See LICENSE and CODE_LICENSE.