ROSETok

Robust Speech Tokenizer · NeurIPS 2026

Paper · Code · Audio demo

Speech Tokenizers are Vulnerable
Transferable Semantic Attack and Robust Tokenizer

ROSETok combines noisy training and semantic distillation to preserve speech reconstruction and semantic representations under noise. The model uses a DAC-based encoder and decoder with residual vector quantization.

Released checkpoints

Checkpoint	File	Description
ROSETok-Base	rosetok_base.pth	Pretrained on English and Chinese speech with Gaussian-noise augmentation and semantic distillation.
ROSETok-AAF	rosetok_aaf.pth	Attack-aware fine-tuning of ROSETok-Base using T-SemAttack examples.

Both files contain PyTorch checkpoint dictionaries with metadata and state_dict. They include the semantic teacher used during training, giving each checkpoint a size of approximately 2.86 GB. Inference uses the encoder.*, quantizer.*, and decoder.* parameters, as shown below.

Model details

Property	Value
Input	Mono speech at 16 kHz
Encoder strides	2, 4, 5, 8
Hop size	320 samples, equivalent to 50 frames per second
Encoder feature dimension	1280
Residual codebooks	9
Entries per codebook	1024
Codebook dimension	8
Training languages	English and Chinese
Semantic teacher	Whisper large-v3

ROSETok-Base is trained on approximately 96.7k hours of English and Chinese speech from Emilia. ROSETok-AAF uses 2,000 LibriSpeech utterances perturbed with T-SemAttack for the second training stage.

Download

hf download wxzyd123/ROSETok rosetok_base.pth rosetok_aaf.pth --local-dir checkpoints/rosetok

Inference

Install the model implementation and its dependencies.

git clone https://github.com/wxzyd123/T-SemAttack.git
cd T-SemAttack
conda create -n rosetok python=3.10 -y
conda activate rosetok
python -m pip install -r requirements.txt

Run this example from the cloned repository root.

import torch
import torchaudio
import soundfile as sf
import yaml
from huggingface_hub import hf_hub_download

from rosetok import ROSETok

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

with open("config/rosetok_16khz.yml", encoding="utf-8") as f:
    config = yaml.safe_load(f)
model_config = {
    key.removeprefix("ROSETok."): value
    for key, value in config.items()
    if key.startswith("ROSETok.")
}
model = ROSETok(**model_config)

checkpoint_path = hf_hub_download(
    repo_id="wxzyd123/ROSETok",
    filename="rosetok_base.pth",  # Use rosetok_aaf.pth for ROSETok-AAF.
)
checkpoint = torch.load(
    checkpoint_path, map_location="cpu", weights_only=True, mmap=True
)
# The released checkpoints also contain the training-time semantic teacher.
codec_state = {
    key: value
    for key, value in checkpoint["state_dict"].items()
    if key.startswith(("encoder.", "quantizer.", "decoder."))
}
model.load_state_dict(codec_state, strict=True)
model = model.to(device).eval()

audio, sample_rate = sf.read(
    "examples/0_train_117.wav", dtype="float32", always_2d=True
)
wave = torch.from_numpy(audio.mean(axis=1)).unsqueeze(0)
wave = torchaudio.functional.resample(wave, sample_rate, 16000).to(device)

with torch.inference_mode():
    result = model(wave, sample_rate=16000)

sf.write("reconstructed.wav", result["audio"][0, 0].cpu().numpy(), 16000)
print("Discrete codes", result["codes"].shape)
print("Semantic features", result["encoded"].shape)

The example selects the inference parameters and loads them with strict=True. Switch the checkpoint filename to rosetok_aaf.pth to use the attack-aware variant.

The output dictionary includes reconstructed audio, pre-quantization features in encoded, quantized features in z, discrete codes, projected latents, and the commitment and codebook losses.

The decoder's stride-5 padding can make reconstruction up to 8 samples shorter than the input (0.5 ms at 16 kHz). For sample-aligned comparisons, trim the reference to the returned audio length.

Intended use and limitations

ROSETok supports speech reconstruction, discrete speech tokenization, and research on robust speech representations. Its reported robustness is evaluated under the noise levels, attacks, datasets, and downstream models described in the paper. Performance on other languages, acoustic conditions, and attack families requires separate evaluation.

Use the model for authorized research and evaluation. The checkpoints are released under the MIT license.