wavjepa

Model Details

The WavJEPA framework comprises a waveform encoder, context encoder, target encoder and a predictor. WavJEPA’s objective is to predict latent representation of various targets blocks based on a single context block extracted from the same sound wave. As waveform encoder, we use the feature encoder of Wav2Vec 2.0, which is composed of stacked temporal convolution layers (Baevski et al., 2020). Similar to the original I-JEPA architecture (Assran et al., 2023), a Vision Transformer (ViT) (Dosovitskiy et al., 2021) is used for the target encoder, context encoder and predictor.

Model Description

WavJEPA is the first framework applying semantic learning to general-purpose audio representations in the time domain, surpassing state-of-the-art time-domain approaches on the HEAR (Turian et al., 2022) and ARCH (La Quatra et al., 2024) benchmark suites while requiring only a fraction of the computational resources. WavJEPA leverages high-level semantic representation learning to tackle the shortcomings of representation learning at the speech unit or token level. We show that this approach substantially outperforms state-of-the-art time-domain audio foundation models across a wide variety of downstream benchmark tasks, while requiring considerably fewer computational resources. Additionally, we address the degraded performance of time-domain models in real-world sound scenes with WavJEPA-Nat, a multi-channel extension of the WavJEPA framework trained on simulated real-world sound scenes. Evaluation on Nat-HEAR (Yuksel et al., 2025), a naturalistic version of the HEAR benchmark suite, demonstrates that WavJEPA-Nat exceeds the robustness of other time-domain foundation models to noise and reverberation.

• Developed by: Goksenin Yuksel, goksenin.yuksel@ru.nl
• Model type: Transformers, Audio Foundation Models, Raw Waveform Models
• Language(s) (NLP): WavJEPA and WavJEPA-Nat support all languages, but mainly English.
• License: MIT

Model Sources

• Repository: https://github.com/labhamlet/wavjepa
• Paper: https://arxiv.org/abs/2509.23238

Uses

WavJEPA can be used as a powerful feature extractor for downstream tasks such as enviromental sound classification, speech recognition, speaker counting etc. Later, training a linear head on top of these extracted features would yield a fine-tuned audio scene analysis model.

How to Get Started with the Model

from transformers import AutoModel, AutoFeatureExtractor

model = AutoModel.from_pretrained("labhamlet/wavjepa-base", trust_remote_code=True)
extractor = AutoFeatureExtractor.from_pretrained("labhamlet/wavjepa-base", trust_remote_code=True)

audio = torch.zeros([1,160000])
extracted = extractor(audio, return_tensors="pt")
audio_feature = extracted['input_values']
print(model(audio_feature).shape)

Training Details

Training Data

We train WavJEPA on the unbalanced training set of AudioSet, which consists of 1.74 million 10-second sound clips scraped from YouTube (Gemmeke et al., 2017).

Training Procedure

Each sound clip was resampled to 16 kHz and mean centered to enforce equal loudness across sound clips. We then randomly sampled 8 sections of 2 s from each sound clip, effectively increasing the batch size by a factor of 8 in a computationally efficient manner. Finally, each instance is instance normalized (Ulyanov et al., 2017). The waveform encoder converts each 2 s instance into an embedding w 200×768, effectively resampling the audio to 100 Hz with a stride of 10 ms and a receptive field size of 12.5 ms

We sampled starting indices for the context block with p = 0.065 and for target blocks with p = 0.025. We set M to 10 for both context block and target block . To update the target encoder parameters ∆, we linearly increased τ from τ0 = 0.999 to τe = 0.99999 over the first 100,000 steps, after which τ was kept constant. We used K = 8 for the top K averaging. We trained WavJEPA for 375,000 steps using a batch size of 32 on two NVIDIA H100 94 GB GPUs. Given our in-batch sampling factor of 8, we boost our effective batch size to 256. We use the AdamW optimizer (Loshchilov & Hutter, 2019) with a weight decay coefficient λw = 0.04. The learning rate schedule follows a cosine decay with linear warm-up over 100,000 steps, reaching a peak learning rate of 2 × 10−4 before decaying to zero

Preprocessing

RMS Normalization was applied to audio clips to get all of them in the same loudness levels, and later instance normalization is applied.

Training Hyperparameters

• Training regime:: WavJEPA and WavJEPA-Nat were trained with mixed precision, torch.compile and flash attention.

Evaluation

We evaluate WavJEPA and other state-of-the-art models on the HEAR and ARCH benchmark task suite, which presents a wide range of tasks to evaluate the downstream performance of audio representation models (Turian et al., 2022).

Testing Data, Factors & Metrics

Testing Data

HEAR: The aim of the HEAR benchmark is to develop a general-purpose audio representation that provides a strong basis for learning in a wide variety of tasks and scenarios. HEAR evaluates audio representations using a benchmark suite across a variety of domains, including speech, environmental sound, and music. HEAR was launched as a NeurIPS 2021 shared challenge. It still remai