Overview of the FuseCodec speech tokenization framework. Input speech x is encoded into latent features Z, then quantized into discrete tokens Q(1:K) via residual vector quantization (RVQ). To enrich these tokens, we incorporate semantic (Si, Ŝ) and contextual (Ci, Ĉ, C* ) representations from frozen pre-trained models. Global vectors Ŝ and Ĉ are formed via mean pooling and [CLS] selection, respectively.

We propose a speech tokenization framework with three different strategies that enrich discrete speech representations with unified and aligned semantic and contextual information:

• FuseCodec-Fusion (Latent Representation Fusion): integrates semantic (Ŝ) and contextual (Ĉ) embeddings into the encoder’s latent space through cross-modal attention and additive fusion, resulting in more robust and coherent latent features Z′.
• FuseCodec-Distill (Global Semantic-Contextual Supervision): globally pooled semantic and contextual vectors supervise each quantized token Q(1) across time, enabling temporally consistent and globally informed representation learning.
• FuseCodec-ContextAlign (Temporally Aligned Contextual Supervision): dynamically aligns contextual embeddings {Ci} with RVQ outputs via a windowed matching algorithm, enforcing timestep-level similarity supervision and enhancing fine-grained cross-modal alignment.

News & Updates

• [2025-09-16] Paper preprint released! Check out FuseCodec on arXiv.
• [2025-09-16] We release three model variations: FuseCodec-Fusion, FuseCodec-Distill, and FuseCodec-ContextAlign.
• [2025-09-16] Released the official implementation of FuseCodec.

Pre-trained Models

Model	Download Link
FuseCodec-Fusion	Download
FuseCodec-Distill	Download
FuseCodec-ContextAlign	Download

Setup

The project environment can be prepared in one of the following ways:

• Install the Python dependencies listed in requirements.txt
• Or build the Docker image using the provided Dockerfile

Training Pipeline

The FuseCodec workflow has three main stages: Dataset Preprocessing, Representation Extraction, and Training. Fusion, Distill and ContextAlign share the same pipeline; below are what each stage does and the main arguments (defaults shown). Only arguments that differ between runs are called out.

1. Dataset Preprocessing

Load audio from a Hugging Face dataset or a local folder, resample everything to 16 kHz, cut into fixed-length segments (pad if shorter) and save processed segments into processed_dataset/. This stage prepares input data for representation extraction.

Arguments:

• --dataset: Hugging Face dataset identifier (e.g., "username/dataset_name").
• --split: Dataset split to use (e.g., train).
• --ratio: Fraction of the dataset to use (1.0 = full dataset; smaller values for debugging).
• --segment_length: Target clip length in seconds (default: 3.0).
• --og_directory: (optional) path to a local folder of audio files instead of a dataset.
• --output_dir: Output directory for processed clips (default: processed_dataset).

2. Representation Extraction

Extracts multimodal representations: Semantic embeddings from a Speech Model (e.g., HuBERT), Transcriptions from a pretrained ASR model (e.g., Wav2Vec2), and Contextual embeddings from a Language Model (e.g., BERT).

Arguments:

• --config: Path to configuration file (JSON).
• --audio_dir: Directory of processed audio (default: processed_dataset).
• --exts: Allowed audio formats (comma-separated, e.g., flac,wav).
• --split_seed: Random seed for train/validation split.
• --valid_set_size: Validation set size (ratio or absolute count, e.g., 0.00400).
• --hubert_model_path: Hugging Face path of the pretrained speech model (e.g., facebook/hubert-base-ls960).
• --hubert_model_layer: Embedding layer(s) to use (e.g., avg or 12).
• --hubert_token_typ: How to broadcast speech token embeddings across time (repeat).
• --stt_model_path: Hugging Face path of the pretrained ASR model (e.g., facebook/wav2vec2-base-960h).
• --llm_model_path: Hugging Face path of the pretrained language model (e.g., bert-base-uncased).
• --llm_model_layer: Embedding layer(s) to use (e.g., avg or 12).
• --llm_token_typ: How to broadcast language token embeddings across time (cls, repeat).
• --rep_typ: Which teacher representations to extract (hubert, llm, or combined for fusion/distillation).

