sw2v_120k

Model Card for SW2V-120k

SW2V is a pure Transformer decoder-based streaming speech representation model, trained via distillation of W2V-BERT 2.0. It provides the self-supervised targets used by JHCodec's self-supervised representation reconstruction (SSRR) loss, with a zero-lookahead architecture suitable for real-time streaming.

• Paper: Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec
• GitHub Repository: https://github.com/jhcodec843/jhcodec
• Demo: https://jhcodec843.github.io/jhcodec/
• License: MIT

Model Details

This checkpoint corresponds to the paper's SW2V model at 120k steps (sw2v_120000.pt). To enhance noise robustness for future applications, we incorporated noise augmentation during its training. For the checkpoint used in the paper's main experiments, see jhcodec/sw2v_60k.

The model operates on 16 kHz mono audio in frames of 320 samples (20 ms, config.w2v.mlp_in.in_features), so the input length must be a multiple of 320. It returns one 1024-dim representation per frame.

Requirements

• Python >= 3.10
• PyTorch/TorchAudio with CUDA support (tested with torch==2.6.0+cu124 and torch==2.9.1+cu128)
• omegaconf==2.3.0
• Flash-Attention (required for the reported performance; tested with flash-attn==2.7.4.post1 and flash-attn==2.8.3)
• huggingface_hub — only if you auto-download the official checkpoint

Note: Running on CPU currently leads to degraded quality.

Usage

Offline (whole utterance at once)

import torch
import torch.nn.functional as F
import torchaudio
from jhcodec.utils import load_pretrained_sw2v

DEVICE = 'cuda'
SAMPLE_RATE = 16000
FRAME_SIZE = 320       # 20 ms hop; input length must be a multiple of this

w2v = load_pretrained_sw2v(repo_id='jhcodec/sw2v_120k').to(DEVICE).eval()

x, sr = torchaudio.load('input.wav')
if sr != SAMPLE_RATE:
    x = torchaudio.transforms.Resample(sr, SAMPLE_RATE)(x)
x = x[0, :].view(1, -1).to(DEVICE)               # [1, T], mono
if x.shape[1] % FRAME_SIZE != 0:
    x = F.pad(x, (0, FRAME_SIZE - x.shape[1] % FRAME_SIZE))

# encode is already decorated with @torch.no_grad()
features, _ = w2v.encode(x, inference_cache=None)  # [1, T//320, 1024]

Streaming (frame by frame)

Pass the returned inference_cache back in on every call, starting from None.

cache = None
feats = []
for i in range(0, x.shape[1], FRAME_SIZE):
    frame_feat, cache = w2v.encode(x[:, i:i + FRAME_SIZE], inference_cache=cache)
    feats.append(frame_feat)                      # each [1, 1, 1024]
features = torch.cat(feats, dim=1)                # [1, T//320, 1024]

Loading a local checkpoint

import omegaconf
import jhcodec.utils as utils
from jhcodec.model.sw2v import AudioEncoder

config = omegaconf.OmegaConf.load('config.json')
w2v = AudioEncoder(config.w2v, training=False)
utils.load_checkpoint(w2v, None, None, 'sw2v_120000.pt', strict_model=True)
w2v = w2v.to(DEVICE).eval()

For CUDA-graph per-frame streaming (AudioEncoderCudaGraph in jhcodec/model/sw2v_cudagraph.py, whose state_dict is identical to AudioEncoder), see the GitHub repository README.

Intended Use

• Streaming speech representation extraction for real-time systems
• Self-supervised targets for neural audio codec training (SSRR loss)
• Research into speech representation distillation
• Neural front-end for speech recognition or synthesis pipelines

Out-of-Scope Use

• Any malicious, deceptive, or privacy-violating applications

Training Details

Distilled from W2V-BERT 2.0; please refer to the GitHub repository README.

Citation

@article{jhcodec2026,
  title={Reconstruct! Don't Encode: Self-Supervised Representation Reconstruction Loss for High-Intelligibility and Low-Latency Streaming Neural Audio Codec},
  author={Anonymous},
  journal={arXiv preprint arXiv:2603.05887},
  year={2026}
}

Authors

Anonymous, Submitted to Interspeech 2026