DinoSR (base, LibriSpeech 960h) — 🤗 Transformers port

This is a conversion of the official DinoSR fairseq checkpoint to 🤗 Transformers.

DinoSR: Self-Distillation and Online Clustering for Self-supervised Speech Representation Learning Alexander H. Liu, Heng-Jui Chang, Michael Auli, Wei-Ning Hsu, James R. Glass — NeurIPS 2023 paper · code

DinoSR combines masked language modelling, self-distillation and online clustering: an EMA teacher encodes unmasked audio, its top-8 layer outputs are quantised by 8 online k-means codebooks (256 clusters each), and the student is trained to predict those discrete units at masked positions. The result is a strong speech encoder whose codebooks give very good phone-discriminative discrete units.

Architecturally DinoSR is identical to data2vec-audio (7-layer conv feature extractor, 12-layer post-LN transformer, stacked convolutional positional embedding), so the port reuses Data2VecAudioModel internally.

What is in this repo

path	contents	how to load
. (root)	student + EMA teacher + codebooks + pre-training heads	AutoModel.from_pretrained(..., trust_remote_code=True)
student/	student encoder only, plain data2vec-audio format	Data2VecAudioModel.from_pretrained(..., subfolder="student")

Usage

Features (no remote code needed)

import torch, soundfile as sf
from transformers import AutoFeatureExtractor, Data2VecAudioModel

fe = AutoFeatureExtractor.from_pretrained("MohammadJRanjbar/DinoSR")
model = Data2VecAudioModel.from_pretrained("MohammadJRanjbar/DinoSR", subfolder="student").eval()

wav, sr = sf.read("sample.wav", dtype="float32")  # 16 kHz mono
inputs = fe(wav, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
    out = model(**inputs, output_hidden_states=True)

out.last_hidden_state          # (batch, frames, 768), 50 frames/second
out.hidden_states[7]           # any intermediate layer

Discrete units (teacher + online-clustering codebooks)

import torch, soundfile as sf
from transformers import AutoFeatureExtractor, AutoModel

fe = AutoFeatureExtractor.from_pretrained("MohammadJRanjbar/DinoSR")
model = AutoModel.from_pretrained("MohammadJRanjbar/DinoSR", trust_remote_code=True).eval()

wav, sr = sf.read("sample.wav", dtype="float32")
inputs = fe(wav, sampling_rate=16000, return_tensors="pt")

units = model.extract_units(**inputs).units   # (batch, frames, 8) cluster ids in [0, 256)
model.codebook_layers                         # (4, 5, 6, 7, 8, 9, 10, 11) — layer per codebook
model.codebooks.shape                         # (8, 256, 768) — the k-means centroids

model(**inputs).last_hidden_state             # `forward` runs the student encoder

Codebook k quantises the feed-forward output of teacher layer model.codebook_layers[k]. The upper-middle layers are usually the most phone-discriminative, so units[..., -2] (layer 10) is a reasonable starting point, but which codebook works best is task-dependent.

Note that the teacher targets are instance-normalised over the time axis, exactly as during pre-training, so units for an utterance depend on the whole utterance: batch with padding gives slightly different units than running each utterance alone. Use batch size 1 (or group by length) when you need reproducible units.

Fine-tuning

student/ is a drop-in data2vec-audio encoder, so the usual heads work:

from transformers import Data2VecAudioForCTC
model = Data2VecAudioForCTC.from_pretrained(
    "MohammadJRanjbar/DinoSR", subfolder="student", vocab_size=32, ctc_loss_reduction="mean",
)

The checkpoint is pre-trained only (no ASR fine-tuning). Inputs must be 16 kHz and zero-mean/unit-variance normalised — AutoFeatureExtractor does this for you.

Conversion & verification

Converted with convert_dinosr_to_hf.py (included in this repo, together with test_equivalence.py) from the official checkpoint (https://data.csail.mit.edu/placesaudio/dinosr/dinosr.ckpt, 400k updates on LibriSpeech-960). The port was verified against the original fairseq model on the same waveform, single and padded-batch:

check	max abs. difference
student last hidden state	3.1e-06
all 12 student layer outputs	≤ 3.1e-06
teacher targets (top-8 layers)	1.1e-05
pre-training head logits	1.1e-05
discrete units	100 % identical

Differences are float32 round-off. One caveat: the original pipeline normalises waveforms with F.layer_norm (variance epsilon 1e-5) whereas Wav2Vec2FeatureExtractor uses 1e-7, a ~0.5 % difference in input scale; feed torch.nn.functional.layer_norm(wav, wav.shape) yourself if you need bit-comparable inputs.

Citation

@inproceedings{liu2023dinosr,
  title     = {DinoSR: Self-Distillation and Online Clustering for Self-supervised Speech Representation Learning},
  author    = {Liu, Alexander H. and Chang, Heng-Jui and Auli, Michael and Hsu, Wei-Ning and Glass, James R.},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2023}
}

Weights are redistributed from the original release; please refer to the upstream repository for licensing terms.