labhamlet/speech-jepa-base

Speech-JEPA Base

Code and recipes

The base model pretrained on 16 kHz sampled speech audio (LibriSpeech, 960 h). When using the model make sure that your speech input is also sampled at 16 kHz.

Note: This model does not have a tokenizer as it was pretrained on audio alone. In order to use this model for speech recognition, a tokenizer should be created and the model should be fine-tuned on labeled text data. Fine-tuned checkpoints for the Libri-light splits are listed below.

Model description

Speech-JEPA is a Joint-Embedding Predictive Architecture for speech. A context encoder sees a partially masked sequence of convolutional speech features and a light-weight convolutional decoder predicts the latent representations of the masked spans produced by an exponential-moving-average teacher (instance-normalised average of the top-8 teacher layers, as in data2vec 2.0). There is no reconstruction of audio and no quantised targets.

Parameters	89.7 M
Feature encoder	7-layer 1-D CNN (wav2vec 2.0 layout), 320× downsampling, 20 ms frames
Transformer	12 layers, 768-d, 12 heads, rotary position embeddings, post-norm
Pretraining data	LibriSpeech 960 h
Pretraining updates	370000

jepa_pretraining_extras.safetensors additionally holds the JEPA decoder, mask token and the EMA teacher encoder so that pre-training can be resumed with the original code.

Usage

import torch
from datasets import load_dataset
from transformers import AutoModel, Wav2Vec2FeatureExtractor

feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("labhamlet/speech-jepa-base")
model = AutoModel.from_pretrained("labhamlet/speech-jepa-base", trust_remote_code=True)

ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
inputs = feature_extractor(ds[0]["audio"]["array"], sampling_rate=16000, return_tensors="pt")

with torch.no_grad():
    hidden_states = model(**inputs).last_hidden_state   # (1, frames, 768), one frame every 20 ms

For batched input pass padding=True to the feature extractor; the returned attention_mask is used by the model to ignore padded samples. output_hidden_states=True returns the input embedding and the output of every Transformer layer.

The original training code can also load this repository:

from wavjepa.hub import load_pretrained      # github.com/labhamlet/SpeechJEPA
model = load_pretrained("labhamlet/speech-jepa-base")          # wavjepa.jepa_d2v2.JEPA, incl. decoder and EMA teacher

Fine-tuned checkpoints (CTC, LibriSpeech character vocabulary)

Labelled data	Repository
Libri-light 10 min	labhamlet/speech-jepa-base-asr-10min
Libri-light 1 h	labhamlet/speech-jepa-base-asr-1h
Libri-light 10 h	labhamlet/speech-jepa-base-asr-10h
LibriSpeech train-clean-100	labhamlet/speech-jepa-base-asr-100h

Files

• model.safetensors, config.json — SpeechJEPAModel weights and configuration
• modeling_speech_jepa.py, configuration_speech_jepa.py — model code (trust_remote_code=True)
• preprocessor_config.json — Wav2Vec2FeatureExtractor (16 kHz, per-utterance normalisation)
• jepa_pretraining_extras.safetensors — JEPA decoder, mask token and EMA teacher (original key names)
• training_config.json — hyper-parameters of the pretraining run

License

BSD-3-Clause.