feat(ssl): add ViP-VL self-supervised pretraining for ChunkFormer - #41
Merged
Conversation
…coder Adds a clean, isolated BEST-RQ (random-projection quantizer) SSL pretraining path for the ChunkFormer encoder, plus a finetune-ready encoder export tool. - chunkformer/ssl: BestRQ model + quantizer/loss/utils modules - init_model/init_dataset/train: register 'bestrq' model type (no tokenizer); behavior-preserving for asr_model/transducer/classification - encoder: optional LayerDrop (encoder_layerdrop, training-only, off by default) - dataset/processor: optional random feature cropping via crop_conf (no-op when unset) - executor.cv: generic scalar-metric accumulation (handles SSL python-float metrics) - train_utils: make deepspeed import optional; vocab_size only for ASR - examples/ssl/bestrq: clean recipe (conf from released checkpoint, sanitized data paths) - tools/export_ssl_encoder.py: export encoder-only checkpoint with sanitized config - tools/verify_bestrq_parity.py: deterministic forward-parity check Verified: forward output is bit-for-bit identical to the kl/add_wav2vec2_ssl branch on avg_50.pt; ASR/RNN-T/classification/BestRQ all still build; exported encoder loads with 0 missing encoder.* keys.
There was a problem hiding this comment.
Pull request overview
This PR introduces a BEST-RQ self-supervised learning (SSL) pretraining path for the ChunkFormer encoder, including a dedicated recipe, encoder-only export tooling for downstream finetuning, and minimal integration into existing training/dataset/model initialization flows.
Changes:
- Added BEST-RQ SSL model implementation and supporting quantizer/loss/util modules under
chunkformer/ssl/. - Added an end-to-end BEST-RQ recipe under
examples/ssl/bestrq/, plus tools for parity verification and encoder-only checkpoint export. - Integrated SSL mode into existing initialization/training/dataset utilities (tokenizer-less path, optional cropping, optional LayerDrop, generic CV metric aggregation, optional DeepSpeed import).
Reviewed changes
Copilot reviewed 23 out of 25 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/verify_bestrq_parity.py | Adds a deterministic forward-pass parity checker for BEST-RQ checkpoints. |
| tools/export_ssl_encoder.py | Adds a utility to export an encoder-only checkpoint bundle + sanitized config/README. |
| examples/ssl/bestrq/run.sh | Adds a staged BEST-RQ pretraining pipeline script (train/avg/export/upload). |
| examples/ssl/bestrq/README.md | Documents BEST-RQ recipe usage, data layout, export, and parity check. |
| examples/ssl/bestrq/path.sh | Sets recipe environment variables and paths for running BEST-RQ stages. |
| examples/ssl/bestrq/conf/bestrq.yaml | Adds a BEST-RQ training configuration (encoder/model/dataset settings). |
| chunkformer/utils/train_utils.py | Makes DeepSpeed optional and adjusts vocab sizing for tokenizer-less modes. |
| chunkformer/utils/mask.py | Adds compute_mask_indices used for BEST-RQ masking behavior. |
| chunkformer/utils/init_model.py | Registers bestrq model type and constructs the BestRQ wrapper model. |
| chunkformer/utils/init_dataset.py | Routes SSL/classification through the tokenizer-less dataset pipeline. |
| chunkformer/utils/executor.py | Generalizes CV scalar metric accumulation (supports SSL scalar metrics). |
| chunkformer/ssl/modules/utils.py | Adds small tensor helper utilities used by BEST-RQ modules. |
| chunkformer/ssl/modules/quantizer.py | Adds random-projection quantizers used by BEST-RQ pretraining. |
| chunkformer/ssl/modules/loss.py | Adds BEST-RQ MLM-style NLL loss wrapper. |
| chunkformer/ssl/modules/init.py | Adds/defines SSL modules package entry point. |
| chunkformer/ssl/bestrq/model.py | Implements the BEST-RQ wrapper model (masking, quantizer targets, loss/metrics). |
| chunkformer/ssl/bestrq/init.py | Exports BestRQ from the BEST-RQ package. |
| chunkformer/ssl/init.py | Adds package-level documentation for SSL support. |
| chunkformer/modules/subsampling.py | Adds stacking() helper used to build BEST-RQ quantizer targets. |
| chunkformer/modules/encoder.py | Adds optional LayerDrop support and cache-handling changes in forward_layers. |
| chunkformer/dataset/processor.py | Adds optional random feature cropping and guards deprecated torchaudio API usage. |
| chunkformer/dataset/dataset.py | Wires crop_conf into batching/padding and dynamic batch window sizing. |
| chunkformer/bin/train.py | Skips tokenizer initialization for bestrq model type. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+558
to
+563
| if crop_conf is not None: | ||
| crop_length = crop_conf.get("crop_length", "min") | ||
| if crop_length == "min": | ||
| crop_length = min([x["feat"].size(0) for x in sample]) | ||
| crop_length = min(crop_length, crop_conf.get("max_crop_length")) | ||
| sample = [crop(s, crop_length=crop_length) for s in sample] |
Comment on lines
+631
to
+641
| if self.crop_conf is not None: | ||
| # Account for cropping when estimating the padded batch size so the | ||
| # frame budget reflects the post-crop sequence length. | ||
| self.shortest_frames = min(self.shortest_frames, new_sample_frames) | ||
| self.longest_frames = max(self.longest_frames, new_sample_frames) | ||
| if self.crop_length == "min": | ||
| shortest_frames = min(self.shortest_frames, self.max_crop_length) | ||
| frames_after_padding = shortest_frames * (buffer_size + 1) | ||
| else: | ||
| longest_frames = min(self.longest_frames, self.crop_length) | ||
| frames_after_padding = longest_frames * (buffer_size + 1) |
Comment on lines
+294
to
+303
| # LayerDrop (https://arxiv.org/abs/1909.11556): randomly skip whole | ||
| # encoder layers during training. Disabled when encoder_layerdrop=0 | ||
| # (default), so inference/finetuning behaviour is unchanged. | ||
| if self.training and self.encoder_layerdrop > 0.0: | ||
| if random.random() < self.encoder_layerdrop: | ||
| new_att_cache = att_cache[idx] if att_cache.size(0) > 0 else torch.zeros((0, 0, 0, 0)) | ||
| new_cnn_cache = cnn_cache[idx] if cnn_cache.size(0) > 0 else torch.zeros((0, 0, 0)) | ||
| r_att_cache.append(new_att_cache) | ||
| r_cnn_cache.append(new_cnn_cache) | ||
| continue |
Comment on lines
+304
to
+320
| xs, _, new_att_cache, new_cnn_cache = layer( | ||
| xs, | ||
| masks, | ||
| pos_emb, | ||
| mask_pad, | ||
| att_cache=att_cache[idx] if att_cache.size(0) > 0 else att_cache, | ||
| cnn_cache=cnn_cache[idx] if cnn_cache.size(0) > 0 else cnn_cache, | ||
| att_cache=att_cache[idx] if att_cache.size(0) > 0 else torch.zeros((0, 0, 0, 0)), | ||
| cnn_cache=cnn_cache[idx] if cnn_cache.size(0) > 0 else torch.zeros((0, 0, 0)), | ||
| chunk_size=chunk_size, | ||
| left_context_size=left_context_size, | ||
| right_context_size=right_context_size, | ||
| ) | ||
| # During training/SSL the caches are empty; force empty placeholders | ||
| # so torch.stack works even when some layers are dropped. | ||
| if att_cache.size(0) == 0: | ||
| new_att_cache = torch.zeros((0, 0, 0, 0)) | ||
| if cnn_cache.size(0) == 0: | ||
| new_cnn_cache = torch.zeros((0, 0, 0)) |
Comment on lines
+132
to
+137
| if self.dist_fn == "cosine": | ||
| # (B, T, num_books, code_dim) -> (B, T, num_books, num_classes) | ||
| xid = torch.einsum('btdh,dch->btdc', x, self.codebooks) | ||
| # (B, T, num_books, num_classes) -> (B, T, num_books) | ||
| xid = xid.max(dim=-1)[1] | ||
| elif self.dist_fn == "l2": |
Comment on lines
+537
to
+546
| for i, mask_idc in enumerate(mask_idcs): | ||
| if target_len is not None and len(mask_idc) > target_len: | ||
| mask_idc = rng.choice(mask_idc, target_len, replace=False) | ||
|
|
||
| mask[i, mask_idc] = True | ||
|
|
||
| if target_len is not None and len(mask_idc) < target_len: | ||
| unmasked = np.flatnonzero(~mask[i]) | ||
| to_mask = rng.choice(unmasked, target_len - len(mask_idc), replace=False) | ||
| mask[i, to_mask] = True |
Comment on lines
+11
to
+22
| def __init__( | ||
| self, | ||
| mask_threshold: float = 0.8, | ||
| ): | ||
| super().__init__() | ||
| self.nll_loss = torch.nn.NLLLoss() | ||
| self.mask_threshold = mask_threshold | ||
|
|
||
| def __call__(self, logits, targets): | ||
| loss = self.nll_loss(logits, targets) | ||
| loss = torch.mean(loss) | ||
| return loss |
Comment on lines
+170
to
+173
| # 2. sanitised config | ||
| with open(args.config, "r") as fin: | ||
| config = yaml.load(fin, Loader=yaml.FullLoader) | ||
| clean_config = sanitize_config(config) |
Comment on lines
+41
to
+44
| def build_model(config_path: str): | ||
| with open(config_path, "r") as fin: | ||
| configs = yaml.load(fin, Loader=yaml.FullLoader) | ||
| # Skip CMVN so we do not depend on an external global_cmvn file. The CMVN |
Comment on lines
226
to
229
| # DeepSpeed automaticly add '--deepspeed' and '--deepspeed_config' to parser | ||
| parser = deepspeed.add_config_arguments(parser) | ||
| if DEEPSPEED_AVAILABLE: | ||
| parser = deepspeed.add_config_arguments(parser) | ||
| return parser |
- processor: guard crop_conf 'max_crop_length' None in padding() and DynamicBatchWindow - encoder: device-aware empty cache placeholders for LayerDrop (avoid CPU/GPU mismatch) - quantizer: remove cosine distance (l2-only, the one we use); fixes the broken cosine branch - loss: remove unused mask_threshold parameter - train_utils: clear RuntimeError when train_engine=deepspeed but DeepSpeed unavailable - tools: use yaml.safe_load instead of yaml.load(FullLoader) - lint: black/isort/flake8 clean (wrap long lines, drop placeholder-less f-strings) Parity with avg_50.pt preserved bit-for-bit; ASR/classification/BestRQ still build.
feature_extractor returns 4 tensors (xs, pos_emb, masks, xs_norm) but was annotated as a 2-tuple, causing a return-value error plus a cascade of unpack / 'cannot determine type' errors in BestRQ.forward under mypy.
Comment on lines
+532
to
+542
| target_len = None | ||
| if require_same_masks: | ||
| if add_masks: | ||
| target_len = max([len(m) for m in mask_idcs]) | ||
| else: | ||
| target_len = min([len(m) for m in mask_idcs]) | ||
|
|
||
| for i, mask_idc in enumerate(mask_idcs): | ||
| if target_len is not None and len(mask_idc) > target_len: | ||
| mask_idc = rng.choice(mask_idc, target_len, replace=False) | ||
|
|
Comment on lines
+441
to
+443
| if padding_mask is not None: | ||
| sz = all_sz - padding_mask[i].long().sum().item() | ||
| assert sz >= 0, sz |
| uniform = sample from uniform distribution [mask_other, mask_length*2] | ||
| normal = sample from normal distribution with mean mask_length and stdev | ||
| mask_other. mask is min 1 element | ||
| poisson = sample from possion distribution with lambda = mask length |
Comment on lines
+101
to
+105
| # (num_books, num_classes, hid_dim) | ||
| codebooks = torch.randn(self.num_books, self.num_classes, self.code_dim).double() | ||
| torch.nn.init.normal_(codebooks, mean=0, std=1) | ||
| codebooks = F.normalize(codebooks, dim=-1) | ||
| self.codebooks = nn.Parameter(codebooks) |
Comment on lines
+33
to
+36
| # NOTE: sox_utils.set_buffer_size is deprecated and was removed in newer | ||
| # torchaudio releases. Guard it so import does not fail across versions. | ||
| if hasattr(torchaudio.utils, "sox_utils"): | ||
| torchaudio.utils.sox_utils.set_buffer_size(16500) |
Comment on lines
+242
to
+245
| logit = x[mask_indices] | ||
| logit = self.final_proj(logit) | ||
| logit = logit.reshape(-1, self.quantizer.num_books, self.quantizer.num_classes) | ||
| logit = logit.log_softmax(-1) |
Comment on lines
+1
to
+5
| # Copyright (c) Facebook, Inc. and its affiliates. | ||
| # | ||
| # This source code is licensed under the MIT license found in the | ||
| # LICENSE file in the root directory of this source tree. | ||
|
|
Comment on lines
+87
to
+97
| readme = f"""--- | ||
| tags: | ||
| - speech | ||
| - self-supervised-learning | ||
| - best-rq | ||
| - chunkformer | ||
| - pretrained-encoder | ||
| - pytorch | ||
| license: apache-2.0 | ||
| library_name: transformers | ||
| --- |
khanld
changed the title
Rename the public-facing model/method/code from BEST-RQ to ViP-VL: - class BestRQ -> ViPVL; package chunkformer/ssl/bestrq -> chunkformer/ssl/vipvl - model type 'bestrq' -> 'vipvl' (kept as a backward-compatible alias so existing checkpoints/configs still load) - recipe examples/ssl/bestrq -> examples/ssl/vipvl; conf/bestrq.yaml -> conf/vipvl.yaml - tools/verify_bestrq_parity.py -> tools/verify_vipvl_parity.py - docstrings/README rebranded to ViP-VL; 'BEST-RQ' kept only as a citation of the underlying random-projection-quantizer technique (arXiv:2202.01855) Parity with avg_50.pt unchanged (verified via the bestrq alias); vipvl and bestrq both build; black/isort/flake8 clean.
khanld
changed the title
Untrack the dev parity script (git rm --cached) and gitignore it so it stays local. Remove the recipe README's reproducibility section that referenced it.
Comment on lines
+6
to
+15
| from typing import Tuple | ||
|
|
||
| import torch | ||
| import torch.nn as nn | ||
|
|
||
| from chunkformer.ssl.modules.loss import MLMLoss | ||
| from chunkformer.ssl.modules.quantizer import RandomProjectionVectorQuantizer | ||
| from chunkformer.ssl.modules.utils import index_put | ||
| from chunkformer.utils.mask import compute_mask_indices, make_pad_mask | ||
|
|
Comment on lines
+539
to
+554
| for i, mask_idc in enumerate(mask_idcs): | ||
| if target_len is not None and len(mask_idc) > target_len: | ||
| mask_idc = rng.choice(mask_idc, target_len, replace=False) | ||
|
|
||
| mask[i, mask_idc] = True | ||
|
|
||
| if target_len is not None and len(mask_idc) < target_len: | ||
| unmasked = np.flatnonzero(~mask[i]) | ||
| to_mask = rng.choice(unmasked, target_len - len(mask_idc), replace=False) | ||
| mask[i, to_mask] = True | ||
|
|
||
| if mask_dropout > 0: | ||
| masked = np.flatnonzero(mask[i]) | ||
| num_holes = np.rint(len(masked) * mask_dropout).astype(int) | ||
| to_drop = rng.choice(masked, num_holes, replace=False) | ||
| mask[i, to_drop] = False |
Comment on lines
+101
to
+105
| # (num_books, num_classes, hid_dim) | ||
| codebooks = torch.randn(self.num_books, self.num_classes, self.code_dim).double() | ||
| torch.nn.init.normal_(codebooks, mean=0, std=1) | ||
| codebooks = F.normalize(codebooks, dim=-1) | ||
| self.codebooks = nn.Parameter(codebooks) |
|
|
||
| loss = self.criterion(logit, y) | ||
|
|
||
| mask_percentile = (mask_indices.sum(-1).float() / x_mask.float().sum(-1)).mean().item() |
Comment on lines
+218
to
+224
| unmasked_features = self.stacking(unmask_xs) | ||
| unmasked_features = unmasked_features.reshape( | ||
| unmasked_features.size(0), unmasked_features.size(1), -1 | ||
| ) | ||
| # B, T, L | ||
| mask_indices = self.stacking(mask_indices) | ||
| # B, T |
Comment on lines
+33
to
+36
| # NOTE: sox_utils.set_buffer_size is deprecated and was removed in newer | ||
| # torchaudio releases. Guard it so import does not fail across versions. | ||
| if hasattr(torchaudio.utils, "sox_utils"): | ||
| torchaudio.utils.sox_utils.set_buffer_size(16500) |
Checkpoint loading configs were migrated to 'model: vipvl', so the backward- compatible 'bestrq' alias is no longer needed. init_model/train now recognise only 'vipvl' for the SSL path. Verified: the released avg_50.pt loads under model: vipvl with identical parity metrics; weights file byte-identical.
- push_model_hf.py now detects the 'vipvl' (and legacy 'bestrq') model type and emits a paper-based ViP-VL self-supervised encoder model card instead of mislabelling it as an ASR model. - create_repository uses create_repo(exist_ok=True) instead of a repo_info() existence probe, which reported private repos as missing and broke uploads.
- Title now spells out the ViP-VL acronym in bold. - States INTERSPEECH 2026 acceptance (badge + intro + citation). - Intro highlights the pretrained ViP-VL model; drops the encoder-only warning. - Architecture table adds the 'Encoder: ChunkFormer' row.
load_checkpoint and load_trained_modules now resolve the path via resolve_checkpoint_path: local file -> local dir/pytorch_model.pt -> download pytorch_model.pt from the Hugging Face Hub. Lets users pass checkpoint=khanhld/vip-vl-base-vie directly. Fully backward compatible with existing local paths. Docs updated (recipe README, model card, export tool).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements the self-supervised pretraining method of ViP-VL (ViP-VL: Vietnamese
Self-supervised Speech Pretraining Model with Vector-Quantization Learning, submitted
to Interspeech 2026) on top of the ChunkFormer encoder.
ViP-VL extends the random-projection-quantizer technique (BEST-RQ, arXiv:2202.01855)
to an aggressive 8× subsampling ChunkFormer backbone and fixes the synchronization
between the masking manifold and the encoder's subsampling rate. This PR ports only the
code needed for that pretraining path — no unrelated experiment cruft — and keeps the
existing ASR / RNN-T / classification training fully intact.
Paper method → code
chunkformer/ssl/vipvl/model.py(
ViPVL) wraps the existingChunkFormerEncoder. Config matches the paper's 78M setup:12 blocks, 512 output size, 8 heads, 2048 FFN, CNN kernel 15, 8× temporal subsampling
(
examples/ssl/vipvl/conf/vipvl.yaml).Subsampling.stacking()stackseach output frame's full receptive field (
size = reverse_calc_length(1),step = subsampling_rate),i.e. the 15-frame window / stride-8 alignment derived from the 3-stacked conv stem
(kernel 3, stride 2). Concatenated stacking feeds the random-projection quantizer targets.
before subsampling; a subsampled frame is labelled "masked" iff ≥ 80 % of its
constituent pre-subsampled frames are masked (
mask_indices >= 0.8inmodel.py).chunkformer/ssl/modules/quantizer.py:frozen Xavier-initialized projection, codebook
c_i ~ N(0, I), L2 nearest-neighbour onL2-normalized vectors; input mean-variance normalization via CMVN to prevent codebook
collapse. MLM negative-log-likelihood loss over masked positions (
modules/loss.py).What's added
chunkformer/ssl/vipvl/—ViPVLmodel + sharedquantizer/loss/utilsmodules.examples/ssl/vipvl/— clean recipe:conf/vipvl.yaml,run.sh,path.sh,README.md, andtools/chunkformersymlinks.tools/export_ssl_encoder.py— exports an encoder-only checkpoint (drops SSL heads)with a sanitized
config.yamlfor downstream finetuning / HF upload.Integration (behavior-preserving)
init_model/init_dataset/train.py: register thevipvlmodel type (no tokenizer).ASR / RNN-T / classification paths are unchanged.
subsampling: adds thestacking()helper used to build quantizer targets.encoder: optional LayerDrop (encoder_layerdrop), training-only, off by default.dataset/processor: optional random feature cropping viacrop_conf(no-op when unset).executor.cv: generic scalar-metric accumulation (handles SSL python-float metrics likecorr).train_utils:deepspeedimport made optional;vocab_sizeset only for ASR.Verification
ASRModel,Transducer,SpeechClassificationModel, andViPVLall build.avg_50.ptloads undermodel: vipvlwith identicalforward-parity metrics (weights byte-identical; only the config's
model:field was migrated).encoder.*keys.pre-commit(black / isort / flake8 / mypy) passes.