3. Training

Trains the codec using the extracted features. Three Fusecodec variants are supported: Fusion, Distill, and ContextAlign (choice controlled via --teacher).

Arguments:

• --config: Path to the training configuration file (e.g., config.json).
• --epochs: Number of training epochs (e.g., 100; can be reduced for debugging).
• --teacher: Supervision type and corresponding script:
  • Fusion: combined → run trainer_fusion.py
  • Distill: combined → run trainer_distill.py
  • ContextAlign: llm_align → run trainer_distill.py
• --layer: RVQ layers for fusion or supervision:
  • Fusion: 0,1,2,3,4,5,6,7
  • Distill / ContextAlign: 0
• --fusion: (Fusion only) Multimodal fusion strategy (sum or concat).
• --attention_stage: Attention mechanism (no, self, cross).
• --dropout: (Fusion only) Dropout probability for modality dropping (default: 0.1).
• --lock: Whether to restrict RVQ layers initially (unlocked = all layers used).
• --exp_dir: Directory to save checkpoints (default: saved_files).
• --continue_train: Continue training from the latest checkpoint if available.

Trainer scripts:

• Fusion: trainer_fusion.py
• Distill / ContextAlign: trainer_distill.py

Note: For an in-depth discussion of the pipeline, methodology, and design decisions, please refer to our paper.

4. Inference

Once training is complete, FuseCodec models can be used for encoding/decoding speech.
You can switch between Fusion and Distill/ContextAlign by changing the model import.

Example Usage:

import os
import torch
import torchaudio

# Choose one depending on variant:
# from codec.model_fusion import Model   # For FuseCodec-Fusion
# from codec.model_distill import Model  # For FuseCodec-Distill / FuseCodec-ContextAlign
from codec.model_fusion import Model     # ← default

# Paths
CHECKPOINT = "FuseCodec_Fusion.pt"
CONFIG = "config.json"
INPUT_DIR = "processed_dataset"
OUTPUT_DIR = "reconstructed"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

os.makedirs(OUTPUT_DIR, exist_ok=True)

model = Model.load_from_checkpoint(CONFIG, CHECKPOINT)
model.eval().to(DEVICE)

for filename in os.listdir(INPUT_DIR):
    if filename.endswith(".flac") or filename.endswith(".wav"):
        filepath = os.path.join(INPUT_DIR, filename)
        wav, sr = torchaudio.load(filepath)

        # Convert stereo → mono
        if wav.shape[0] > 1:
            wav = wav[:1, :]

        # Resample if needed
        if sr != model.sample_rate:
            wav = torchaudio.functional.resample(wav, sr, model.sample_rate)

        wav = wav.unsqueeze(0).to(DEVICE)

        # Encode → Quantize → Decode
        with torch.no_grad():
            codes = model.encode(wav) # Quantized codebook indices
            fake_wav = model.decode(codes).squeeze(0) # Reconstructed speech

        # Save output
        out_path = os.path.join(OUTPUT_DIR, "fake_" + filename)
        torchaudio.save(out_path, fake_wav.cpu().detach(), model.sample_rate, format="FLAC")
        print(f"Saved: {out_path}")

Citation

If you find FuseCodec useful in your research, please cite:

@misc{ahasan2025fusecodecsemanticcontextualfusionsupervision,
      title={FuseCodec: Semantic-Contextual Fusion and Supervision for Neural Codecs}, 
      author={Md Mubtasim Ahasan and Rafat Hasan Khan and Tasnim Mohiuddin and Aman Chadha and Tariq Iqbal and M Ashraful Amin and Amin Ahsan Ali and Md Mofijul Islam and A K M Mahbubur Rahman},
      year={2025},
      eprint={2509.11425},
      archivePrefix={arXiv},
      primaryClass={cs.SD},
      url={https://arxiv.org/abs/2509.11425}, 
}

License

This project is licensed under the Apache 2.0 License - see the LICENSE file for details.